code: def _batch_norm_new_params(input_shape, rng, axis=(0, 1, 2),
                             center=True, scale=True, **kwargs):
    
    del rng, kwargs
    axis = (axis,) if np.isscalar(axis) else axis
    shape = tuple(d for i, d in enumerate(input_shape) if i not in axis)
    beta = np.zeros(shape, dtype='float32') if center else ()
    gamma = np.ones(shape, dtype='float32') if scale else ()
    return (beta, gamma)
docs: Helper to initialize batch norm params.