code: def testOneHotCustomDtype(self):
      actual = nn.one_hot(jnp.array([0, 1, 2]), 3, dtype=jnp.bool_)
      expected = jnp.array(
          [
              [True, False, False],
              [False, True, False],
              [False, False, True],
          ]
      )
      self.assertAllClose(actual, expected)
  
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 32 for calling the NNFunctionsTest obj's underlying member method code with example usage: obj.testOneHotCustomDtype() without return types