src/transformers/models/jukebox/modeling_jukebox.py/JukeboxLayerNorm/__init__
class JukeboxLayerNorm:    def __init__(self, normalized_shape, eps=1e-5, elementwise_affine=True):
        super().__init__(normalized_shape, eps=eps, elementwise_affine=elementwise_affine)
        self.width = np.prod(normalized_shape)
        self.max_numel = 65535 * self.width