code: def _one_hot(
      x: Any, num_classes: int, *, dtype: Any, axis: int | AxisName
  ) -> Array:
      num_classes = core.concrete_dim_or_error(
          num_classes,
          "The error arose in jax.nn.one_hot argument `num_classes`.",
      )
      dtype = dtypes.canonicalize_dtype(dtype)
      x_arr = jnp.asarray(x)
      try:
          output_pos_axis = util.canonicalize_axis(axis, x_arr.ndim + 1)
      except TypeError:
          axis_size = lax.psum(1, axis)
          if num_classes != axis_size:
              raise ValueError(
                  f"Expected num_classes to match the size of axis {axis}, "
                  f"but {num_classes} != {axis_size}"
              ) from None
          axis_idx = lax.axis_index(axis)
          return jnp.asarray(x_arr == axis_idx, dtype=dtype)
      axis = operator.index(axis)  # type: ignore[arg-type]
      lhs = lax.expand_dims(x_arr, (axis,))
      rhs_shape = [1] * x_arr.ndim
      rhs_shape.insert(output_pos_axis, num_classes)
      rhs = lax.broadcasted_iota(
          x_arr.dtype, rhs_shape, output_pos_axis
      )
      return jnp.asarray(lhs == rhs, dtype=dtype)
  
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 21 for calling a global function with example usage: _one_hot(x, num_classes) and returns: jnp, jnp