code: def layer_norm_forward_kernel(
      x_ref,
      weight_ref,
      bias_ref,  # Input arrays
      o_ref,
      mean_ref=None,
      rstd_ref=None,  # Output arrays
      *,
      eps: float,
      block_size: int,
  ):
      n_col = x_ref.shape[0]
  
      def mean_body(i, acc_ref):
          col_idx = i * block_size + jnp.arange(block_size)
          mask = col_idx < n_col
          a = pl.load(
              x_ref,
              (col_idx,),
              mask=mask,
              other=0.0,
              eviction_policy="evict_last",
          ).astype(jnp.float32)
          acc_ref[:] += a
  
      mean = (
          for_loop(
              pl.cdiv(n_col, block_size),
              mean_body,
              jnp.zeros(block_size),
          ).sum()
          / n_col
      )
  
      def var_body(i, acc_ref):
          col_idx = i * block_size + jnp.arange(block_size)
          mask = col_idx < n_col
          a = pl.load(
              x_ref,
              (col_idx,),
              mask=mask,
              other=0.0,
              eviction_policy="evict_last",
          ).astype(jnp.float32)
          a = jnp.where(mask, a - mean, 0.0)
          acc_ref[:] += a * a
  
      var = (
          for_loop(
              pl.cdiv(n_col, block_size),
              var_body,
              jnp.zeros(block_size),
          ).sum()
          / n_col
      )
      rstd = 1 / jnp.sqrt(var + eps)
      if mean_ref is not None:
          mean_ref[...] = mean.astype(mean_ref.dtype)
      if rstd_ref is not None:
          rstd_ref[...] = rstd.astype(rstd_ref.dtype)
  
      def body(i, _):
          col_idx = i * block_size + jnp.arange(block_size)
          mask = col_idx < n_col
          weight = pl.load(weight_ref, (col_idx,), mask=mask)
          bias = pl.load(bias_ref, (col_idx,), mask=mask)
          x = pl.load(
              x_ref,
              (col_idx,),
              mask=mask,
              other=0.0,
              eviction_policy="evict_first",
          ).astype(jnp.float32)
          out = (x - mean) * rstd * weight + bias
          pl.store(
              o_ref, (col_idx,), out.astype(o_ref.dtype), mask=mask
          )
  
      for_loop(pl.cdiv(n_col, block_size), body, ())
  
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 1 for calling a global function with example usage: layer_norm_forward_kernel(x_ref, weight_ref, bias_ref, o_ref, mean_ref, rstd_ref) without return types