code: def test_onehot_shape(self):
      indices = jnp.arange(24, dtype=jnp.float32).reshape([2, 3, 4])
      num_classes = 24
      out = basic.one_hot(indices, num_classes=num_classes)
      self.assertEqual(out.shape, (2, 3, 4, num_classes))
  
imports: from absl.testing import absltest
  from absl.testing import parameterized
  from haiku._src import base
  from haiku._src import basic
  from haiku._src import test_utils
  from haiku._src import transform
  import jax
  from jax import random
  import jax.numpy
  import numpy
label_desc: Use image node_id 1 for calling the BasicTest obj's underlying member method code with example usage: obj.test_onehot_shape() without return types