code: def rms_norm_forward(
      x,
      weight,
      bias,
      num_warps: int | None = None,
      num_stages: int | None = 3,
      eps: float = 1e-5,
      backward_pass_impl: str = "triton",
      interpret: bool = False,
  ):
      del num_stages
      del backward_pass_impl
      n = x.shape[-1]
      # Triton heuristics
      # Less than 64KB per feature: enqueue fused kernel
      max_fused_size = 65536 // x.dtype.itemsize
      block_size = min(max_fused_size, pl.next_power_of_2(n))
      block_size = min(max(block_size, 128), 4096)
      num_warps = min(max(block_size // 256, 1), 8)
  
      kernel = functools.partial(
          rms_norm_forward_kernel, eps=eps, block_size=block_size
      )
      out_shape = [
          jax.ShapeDtypeStruct(shape=(n,), dtype=x.dtype),
          jax.ShapeDtypeStruct(shape=(), dtype=x.dtype),
      ]
      method = pl.pallas_call(
          kernel,
          num_warps=num_warps,
          grid=(),
          out_shape=out_shape,
          debug=False,
          interpret=interpret,
          name="rms_forward",
      )
  
      method = jax.vmap(
          jax.vmap(method, in_axes=(0, None, None)),
          in_axes=(0, None, None),
      )
      out, rstd = method(x, weight, bias)
      return out, (x, weight, bias, rstd)
  
file_path: jax/jax/experimental/pallas/ops/rms_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 2 for calling a global function with example usage: rms_norm_forward(x, weight, bias, num_warps, num_stages, eps, backward_pass_impl, interpret) and returns: out,  x,  weight,  bias,  rstd