Package: wrap-torch2jax
NOTE:wrap_torch2jaxis a pip alias fortorch2jaxtorch2jaxDocumentationThis package is designed to facilitate no-copy PyTorch calling from JAX under
both eager execution and JIT. It leverages the JAX C++ extension interface,
enabling operations on both CPU and GPU platforms. Moreover, it allows for
executing arbitrary PyTorch code from JAX under eager execution and JIT.The intended application is efficiently running existing PyTorch code (like ML
models) in JAX applications with very low overhead.This project was inspired by the jax2torch repositoryhttps://github.com/lucidrains/jax2torchand has been made possible due to an amazing tutorial on extending JAXhttps://github.com/dfm/extending-jax.
Comprehensive JAX documentationhttps://github.com/google/jaxalso
significantly contributed to this work.Although I am unsure this functionality could be achieved without C++/CUDA, the
C++ compilation is efficiently done using PyTorch's portable CUDA & C++
compilation features, requiring minimal configuration.Install$pipinstallgit+https://github.com/rdyro/torch2jax.gitUsageWith a single outputimporttorchimportjaxfromjaximportnumpyasjnpfromwrap_torch2jaximporttorch2jax# this converts a Python function to JAXfromwrap_torch2jaximportSize,dtype_t2j# this is torch.Size, a tuple-like shape representationdeftorch_fn(a,b):returna+bshape=(10,2)a,b=torch.randn(shape),torch.randn(shape)jax_fn=torch2jax(torch_fn,a,b)# without output_shapes, torch_fn **will be evaluated once**jax_fn=torch2jax(torch_fn,a,b,output_shapes=Size(a.shape))# torch_fn will NOT be evaluated# you can specify the whole input and output structure without instantiating the tensors# torch_fn will NOT be evaluatedjax_fn=torch2jax(torch_fn,jax.ShapeDtypeStruct(a.shape,dtype_t2j(a.dtype)),jax.ShapeDtypeStruct(b.shape,dtype_t2j(b.dtype)),output_shapes=jax.ShapeDtypeStruct(a.shape,dtype_t2j(a.dtype)),)prngkey=jax.random.PRNGKey(0)device=jax.devices("cuda")[0]# both CPU and CUDA are supporteda=jax.device_put(jax.random.normal(prngkey,shape),device)b=jax.device_put(jax.random.normal(prngkey,shape),device)# call the no-copy torch functionout=jax_fn(a,b)# call the no-copy torch function **under JIT**out=jax.jit(jax_fn)(a,b)With a multiple outputsdeftorch_fn(a,b):layer=torch.nn.Linear(2,20).to(a)returna+b,torch.norm(a),layer(a*b)shape=(10,2)a,b=torch.randn(shape),torch.randn(shape)jax_fn=torch2jax(torch_fn,a,b)# with example argumetnsprngkey=jax.random.PRNGKey(0)device=jax.devices("cuda")[0]a=jax.device_put(jax.random.normal(prngkey,shape),device)b=jax.device_put(jax.random.normal(prngkey,shape),device)# call the no-copy torch functionx,y,z=jax_fn(a,b)# call the no-copy torch function **under JIT**x,y,z=jax.jit(jax_fn)(a,b)For a more advanced discussion on different ways of specifying input/output
specification of the wrapped function, take a look at:input_output_specification.ipynbnotebook in theexamplesfolder.Automatically defining gradientsAutomatic reverse-mode gradient definitions are now supported for wrapped
pytorch functions with the methodtorch2jax_with_vjpimporttorchimportjaxfromjaximportnumpyasjnpimportnumpyasnpfromwrap_torch2jaximporttorch2jax_with_vjpdeftorch_fn(a,b):returntorch.nn.MSELoss()(a,b)shape=(6,)xt,yt=torch.randn(shape),torch.randn(shape)# `depth` determines how many times the function can be differentiatedjax_fn=torch2jax_with_vjp(torch_fn,xt,yt,depth=2)# we can now differentiate the function (derivatives are taken using PyTorch autodiff)g_fn=jax.grad(jax_fn,argnums=(0,1))x,y=jnp.array(np.random.randn(*shape)),jnp.array(np.random.randn(*shape))print(g_fn(x,y))# JIT works tooprint(jax.jit(g_fn)(x,y))Caveats:jax.hessian(f)will not work sincetorch2jaxuses forward differentiation, but
the same functionality can be achieved usingjax.jacobian(jax.jacobian(f))input shapes are fixed for one wrapped function and cannot change, usetorch2jax_with_vjp/torch2jaxagain if you need to alter the input shapesin line with JAX philosphy, PyTorch functions must be non-mutable,torch.funchas a good description
of how to convert e.g., PyTorch models, to non-mutable formulationDealing with Changing ShapesYou can deal with changing input shapes by callingtorch2jax(andtorch2jax_with_vjp) in the JAX function, both under JIT and eagerly!@jax.jitdefcompute(a,b,c):d=torch2jax_with_vjp(torch_fn,jax.ShapeDtypeStruct(a.shape,dtype_t2j(a.dtype)),jax.ShapeDtypeStruct(b.shape,dtype_t2j(b.dtype)),output_shapes=jax.ShapeDtypeStruct(a.shape,dtype_t2j(a.dtype)),)(a,b)returnd-cprint(compute(a,b,a))Timing Comparison vspure_callbackThis package achieves a much better performance when calling PyTorch code from
JAX because it does not copy its input arguments and does not move CUDA data off
the GPU.Current Limitations oftorch2jaxcompilation happens on module import and can take 1-2 minutes (it will be cached afterwards)in the Pytorch function all arguments must be tensors, all outputs must be tensorsall arguments must be on the same device and of the same datatype, either float32 or float64an input/output shape (e.g.output_shapes=kw argument) representations (for
flexibility in input and output structure) must be wrapped intorch.Sizeorjax.ShapeDtypeStructthe current implementation does not support batching, that's on the roadmapthe current implementation does not define the VJP rule, in current design, this has to be done in
PythonChangelogno version changeadded helper scriptinstall_package_aliased.pyto automatically install
the package with a different name (to avoid a name conflict)version 0.4.7support for newest JAX (0.4.17) with backwards compatibility maintainedcompilation now delegated to python version subfolders for multi-python systemsversion 0.4.6bug-fix: cuda stream is now synchronized before and after a torch call explicitly to
avoid reading unwritten dataversion 0.4.5torch2jax_with_vjpnow automatically selectsuse_torch_vjp=Falseif theTruefailsbug-fix: cuda stream is now synchronized after a torch call explicitly to
avoid reading unwritten dataversion 0.4.4introduced ause_torch_vjp(defaulting to True) flag intorch2jax_with_vjpwhich
can be set to False to use the oldtorch.autograd.gradfor taking
gradients, it is the slower method, but is more compatibleversion 0.4.3added a note in README about specifying input/output structure without instantiating dataversion 0.4.2addedexamples/input_output_specification.ipynbshowing how input/output
structure can be specifiedversion 0.4.1bug-fix: intorch2jax_with_vjp, nondiff arguments were erroneously memorizedversion 0.4.0added batching (vmap support) usingtorch.vmap, this makesjax.jacobianworkrobustified support for gradientsadded mixed type arguments, including support for float16, float32, float64 and integer typesremoved unnecessary torch function calls in defining gradientsadded an example of wrapping a BERT model in JAX (with weights modified from JAX),examples/bert_from_jax.ipynbversion 0.3.0added a beta-version of a new wrapping methodtorch2jax_with_vjpwhich
allows recursively defining reverse-mode gradients for the wrapped torch
function that works in JAX both normally and under JITversion 0.2.0arbitrary input and output structure is now allowedremoved the restriction on the number of arguments or their maximum dimensionold interface is available viatorch2jax.compat.torch2jaxversion 0.1.2full CPU only version support, selected viatorch.cuda.is_available()bug-fix: compilation should now cache properlyversion 0.1.1bug-fix: functions do not get overwritten, manual fn id parameter replaced with automatic id generationcompilation caching is now betterversion 0.1.0first working version of the packageRoadmapcall PyTorch functions on JAX data without input data copycall PyTorch functions on JAX data without input data copy under jitsupport both GPU and CPU(feature) support partial CPU building on systems without CUDA(user-friendly) support functions with a single output (return a single output, not a tuple)(user-friendly) support arbitrary argument input and output structure (use pytrees on the
Python side)(feature) support batching (e.g., support forjax.vmap)(feature) support integer input/output types(feature) support mixed-precision arguments in inputs/outputs(feature) support defining VJP for the wrapped function (import the experimental functionality
fromjit-JAXFriendlyInterface)(tests) test how well device mapping works on multiple GPUs(tests) setup automatic tests for multiple versions of Python, PyTorch and JAX(feature) look into supporting in-place functions (support for output without copy)(feature) support TPURelated WorkOur Python package wraps PyTorch code as-is (so custom code and mutating code
will work!), but if you're looking for an automatic way to transcribe a
supported subset of PyTorch code to JAX, take a look athttps://github.com/samuela/torch2jax/tree/main.We realize that two packages named the same is not ideal. As we work towards a
solution, here's a stop-gap solution. We offer a helper script to install the
package with an alias name, installing our package using pip under a different
name.$ git clone https://github.com/rdyro/torch2jax.git- clone this repo$ python3 install_package_aliased.py new_name_torch2jax --install --test- install and test this package under the namenew_name_torch2jaxyou can now use this package under the namenew_name_torch2jax