areiner222
commented
10 months ago I found that keras_core (I believe soon to be Keras 3) provides a workaround for my example. In particular, the stateless_call method of keras_core layers makes it possible. Here is a contrived example:
import os
os.environ["KERAS_BACKEND"] = "tensorflow"
import jax
from jax.experimental import jax2tf
import tensorflow as tf
import keras_core
class Jax2TFLayer(keras_core.layers.Layer):
def __init__(self, units, *args, **kwargs):
super().__init__(*args, **kwargs)
self._layer = keras_core.layers.Dense(units)
def build(self, input_shape):
super().build(input_shape)
self._layer.build(input_shape)
def _call_tf(self, x, trainable_variables, non_trainable_variables):
return self._layer.stateless_call(trainable_variables, non_trainable_variables, x)[0]
def _call_jax(self, x, trainable_variables, non_trainable_variables):
length = 100
def f(x, i):
out = jax2tf.call_tf(self._call_tf)(x, trainable_variables, non_trainable_variables)
return x, out
return jax.lax.scan(f, x, jp.arange(length), length=length)
@tf.function(autograph=False, jit_compile=True)
def call(self, x):
return jax2tf.convert(self._call_jax)(
x, [v.value for v in self._layer.trainable_variables], [v.value for v in self._layer.non_trainable_variables]
)[1]
inp = tf.random.uniform((10, 16))
out = layer(inp)This can work great for my use case! but can't use it just yet as I understand polymorphism in call_tf doesn't work with native_serialization which I currently require
gnecula
maxwillzq
Edgeworth
Description
Having trouble nesting a call_tf inside a converted jax function when using a stateful transform / tf.Variable.
Error is:
What jax/jaxlib version are you using?
jax v0.4.17, jaxlib v0.4.17, tensorflow v2.14.0
Which accelerator(s) are you using?
CPU
Additional system info
Apple M1 Max
NVIDIA GPU info
No response