Repository: google/jax
Issue: jax2tf does not preserve argument names

When converting a function from JAX to TensorFlow using experimental `jax2tf.convert`, the function arguments are not preserved, thus it is not possible to call the code with named instead of positional arguments.

For example, the following snippet will break:

```python
    def f_jax(x):
        return jnp.sum(x)
    
    f_tf = jax2tf.convert(f_jax)

    x = jnp.zeros(3)
    print(f_tf(x=x))
```

```
Traceback (most recent call last):
  File "test.py", line 8, in <module>
    print(f_tf(x=x))
TypeError: converted_fun() got an unexpected keyword argument 'x'
```

My workaround so far has been to wrap the converted function around another function, and then proceed to export the saved model. 

```python
    def f_jax(x):
        return jnp.sum(x)
    
    f_tf = jax2tf.convert(f_jax)
    f_wrap = lambda x: f_tf(x)

    x = jnp.zeros(3)
    print(f_wrap(x=x))
```

Since jax2tf is experimental, I expected to find some issues. However, this one is not listed as a [known issue](https://github.com/google/jax/tree/master/jax/experimental/jax2tf#known-issues). An important use case for this is to allow models with multiple inputs to be called via Tensorflow Serving.

Changed files:
jax/experimental/jax2tf/jax2tf.py, jax/experimental/jax2tf/tests/jax2tf_test.py

diff --git a/jax/experimental/jax2tf/jax2tf.py b/jax/experimental/jax2tf/jax2tf.py
index fafea2703..dab5b32db 100644
--- a/jax/experimental/jax2tf/jax2tf.py
+++ b/jax/experimental/jax2tf/jax2tf.py
@@ -170,6 +170,9 @@ def convert(fun: Callable, *,
       The conversion fails if it cannot ensure that the it would produce the same
       sequence of TF ops for any non-zero values of the dimension variables.
 
+      polymorphic_shapes are only supported for positional arguments; shape
+      polymorphism is not supported for keyword arguments.
+
       See [the README](https://github.com/google/jax/blob/master/jax/experimental/jax2tf/README.md#shape-polymorphic-conversion)
       for more details.
 
@@ -192,7 +195,7 @@ def convert(fun: Callable, *,
   """
   api._check_callable(fun)
 
-  def converted_fun(*args: TfVal) -> TfVal:
+  def converted_fun(*args: TfVal, **kwargs: TfVal) -> TfVal:
     # TODO: is there a better way to check if we are inside a transformation?
     if not core.trace_state_clean():
       raise ValueError("convert must be used outside all JAX transformations."
@@ -204,21 +207,23 @@ def convert(fun: Callable, *,
                "be NumPy array, scalar, tf.Variable, or tf.Tensor")
         raise TypeError(msg)
     tree_util.tree_map(check_arg, args)
+    tree_util.tree_map(check_arg, list(kwargs.values()))
 
     # Name input tensors
     args = tuple(
         tree_util.tree_map(lambda x, i=i: tf.identity(x, f"jax2tf_arg_{i}"), a)  # type: ignore
         for i, a in enumerate(args))
+    kwargs = {k: tf.identity(v, f"jax2tf_arg_{k}") for k, v in kwargs.items()}
 
     # This function may take pytrees of TfVals. We can only set
     # tf.custom_gradient on functions that take a flat argument list.
-    args_flat, in_tree = tree_util.tree_flatten((args, {}))
+    args_flat, in_tree = tree_util.tree_flatten((args, kwargs))
 
     if polymorphic_shapes is None:
       polymorphic_shapes_ = (None,) * len(args)
     else:
       if not isinstance(polymorphic_shapes, Sequence) or len(args) != len(polymorphic_shapes):
-        msg = ("polymorphic_shapes must be a sequence with the same length as the argument list "
+        msg = ("polymorphic_shapes must be a sequence with the same length as the positional argument list "
                f"({len(args)}). Got polymorphic_shapes={polymorphic_shapes}.")
         raise TypeError(msg)
       polymorphic_shapes_ = tuple(polymorphic_shapes)
@@ -227,6 +232,9 @@ def convert(fun: Callable, *,
     polymorphic_shapes_flat = tuple(api_util.flatten_axes("jax2tf.convert polymorphic_shapes",
                                                           in_tree.children()[0],
                                                           polymorphic_shapes_))
+    # Add kwargs shapes.
+    polymorphic_shapes_flat = polymorphic_shapes_flat + tuple(
+      (None,) * (len(args_flat) - len(polymorphic_shapes_flat)))
 
     # Construct the abstract values for the flat arguments, possibly based on
     # the input shapes and the polymorphic_shapes if given. May create new shape
diff --git a/jax/experimental/jax2tf/tests/jax2tf_test.py b/jax/experimental/jax2tf/tests/jax2tf_test.py
index e3d8cee59..296805f21 100644
--- a/jax/experimental/jax2tf/tests/jax2tf_test.py
+++ b/jax/experimental/jax2tf/tests/jax2tf_test.py
@@ -517,6 +517,15 @@ class Jax2TfTest(tf_test_util.JaxToTfTestCase):
         tf_fn_array(np.array([3, 4, 5])), np.array([4.5, 10, 17.5],
                                                    jnp.bfloat16))
 
+  def test_kwargs(self):
+    # Re: https://github.com/google/jax/issues/6791
+    def f_jax(*, x):
+      return jnp.sum(x)
+    f_tf = jax2tf.convert(f_jax)
+    self.assertAllClose(
+      f_tf(x=np.zeros(3, dtype=np.float32)),  # Call with kwargs.
+      np.zeros((), dtype=np.float32))
+
   def test_enable_xla(self):
     # Tests that enable_xla flag is properly scoped to a conversion.
     def fun(x):