Compatibility issue between `add_loss` and `TimeDistributed` when using JAX backend

[AmedeoBiolatti]

When using JAX backend with jit enabled, I am getting a TracerBoolConversionError when trying to use a custom layer with add_loss inside a TimeDistributed wrapper inside a multi-input functional model.

This is a minimal example that raises an error. Running the code twice will make the error go away

import os

os.environ["KERAS_BACKEND"] = "jax"

import keras
import numpy as np


class CustomLayer(keras.layers.Layer):
    def __init__(self, ):
        super().__init__()

    def call(self, x):
        self.add_loss(1e-6 * keras.ops.mean(x ** 2))
        return x

    def compute_output_shape(self, input_shape):
        return input_shape


X = keras.layers.Input(shape=(3, 2))
Z = keras.layers.Input(shape=(3, 2))

encoder = keras.layers.TimeDistributed(CustomLayer())
E = encoder(X) + encoder(Z)
Y = keras.layers.Dense(1)(E)

model = keras.Model(
    inputs=[X, Z],
    outputs=Y
)

model.compile(
    loss="mse",
    optimizer="adam",
    jit_compile=True
)
h = model.fit(
    [np.random.normal(size=(100, 3, 2)), np.random.normal(size=(100, 3, 2))],
    np.random.normal(size=(100, 3, 1))
)

I tested in the following configurations
keras: 3.8.0, 3.11.2
jax: 0.5.2, 0.7.1

[mehtamansi29]

Hi @AmedeoBiolatti -

When you apply jit_compile=true,the error TracerBoolConversionError indicates that JAX's JIT compiler is encountering a Python control flow operation gives concrete boolean valu. Since tracers represent abstract values without concrete numbers. More detail about TracerBoolConversionError you can find here.
During tracing based on boolean condition it can't decide the result, because of that getting error.

When TimeDistributed calls CustomLayer, it creates a JAX Tracer for the loss value via add_loss method. Then it uses this tracer in a context that requires a concrete boolean value, which is not available during the JIT compilation phase.

You can add loss outside TimeDistributed layer like this:

import os
os.environ["KERAS_BACKEND"] = "jax"

import keras
import numpy as np

class CustomLayer(keras.layers.Layer):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def call(self, x):
        return x

    def compute_output_shape(self, input_shape):
        return input_shape

class CustomModel(keras.Model):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.encoder = keras.layers.TimeDistributed(CustomLayer())
        self.dense = keras.layers.Dense(1)

    def call(self, inputs):
        X, Z = inputs
        activations_x = self.encoder(X)
        activations_z = self.encoder(Z)
        E = activations_x + activations_z
        Y = self.dense(E)
        loss_x = 1e-6 * keras.ops.mean(keras.ops.square(activations_x))
        loss_z = 1e-6 * keras.ops.mean(keras.ops.square(activations_z))
        self.add_loss(loss_x)
        self.add_loss(loss_z)

        return Y

model = CustomModel()

model.compile(
    loss="mse",
    optimizer="adam",
    jit_compile=True
)

print("Model compiled successfully.")

h = model.fit(
    [np.random.normal(size=(100, 3, 2)), np.random.normal(size=(100, 3, 2))],
    np.random.normal(size=(100, 3, 1))
)

print("\nModel trained successfully.")

[AmedeoBiolatti]

Thanks for the reply! I am already using a similar workaround, but I think that something strange is going on and it could be possible to improve the user experience.

What I think is strange

• the issue only occurs when using two inputs for the functional model
• the loss code in CustomLayer doesn't seem to require the convertion of a traced array to a bool, at least not in a trivial way
• the issue goes away running the code twice