code: def rms_norm(cfg, scale, x):
      x = jnp.asarray(x, jnp.float32)
      mean2 = jnp.mean(jax.lax.square(x), axis=-1, keepdims=True)
      y = jnp.asarray(x * jax.lax.rsqrt(mean2 + cfg.epsilon), cfg.dtype)
      return y * jnp.asarray(scale, cfg.dtype)
  
imports: import dataclasses
  import typing
  import jax
  import jax.numpy
  import numpy
  from jax.sharding import PartitionSpec
  from flax.experimental import nnx
label_desc: Use image node_id 7 for calling a global function with example usage: rms_norm(cfg, scale, x) and returns: unknown