code: def test_one_hot(self):
      num_classes = 3
      indices = jnp.array([[[1.0, 2.0, 3.0], [1.0, 2.0, 2.0]]])
      expected_result = jnp.array(
          [
              [
                  [[0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 0.0]],
                  [[0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 1.0]],
              ]
          ]
      )
      result = base.one_hot(indices, num_classes)
      np.testing.assert_array_almost_equal(result, expected_result)
  
file_path: rlax/rlax/_src/base_test.py
imports: from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  import jax.numpy
  import numpy
  from rlax._src import base
label_desc: Use image node_id 1 for calling the OneHotTest obj's underlying member method code with example usage: obj.test_one_hot() without return types