code: def one_hot(
      x: Any,
      num_classes: int,
      *,
      dtype: Any = jnp.float_,
      axis: int | AxisName = -1,
  ) -> Array:
      """One-hot encodes the given indices.
  
      Each index in the input ``x`` is encoded as a vector of zeros of length
      ``num_classes`` with the element at ``index`` set to one::
  
        >>> jax.nn.one_hot(jnp.array([0, 1, 2]), 3)
        Array([[1., 0., 0.],
               [0., 1., 0.],
               [0., 0., 1.]], dtype=float32)
  
      Indices outside the range [0, num_classes) will be encoded as zeros::
  
        >>> jax.nn.one_hot(jnp.array([-1, 3]), 3)
        Array([[0., 0., 0.],
               [0., 0., 0.]], dtype=float32)
  
      Args:
        x: A tensor of indices.
        num_classes: Number of classes in the one-hot dimension.
        dtype: optional, a float dtype for the returned values (default :obj:`jnp.float_`).
        axis: the axis or axes along which the function should be
          computed.
      """
      num_classes = core.concrete_dim_or_error(
          num_classes,
          "The error arose in jax.nn.one_hot argument `num_classes`.",
      )
      return _one_hot(x, num_classes, dtype=dtype, axis=axis)
  
imports: from __future__ import annotations
  from functools import partial
  import operator
  import numpy
  from typing import Any
  import jax
  import jax.numpy
  from jax import custom_jvp
  from jax import lax
  from jax._src import config
  from jax._src import core
  from jax._src import dtypes
  from jax._src import util
  from jax._src.core import AxisName
  from jax._src.numpy import util
  from jax._src.typing import Array, ArrayLike
  from jax._src.ops.special import logsumexp
label_desc: Use image node_id 22 for calling a global function with example usage: one_hot(x, num_classes) and returns: _one_hot