currentframe is only available in CPython

[behrica]

https://github.com/google/trax/blob/b8acc0e8f5c52a5ac25fd98a7a196f13b5fc1f69/trax/layers/base.py#L127

See here https://docs.python.org/3/library/inspect.html#inspect.currentframe