Custom Model Data Cardinality Check Ambiguous

[whitead]

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): No
• OS Platform and Distribution: Linux Ubuntu 16.04
• TensorFlow installed from (source or binary): binary
• TensorFlow version (use command below): 2.3.2/2.4.1
• Python version: 3.7.8

Describe the current behavior

Training Keras custom model when inputs are lists fail data cardinality check if leading axis is not batch. The code below works in 2.3.2 but not 2.4.1. An example of why you would not want a leading axis as batch is a per-batch weight or dynamic parameter (see example code).

Describe the expected behavior

This worked in 2.3.2 but was changed. I read the release notes and cannot seem to understand why this behavior changed. I am unsure how to fix. Either a flag to remove this check, a way to change the input spec of model call, or clarification on how to allow inputs that do not have a leading batch axis would help me fix this problem.

Standalone code to reproduce the issue

import tensorflow as tf
import numpy as np

class MyModel(tf.keras.Model):    
    def call(self, inputs, training):
        # make a model with a batched input (x) and per-batch tensor (s)
        x = inputs[0]
        s = inputs[1]
        return tf.reduce_mean(x * s, axis=1)
    
m = MyModel()
x = np.ones((10, 2))
s = np.ones(1)
y = np.ones(10)

m.compile('sgd', loss='mean_squared_error')

# works as call
m([x, s])

# fails on TF 2.4.1
# succeeds on TF 2.3.2
m.train_on_batch(x=[x, s], y=y)

** Error Message**

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-5-1e85b8f34d3e> in <module>
     21 # fails on 2.4.1
     22 # succeeds on 2.3.2
---> 23 m.train_on_batch(x=[x, s], y=y)

~/miniconda3/envs/mmm/lib/python3.7/site-packages/tensorflow/python/keras/engine/training.py in train_on_batch(self, x, y, sample_weight, class_weight, reset_metrics, return_dict)
   1723       iterator = data_adapter.single_batch_iterator(self.distribute_strategy, x,
   1724                                                     y, sample_weight,
-> 1725                                                     class_weight)
   1726       self.train_function = self.make_train_function()
   1727       logs = self.train_function(iterator)

~/miniconda3/envs/mmm/lib/python3.7/site-packages/tensorflow/python/keras/engine/data_adapter.py in single_batch_iterator(strategy, x, y, sample_weight, class_weight)
   1511     data = (x, y, sample_weight)
   1512 
-> 1513   _check_data_cardinality(data)
   1514   dataset = dataset_ops.DatasetV2.from_tensors(data)
   1515   if class_weight:

~/miniconda3/envs/mmm/lib/python3.7/site-packages/tensorflow/python/keras/engine/data_adapter.py in _check_data_cardinality(data)
   1527           label, ", ".join(str(i.shape[0]) for i in nest.flatten(single_data)))
   1528     msg += "Make sure all arrays contain the same number of samples."
-> 1529     raise ValueError(msg)
   1530 
   1531 

ValueError: Data cardinality is ambiguous:
  x sizes: 10, 1
  y sizes: 10
Make sure all arrays contain the same number of samples.

[Saduf2019]

I have replicated the issue on tf 2.4 and nightly, please find the gist here.

[Saduf2019]

@whitead
I ran the code shared on tf 2.3 and face a different error whereas the issue description says the code works fine on tf 2.3, please find the gist here.

[whitead]

Sorry about that @Saduf2019. I've edited the post for the shape error. I've attached output for 2.3.2 below and here is a gist

[Saduf2019]

@whitead
You can safely ignore the warning, to suppress the warnings please use.

import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import tensorflow as tf

[whitead]

Hi @Saduf2019 that was the message from 2.3.2. I will repost the error message for 2.4.1 below. My issue is that the code works in 2.3.2 but not in 2.4.1 and I do not know how to fix.

[saikumarchalla]

Issue still exists in TF 2.5. Please find the gist here.Thanks!