adapter for arraylike objects e g tfjax tensors numpy arrays def init self x ynone sampleweightnone batchsizenone stepsnone shufflefalse classweightnone if not canconvertarraysx y sampleweight raise valueerror expected all elements of x to be arraylike freceived invalid types xx x y sampleweight converttoarraysx y sampleweight if sampleweight is not none if classweight is not none raise valueerror you cannot classweight and sampleweight at the same time if tree isnestedy if isinstancesampleweight np ndarray issamplewise lensampleweight shape 1 or lensampleweight shape 2 and sampleweight shape1 1 if not issamplewise raise valueerror for a model with multiple outputs when providing a single sampleweight array it should only have one scalar score per sample i e shape numsamples if you want to use nonscalar sample weights pass a sampleweight argument with one array per model output replicate the same sampleweight array on all outputs sampleweight tree mapstructure lambda sampleweight y else try tree assertsamestructurey sampleweight except valueerror raise valueerror you should provide one sampleweight array per output in y the two structures did not match n f y yn f sampleweight sampleweightn if classweight is not none if tree isnestedy raise valueerror classweight is only supported for models with a single output sampleweight dataadapterutils classweighttosampleweights y classweight inputs dataadapterutils packxysampleweightx y sampleweight dataadapterutils checkdatacardinalityinputs numsamples seti shape0 for i in tree flatteninputs pop self numsamples numsamples self inputs inputs if batchsize is not passed but steps is calculate from the input data defaults to 32 for backwards compatibility if not batchsize batchsize intmath ceilnumsamples steps if steps else 32 self size intmath ceilnumsamples batchsize self batchsize batchsize self partialbatchsize numsamples batchsize self shuffle shuffle def getnumpyiteratorself inputs self inputs if self shuffle inputs dataadapterutils syncshuffle inputs numsamplesself numsamples for i in rangeself size start stop i self batchsize i 1 self batchsize yield tree mapstructurelambda x xstart stop inputs def gettfdatasetself from keras utils moduleutils import tensorflow as tf inputs self inputs shuffle self shuffle batchsize self batchsize numsamples self numsamples numfullbatches intself numsamples batchsize vectorized version of shuffle this is a performance improvement over using fromtensorslices the indices of the data are shuffled and batched and these indices are then zipped with the data and used to extract a batch of the data at each step the performance improvements here come from 1 vectorized batch using gather 2 parallelized map 3 pipelined permutation generation 4 optimized permutation batching 5 disabled static optimizations indicesdataset tf data dataset range1 def permutation it turns out to be more performant to make a new set of indices rather than reusing the same range tensor presumably because of buffer forwarding indices tf rangenumsamples dtypetf int64 if shuffle and shuffle batch indices tf random shuffleindices return indices we prefetch a single element computing large permutations can take quite a while so we don t want to wait for prefetching over an epoch boundary to trigger the next permutation on the other hand too many simultaneous shuffles can contend on a hardware level and degrade all performance indicesdataset indicesdataset mappermutation prefetch1 def slicebatchindicesindices numinfullbatch numfullbatches batchsize firstkindices tf sliceindices 0 numinfullbatch firstkindices tf reshape firstkindices numfullbatches batchsize flatdataset tf data dataset fromtensorslicesfirstkindices if self partialbatchsize indexremainder tf data dataset fromtensors tf slice indices numinfullbatch self partialbatchsize flatdataset flatdataset concatenateindexremainder return flatdataset def sliceinputsindicesdataset inputs dataset tf data dataset zip indicesdataset tf data dataset fromtensorsinputs repeat def grabbatchi data return tree mapstructure lambda d tf gatherd i axis0 data dataset dataset map grabbatch numparallelcallstf data autotune default optimizations are disabled to avoid the overhead of unnecessary input pipeline graph serialization deserialization options tf data options options experimentaloptimization applydefaultoptimizations false if self shuffle options experimentalexternalstatepolicy tf data experimental externalstatepolicy ignore dataset dataset withoptionsoptions return dataset indicesdataset indicesdataset flatmapslicebatchindices dataset sliceinputsindicesdataset inputs if shuffle batch def shufflebatchbatch return tree mapstructuretf random shuffle batch dataset dataset mapshufflebatch options tf data options options experimentaldistribute autoshardpolicy tf data experimental autoshardpolicy data dataset dataset withoptionsoptions return dataset prefetchtf data autotune property def numbatchesself return self size property def batchsizeself return self batchsize property def haspartialbatchself return self partialbatchsize 0 property def partialbatchsizeself return self partialbatchsize or none def canconvertarraysarrays def canconvertsinglearrayx isnone x is none knowntype isinstancex dataadapterutils arraytypes convertabletype hasattrx array return isnone or knowntype or convertabletype return all tree flattentree mapstructurecanconvertsinglearray arrays def converttoarraysarrays def convertsinglearrayx if x is none return x if pandas is not none if isinstancex pandas series x np expanddimsx tonumpy axis1 elif isinstancex pandas dataframe x x tonumpy if istfraggedtensorx from keras utils moduleutils import tensorflow as tf convert floats to floatx if backend isfloatdtypex dtype and not backend standardizedtypex dtype backend floatx x tf castx backend floatx return x if not isinstancex np ndarray using array should handle tf tensor jax np ndarray torch tensor as well as any other tensorlike object that has added numpy support if hasattrx array x backend converttonumpyx else raise valueerror expected a numpy array tf tensor tf raggedtensor jax np ndarray torch tensor pandas dataframe or pandas series received invalid input fx of type typex if x dtype object return x convert floats to floatx if backend isfloatdtypex dtype and not backend standardizedtypex dtype backend floatx x x astypebackend floatx return x arrays tree mapstructureconvertsinglearray arrays return liststotuplesarrays def istfraggedtensorx return x class name raggedtensor adapter for array like objects e g tf jax tensors numpy arrays replicate the same sample_weight array on all outputs if batch_size is not passed but steps is calculate from the input data defaults to 32 for backwards compatibility vectorized version of shuffle this is a performance improvement over using from_tensor_slices the indices of the data are shuffled and batched and these indices are then zipped with the data and used to extract a batch of the data at each step the performance improvements here come from 1 vectorized batch using gather 2 parallelized map 3 pipelined permutation generation 4 optimized permutation batching 5 disabled static optimizations it turns out to be more performant to make a new set of indices rather than reusing the same range tensor presumably because of buffer forwarding we prefetch a single element computing large permutations can take quite a while so we don t want to wait for prefetching over an epoch boundary to trigger the next permutation on the other hand too many simultaneous shuffles can contend on a hardware level and degrade all performance convert a tensor of indices into a dataset of batched indices this step can be accomplished in several ways the most natural is to slice the tensor in a dataset map with a condition on the upper index to handle the partial batch however it turns out that coercing the tensor into a shape which is divisible by the batch size and handling the last partial batch separately allows for a much more favorable memory access pattern and improved performance args indices tensor which determines the data order for an entire epoch returns a dataset of batched indices slice inputs into a dataset of batches given a dataset of batch indices and the unsliced inputs this step slices the inputs in a parallelized fashion and produces a dataset of input batches args indices_dataset a dataset of batched indices inputs a python data structure that contains the inputs targets and possibly sample weights returns a dataset of input batches matching the batch indices default optimizations are disabled to avoid the overhead of unnecessary input pipeline graph serialization deserialization check if array like inputs can be handled by arraydataadapter args inputs structure of tensor s numpy arrays or tensor like returns true if arrays can be handled by arraydataadapter false otherwise process array like inputs this function converts tf tensors to numpy arrays converts pandas series to np ndarray converts list s to tuple s for tf data support args inputs structure of tensor s numpy arrays or tensor like returns structure of numpy ndarray s convert floats to floatx using __array__ should handle tf tensor jax np ndarray torch tensor as well as any other tensor like object that has added numpy support convert floats to floatx
import math

import numpy as np
import tree

from keras import backend
from keras.trainers.data_adapters import data_adapter_utils
from keras.trainers.data_adapters.data_adapter import DataAdapter
from keras.utils.nest import lists_to_tuples

try:
    import pandas
except ImportError:
    pandas = None


class ArrayDataAdapter(DataAdapter):
    

    def __init__(
        self,
        x,
        y=None,
        sample_weight=None,
        batch_size=None,
        steps=None,
        shuffle=False,
        class_weight=None,
    ):
        if not can_convert_arrays((x, y, sample_weight)):
            raise ValueError(
                "Expected all elements of `x` to be array-like. "
                f"Received invalid types: x={x}"
            )

        x, y, sample_weight = convert_to_arrays((x, y, sample_weight))
        if sample_weight is not None:
            if class_weight is not None:
                raise ValueError(
                    "You cannot `class_weight` and `sample_weight` "
                    "at the same time."
                )
            if tree.is_nested(y):
                if isinstance(sample_weight, np.ndarray):
                    is_samplewise = len(sample_weight.shape) == 1 or (
                        len(sample_weight.shape) == 2
                        and sample_weight.shape[1] == 1
                    )
                    if not is_samplewise:
                        raise ValueError(
                            "For a model with multiple outputs, when providing "
                            "a single `sample_weight` array, it should only "
                            "have one scalar score per sample "
                            "(i.e. shape `(num_samples,)`). If you want to use "
                            "non-scalar sample weights, pass a `sample_weight` "
                            "argument with one array per model output."
                        )
                    
                    sample_weight = tree.map_structure(
                        lambda _: sample_weight, y
                    )
                else:
                    try:
                        tree.assert_same_structure(y, sample_weight)
                    except ValueError:
                        raise ValueError(
                            "You should provide one `sample_weight` array per "
                            "output in `y`. The two structures did not match:\n"
                            f"- y: {y}\n"
                            f"- sample_weight: {sample_weight}\n"
                        )
        if class_weight is not None:
            if tree.is_nested(y):
                raise ValueError(
                    "`class_weight` is only supported for Models with a single "
                    "output."
                )
            sample_weight = data_adapter_utils.class_weight_to_sample_weights(
                y, class_weight
            )

        inputs = data_adapter_utils.pack_x_y_sample_weight(x, y, sample_weight)

        data_adapter_utils.check_data_cardinality(inputs)
        num_samples = set(i.shape[0] for i in tree.flatten(inputs)).pop()
        self._num_samples = num_samples
        self._inputs = inputs

        
        
        if not batch_size:
            batch_size = int(math.ceil(num_samples / steps)) if steps else 32

        self._size = int(math.ceil(num_samples / batch_size))
        self._batch_size = batch_size
        self._partial_batch_size = num_samples % batch_size
        self._shuffle = shuffle

    def get_numpy_iterator(self):
        inputs = self._inputs
        if self._shuffle:
            inputs = data_adapter_utils.sync_shuffle(
                inputs, num_samples=self._num_samples
            )
        for i in range(self._size):
            start, stop = i * self._batch_size, (i + 1) * self._batch_size
            yield tree.map_structure(lambda x: x[start:stop], inputs)

    def get_tf_dataset(self):
        from keras.utils.module_utils import tensorflow as tf

        inputs = self._inputs
        shuffle = self._shuffle
        batch_size = self._batch_size
        num_samples = self._num_samples
        num_full_batches = int(self._num_samples // batch_size)

        
        
        
        
        
        
        
        
        
        

        indices_dataset = tf.data.Dataset.range(1)

        def permutation(_):
            
            
            
            indices = tf.range(num_samples, dtype=tf.int64)
            if shuffle and shuffle != "batch":
                indices = tf.random.shuffle(indices)
            return indices

        
        
        
        
        
        indices_dataset = indices_dataset.map(permutation).prefetch(1)

        def slice_batch_indices(indices):
            
            num_in_full_batch = num_full_batches * batch_size
            first_k_indices = tf.slice(indices, [0], [num_in_full_batch])
            first_k_indices = tf.reshape(
                first_k_indices, [num_full_batches, batch_size]
            )

            flat_dataset = tf.data.Dataset.from_tensor_slices(first_k_indices)
            if self._partial_batch_size:
                index_remainder = tf.data.Dataset.from_tensors(
                    tf.slice(
                        indices, [num_in_full_batch], [self._partial_batch_size]
                    )
                )
                flat_dataset = flat_dataset.concatenate(index_remainder)

            return flat_dataset

        def slice_inputs(indices_dataset, inputs):
            
            dataset = tf.data.Dataset.zip(
                (indices_dataset, tf.data.Dataset.from_tensors(inputs).repeat())
            )

            def grab_batch(i, data):
                return tree.map_structure(
                    lambda d: tf.gather(d, i, axis=0), data
                )

            dataset = dataset.map(
                grab_batch, num_parallel_calls=tf.data.AUTOTUNE
            )

            
            
            options = tf.data.Options()
            options.experimental_optimization.apply_default_optimizations = (
                False
            )
            if self._shuffle:
                options.experimental_external_state_policy = (
                    tf.data.experimental.ExternalStatePolicy.IGNORE
                )
            dataset = dataset.with_options(options)
            return dataset

        indices_dataset = indices_dataset.flat_map(slice_batch_indices)

        dataset = slice_inputs(indices_dataset, inputs)

        if shuffle == "batch":

            def shuffle_batch(*batch):
                return tree.map_structure(tf.random.shuffle, batch)

            dataset = dataset.map(shuffle_batch)

        options = tf.data.Options()
        options.experimental_distribute.auto_shard_policy = (
            tf.data.experimental.AutoShardPolicy.DATA
        )
        dataset = dataset.with_options(options)
        return dataset.prefetch(tf.data.AUTOTUNE)

    @property
    def num_batches(self):
        return self._size

    @property
    def batch_size(self):
        return self._batch_size

    @property
    def has_partial_batch(self):
        return self._partial_batch_size > 0

    @property
    def partial_batch_size(self):
        return self._partial_batch_size or None


def can_convert_arrays(arrays):
    

    def can_convert_single_array(x):
        is_none = x is None
        known_type = isinstance(x, data_adapter_utils.ARRAY_TYPES)
        convertable_type = hasattr(x, "__array__")
        return is_none or known_type or convertable_type

    return all(
        tree.flatten(tree.map_structure(can_convert_single_array, arrays))
    )


def convert_to_arrays(arrays):
    

    def convert_single_array(x):
        if x is None:
            return x
        if pandas is not None:
            if isinstance(x, pandas.Series):
                x = np.expand_dims(x.to_numpy(), axis=-1)
            elif isinstance(x, pandas.DataFrame):
                x = x.to_numpy()
        if is_tf_ragged_tensor(x):
            from keras.utils.module_utils import tensorflow as tf

            
            if (
                backend.is_float_dtype(x.dtype)
                and not backend.standardize_dtype(x.dtype) == backend.floatx()
            ):
                x = tf.cast(x, backend.floatx())
            return x
        if not isinstance(x, np.ndarray):
            
            
            
            if hasattr(x, "__array__"):
                x = backend.convert_to_numpy(x)
            else:
                raise ValueError(
                    "Expected a NumPy array, tf.Tensor, tf.RaggedTensor, "
                    "jax.np.ndarray, torch.Tensor, Pandas Dataframe, or "
                    "Pandas Series. Received invalid input: "
                    f"{x} (of type {type(x)})"
                )
        if x.dtype == object:
            return x
        
        if (
            backend.is_float_dtype(x.dtype)
            and not backend.standardize_dtype(x.dtype) == backend.floatx()
        ):
            x = x.astype(backend.floatx())
        return x

    arrays = tree.map_structure(convert_single_array, arrays)
    return lists_to_tuples(arrays)


def is_tf_ragged_tensor(x):
    return x.__class__.__name__ == "RaggedTensor"