keras/src/layers/normalization/layer_normalization.py/LayerNormalization/_broadcast
class LayerNormalization:    def call(self, inputs):
        inputs = ops.cast(inputs, self.compute_dtype)
        # Compute the axes along which to reduce the mean / variance
        input_shape = inputs.shape
        ndims = len(input_shape)

        # Broadcasting only necessary for norm when the axis is not just
        # the last dimension
        broadcast_shape = [1] * ndims
        for dim in self.axis:
            broadcast_shape[dim] = input_shape[dim]

        def _broadcast(v):
            if (
                v is not None
                and len(v.shape) != ndims
                and self.axis != [ndims - 1]
            ):
                return ops.reshape(v, broadcast_shape)
            return v

        input_dtype = inputs.dtype
        if input_dtype in ("float16", "bfloat16") and self.dtype == "float32":
            # If mixed precision is used, cast inputs to float32 so that
            # this is at least as numerically stable as the fused version.
            inputs = ops.cast(inputs, "float32")

        if self.rms_scaling:
            # Calculate outputs with only variance and gamma if rms scaling
            # is enabled
            # Calculate the variance along self.axis (layer activations).
            variance = ops.var(inputs, axis=self.axis, keepdims=True)
            inv = ops.rsqrt(variance + self.epsilon)

            outputs = (
                inputs * inv * ops.cast(_broadcast(self.gamma), inputs.dtype)
            )
        else:
            # Calculate the mean & variance along self.axis (layer activations).
            mean, variance = ops.moments(inputs, axes=self.axis, keepdims=True)
            gamma, beta = _broadcast(self.gamma), _broadcast(self.beta)

            inv = ops.rsqrt(variance + self.epsilon)
            if gamma is not None:
                gamma = ops.cast(gamma, inputs.dtype)
                inv = inv * gamma

            res = -mean * inv
            if beta is not None:
                beta = ops.cast(beta, inputs.dtype)
                res = res + beta

            outputs = inputs * inv + res

        return ops.cast(outputs, input_dtype)