Support tf2jax.convert in TensorFlow tracing contexts.

Adds a new configuration option `skip_variables_evaluation_inside_tf_tracing` which users can set to True to enable `tf2jax.convert` to work in TensorFlow 2 tracing contexts. Fore example, when exporting a JAX model as a TF saved model, where that JAX model uses `tf2jax.convert` internally.

PiperOrigin-RevId: 825829137

Author: rryan

tf2jax/_src/config.py

@@ -39,6 +39,11 @@
     # converting `tensorflow_probability`'s `HiddenMarkovModel.posterior_mode`
     # to JAX. Here we allow to disable it.
     disable_assert_in_tensor_list_get_item=False,
+    # If True, avoids evaluating variables in a TF tracing context. Set to False
+    # when using tf2jax.convert from within a tf.function in TF2 mode (e.g. when
+    # exporting a JAX function using tf2jax.convert internally as a TensorFlow
+    # saved model).
+    skip_variables_evaluation_inside_tf_tracing=False,
 )
 
 

tf2jax/_src/tf2jax.py

@@ -537,7 +537,7 @@ def maybe_tensor_to_spec(v):
     exp_args, exp_kwargs = structured_inputs
     if exp_args:
       raise ValueError("If function_spec is None then only keyword arguments "
-                       f"are expectd, found args={exp_args} in structure.")
+                       f"are expected, found args={exp_args} in structure.")
     parameters = tuple([
         # TODO(b/266552275) Remove temporary fix for TF-Hub.
         inspect.Parameter(
@@ -999,7 +999,7 @@ def _convert(
     constants: A mapping from tensor names to constant values. The keys are a
       subset of captured_input_names.
     library: A mapping from function names to Callable. This is non-empty on
-      recurisve calls if the FunctionDefLibrary in the GraphDef is non-empty.
+      recursive calls if the FunctionDefLibrary in the GraphDef is non-empty.
 
   Returns:
     A tuple: the first is a Jax functions that takes a flat parameter dict and a
@@ -1068,7 +1068,7 @@ def _convert(
     raise ValueError(err_message)
 
   # Extract variables.
-  if tf.executing_eagerly():
+  if tf.executing_eagerly() or tf.inside_function():
     # Uniqueify variables with identical names.
     variables_tf = {}
     var_name_by_ref = {}
@@ -1082,11 +1082,23 @@ def _convert(
       variables_tf[var_name] = v
       var_name_by_ref[v.ref()] = var_name
 
-    variables = {
-        k: Variable(v.numpy(), v.trainable, v.name)
-        for k, v in variables_tf.items()
-    }
-  else:
+    if tf.inside_function():
+      # We cannot evaluate variables inside a tf.function.
+      variables = {}
+      if (
+          not config.get_config("skip_variables_evaluation_inside_tf_tracing")
+          and variable_map
+      ):
+        raise ValueError(
+            "Unable to to evaluate variables inside a TF tracing context, "
+            "and `skip_variables_evaluation_inside_tf_tracing` is False."
+        )
+    else:
+      variables = {
+          k: Variable(v.numpy(), v.trainable, v.name)
+          for k, v in variables_tf.items()
+      }
+  else:  # We are in TF1 mode
     variables_tf = {_parse_input(v.name): v for _, v in variable_map.items()}
     var_name_by_ref = {
         v.ref(): _parse_input(v.name) for v in variable_map.values()
@@ -1103,7 +1115,10 @@ def _convert(
 
   assert len(variable_map) == len(variables_tf)
   assert len(variable_map) == len(var_name_by_ref)
-  assert len(variable_map) == len(variables)
+  if not tf.inside_function() or not config.get_config(
+      "skip_variables_evaluation_inside_tf_tracing"
+  ):
+    assert len(variable_map) == len(variables)
 
   var_by_node = {k: var_name_by_ref[v.ref()] for k, v in variable_map.items()}
   node_by_var = {v: k for k, v in var_by_node.items()}

tf2jax/_src/tf2jax_test.py

@@ -17,13 +17,13 @@
 import inspect
 
 from absl.testing import parameterized
-
 import chex
+import flax
 from flax import linen as nn
 import jax
+from jax.experimental import jax2tf
 import jax.numpy as jnp
 import numpy as np
-
 import tensorflow as tf
 from tf2jax._src import config
 from tf2jax._src import tf2jax
@@ -63,6 +63,34 @@ def __call__(self, x):
     return x
 
 
+class FlaxTestDense(nn.Module):
+  """A Flax module that wraps a TestDense saved model."""
+
+  saved_model_path: str
+  input_dim: int
+
+  def setup(self):
+    self.m = tf.saved_model.load(self.saved_model_path)
+    if self.is_initializing():
+      _, params = tf2jax.convert(
+          tf.function(self.m.__call__),
+          tf.TensorSpec([None, self.input_dim], tf.float32),
+      )
+      params = flax.traverse_util.unflatten_dict(params, sep="/")
+      self.params = self.param("saved_model", lambda _: params)
+    else:
+      self.params = self.param("saved_model", lambda _: None)
+
+  def __call__(self, x: jax.Array) -> jax.Array:
+    fn, _ = tf2jax.convert(
+        tf.function(self.m.__call__),
+        tf.TensorSpec(x.shape, tf.float32),
+    )
+    params = flax.traverse_util.flatten_dict(self.params, sep="/")
+    y, _ = fn(params, x)
+    return y
+
+
 class FeaturesTest(tf.test.TestCase, parameterized.TestCase):
 
   def _setup_saved_model(self, *inputs):
@@ -494,6 +522,71 @@ def tf_func(x):
     ):
       self.variant(jax_func)({"aaa": jax_params["aaa"]}, np_inputs)
 
+  def test_export_saved_model_export_jax_module(self):
+    input_dim = 5
+    l = TestDense(input_dim=input_dim, output_size=5)
+    l.__call__ = tf.function(
+        l.__call__, input_signature=[tf.TensorSpec((None, input_dim))]
+    )
+
+    x_tf = tf.ones((2, input_dim))
+    y_tf = l(x_tf)
+    kernel_tf, bias_tf = l.variables
+    tf_export_dir = self.create_tempdir()
+    jax_export_dir = self.create_tempdir()
+    tf.saved_model.save(l, tf_export_dir)
+
+    module = FlaxTestDense(tf_export_dir, input_dim=input_dim)
+    x_jax = jnp.ones((2, input_dim))
+    y_jax, jax_params = module.init_with_output(jax.random.key(0), x_jax)
+
+    kernel_jax, bias_jax = jax.tree.leaves(jax_params)
+
+    np.testing.assert_array_equal(kernel_tf.numpy(), kernel_jax)
+    np.testing.assert_array_equal(bias_tf.numpy(), bias_jax)
+    np.testing.assert_allclose(y_tf.numpy(), y_jax, atol=1e-6, rtol=1e-6)
+
+    def model_fn(params, inputs):  # The JAX model function to export.
+      fm = FlaxTestDense(tf_export_dir, input_dim=input_dim)
+      return fm.apply(params, inputs)
+
+    class ExportWrapper(tf.Module):
+
+      def __init__(self, params):
+        super().__init__()
+        self._tf_fn = jax2tf.convert(model_fn)
+        self._params = params
+
+      @tf.function(
+          input_signature=[
+              tf.TensorSpec(shape=[1, input_dim], dtype=tf.float32)
+          ]
+      )
+      def __call__(self, x):
+        return self._tf_fn(self._params, x)
+
+    with config.override_config(
+        "skip_variables_evaluation_inside_tf_tracing", False
+    ), self.assertRaises(ValueError):
+      tf.saved_model.save(
+          ExportWrapper(jax_params),
+          jax_export_dir,
+          options=tf.saved_model.SaveOptions(
+              experimental_custom_gradients=False
+          ),
+      )
+
+    with config.override_config(
+        "skip_variables_evaluation_inside_tf_tracing", True
+    ):
+      tf.saved_model.save(
+          ExportWrapper(jax_params),
+          jax_export_dir,
+          options=tf.saved_model.SaveOptions(
+              experimental_custom_gradients=False
+          ),
+      )
+
 
 if __name__ == "__main__":
   tf.test.main()