code: def one_hot(indices, num_classes, dtype=jnp.float32):
      """Returns a one-hot version of indices.
  
      Args:
        indices: A tensor of indices.
        num_classes: Number of classes in the one-hot dimension.
        dtype: The dtype.
  
      Returns:
        The one-hot tensor. If indices' shape is [A, B, ...], shape is
        [A, B, ..., num_classes].
      """
      labels = jnp.arange(num_classes)
      for _ in range(indices.ndim):
          labels = jnp.expand_dims(labels, axis=0)
      return jnp.array(indices[..., jnp.newaxis] == labels, dtype=dtype)
  
file_path: rlax/rlax/_src/base.py
imports: from typing import Optional, Sequence, Union
  import chex
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 2 for calling a global function with example usage: one_hot(indices, num_classes, dtype) and returns: jnp