Problems with Masking and Concatenate Layer

TensorFlow
1

I am trying to use masking in combination with a Concatenate layer but keep getting an

Failed to convert elements of (None, 3, 2) to Tensor. Consider casting elements to a supported type. See https://www.tensorflow.org/api_docs/python/tf/dtypes for supported TF dtypes.

error. Below you can find a minimal working example of the problem.

import tensorflow as tf
import numpy as np
import tensorflow.keras as keras
from keras.src import ops


jets = np.ones((100,3,2))

leptons = np.ones((100,1))
labels = np.zeros((100,1))
globals = np.ones((100,1))
for i in range(100):
    if i % 2 == 0:
        jets[i,2,0] = 0
        jets[i,2,1] = 0
        labels[i,0] = 1

jet_input = keras.Input(shape=(3,2), name='jet_input')
jet_masked = keras.layers.Masking(mask_value=0.0)(jet_input)
lepton_input = keras.Input(shape=(1,), name='lepton_input')
global_input = keras.Input(shape=(1,), name='global_input')

lepton_tiled = keras.layers.RepeatVector(3)(lepton_input)
global_tiled = keras.layers.RepeatVector(3)(global_input)

lepton_jet_concat = keras.layers.Concatenate()([jet_masked, lepton_tiled,global_tiled])

lepton_jet_lstm = keras.layers.LSTM(2, return_sequences=True)(lepton_jet_concat)
lepton_jet_lstm = keras.layers.Flatten()(lepton_jet_lstm)
lepton_jet_output = keras.layers.Dense(1, activation='sigmoid', name='lepton_jet_output')(lepton_jet_lstm)
model = keras.Model(inputs=[jet_input, lepton_input,global_input], outputs=[lepton_jet_output])
model.summary()
model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])

model.fit([jets, leptons,globals], labels, epochs=100, batch_size=32, verbose=1)