keras/src/ops/nn.py/Moments/compute_output_spec
class Moments:    def compute_output_spec(self, x):
        return (
            KerasTensor(
                reduce_shape(x.shape, axis=self.axes, keepdims=self.keepdims),
                dtype=x.dtype,
            ),
            KerasTensor(
                reduce_shape(x.shape, axis=self.axes, keepdims=self.keepdims),
                dtype=x.dtype,
            ),
        )