a preprocessing layer which buckets continuous features by ranges this layer will place each element of its input data into one of several contiguous ranges and output an integer index indicating which range each element was placed in note this layer is safe to use inside a tf data pipeline independently of which backend you re using input shape any array of dimension 2 or higher output shape same as input shape arguments binboundaries a list of bin boundaries the leftmost and rightmost bins will always extend to inf and inf so binboundaries0 1 2 generates bins inf 0 0 1 1 2 and 2 inf if this option is set adapt should not be called numbins the integer number of bins to compute if this option is set adapt should be called to learn the bin boundaries epsilon error tolerance typically a small fraction close to zero e g 0 01 higher values of epsilon increase the quantile approximation and hence result in more unequal buckets but could improve performance and resource consumption outputmode specification for the output of the layer values can be int onehot multihot or count configuring the layer as follows int return the discretized bin indices directly onehot encodes each individual element in the input into an array the same size as numbins containing a 1 at the input s bin index if the last dimension is size 1 will encode on that dimension if the last dimension is not size 1 will append a new dimension for the encoded output multihot encodes each sample in the input into a single array the same size as numbins containing a 1 for each bin index index present in the sample treats the last dimension as the sample dimension if input shape is samplelength output shape will be numtokens count as multihot but the int array contains a count of the number of times the bin index appeared in the sample defaults to int sparse boolean only applicable to onehot multihot and count output modes only supported with tensorflow backend if true returns a sparsetensor instead of a dense tensor defaults to false examples discretize float values based on provided buckets input np array1 5 1 0 3 4 5 0 0 3 0 1 3 0 0 layer discretizationbinboundaries0 1 2 layerinput array0 2 3 1 1 3 2 1 discretize float values based on a number of buckets to compute input np array1 5 1 0 3 4 5 0 0 3 0 1 3 0 0 layer discretizationnumbins4 epsilon0 01 layer adaptinput layerinput array0 2 3 2 1 3 3 1 computes bin boundaries from quantiles in a input dataset calling adapt on a discretization layer is an alternative to passing in a binboundaries argument during construction a discretization layer should always be either adapted over a dataset or passed binboundaries during adapt the layer will estimate the quantile boundaries of the input dataset the number of quantiles can be controlled via the numbins argument and the error tolerance for quantile boundaries can be controlled via the epsilon argument arguments data the data to train on it can be passed either as a batched tf data dataset or as a numpy array steps integer or none total number of steps batches of samples to process if data is a tf data dataset and steps is none adapt will run until the input dataset is exhausted when passing an infinitely repeating dataset you must specify the steps argument this argument is not supported with array inputs or list inputs legacy format case reduce a 1d sequence of values to a summary this algorithm is based on numpy quantiles but modified to allow for intermediate steps between multiple data sets it first finds the target number of bins as the reciprocal of epsilon and then takes the individual values spaced at appropriate intervals to arrive at that target the final step is to return the corresponding counts between those values if the target numbins is larger than the size of values the whole array is returned with weights of 1 args values 1d np ndarray to be summarized epsilon a float32 that determines the approximate desired precision returns a 2d np ndarray that is a summary of the inputs first column is the interpolated partition values the second is the weights counts weighted merge sort of summaries given two summaries of distinct data this function merges and compresses them to stay within epsilon error tolerance args prevsummary 2d np ndarray summary to be merged with nextsummary nextsummary 2d np ndarray summary to be merged with prevsummary epsilon a float that determines the approxmiate desired precision returns a 2d np ndarray that is a merged summary first column is the interpolated partition values the second is the weights counts compress a summary to within epsilon accuracy the compression step is needed to keep the summary sizes small after merging and also used to return the final target boundaries it finds the new bins based on interpolating cumulative weight percentages from the large summary taking the difference of the cumulative weights from the previous bin s cumulative weight will give the new weight for that bin args summary 2d np ndarray summary to be compressed epsilon a float32 that determines the approxmiate desired precision returns a 2d np ndarray that is a compressed summary first column is the interpolated partition values the second is the weights counts a preprocessing layer which buckets continuous features by ranges this layer will place each element of its input data into one of several contiguous ranges and output an integer index indicating which range each element was placed in note this layer is safe to use inside a tf data pipeline independently of which backend you re using input shape any array of dimension 2 or higher output shape same as input shape arguments bin_boundaries a list of bin boundaries the leftmost and rightmost bins will always extend to inf and inf so bin_boundaries 0 1 2 generates bins inf 0 0 1 1 2 and 2 inf if this option is set adapt should not be called num_bins the integer number of bins to compute if this option is set adapt should be called to learn the bin boundaries epsilon error tolerance typically a small fraction close to zero e g 0 01 higher values of epsilon increase the quantile approximation and hence result in more unequal buckets but could improve performance and resource consumption output_mode specification for the output of the layer values can be int one_hot multi_hot or count configuring the layer as follows int return the discretized bin indices directly one_hot encodes each individual element in the input into an array the same size as num_bins containing a 1 at the input s bin index if the last dimension is size 1 will encode on that dimension if the last dimension is not size 1 will append a new dimension for the encoded output multi_hot encodes each sample in the input into a single array the same size as num_bins containing a 1 for each bin index index present in the sample treats the last dimension as the sample dimension if input shape is sample_length output shape will be num_tokens count as multi_hot but the int array contains a count of the number of times the bin index appeared in the sample defaults to int sparse boolean only applicable to one_hot multi_hot and count output modes only supported with tensorflow backend if true returns a sparsetensor instead of a dense tensor defaults to false examples discretize float values based on provided buckets input np array 1 5 1 0 3 4 5 0 0 3 0 1 3 0 0 layer discretization bin_boundaries 0 1 2 layer input array 0 2 3 1 1 3 2 1 discretize float values based on a number of buckets to compute input np array 1 5 1 0 3 4 5 0 0 3 0 1 3 0 0 layer discretization num_bins 4 epsilon 0 01 layer adapt input layer input array 0 2 3 2 1 3 3 1 computes bin boundaries from quantiles in a input dataset calling adapt on a discretization layer is an alternative to passing in a bin_boundaries argument during construction a discretization layer should always be either adapted over a dataset or passed bin_boundaries during adapt the layer will estimate the quantile boundaries of the input dataset the number of quantiles can be controlled via the num_bins argument and the error tolerance for quantile boundaries can be controlled via the epsilon argument arguments data the data to train on it can be passed either as a batched tf data dataset or as a numpy array steps integer or none total number of steps batches of samples to process if data is a tf data dataset and steps is none adapt will run until the input dataset is exhausted when passing an infinitely repeating dataset you must specify the steps argument this argument is not supported with array inputs or list inputs legacy format case reduce a 1d sequence of values to a summary this algorithm is based on numpy quantiles but modified to allow for intermediate steps between multiple data sets it first finds the target number of bins as the reciprocal of epsilon and then takes the individual values spaced at appropriate intervals to arrive at that target the final step is to return the corresponding counts between those values if the target num_bins is larger than the size of values the whole array is returned with weights of 1 args values 1d np ndarray to be summarized epsilon a float32 that determines the approximate desired precision returns a 2d np ndarray that is a summary of the inputs first column is the interpolated partition values the second is the weights counts weighted merge sort of summaries given two summaries of distinct data this function merges and compresses them to stay within epsilon error tolerance args prev_summary 2d np ndarray summary to be merged with next_summary next_summary 2d np ndarray summary to be merged with prev_summary epsilon a float that determines the approxmiate desired precision returns a 2 d np ndarray that is a merged summary first column is the interpolated partition values the second is the weights counts compress a summary to within epsilon accuracy the compression step is needed to keep the summary sizes small after merging and also used to return the final target boundaries it finds the new bins based on interpolating cumulative weight percentages from the large summary taking the difference of the cumulative weights from the previous bin s cumulative weight will give the new weight for that bin args summary 2d np ndarray summary to be compressed epsilon a float32 that determines the approxmiate desired precision returns a 2d np ndarray that is a compressed summary first column is the interpolated partition values the second is the weights counts
import numpy as np

from keras import backend
from keras.api_export import keras_export
from keras.layers.preprocessing.tf_data_layer import TFDataLayer
from keras.utils import argument_validation
from keras.utils import numerical_utils
from keras.utils.module_utils import tensorflow as tf


@keras_export("keras.layers.Discretization")
class Discretization(TFDataLayer):
    

    def __init__(
        self,
        bin_boundaries=None,
        num_bins=None,
        epsilon=0.01,
        output_mode="int",
        sparse=False,
        dtype=None,
        name=None,
    ):
        if dtype is None:
            dtype = "int64" if output_mode == "int" else backend.floatx()

        super().__init__(name=name, dtype=dtype)

        if sparse and backend.backend() != "tensorflow":
            raise ValueError(
                "`sparse` can only be set to True with the "
                "TensorFlow backend."
            )
        if sparse and output_mode == "int":
            raise ValueError(
                "`sparse` may only be true if `output_mode` is "
                "`'one_hot'`, `'multi_hot'`, or `'count'`. "
                f"Received: sparse={sparse} and "
                f"output_mode={output_mode}"
            )

        argument_validation.validate_string_arg(
            output_mode,
            allowable_strings=(
                "int",
                "one_hot",
                "multi_hot",
                "count",
            ),
            caller_name=self.__class__.__name__,
            arg_name="output_mode",
        )

        if num_bins is not None and num_bins < 0:
            raise ValueError(
                "`num_bins` must be greater than or equal to 0. "
                f"Received: `num_bins={num_bins}`"
            )
        if num_bins is not None and bin_boundaries is not None:
            if len(bin_boundaries) != num_bins - 1:
                raise ValueError(
                    "Both `num_bins` and `bin_boundaries` should not be "
                    f"set. Received: `num_bins={num_bins}` and "
                    f"`bin_boundaries={bin_boundaries}`"
                )

        self.input_bin_boundaries = bin_boundaries
        self.bin_boundaries = (
            bin_boundaries if bin_boundaries is not None else []
        )
        self.num_bins = num_bins
        self.epsilon = epsilon
        self.output_mode = output_mode
        self.sparse = sparse

        if self.bin_boundaries:
            self.built = True
            self.summary = None
        else:
            self.summary = np.array([[], []], dtype="float32")

    def build(self, input_shape=None):
        self.built = True

    @property
    def input_dtype(self):
        return backend.floatx()

    def adapt(self, data, steps=None):
        
        if self.input_bin_boundaries is not None:
            raise ValueError(
                "Cannot adapt a Discretization layer that has been initialized "
                "with `bin_boundaries`, use `num_bins` instead."
            )
        self.reset_state()
        if isinstance(data, tf.data.Dataset):
            if steps is not None:
                data = data.take(steps)
            for batch in data:
                self.update_state(batch)
        else:
            self.update_state(data)
        self.finalize_state()

    def update_state(self, data):
        data = np.array(data).astype("float32")
        summary = summarize(data, self.epsilon)
        self.summary = merge_summaries(summary, self.summary, self.epsilon)

    def finalize_state(self):
        if self.input_bin_boundaries is not None:
            return
        self.bin_boundaries = get_bin_boundaries(
            self.summary, self.num_bins
        ).tolist()

    def reset_state(self):
        if self.input_bin_boundaries is not None:
            return
        self.summary = np.array([[], []], dtype="float32")

    def compute_output_spec(self, inputs):
        return backend.KerasTensor(shape=inputs.shape, dtype=self.compute_dtype)

    def load_own_variables(self, store):
        if len(store) == 1:
            
            self.summary = store["0"]
        return

    def call(self, inputs):
        indices = self.backend.numpy.digitize(inputs, self.bin_boundaries)
        outputs = numerical_utils.encode_categorical_inputs(
            indices,
            output_mode=self.output_mode,
            depth=len(self.bin_boundaries) + 1,
            dtype=self.compute_dtype,
            backend_module=self.backend,
        )
        return outputs

    def get_config(self):
        return {
            "bin_boundaries": self.bin_boundaries,
            "num_bins": self.num_bins,
            "epsilon": self.epsilon,
            "output_mode": self.output_mode,
            "sparse": self.sparse,
            "name": self.name,
            "dtype": self.dtype,
        }


def summarize(values, epsilon):
    
    values = np.reshape(values, [-1])
    values = np.sort(values)
    elements = np.size(values)
    num_buckets = 1.0 / epsilon
    increment = elements / num_buckets
    start = increment
    step = max(increment, 1)
    boundaries = values[int(start) :: int(step)]
    weights = np.ones_like(boundaries)
    weights = weights * step
    return np.stack([boundaries, weights])


def merge_summaries(prev_summary, next_summary, epsilon):
    
    merged = np.concatenate((prev_summary, next_summary), axis=1)
    merged = np.take(merged, np.argsort(merged[0]), axis=1)
    return compress_summary(merged, epsilon)


def get_bin_boundaries(summary, num_bins):
    return compress_summary(summary, 1.0 / num_bins)[0, :-1]


def compress_summary(summary, epsilon):
    
    if summary.shape[1] * epsilon < 1:
        return summary

    percents = epsilon + np.arange(0.0, 1.0, epsilon)
    cum_weights = summary[1].cumsum()
    cum_weight_percents = cum_weights / cum_weights[-1]
    new_bins = np.interp(percents, cum_weight_percents, summary[0])
    cum_weights = np.interp(percents, cum_weight_percents, cum_weights)
    new_weights = cum_weights - np.concatenate(
        (np.array([0]), cum_weights[:-1])
    )
    summary = np.stack((new_bins, new_weights))
    return summary.astype("float32")