def METHOD_NAME(self):
    # With fixed bins
    layer = layers.Discretization(bin_boundaries=[0.0, 0.35, 0.5, 1.0])
    x = np.array([[-1.0, 0.0, 0.1, 0.2, 0.4, 0.5, 1.0, 1.2, 0.98]])
    self.assertAllClose(layer(x), np.array([[0, 1, 1, 1, 2, 3, 4, 4, 3]]))
    ds = tf_data.Dataset.from_tensor_slices(x).batch(1).map(layer)
    for output in ds.take(1):
        output = output.numpy()
    self.assertAllClose(output, np.array([[0, 1, 1, 1, 2, 3, 4, 4, 3]]))
    # With adapt flow
    layer = layers.Discretization(num_bins=4)
    layer.adapt(
        np.random.random((32, 3)),
    )
    x = np.array([[0.0, 0.1, 0.3]])
    ds = tf_data.Dataset.from_tensor_slices(x).batch(1).map(layer)
    for output in ds.take(1):
        output.numpy()