Package: tf2jax
TF2JAXTF2JAX is an experimental library for convertingTensorFlowfunctions/graphs
toJAXfunctions.Specifically, it aims to transform atf.function, e.g.@tf.functiondeftf_fn(x):returntf.sin(tf.cos(x))to a python function equivalent to the following JAX code.defjax_fn(x):returnjnp.sin(jnp.cos(x))Users are able to apply additional JAX transforms (e.g.jit,grad,vmap,make_jaxpr, etc.) to the converted function as they would any other code
written in JAX.[TOC]InstallationYou can install the latest released version of TF2JAX from PyPI via:pipinstalltf2jaxor you can install the latest development version from GitHub:pipinstallgit+https://github.com/deepmind/tf2jax.gitMotivationsTF2JAX enables existing TensorFlow functions and models (includingSavedModelandTensorFlow Hub) to be reused
and/or fine-tuned within JAX codebases. The conversion process is transparent
to the users, which is useful for debugging and introspection.This also provide a pathway for JAX users to integrate JAX functions serialized
viajax2tf.convert, back into their existing JAX codebases.Seesectionat the end for comparison with an alternative
approach provided byjax2tf.call_tf.DisclaimerThis is experimental code with potentially unstable API, and there are no
guarantees for using it at this point in time. We highly recommend you
thoroughly test the resulting JAX functions to ensure they meet your
requirements.Quick startThe rest of this document assumes the following imports:importjaximportjax.numpyasjnpimportnumpyasnpimporttensorflowastf# Assumes this is v2.importtf2jaxAn example using theconvertAPI and the Sonnet v2 MLP.importsonnet.v2assntmodel=snt.nets.MLP((64,10,))@tf.functiondefforward(x):returnmodel(x)x=np.random.normal(size=(128,16)).astype(np.float32)# TF -> JAX, jax_params are the network parameters of the MLPjax_func,jax_params=tf2jax.convert(forward,np.zeros_like(x))# Call JAX, also return updated jax_params (e.g. variable, batchnorm stats)jax_outputs,jax_params=jax_func(jax_params,x)tf2jax.converthas the signatureconvert(fn: tf.Function, *args, **kwargs),
wherefn(*args, **kwargs)is used to trace the functionfnand generates the
correspondingtf.GraphDef. Thezeros_likeis not necessary, only used here
to demonstrate the JAX function is not memorizing the outputs.Example with a pure functionIf your function is pure, i.e. it does not capture any variables, then you can
drop the parameters from the inputs and outputs of the converted function withtf2jax.convert_functional.@tf.functiondefforward(x):returntf.sin(tf.cos(x))jax_func=tf2jax.convert_functional(forward,np.zeros_like(x))jax_outputs=jax_func(x)Randomness and PRNG KeysA TensorFlow function that make use of random ops will be converted to a JAX
function that takes a PRNG key as a keyword-only argument. TF2JAX will
complain loudly if a PRNG key is required but not provided.jax_outputs,jax_params=jax_func(jax_params,x,rng=jax.random.PRNGKey(42))Custom GradientCustom gradient support is highly experimental, please report any errors.@tf.function@tf.custom_gradientdefforward(x):e=tf.exp(x)defgrad(dy):returndy*tf.sin(x)+e# # This is deliberately the wrong gradient.returntf.reduce_sum(e),gradwithtf2jax.override_config("convert_custom_gradient",True):jax_func=tf2jax.convert_functional(forward,np.zeros_like(x))jax_grads=jax.grad(jax_func)(x)Support for Serialization FormatsSavedModelSavedModelis the preferred
format for serializing TF2 functions.model=tf.Module()model.f=forwardmodel.f(x)# Dummy call.tf.saved_model.save(model,"/tmp/blah")restored=tf.saved_model.load("/tmp/blah")jax_func,jax_params=tf2jax.convert(restored.f,np.zeros_like(x))If the restored function has an unambiguous signature, i.e. it was only traced
once prior to export. Then TF2JAX can convert the function directly from its
GraphDef without tracing it again.jax_func,jax_params=tf2jax.convert_from_restored(restored.f)TF-HubThe (legacy, TF1)TF-Hubformat is supported with minor boilerplate.importtensorflow_hubashubhub_model=hub.load("/tmp/blah")jax_func,jax_params=tf2jax.convert(tf.function(hub_model),tf.zeros_like(x))jax_outputs,updated_jax_params=jax_func(jax_params,x)JAX to TensorFlow and back again.tf2jax.convert_functionalcan convert the outputs ofjax2tf.convertback
into JAX code.# Some JAX function.defforward(*inputs):...# JAX -> TFtf_func=jax2tf.convert(forward)# JAX -> TF -> JAXjax_func=tf2jax.convert_functional(tf.function(tf_func),*tree.map_structure(np.zeros_like,inputs))# JAX -> TF -> SavedModel -> TFmodel=tf.Module()model.f=tf.function(tf_func)model.f(*tree.map_structure(tf.zeros_like,inputs))# Dummy call.tf.saved_model.save(model,"/tmp/blah")restored=tf.saved_model.load("/tmp/blah")# JAX -> TF -> SavedModel -> TF -> JAXjax_too_func=tf2jax.convert_functional(restored.f,*tree.map_structure(np.zeros_like,inputs))Additional ConfigurationThe behaviour of TF2JAX can be configured globally viatf2jax.update_config,
or configured locally via the context managertf2jax.override_config.Strict shape and dtype checkingBy default, TF2JAX will assert that the input shapes to the converted function
are compatible with the input shapes of the original function. This is because
some functions have shape dependent behaviours that will silently return the
incorrect outputs after conversion, e.g. some batchnorm implementation.jax_func=tf2jax.convert_functional(forward,np.zeros((10,5),np.float32))# This will raise an error.jax_func(np.zeros((20,5),np.float32))# This will not.withtf2jax.override_config("strict_shape_check",False):jax_func(np.zeros((20,5),np.float32))The input dtypes are not currently checked but this may change in the future.jax_func=tf2jax.convert_functional(forward,np.zeros((10,5),np.float32))# This will not raise an error.jax_func(np.zeros((20,5),np.int32))# This will.withtf2jax.override_config("strict_dtype_check",True):jax_func(np.zeros((20,5),np.int32))Convert constants to bfloat16TF2JAX allows users to trace the converted function with parameters and inputs
of different precision than the original function, e.g.bfloat16instead offloat32. To aid this, the configurationforce_const_float32_to_bfloat16andforce_const_float64_to_bfloat16can be used to force float constants in
the original function intobfloat16precision, to avoid accidental type
promotion.@tf.functiondefforward(x):returnx+tf.constant(3.14,dtype=tf.float32)withtf2jax.override_config("force_const_float32_to_bfloat16",True):jax_func=tf2jax.convert_functional(forward,np.zeros_like(x))jax_bf16_outputs=jax_func(jnp.asarray(x,jnp.bfloat16))Disable PreventGradientIfjax2tf.convert(..., with_gradient=False)is used to produce the initial TF
function (and possibly exported as SavedModel), then TF2JAX will respect the
insertedtf.raw_ops.PreventGradientops and raiseLookupErrorwhen computing
gradients.This can be disabled by setting the configurationraise_on_prevent_gradientto
false (default is true), so that TF2JAX will only log a warning but otherwise
allow the gradient to be computed as though thePreventGradientops were not
present.@tf.functiondefprevent(x):returntf.raw_ops.PreventGradient(input=x*x,message="prevented")jax_func=tf2jax.convert_functional(prevent,0.0)jax.grad(jax_func)(3.14)# Raise LookupError.withtf2jax.config.override_config("raise_on_prevent_gradient",False):jax_func=tf2jax.convert_functional(prevent,0.0)g=jax.grad(jax_func)(3.14)# Returns 6.28.Infer Cumulative ReductionsIf theinfer_cumulative_reduction_from_jax2tfflag is true (default) then
TF2JAX will attempt to infercummax,cummin,cumprodandcumsumoperations fromreduce_windowoperations generated by JAX2TF. This provides
better performance becausereduce_windowimplementation of these ops have
O(N^2) complexity on CPU and GPU backends, and can suffer from long compilation
times due to aggressive constant folding.Seejax2tf_cumulative_reductionfor more context.JAX2TF Native Serialization and XlaCallModule.From JAX v0.4.7 and onward,jax2tf.convertpreferred mode of operation (soon
to be default) isnative serializationin which the target function is
lowered toStableHLOand wrapped in a single TensorFlow op,XlaCallModule.The new native serialization format will more faithfully reproduce the semantics
of the target function, at the cost of some reduced flexibility for downstream
processing as the computation graph is no longer exposed as a tf.Graph.XlaCallModuleis supported by TF2JAX from v0.3.4 and onward.However as this makes use of a custom JAX primitive that aims to encapsulateStableHLOpayload found inXlaCallModule, it does not possess JAX rules for
transformations such as (but not limited to) batching and differentiation.Differentiation: first order derivative of serialized function is
still supported through custom gradients requested at serialization time withjax2tf.convert(..., with_gradient=True). This is the default behaviour.Batching:jax.vmapwill fail, though users may be able to naively
replicate the desired behavior withjax.lax.map, albeit with poorer
performance.Platform SpecificityNatively serialized JAX programs are platform specific (link). Executing a natively
serialized program on platforms other than the one for which it was lowered,
would raise a ValueError, e.g.:ValueError:Unsupportedbackend:`cpu`notin`('tpu',)`.This matches the behaviour ofXlaCallModule.Users can disable this check via a config flag, but the resulting program may
be slower or fail to execute completely.withtf2jax.override_config("xlacallmodule_strict_checks",False):jax_func(np.zeros((20,5),np.float32))LimitationsCurrently, only a subset of TensorFlow ops are supported, and not necessarily
all functionalities are supported for some ops. The code will fail fast. Support
for additional TensorFlow ops are added on a as-needed basis. Please submit your
requests via Github issues or send in your pull requests.There will likely to be some cases where the resulting JAX code is not
equivalent to the TensorFlow code, both in terms of performance and numerical
outputs. The goal is to minimise differences in the latter for model endpoints,
ahead of improving performance.TF2 control flows are supported with some limitations, e.g. for while loops,
thecondandbodyfunctions cannot have side effects such as assigning to
variables.TF1 control flows are not supported.Alternativesjax2tf.call_tfjax2tfnow also offers the experimentalcall_tffunction which allows JAX to
call TensorFlow functions. For compiled code, this works by staging out
TensorFlow functions to XLA.From thejax2tf documentation, as of 2022-07-22:The functioncall_tfallows JAX functions to call TensorFlow functions.
These functions can be called anywhere in a JAX computation, including in
staging contextsjax.jit,jax.pmap,jax.xmap, or inside JAX's
control-flow primitives. In non-staging contexts, the TensorFlow function is
called in eager mode. For now, only reverse-mode autodiff is supported for
these functions (no forward-mode autodiff, norvmap).The advantage ofcall_tfis that it implicitly covers all TensorFlow ops and
supportscustom_gradientby deferring to TensorFlow during eager execution and
to XLA for compiled code.The disadvantage is that it only supports a limited set of JAX transforms
(jit,grad,pmap,remat) and otherwise appears as a "black box" to
JAX (e.g.vmapis not supported, nor custom transforms). A TensorFlow function
must be compileable to XLA if it is to be jitted aftercall_tf.Citing TF2JAXThis repository is part of theDeepMind JAX Ecosystem, to cite TF2JAX please
use theDeepMind JAX Ecosystem citation.ContributingWe are happy to receive pull requests that improve our coverage of TensorFlow
ops.