model.predict(...) fails on a multi-input model only with the jax backend

[keras-team]

model.predict(...) fails on a multi-input model only with the jax backend

[fchollet]

Thanks for the report! This is now fixed.

"Reproducible bugs are candy" -Salvatore Sanfilippo

[ageron]

Note: if I use named inputs instead of a tuple, it works fine:

import keras_core as keras
import numpy as np

class MultiInputModel(keras.Model):
    def call(self, inputs):
        a = inputs["input_a"]
        b = inputs["input_b"]
        return a + b

model = MultiInputModel()
x1, x2 = np.random.rand(2, 3, 4)
y_pred1 = model({"input_a": x1, "input_b": x2}) # still works fine
y_pred2 = model.predict({"input_a": x1, "input_b": x2}) # now works too