def adapt(self, data):
        """Computes the mean and variance of values in a dataset.

        Calling `adapt()` on a `Normalization` layer is an alternative to
        passing in `mean` and `variance` arguments during layer construction. A
        `Normalization` layer should always either be adapted over a dataset or
        passed `mean` and `variance`.

        During `adapt()`, the layer will compute a `mean` and `variance`
        separately for each position in each axis specified by the `axis`
        argument. To calculate a single `mean` and `variance` over the input
        data, simply pass `axis=None` to the layer.

        Arg:
            data: The data to train on. It can be passed either as a
                `tf.data.Dataset`, as a NumPy array, or as a backend-native
                eager tensor.
                If a dataset, *it must be batched*. Keras will assume that the
                data is batched, and if that assumption doesn't hold, the mean
                and variance may be incorrectly computed.
        """
        if isinstance(data, np.ndarray) or backend.is_tensor(data):
            input_shape = data.shape
        elif isinstance(data, tf.data.Dataset):
            input_shape = tuple(data.element_spec.shape)
            if len(input_shape) == 1:
                # Batch dataset if it isn't batched
                data = data.batch(128)
            input_shape = tuple(data.element_spec.shape)

        if not self.built:
            self.build(input_shape)
        else:
            for d in self._keep_axis:
                if input_shape[d] != self._build_input_shape[d]:
                    raise ValueError(
                        "The layer was built with "
                        f"input_shape={self._build_input_shape}, "
                        "but adapt() is being called with data with "
                        f"an incompatible shape, data.shape={input_shape}"
                    )

        if isinstance(data, np.ndarray):
            total_mean = np.mean(data, axis=self._reduce_axis)
            total_var = np.var(data, axis=self._reduce_axis)
        elif backend.is_tensor(data):
            total_mean = ops.mean(data, axis=self._reduce_axis)
            total_var = ops.var(data, axis=self._reduce_axis)
        elif isinstance(data, tf.data.Dataset):
            total_mean = ops.zeros(self._mean_and_var_shape)
            total_var = ops.zeros(self._mean_and_var_shape)
            total_count = 0
            for batch in data:
                batch = backend.convert_to_tensor(
                    batch, dtype=self.compute_dtype
                )
                batch_mean = ops.mean(batch, axis=self._reduce_axis)
                batch_var = ops.var(batch, axis=self._reduce_axis)
                if self._reduce_axis:
                    batch_reduce_shape = (
                        batch.shape[d] for d in self._reduce_axis
                    )
                    batch_count = math.prod(batch_reduce_shape)
                else:
                    batch_count = 1

                total_count += batch_count
                batch_weight = float(batch_count) / total_count
                existing_weight = 1.0 - batch_weight
                new_total_mean = (
                    total_mean * existing_weight + batch_mean * batch_weight
                )
                # The variance is computed using the lack-of-fit sum of squares
                # formula (see
                # https://en.wikipedia.org/wiki/Lack-of-fit_sum_of_squares).
                total_var = (
                    total_var + (total_mean - new_total_mean) ** 2
                ) * existing_weight + (
                    batch_var + (batch_mean - new_total_mean) ** 2
                ) * batch_weight
                total_mean = new_total_mean
        else:
            raise NotImplementedError(f"Unsupported data type: {type(data)}")

        self.adapt_mean.assign(total_mean)
        self.adapt_variance.assign(total_var)
        self.finalize_state()
Computes the mean and variance of values in a dataset.

        Calling `adapt()` on a `Normalization` layer is an alternative to
        passing in `mean` and `variance` arguments during layer construction. A
        `Normalization` layer should always either be adapted over a dataset or
        passed `mean` and `variance`.

        During `adapt()`, the layer will compute a `mean` and `variance`
        separately for each position in each axis specified by the `axis`
        argument. To calculate a single `mean` and `variance` over the input
        data, simply pass `axis=None` to the layer.

        Arg:
            data: The data to train on. It can be passed either as a
                `tf.data.Dataset`, as a NumPy array, or as a backend-native
                eager tensor.
                If a dataset, *it must be batched*. Keras will assume that the
                data is batched, and if that assumption doesn't hold, the mean
                and variance may be incorrectly computed.
        
keras/src/layers/preprocessing/normalization.py