def METHOD_NAME(self):
    layer = layers.LayerNormalization(
        dtype="float32",
        rms_scaling=True,
        gamma_initializer="ones",
    )
    inputs = np.arange(5).astype("float32")[None, :]
    out = layer(inputs)
    self.assertAllClose(out, [[0.0, 0.70693, 1.41386, 2.12079, 2.82772]])