How to use custom_vjp as the inputs of jitted function when using call_tf ?

[hejujie]

I want to use a TensorFlow function in JAX with jit mode, so I choose jax2tf.call_tf to translate my TensorFlow function and pass it as an input of the jitted function.
It get an error:
TypeError: Argument '<jax._src.custom_derivatives.custom_vjp object at 0x7f2c15350990>' of type <class 'jax._src.custom_derivatives.custom_vjp'> is not a valid JAX type.
How can I transform custom_vjp to a valid JAX type, or how can I use the tf function in a jitted function ?

from jax.experimental import jax2tf
import jax
import tensorflow as tf


def cos_tf(x):
    return tf.math.cos(x)

jax_cos = jax2tf.call_tf(cos_tf)
print(type(jax_cos))

def func(recurrent_net):
    "do something with recurrent_net"
    return None

jax_func = jax.jit(func)
jax_func(jax_cos)

<class 'jax._src.custom_derivatives.custom_vjp'>
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
Traceback (most recent call last):
  File "test.py", line 17, in <module>
    jax_func(jax_cos)
  File "/root/miniconda3/lib/python3.7/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/root/miniconda3/lib/python3.7/site-packages/jax/_src/api.py", line 430, in cache_miss
    _check_arg(arg)
  File "/root/miniconda3/lib/python3.7/site-packages/jax/_src/api.py", line 2816, in _check_arg
    raise TypeError(f"Argument '{arg}' of type {type(arg)} is not a valid JAX type.")
jax._src.traceback_util.UnfilteredStackTrace: TypeError: Argument '<jax._src.custom_derivatives.custom_vjp object at 0x7f2c15350990>' of type <class 'jax._src.custom_derivatives.custom_vjp'> is not a valid JAX type.

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "test.py", line 17, in <module>
    jax_func(jax_cos)
TypeError: Argument '<jax._src.custom_derivatives.custom_vjp object at 0x7f2c15350990>' of type <class 'jax._src.custom_derivatives.custom_vjp'> is not a valid JAX type.

[YouJiacheng]

You can try to use jax.tree_util.Partial to wrap your function.

Passing zero arguments to Partial effectively wraps the original function, making it a valid argument in JAX transformed functions

Alternatively, you can use static_argnums or static_argnames in jit, i.e. use jax_func = jax.jit(func, static_argnums=0)

Or you can use closure:

from jax.experimental import jax2tf
import jax
import tensorflow as tf


def cos_tf(x):
    return tf.math.cos(x)

jax_cos = jax2tf.call_tf(cos_tf)

def make_jax_func(recurrent_net):
    def func(x): # I guess you actually need some inputs here, not just recurrent_net
        return recurrent_net(x) + x # do sth with recurrent_net
    return jax.jit(func)

jax_func = make_jax_func(jax_cos)
jax_func(1.0)

BTW, <class 'jax._src.custom_derivatives.custom_vjp'> actually refer to "functions with custom vjp".

[hejujie]

Thank you so much!