Expected behavior

When I create a custom PyDataset and then pass it into the Normalization adapt method, the mean and variance should be computed successfully, as for a TensorFlow dataset.

Actual behavior

Method adapt does not handle the PyDataset like inputs. 

Traceback (most recent call last):
  File "/home/user/project/keras_bug_sample.py", line 14, in <module>
    normalizer.adapt(CustomDataset)
  File "/home/user/.pyenv/versions/3.12.8/lib/python3.12/site-packages/keras/src/layers/preprocessing/normalization.py", line 230, in adapt
    self.build(input_shape)
               ^^^^^^^^^^^
UnboundLocalError: cannot access local variable 'input_shape' where it is not associated with a value

Steps to reproduce

import keras

class CustomDataset(keras.utils.PyDataset):
    def __len__(self):
        return 100

    def __getitem__(self, idx):
        # Generate dummy data
        x = np.random.rand(32, 32, 3)
        y = np.random.randint(0, 10, size=(1,))
        return x, y

normalizer = keras.layers.Normalization()
normalizer.adapt(CustomDataset)

Comment From: sonali-kumari1

Hi @limzikiki - Thanks for reporting this. The error UnboundLocalError: cannot access local variable 'input_shape' where it is not associated with a value occurs because adapt() does not directly support custom datasets like keras.utils.PyDataset. The adapt() method expects data to be in the form of tf.data.Dataset, NumPy array, or a backend-native eager tensor. Since you are passing keras.utils.PyDataset, adapt() is not able to extract the input_shape which leads to the UnboundLocalError. To resolve this you can convert your custom dataset into tf.data.Dataset and ensure that dataset is properly batched to compute the mean and variance correctly.

Comment From: limzikiki

I ended up fixing it by subclassing the Normalization class and introducing support for the PyDataset type inputs. I was wondering if the Keras community would also benefit from this feature. Basically, what I did in my solution is that I fetched the first batch, and then, based on its shape, I built the normalisation. I don't know if that's a proper solution. If that’s okay, I can open a PR with the changes.

Comment From: VarunS1997

Please go ahead and make a PR, @limzikiki -- thank you!

Comment From: tripathi-genius

Is it open to contributions? I want to contribute.

Comment From: hertschuh

@tripathi-genius

Is it open to contributions? I want to contribute.

Yes!

The idea is to reuse as much of the existing logic from data_adapters.

Maybe you can just use get_adapter and then call get_tf_dataset() on the returned adapter.