code: def embedding(
      weights: JaxArray,
      indices: JaxArray,
      /,
      *,
      max_norm: Optional[int] = None,
      out: Optional[JaxArray] = None,
  ) -> JaxArray:
      ivy.utils.assertions.check_equal(
          len(weights.shape),
          2,
          message="weights must be 2-d",
          as_array=False,
      )
  
      embeddings = jnp.take(weights, indices, axis=0)
      if max_norm is not None:
          norms = jnp.linalg.norm(embeddings, axis=-1, keepdims=True)
          embeddings = jnp.where(
              norms > max_norm,
              embeddings * max_norm / norms,
              embeddings,
          )
          embeddings = jnp.where(
              norms < -max_norm,
              embeddings * -max_norm / norms,
              embeddings,
          )
      return embeddings
  
imports: from typing import Optional, Union, Tuple, List, Literal, Sequence, Callable
  import jax
  import jax.lax
  import jax.numpy
  import math
  import ivy
  from ivy import output_to_native_arrays
  from ivy.functional.backends.jax import JaxArray
  from ivy.functional.backends.jax.random import RNG
  from ivy.functional.ivy.experimental.general import _correct_ivy_callable
  from ivy.functional.ivy.layers import _handle_padding, _validate_max_pool_params, _depth_max_pooling_helper
  from ivy.functional.ivy.experimental.layers import _padding_ceil_mode, _get_size
  from ivy.func_wrapper import with_supported_dtypes
  from ivy.func_wrapper import with_unsupported_dtypes
  from .None import backend_version
  from ivy.functional.backends.jax.experimental.manipulation import _to_nested_tuple
label_desc: Use image node_id 21 for calling a global function with example usage: embedding() and returns: embeddings