<jupyter_start><jupyter_text>Please find torch implementation of this notebook here: https://colab.research.google.com/github/probml/pyprobml/blob/master/notebooks/book1/14/layer_norm_torch.ipynb<jupyter_code>import numpy as np

import jax
import jax.numpy as jnp

try:
    from flax import linen as nn
except ModuleNotFoundError:
    %pip install -qq flax
    from flax import linen as nn
# batch size 3, feature size 2
np.random.seed(42)
X = np.random.normal(size=(2, 3))

print("batch norm")
mu_batch = np.mean(X, axis=0)
sigma_batch = np.std(X, axis=0)
XBN = (X - mu_batch) / sigma_batch
print(XBN)

print("layer norm")
mu_layer = np.expand_dims(np.mean(X, axis=1), axis=1)
sigma_layer = np.expand_dims(np.std(X, axis=1), axis=1)
XLN = (X - mu_layer) / sigma_layer
print(XLN)
X = jnp.float32(X)

rng = jax.random.PRNGKey(42)
bn_rng, ln_rng = jax.random.split(rng)

print("batch norm")
bn = nn.BatchNorm(use_running_average=False, epsilon=1e-6)
bn_params = bn.init(bn_rng, X)
XBN_t, _ = bn.apply(bn_params, X, mutable=["batch_stats"])
print(XBN_t)
assert np.allclose(np.array(XBN_t), XBN, atol=1e-3)

print("layer norm")
ln = nn.LayerNorm()
ln_params = ln.init(ln_rng, X)
XLN_t = ln.apply(ln_params, X)
print(XLN_t)
assert np.allclose(np.array(XLN_t), XLN, atol=1e-3)<jupyter_output>WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
{'fraction_non_alphanumeric': 0.1409058231488138, 'fraction_numerical': 0.018691588785046728, 'mean_word_length': 5.327272727272727, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 4, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '358547', 'license': 'MIT', 'n_tokens_mistral': 609, 'n_tokens_neox': 586, 'n_words': 127, 'path': 'notebooks/book1/14/layer_norm_jax.ipynb', 'repo_name': 'patel-zeel/pyprobml', 'size': '1391'}