def METHOD_NAME(self):
    input_data = np.array([3, 2, 0, 1])
    expected_output = np.array(
        [
            [0, 0, 0, 1],
            [0, 0, 1, 0],
            [1, 0, 0, 0],
            [0, 1, 0, 0],
        ]
    )
    num_tokens = 4
    expected_output_shape = (num_tokens, num_tokens)
    # Test call on layer directly.
    layer = layers.CategoryEncoding(
        num_tokens=num_tokens, output_mode="one_hot"
    )
    output_data = layer(input_data)
    self.assertAllClose(expected_output, output_data)
    self.assertEqual(expected_output_shape, output_data.shape)