lucidrains / jax2torch Goto Github PK

Hello,

My understanding was that Jax preallocates ~90% of the GPU memory, and setting XLA_PYTHON_CLIENT_PREALLOCATE to false is needed when interoping with other frameworks in order to avoid this. However, this isn't mentioned anywhere here. I figure this should probably be explicitly mentioned?

The SMURF repo does set os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false", so I assume it is necessary?

When I run this on PyTorch 2.1 and CUDA 11.8, I get this error:
Traceback (most recent call last):
File "", line 12, in
File "/home/nzubic/anaconda3/envs/events_signals/lib/python3.11/site-packages/jax2torch/jax2torch.py", line 54, in inner
return JaxFun.apply(*bound.arguments.values())
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/nzubic/anaconda3/envs/events_signals/lib/python3.11/site-packages/torch/autograd/function.py", line 542, in apply
raise RuntimeError(
RuntimeError: In order to use an autograd.Function with functorch transforms (vmap, grad, jvp, jacrev, ...), it must override the setup_context staticmethod. For more details, please see https://pytorch.org/docs/master/notes/extending.func.html

see issue posted here:
google/jax#14399

I think it might be a jax bug, not sure if it's jax2torch?