Open areiner222 opened 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
I will take a look at this repro, but in general the call_tf mechanism works best when there are not tf.Variables captured in the called function. It is best to refactor the TF code to pass the variable values in and out of the TF function.
I looked at the original repro.
The problem here is that when call_tf processes func_tf
is notices that func_tf
captures two tf.Variables
, and it tries to read the values of those variables using .numpy()
. This works if the code runs in eager mode, but does not work when the outer code is under tf.function
. You can verify this if you remove tf.function
from your repro.
@maxwillzq Can you please take a look if there is some workaround?
Sure.
and it tries to read the values of those variables using .numpy()
Could you point me where it call .numpy()
in keras lib ? I will try to see if we can remove it.
The call to numpy
happens as part of np.asarray(inp)
here. This works if we are in eager mode bug fails under tf.function
@areiner222 As @gnecula 's suggestion, the best way is to rewrite the tf function so it has on captured tf.Variables. Another choice is https://www.tensorflow.org/api_docs/python/tf/config/run_functions_eagerly but it has performance penalty.
I test it on colab https://shorturl.at/bkwHJ and it works.
By the way, I try to use same way on jax2tf.py but it is much complicated. @gnecula, do you know where is the best location insert this context_manager call ? Thanks
I do not think that we should change jax2tf to force running of TF functions eagerly.
If the user starts with a tf.keras.layer
, which references variables, is it possible to split that function into a set of variable values and a function that takes the variables as inputs?
It would be nice if there were some other way of resolving this. I am loading a saved model (originally JAX code) to run inside more JAX code, which I then want to export into another savedmodel, but it is tricky to use call_tf on the inner savedmodel while also passing in all variables / etc it uses (since they are tf tensors).
For resolving this when functions are not run eagerly, what about modifying _call_tf_lowering to grab the value of variables? This is what I had to do to get my use case working for https://github.com/google/jax/issues/11753
if tf.executing_eagerly():
np_captured_inputs = [np.asarray(inp) for inp in captured_inputs]
else:
if captured_inputs:
with tf.compat.v1.Session(graph=captured_inputs[0].graph) as sess:
sess.run([v.initializer for v in captured_inputs])
np_captured_inputs = sess.run(captured_inputs)
else:
np_captured_inputs = []
captured_ops = tuple(
mlir.ir_constant(inp)
for inp in np_captured_inputs
)
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