code: def one_hot(
      indices: JaxArray,
      depth: int,
      /,
      *,
      on_value: Optional[Number] = None,
      off_value: Optional[Number] = None,
      axis: Optional[int] = None,
      dtype: Optional[jnp.dtype] = None,
      device: jaxlib.xla_extension.Device = None,
      out: Optional[JaxArray] = None,
  ) -> JaxArray:
      on_none = on_value is None
      off_none = off_value is None
  
      if dtype is None:
          if on_none and off_none:
              dtype = jnp.float32
          else:
              if not on_none:
                  dtype = jnp.array(on_value).dtype
              elif not off_none:
                  dtype = jnp.array(off_value).dtype
  
      res = jnp.eye(depth, dtype=dtype)[
          jnp.array(indices, dtype="int64").reshape(-1)
      ]
      res = res.reshape(list(indices.shape) + [depth])
  
      if not on_none and not off_none:
          res = jnp.where(res == 1, on_value, off_value)
  
      if axis is not None:
          res = jnp.moveaxis(res, -1, axis)
  
      return res
  
imports: from numbers import Number
  import numpy
  from typing import Union, Optional, List, Sequence, Tuple
  import jax.dlpack
  import jax.numpy
  import jax._src
  import jaxlib.xla_extension
  import ivy
  from ivy import as_native_dtype
  from ivy.functional.backends.jax import JaxArray
  from ivy.functional.backends.jax.device import dev
  from ivy.functional.ivy.creation import _asarray_to_native_arrays_and_back, _asarray_infer_device, _asarray_infer_dtype, _asarray_handle_nestable, NestedSequence, SupportsBufferProtocol, _asarray_inputs_to_native_shapes
label_desc: Use image node_id 19 for calling a global function with example usage: one_hot() and returns: res