code: def BatchNorm(x, params, axis=(0, 1, 2), epsilon=1e-5,
                center=True, scale=True, **unused_kwargs):
    
    mean = np.mean(x, axis, keepdims=True)
    # Fast but less numerically-stable variance calculation than np.var.
    m1 = np.mean(x**2, axis, keepdims=True)
    var = m1 - mean**2
    z = (x - mean) / np.sqrt(var + epsilon)
  
    # Expand the parameters to have the right axes.
    beta, gamma = params
    # TODO(phawkins): np.expand_dims should accept an axis tuple.
    # (https://github.com/numpy/numpy/issues/12290)
    ed = tuple(None if i in axis else slice(None) for i in range(np.ndim(x)))
    beta = beta[ed]
    gamma = gamma[ed]
  
    # Return the z rescaled by the parameters if requested.
    if center and scale:
      return gamma * z + beta
    if center:
      return z + beta
    if scale:
      return gamma * z
    return z
docs: Layer construction function for a batch normalization layer.