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