code: def layer_norm_reference(x, weight, bias, *, eps: float = 1e-5):
      mean = jnp.mean(x, axis=1)
      mean2 = jnp.mean(jnp.square(x), axis=1)
      var = jnp.maximum(0.0, mean2 - jnp.square(mean))
      y = x - mean[:, None]
      mul = lax.rsqrt(var + eps)
      return y * mul[:, None] * weight[None] + bias[None]
  
file_path: jax/jax/experimental/pallas/ops/layer_norm.py
imports: from __future__ import annotations
  import functools
  from typing import Optional
  import jax
  from jax import lax
  import jax.numpy
  from jax._src.lax.control_flow.for_loop import for_loop
  from jax.experimental import pallas
label_desc: Use image node_id 7 for calling a global function with example usage: layer_norm_reference(x, weight, bias) and returns: unknown