code: def _abs_sq(x):
      """Computes the elementwise square of the absolute value |x|^2."""
      if jnp.iscomplexobj(x):
          return lax.square(lax.real(x)) + lax.square(lax.imag(x))
      else:
          return lax.square(x)
  
file_path: flax/flax/experimental/nnx/nnx/nn/normalization.py
imports: import typing
  import jax
  import jax.numpy
  from jax import lax
  from flax.experimental import nnx
  from flax.experimental.nnx.nnx import flaglib, rnglib
  from flax.experimental.nnx.nnx.module import Module, first_from
  from flax.experimental.nnx.nnx.nn import dtypes, initializers
label_desc: Use image node_id 2 for calling a global function with example usage: _abs_sq(x) and returns: unknown, lax