Split tensorflow BatchDataset for LSTM with multiple inputs

I construct a LSTM model with two inputs: one for categorical variables, one for numerical variables:

model = Model(inputs = [cat_input, num_input], outputs = x, name = "LSTM")

The input data for the LSTM is generated by means of tensorflow.keras.utils.timeseries_dataset_from_array():

input_dataset = timeseries_dataset_from_array(
    df[["cat", "num1", "num2"]], df["target"], sequence_length=n_timesteps, sequence_stride=1, batch_size=20
)

When I directly feed input_dataset into the model, I get the following error: "ValueError: Layer "LSTM" expects 2 input(s), but it received 1 input tensors. Inputs received: [<tf.Tensor 'IteratorGetNext:0' shape=(None, None, 3) dtype=int64>]", because the model expects two inputs and not one.

I can achieve this (a bit ugly) like so:

input_dataset2 = input_dataset.map(lambda x, y: ((x[:,:,0:1], x[:,:,1:3]), y))
model.fit(
    input_dataset2, steps_per_epoch=20, epochs=50, verbose=0, shuffle=True
) # this now works

My question: The solution I found is not very elegant. Is this kind of split also possible with tf.split() or another function?

EDIT: When I try the following:

input_dataset.map(lambda x, y: ((split(value=x, num_or_size_splits=[1, 2], axis = -1)), y))

I get this error: "ValueError: Value [<tf.Tensor 'split:0' shape=(None, None, 1) dtype=int64>, <tf.Tensor 'split:1' shape=(None, None, 2) dtype=int64>] is not convertible to a tensor with dtype <dtype: 'int64'> and shape (2, None, None, None)."

You can use the tf.split() function to split the input tensor along the last axis, then you can pass the split tensors as separate inputs to your model. To fix the "ValueError" you encountered, you need to cast the outputs of tf.split() to tf.int64 or another compatible data type for the input of your model:

input_dataset = timeseries_dataset_from_array(
    df[["cat", "num1", "num2"]], df["target"], sequence_length=n_timesteps, sequence_stride=1, batch_size=20
)

input_dataset2 = input_dataset.map(lambda x, y: (tf.split(x, [1, 2], axis=-1), y))
input_dataset2 = input_dataset2.map(lambda x, y: (list(map(lambda x: tf.cast(x, tf.int64), x)), y))

model.fit(
    input_dataset2, steps_per_epoch=20, epochs=50, verbose=0, shuffle=True
)

I solved it by explicitly adding tuple() - the brackets are otherwise not recognized as tuple:

input_dataset = input_dataset.map(lambda x, y: 
    (
        tuple(
            tf.split(x, [1,2], -1)
        ),
        y
    )
)