[jax2tf] train=True/False argument in jax2tf converted functions

[gnecula]

(posted on behalf of a user)

I was testing jax2tf (works very well btw. I'm able to save a model trained in JAX and continue training in TF) and I could not figure out how to convert and save the fwd pass of a model with a dropout layer.

The fwd pass function has a train=True/False switch to control the behavior of the dropout layer. But when I convert it with jax2tf.convert and call it, I get:

ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected

Do I have to convert this function twice, once with train=True and once with Train=False ? Is that the only way?

[gnecula]

It helps to think of jax2tf.convert having the same limitations as jax.jit. Hence, you have two options, similar to the options you'd have if you want to jit the function in JAX:

• in most cases, when the train parameter is used in Python conditionals, you have to convert the function twice (just like you'd have to jit it twice, either explicitly, or implicitly by using the static_argnums parameter for jax.jit). Otherwise, you will get the ConcretizationError (both from jax.jit and jax2tf.convert). For example,

    def f_jax_1(x, train):
      if train:  # Note the Python conditional
        return x
      else:
        return x * 2.

    # jax.jit(f_jax_1)(2., True)  # --> Concretization error
    print(jax.jit(lambda x: f_jax_1(x, True))(2.))  # --> 2.0
    print(jax.jit(lambda x: f_jax_1(x, False))(2.))  # --> 4.0

    # jax2tf.convert(f_jax_1)(2., True)  # Concretization error
    print(jax2tf.convert(lambda x: f_jax_1(x, True))(2.))  # --> 2.0

• if however you use the train parameter only in numerical computation or with lax.cond, then you can convert the function
  once and pass the train argument to the resulting TF function:

    def f_jax_2(x, train):
      return lax.cond(train, lambda: x, lambda: x * 2.)

    print(jax.jit(f_jax_2)(2., True))
    print(jax2tf.convert(f_jax_2)(2., True))