Batch Shape inside customized train_step of keras subclassed model is None
Hi, for me this is not a bug at all here my explanation:

During the training process in TensorFlow, the batch dimension is usually automatically removed before passing the data to the train_step method. In my opinion and this is part of the internal functionality of the TensorFlow facilitate the deployment of custom models, this make models more flexible and can be used with different amounts of data at each training step.


I have assumed this behavior from the tensor docu and the following argument:
* dynamic_batch: Whether to set the batch sizes of all the returned tf.TensorSpec to None. (Note that when defining functional or Sequential models with tf.keras.Input([...], batch_size=X), the batch size will always be preserved). Defaults to True.

If I'm wrong please make me know since I would also like to learn. Hi @danielreyes9756, you are right, this is the intentional API behaviour. 
I've better identified the issue. I'm trying to sample a random vector from a tensorflow_probability distribution. If we follow the example at the  docs for a [custom train step](https://github.com/keras-team/keras-io/blob/6623a072e7eef484c3defa44dc16d824eac434cb/guides/custom_train_step_in_tensorflow.py#L377C1-L384C10) tf.shape(inputs)[0] is used to get the batch dimension and sample a random vector from keras.backend.random_normal, and in the example below it outputs the right dimensions (None, dim1). 

However if we use a Normal distribution from tensorflow_probability we:
1) Cannot use tf.shape(inputs)[0] as it is a 'tensorflow.python.framework.ops.SymbolicTensor' class, and therfore we must use inputs.shape[0]
2) Cannot sample from a (None, 2) input shape, since we get  the following:
ValueError: Attempt to convert a value (None) with an unsupported type (<class 'NoneType'>) to a Tensor.


I've updated the code example to better show case the issue:

```python
import tensorflow as tf 
import tensorflow_probability.python.distributions as tfd 

class cl(tf.keras.layers.Layer):
    def __init__(self):
        super().__init__()

    def build(self, input_shape):
        super().build(input_shape)
        self.normal = tfd.Normal(loc=[1.] * input_shape[-1], scale=[1.] * input_shape[-1])

    def call(self, inputs, *args):
        batch_shape = inputs.shape[0]
        return self.normal.sample(tf.TensorShape((batch_shape, 2)))


class cm(tf.keras.models.Model):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.normal = cl()

    def call(self, inputs, *args):
        print(f"tf.shape: {tf.shape(inputs)}, inputs.shape:{inputs.shape}")
        sample = self.normal(inputs)
        return sample

    def train_step(self, data):
        B = tf.shape(data)[0]
        B2 = data.shape[0]
        rnd = keras.backend.random_normal(shape=(B, 2))
        print(f"random normal from keras.backend shape: {rnd.shape}")
        ta = tf.TensorArray(dtype=tf.float32, size=100, element_shape=tf.TensorShape((B2, 1)))
        sample = self(data)  # Error
        return {"loss":1}

inputs = tf.random.normal(shape=(100, 10, 2))
c0 = cm()
c0.compile("adam")
c0.fit(tf.constant(inputs),epochs=1)
```

which outputs: 
```console
random normal from keras.backend shape: (None, 2)
tf.shape: Tensor("cm_35/Shape:0", shape=(3,), dtype=int32), inputs.shape:(None, 10, 2)

ValueError: Cannot convert a partially known TensorShape to a Tensor: (None, 2)
```

I think this issue may best be opened on the tfp repo.  Hi @claCase, thanks for the answer and for your effort!

I think it's impressive how you found this!

best regards,
Reyes Hi @danielreyes9756 , thanks for the kind words.
I think I may have found the culprit in the [Tensorflow Probability  Distribution Class](https://github.com/tensorflow/probability/blob/230463ad39a4a089965a9ab1aeb5ac3de7ff398c/tensorflow_probability/python/distributions/distribution.py#L736C1-L739C31). Tensors must be statically defined, and a shape of None type can't be automatically inferred. I don't think there's a solution to this issue, the only solution would be to create a custom fit function... I'm still opening a new issue to get more insights in the tfp repo.  @claCase,
As mentioned above, this looks like the intended behaviour. The batch dimension is automatically removed before passing the data to the **train_step** method. Also it is mentioned in the tf.keras.Sequential official document. [Reference](https://www.tensorflow.org/api_docs/python/tf/keras/Sequential)

Also as you mentioned above, the issue related to tensorflow probabilty will discuss and resolved in this respective [repo](https://github.com/tensorflow/probability/issues).

Thank you!


  Hi @tilakrayal 
I've opened [the issue](https://github.com/tensorflow/probability/issues/1786) under the tfp repo, this issue can be closed, 
Thank you Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/62889">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/62889">No</a>
 I solved the issue by passing the batch_size parameter to the fit function
```python
c0.fit(tf.constant(inputs),epochs=1, batch_size=10)
```
### Issue type

Bug

### Have you reproduced the bug with TensorFlow Nightly?

No

### Source

source

### TensorFlow version

tf 2.15

### Custom code

Yes

### OS platform and distribution

Windows WSL

### Mobile device

_No response_

### Python version

3.11

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

I'm trying to get the full input shape inside the train_step function of a subclassed keras model. When using tf.shape(inputs) will output the rank of the tensor, when using inputs.shape will output (None, d1, d2, ...).  
This issue is similar to #36991 and #38155 .

### Standalone code to reproduce the issue

```python
import tensorflow as tf 

class cm(tf.keras.models.Model):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    @tf.function
    def call(self, inputs, *args):
        print(f"tf.shape: {tf.shape(inputs)}, inputs.shape:{inputs.shape}")
        return inputs

    @tf.function
    def train_step(self, data):
        print(tf.shape(data), data.shape)
        return {"loss":1}

inputs = tf.random.normal(shape=(100, 10, 1))
c0 = cm()
_ = c0(inputs)
c0.compile("adam")
c0.fit(tf.constant(inputs),epochs=1)
```


### Relevant log output

```shell
Inside call function: 
tf.shape: Tensor("Shape:0", shape=(3,), dtype=int32), inputs.shape:(100, 10, 1)

Inside train_step:
Tensor("Shape:0", shape=(3,), dtype=int32) (None, 10, 1)
```
