code: def testOneHotAxis(self):
      expected = jnp.array(
          [[0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [1.0, 0.0, 0.0]]
      ).T
  
      actual = nn.one_hot(jnp.array([1, 2, 0]), 3, axis=0)
      self.assertAllClose(actual, expected, check_dtypes=False)
  
      actual = nn.one_hot(jnp.array([1, 2, 0]), 3, axis=-2)
      self.assertAllClose(actual, expected, check_dtypes=False)
  
imports: import collections
  from functools import partial
  import itertools
  import unittest
  from absl.testing import absltest
  from absl.testing import parameterized
  import scipy.stats
  from jax._src import config
  from jax._src import core
  from jax._src import test_util
  from jax._src import ad_checkpoint
  from jax.test_util import check_grads
  from jax import nn
  from jax import random
  import jax
  import jax.numpy
label_desc: Use image node_id 34 for calling the NNFunctionsTest obj's underlying member method code with example usage: obj.testOneHotAxis() without return types