<commit_before># TF2JAX

![CI status](https://github.com/deepmind/tf2jax/workflows/ci/badge.svg)
![pypi](https://img.shields.io/pypi/v/tf2jax)

TF2JAX is an experimental library for converting [TensorFlow] functions/graphs
to [JAX] functions.

Specifically, it aims to transform a `tf.function`, e.g.

```python
@tf.function
def tf_fn(x):
  return tf.sin(tf.cos(x))
```

to a python function equivalent to the following JAX code.

```python
def jax_fn(x):
  return jnp.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]

## Installation

TF2JAX can be installed with `pip` directly from github, with the following command:

`pip install git+git://github.com/deepmind/tf2jax.git`

or from PyPI:

`pip install tf2jax`

## Motivations

TF2JAX enables existing TensorFlow functions and models (including
[SavedModel](https://www.tensorflow.org/guide/saved_model) and
[TensorFlow Hub](https://www.tensorflow.org/hub/tf1_hub_module)) 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
via `jax2tf.convert`, back into their existing JAX codebases.

See [section](#alternatives) at the end for comparison with an alternative
approach provided by `jax2tf.call_tf`.

## Disclaimer

This 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 start

The rest of this document assumes the following imports:

```python
import jax
import jax.numpy as jnp
import numpy as np
import tensorflow as tf  # Assumes this is v2.
import tf2jax
```

An example using the `convert` API and the Sonnet v2 MLP.

```python
import sonnet.v2 as snt

model = snt.nets.MLP((64, 10,))

@tf.function
def forward(x):
  return model(x)

x = np.random.normal(size=(128, 16)).astype(np.float32)

# TF -> JAX, jax_params are the network parameters of the MLP
jax_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.convert` has the signature `convert(fn: tf.Function, *args, **kwargs)`,
where `fn(*args, **kwargs)` is used to trace the function `fn` and generates the
corresponding `tf.GraphDef`. The `zeros_like` is not necessary, only used here
to demonstrate the JAX function is not memorizing the outputs.

### Example with a pure function

If 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 with
`tf2jax.convert_functional`.

```python
@tf.function
def forward(x):
  return tf.sin(tf.cos(x))

jax_func = tf2jax.convert_functional(forward, np.zeros_like(x))
jax_outputs = jax_func(x)
```

## Randomess and PRNG Keys

A 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.

```python
jax_outputs, jax_params = jax_func(jax_params, x, rng=jax.random.PRNGKey(42))
```

## Custom Gradient

Custom gradient support is highly experimental, please report any errors.

```python
@tf.function
@tf.custom_gradient
def forward(x):
  e = tf.exp(x)
  def grad(dy):
    return dy * tf.sin(x) + e  # # This is deliberately the wrong gradient.
  return tf.reduce_sum(e), grad

with tf2jax.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 Formats

### SavedModel

[SavedModel](https://www.tensorflow.org/guide/saved_model) is the preferred
format for serializing TF2 functions.

```python
model = tf.Module()
model.f = forward
model.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.

```python
jax_func, jax_params = tf2jax.convert_from_restored(restored.f)
```

### TF-Hub

The (legacy, TF1) [TF-Hub](https://www.tensorflow.org/hub/tf1_hub_module)
format is supported with minor boilerplate.

```python
import tensorflow_hub as hub

hub_model = hub.load("/tmp/blah")
jax_func, jax_params = tf2jax.convert(tf.function(hub_model))
jax_outputs, jax_params = jax_func(x)
```

## JAX to TensorFlow and Back Again.

`tf2jax.convert_functional` can convert the outputs of `jax2tf.convert` back
into JAX code.

```python<commit_msg>Minor fix to tf_hub documentation.<commit_after># TF2JAX

![CI status](https://github.com/deepmind/tf2jax/workflows/ci/badge.svg)
![pypi](https://img.shields.io/pypi/v/tf2jax)

TF2JAX is an experimental library for converting [TensorFlow] functions/graphs
to [JAX] functions.

Specifically, it aims to transform a `tf.function`, e.g.

```python
@tf.function
def tf_fn(x):
  return tf.sin(tf.cos(x))
```

to a python function equivalent to the following JAX code.

```python
def jax_fn(x):
  return jnp.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]

## Installation

TF2JAX can be installed with `pip` directly from github, with the following command:

`pip install git+git://github.com/deepmind/tf2jax.git`

or from PyPI:

`pip install tf2jax`

## Motivations

TF2JAX enables existing TensorFlow functions and models (including
[SavedModel](https://www.tensorflow.org/guide/saved_model) and
[TensorFlow Hub](https://www.tensorflow.org/hub/tf1_hub_module)) 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
via `jax2tf.convert`, back into their existing JAX codebases.

See [section](#alternatives) at the end for comparison with an alternative
approach provided by `jax2tf.call_tf`.

## Disclaimer

This 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 start

The rest of this document assumes the following imports:

```python
import jax
import jax.numpy as jnp
import numpy as np
import tensorflow as tf  # Assumes this is v2.
import tf2jax
```

An example using the `convert` API and the Sonnet v2 MLP.

```python
import sonnet.v2 as snt

model = snt.nets.MLP((64, 10,))

@tf.function
def forward(x):
  return model(x)

x = np.random.normal(size=(128, 16)).astype(np.float32)

# TF -> JAX, jax_params are the network parameters of the MLP
jax_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.convert` has the signature `convert(fn: tf.Function, *args, **kwargs)`,
where `fn(*args, **kwargs)` is used to trace the function `fn` and generates the
corresponding `tf.GraphDef`. The `zeros_like` is not necessary, only used here
to demonstrate the JAX function is not memorizing the outputs.

### Example with a pure function

If 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 with
`tf2jax.convert_functional`.

```python
@tf.function
def forward(x):
  return tf.sin(tf.cos(x))

jax_func = tf2jax.convert_functional(forward, np.zeros_like(x))
jax_outputs = jax_func(x)
```

## Randomess and PRNG Keys

A 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.

```python
jax_outputs, jax_params = jax_func(jax_params, x, rng=jax.random.PRNGKey(42))
```

## Custom Gradient

Custom gradient support is highly experimental, please report any errors.

```python
@tf.function
@tf.custom_gradient
def forward(x):
  e = tf.exp(x)
  def grad(dy):
    return dy * tf.sin(x) + e  # # This is deliberately the wrong gradient.
  return tf.reduce_sum(e), grad

with tf2jax.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 Formats

### SavedModel

[SavedModel](https://www.tensorflow.org/guide/saved_model) is the preferred
format for serializing TF2 functions.

```python
model = tf.Module()
model.f = forward
model.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.

```python
jax_func, jax_params = tf2jax.convert_from_restored(restored.f)
```

### TF-Hub

The (legacy, TF1) [TF-Hub](https://www.tensorflow.org/hub/tf1_hub_module)
format is supported with minor boilerplate.

```python
import tensorflow_hub as hub

hub_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_functional` can convert the outputs of `jax2tf.convert` back
into JAX code.

```python
{'fraction_non_alphanumeric': 0.10549147195109929, 'fraction_numerical': 0.00867593414177265, 'mean_word_length': 4.914868804664723, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 12, 'lorem ipsum': 0, 'www.': 8, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2780042', 'n_tokens_mistral': 3442, 'n_tokens_neox': 3315, 'n_words': 1157}