keras/src/metrics/reduction_metrics.py/Mean/update_state
class Mean:    def update_state(self, values, sample_weight=None):
        values, sample_weight = reduce_to_samplewise_values(
            values, sample_weight, reduce_fn=ops.mean, dtype=self.dtype
        )
        self.total.assign_add(ops.sum(values))
        if sample_weight is not None:
            num_samples = ops.sum(sample_weight)
        elif len(values.shape) >= 1:
            num_samples = ops.shape(values)[0]
        else:
            num_samples = 1
        self.count.assign_add(ops.cast(num_samples, dtype=self.dtype))