Repository: mlflow/mlflow

Issue:
[BUG] Wrong batch_size value when fitting keras tf.keras.utils.Sequence
### Issues Policy acknowledgement

- [X] I have read and agree to submit bug reports in accordance with the [issues policy](https://www.github.com/mlflow/mlflow/blob/master/ISSUE_POLICY.md)

### Willingness to contribute

Yes. I would be willing to contribute a fix for this bug with guidance from the MLflow community.

### MLflow version

- Client: 2.2.0



### System information

- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: MacOS 13.2.1
- **Python version**:  3.10.8
- **yarn version, if running the dev UI**: -


### Describe the problem

`batch_size` has the wrong value when using `mlflow.tensorflow.autolog()` and the model is trained on a `tf.keras.utils.Sequence` and the model has more than one input. In this situation the value stored in `batch_size` is the # inputs the model has.

IMO, the issue is in this [line](https://github.com/mlflow/mlflow/blob/0130976ee27f20624d23aeb3696c42275fb6f83a/mlflow/tensorflow/__init__.py#L1184), it's taking as `batch_size` value always the 1st dimension, when it should pick the 2nd dimension when the model has more than one input in order to take the batch length.

### Code to reproduce issue

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```python
import mlflow
import tensorflow as tf
import numpy as np
import math

BATCH_SIZE = 5

x_train = np.random.randint(0, 5, size=(10,2))
y_train = np.array((x_train[:,0] == x_train[:,1]), dtype=int)

class DataSequence(tf.keras.utils.Sequence):

    def __init__(self, x1, x2, y, batch_size):
        self.x1, self.x2, self.y = x1, x2, y
        self.batch_size = batch_size

    def __len__(self):
        return math.ceil(len(self.x1) / self.batch_size)

    def __getitem__(self, idx):
        batch_x1 = self.x1[idx * self.batch_size:(idx + 1) * self.batch_size]
        batch_x2 = self.x2[idx * self.batch_size:(idx + 1) * self.batch_size]
        batch_y = self.y[idx * self.batch_size:(idx + 1) * self.batch_size]
        return [batch_x1, batch_x2], batch_y
    
generator = DataSequence(x_train[:,0], x_train[:,1], y_train, BATCH_SIZE)

input1 = tf.keras.Input(shape=(1,))
input2 = tf.keras.Input(shape=(1,))
concat = tf.keras.layers.Concatenate()([input1, input2])
output = tf.keras.layers.Dense(1, activation='sigmoid')(concat)

model = tf.keras.models.Model(inputs=[input1, input2], outputs=output)
model.summary()

model.compile(optimizer='adam', loss='mse', metrics=['accuracy'])

mlflow.tensorflow.autolog()
model.fit(generator, epochs=1)
mlflow.end_run()

autolog_run = mlflow.last_active_run()

print("Expected batch_size:", BATCH_SIZE)
print("MLflow batch_size:", autolog_run.data.params['batch_size'])
```

### Other info / logs

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
The previous code prints:
```
Expected batch_size: 5
MLflow batch_size: 2
```

I'd like to contribute on a fix, but I'm not sure how to proceed [tf.keras.utils.Sequence](https://github.com/keras-team/keras/blob/v2.11.0/keras/utils/data_utils.py#L445-L520) has no information about the batch size. Probably we should take information about the model trained and check how many inputs it expect. What do you think?


### What component(s) does this bug affect?

- [ ] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [ ] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/recipes`: Recipes, Recipe APIs, Recipe configs, Recipe Templates
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [ ] `area/server-infra`: MLflow Tracking server backend
- [X] `area/tracking`: Tracking Service, tracking client APIs, autologging

### What interface(s) does this bug affect?

- [ ] `area/uiux`: Front-end, user experience, plotting, JavaScript, JavaScript dev server
- [ ] `area/docker`: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
- [ ] `area/sqlalchemy`: Use of SQLAlchemy in the Tracking Service or Model Registry
- [ ] `area/windows`: Windows support

### What language(s) does this bug affect?

- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

### What integration(s) does this bug affect?

- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations


Discussion:
Hi @jaume-ferrarons, when the model has more than one input, does the 2nd dimension **_always_** correspond to the batch size?
> Hi @jaume-ferrarons, when the model has more than one input, does the 2nd dimension **_always_** corresponds to the batch size?

Hi @harupy , I think so but I'm not sure at 100%. 
@jaume-ferrarons Can you confirm that by asking a question in the tensorflow repository?
From reading the Keras documentation, [fit method](https://keras.io/api/models/model_training_apis/), I deduce that yes, the output of the Sequence is **always** a batched `input` for the model. The input x is an _array (or array-like), or a list of arrays (in case the model has multiple inputs)_, no other option is mentioned.

Also it states that the input can be _keras.utils.Sequence returning (inputs, targets) or (inputs, targets, sample_weights)_.

From the [Sequence class](https://keras.io/api/utils/python_utils/#sequence-class), it states: _The method \_\_getitem\_\_ should return a complete batch._.

When connecting points, Sequence has to return a complete batch, in which the inputs should  _array (or array-like), or a list of arrays (in case the model has multiple inputs)_. Then, for a single input models the batch size is the 1st dimension and multi-input the 2nd.
@jaume-ferrarons Makes sense, would you be interested in creating a PR to fix this issue?
@harupy Yes! I'll make a PR 😄 

Resolving patch:
diff --git a/mlflow/tensorflow/__init__.py b/mlflow/tensorflow/__init__.py
--- a/mlflow/tensorflow/__init__.py
+++ b/mlflow/tensorflow/__init__.py
@@ -1174,6 +1174,7 @@ def _patch_implementation(
 
             batch_size = None
             try:
+                is_single_input_model = isinstance(inst.input_shape, tuple)
                 training_data = kwargs["x"] if "x" in kwargs else args[0]
                 if isinstance(training_data, tensorflow.data.Dataset) and hasattr(
                     training_data, "_batch_size"
@@ -1181,10 +1182,16 @@ def _patch_implementation(
                     batch_size = training_data._batch_size.numpy()
                 elif isinstance(training_data, tensorflow.keras.utils.Sequence):
                     first_batch_inputs, _ = training_data[0]
-                    batch_size = len(first_batch_inputs)
+                    if is_single_input_model:
+                        batch_size = len(first_batch_inputs)
+                    else:
+                        batch_size = len(first_batch_inputs[0])
                 elif is_iterator(training_data):
                     peek = next(training_data)
-                    batch_size = len(peek[0])
+                    if is_single_input_model:
+                        batch_size = len(peek[0])
+                    else:
+                        batch_size = len(peek[0][0])
 
                     def __restore_generator(prev_generator):
                         yield peek


Test patch:
diff --git a/tests/tensorflow/test_tensorflow2_autolog.py b/tests/tensorflow/test_tensorflow2_autolog.py
--- a/tests/tensorflow/test_tensorflow2_autolog.py
+++ b/tests/tensorflow/test_tensorflow2_autolog.py
@@ -324,6 +324,85 @@ def test_tf_keras_autolog_implicit_batch_size_works(generate_data, batch_size):
     assert mlflow.last_active_run().data.params["batch_size"] == str(batch_size)
 
 
+def __tf_dataset_multi_input(batch_size):
+    a = tf.data.Dataset.range(1)
+    b = tf.data.Dataset.range(1)
+    c = tf.data.Dataset.range(1)
+    ds = tf.data.Dataset.zip(((a, b), c))
+    return ds.batch(batch_size)
+
+
+class __SequenceMultiInput(tf.keras.utils.Sequence):
+    def __init__(self, batch_size):
+        self.batch_size = batch_size
+
+    def __len__(self):
+        return 10
+
+    def __getitem__(self, idx):
+        return (np.random.rand(self.batch_size), np.random.rand(self.batch_size)), np.random.rand(
+            self.batch_size
+        )
+
+
+def __generator_multi_input(data, target, batch_size):
+    data_batches = np.split(data, data.shape[1] // batch_size, axis=1)
+    target_batches = np.split(target, target.shape[0] // batch_size)
+    for inputs, output in zip(data_batches, target_batches):
+        yield tuple(inputs), output
+
+
+class __GeneratorClassMultiInput:
+    def __init__(self, data, target, batch_size):
+        self.data = data
+        self.target = target
+        self.batch_size = batch_size
+        self.ptr = 0
+
+    def __next__(self):
+        if self.ptr >= len(self.data):
+            raise StopIteration
+        idx = self.ptr % len(self.data)
+        self.ptr += 1
+        return (
+            self.data[idx : idx + self.batch_size, 0],
+            self.data[idx : idx + self.batch_size, 1],
+        ), self.target[idx : idx + self.batch_size]
+
+    def __iter__(self):
+        return self
+
+
+@pytest.mark.parametrize(
+    "generate_data",
+    [
+        __tf_dataset_multi_input,
+        __SequenceMultiInput,
+        functools.partial(__generator_multi_input, np.random.rand(2, 10), np.random.rand(10)),
+        functools.partial(__GeneratorClassMultiInput, np.random.rand(10, 2), np.random.rand(10, 1)),
+    ],
+)
+@pytest.mark.parametrize("batch_size", [5, 10])
+def test_tf_keras_autolog_implicit_batch_size_works_multi_input(generate_data, batch_size):
+    mlflow.tensorflow.autolog()
+
+    input1 = tf.keras.Input(shape=(1,))
+    input2 = tf.keras.Input(shape=(1,))
+    concat = tf.keras.layers.Concatenate()([input1, input2])
+    output = tf.keras.layers.Dense(1, activation="sigmoid")(concat)
+
+    model = tf.keras.models.Model(inputs=[input1, input2], outputs=output)
+    model.compile(loss="mse")
+
+    # 'x' passed as arg
+    model.fit(generate_data(batch_size), verbose=0)
+    assert mlflow.last_active_run().data.params["batch_size"] == str(batch_size)
+
+    # 'x' passed as kwarg
+    model.fit(x=generate_data(batch_size), verbose=0)
+    assert mlflow.last_active_run().data.params["batch_size"] == str(batch_size)
+
+
 @pytest.mark.skipif(
     Version(tf.__version__) < Version("2.1.4"),
     reason="Does not support passing of generator classes as `x` in `fit`",