def testOneHot(self):
    actual = nn.one_hot(jnp.array([0, 1, 2]), 3)
    expected = jnp.array([[1., 0., 0.],
                         [0., 1., 0.],
                         [0., 0., 1.]])
    self.assertAllClose(actual, expected)

    actual = nn.one_hot(jnp.array([1, 2, 0]), 3)
    expected = jnp.array([[0., 1., 0.],
                         [0., 0., 1.],
                         [1., 0., 0.]])
    self.assertAllClose(actual, expected)