keras-hub-nightly 0.15.0.dev20240823171555__py3-none-any.whl

keras_hub/src/layers/preprocessing/multi_segment_packer.py

@@ -0,0 +1,319 @@
+# Copyright 2024 The KerasHub Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.layers.preprocessing.preprocessing_layer import (
+    PreprocessingLayer,
+)
+from keras_hub.src.utils.tensor_utils import convert_to_ragged_batch
+
+try:
+    import tensorflow as tf
+    import tensorflow_text as tf_text
+except ImportError:
+    tf = None
+    tf_text = None
+
+
+@keras_hub_export("keras_hub.layers.MultiSegmentPacker")
+class MultiSegmentPacker(PreprocessingLayer):
+    """Packs multiple sequences into a single fixed width model input.
+
+    This layer packs multiple input sequences into a single fixed width sequence
+    containing start and end delimeters, forming a dense input suitable for a
+    classification task for BERT and BERT-like models.
+
+    Takes as input a tuple of token segments. Each tuple element should contain
+    the tokens for a segment, passed as tensors, `tf.RaggedTensor`s, or lists.
+    For batched input, each element in the tuple of segments should be a list of
+    lists or a rank two tensor. For unbatched inputs, each element should be a
+    list or rank one tensor.
+
+    The layer will process inputs as follows:
+     - Truncate all input segments to fit within `sequence_length` according to
+       the `truncate` strategy.
+     - Concatenate all input segments, adding a single `start_value` at the
+       start of the entire sequence, and multiple `end_value`s at the end of
+       each segment.
+     - Pad the resulting sequence to `sequence_length` using `pad_tokens`.
+     - Calculate a separate tensor of "segment ids", with integer type and the
+       same shape as the packed token output, where each integer index of the
+       segment the token originated from. The segment id of the `start_value`
+       is always 0, and the segment id of each `end_value` is the segment that
+       precedes it.
+
+    Args:
+        sequence_length: int. The desired output length.
+        start_value: int/str/list/tuple. The id(s) or token(s) that are to be
+            placed at the start of each sequence (called "[CLS]" for BERT). The
+            dtype must match the dtype of the input tensors to the layer.
+        end_value: int/str/list/tuple. The id(s) or token(s) that are to be
+            placed at the end of the last input segment (called "[SEP]" for
+            BERT). The dtype must match the dtype of the input tensors to the
+            layer.
+        sep_value: int/str/list/tuple. The id(s) or token(s) that are to be
+            placed at the end of every segment, except the last segment (called
+            "[SEP]" for BERT). If `None`, `end_value` is used. The dtype must
+            match the dtype of the input tensors to the layer.
+        pad_value: int/str. The id or token that is to be placed into the unused
+            positions after the last segment in the sequence
+            (called "[PAD]" for BERT).
+        truncate: str. The algorithm to truncate a list of batched segments to
+            fit a per-example length limit. The value can be either
+            `"round_robin"` or `"waterfall"`:
+            - `"round_robin"`: Available space is assigned one token at a
+                time in a round-robin fashion to the inputs that still need
+                some, until the limit is reached.
+            - `"waterfall"`: The allocation of the budget is done using a
+                "waterfall" algorithm that allocates quota in a
+                left-to-right manner and fills up the buckets until we run
+                out of budget. It support arbitrary number of segments.
+
+    Returns:
+        A tuple with two elements. The first is the dense, packed token
+        sequence. The second is an integer tensor of the same shape, containing
+        the segment ids.
+
+    Examples:
+
+    *Pack a single input for classification.*
+    >>> seq1 = [1, 2, 3, 4]
+    >>> packer = keras_hub.layers.MultiSegmentPacker(
+    ...     sequence_length=8, start_value=101, end_value=102
+    ... )
+    >>> token_ids, segment_ids = packer((seq1,))
+    >>> np.array(token_ids)
+    array([101, 1, 2, 3, 4, 102, 0, 0], dtype=int32)
+    >>> np.array(segment_ids)
+    array([0, 0, 0, 0, 0, 0, 0, 0], dtype=int32)
+
+    *Pack multiple inputs for classification.*
+    >>> seq1 = [1, 2, 3, 4]
+    >>> seq2 = [11, 12, 13, 14]
+    >>> packer = keras_hub.layers.MultiSegmentPacker(
+    ...     sequence_length=8, start_value=101, end_value=102
+    ... )
+    >>> token_ids, segment_ids = packer((seq1, seq2))
+    >>> np.array(token_ids)
+    array([101, 1, 2, 3, 102,  11,  12, 102], dtype=int32)
+    >>> np.array(segment_ids)
+    array([0, 0, 0, 0, 0, 1, 1, 1], dtype=int32)
+
+    *Pack multiple inputs for classification with different sep tokens.*
+    >>> seq1 = [1, 2, 3, 4]
+    >>> seq2 = [11, 12, 13, 14]
+    >>> packer = keras_hub.layers.MultiSegmentPacker(
+    ...     sequence_length=8,
+    ...     start_value=101,
+    ...     end_value=102,
+    ...     sep_value=[102, 102],
+    ... )
+    >>> token_ids, segment_ids = packer((seq1, seq2))
+    >>> np.array(token_ids)
+    array([101,   1,   2, 102, 102,  11,  12, 102], dtype=int32)
+    >>> np.array(segment_ids)
+    array([0, 0, 0, 0, 0, 1, 1, 1], dtype=int32)
+
+    Reference:
+        [Devlin et al., 2018](https://arxiv.org/abs/1810.04805).
+    """
+
+    def __init__(
+        self,
+        sequence_length,
+        start_value,
+        end_value,
+        sep_value=None,
+        pad_value=None,
+        truncate="round_robin",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.sequence_length = sequence_length
+        if truncate not in ("round_robin", "waterfall"):
+            raise ValueError(
+                "Only 'round_robin' and 'waterfall' algorithms are "
+                "supported. Received %s" % truncate
+            )
+        self.truncate = truncate
+
+        # Maintain private copies of start/end values for config purposes.
+        self._start_value = start_value
+        self._sep_value = sep_value
+        self._end_value = end_value
+
+        def check_special_value_type(value, value_name):
+            if isinstance(value, (int, str)):
+                return [value]
+            if value and not isinstance(value, (list, tuple)):
+                raise ValueError(
+                    f"{value_name} should be of type int/str/list/tuple."
+                    f"Received type: `{type(value)}`."
+                )
+            return value
+
+        start_value = check_special_value_type(start_value, "start_value")
+        if sep_value is None:
+            sep_value = end_value
+        sep_value = check_special_value_type(sep_value, "sep_value")
+        end_value = check_special_value_type(end_value, "end_value")
+
+        self.start_value = start_value
+        self.sep_value = sep_value
+        self.end_value = end_value
+
+        self.pad_value = pad_value
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "sequence_length": self.sequence_length,
+                "start_value": self._start_value,
+                "end_value": self._end_value,
+                "sep_value": self._sep_value,
+                "pad_value": self.pad_value,
+                "truncate": self.truncate,
+            }
+        )
+        return config
+
+    def _sanitize_inputs(self, inputs):
+        """Force inputs to a list of rank 2 ragged tensors."""
+        # Sanitize inputs.
+        if not isinstance(inputs, (list, tuple)):
+            inputs = (inputs,)
+        if not inputs:
+            raise ValueError(
+                "At least one input is required for packing. "
+                f"Received: `inputs={inputs}`"
+            )
+        inputs, unbatched_list, _ = list(
+            zip(*(convert_to_ragged_batch(x) for x in inputs))
+        )
+        if len(set(unbatched_list)) != 1:
+            ranks = [1 if unbatched else 2 for unbatched in unbatched_list]
+            raise ValueError(
+                "All inputs for packing must have the same rank. "
+                f"Received: `inputs={inputs}` with ranks {ranks}"
+            )
+        return inputs, unbatched_list[0]
+
+    def _trim_inputs(self, inputs):
+        """Trim inputs to desired length."""
+        num_segments = len(inputs)
+        num_special_tokens = (
+            len(self.start_value)
+            + (num_segments - 1) * len(self.sep_value)
+            + len(self.end_value)
+        )
+        if self.truncate == "round_robin":
+            return tf_text.RoundRobinTrimmer(
+                self.sequence_length - num_special_tokens
+            ).trim(inputs)
+        elif self.truncate == "waterfall":
+            return tf_text.WaterfallTrimmer(
+                self.sequence_length - num_special_tokens
+            ).trim(inputs)
+        else:
+            raise ValueError("Unsupported truncate: %s" % self.truncate)
+
+    def _combine_inputs(
+        self,
+        segments,
+        add_start_value=True,
+        add_end_value=True,
+    ):
+        """Combine inputs with start and end values added."""
+        dtype = segments[0].dtype
+        batch_size = segments[0].nrows()
+        start_value = tf.convert_to_tensor(self.start_value, dtype=dtype)
+        sep_value = tf.convert_to_tensor(self.sep_value, dtype=dtype)
+        end_value = tf.convert_to_tensor(self.end_value, dtype=dtype)
+
+        start_columns = tf.repeat(
+            start_value[tf.newaxis, :], repeats=batch_size, axis=0
+        )
+        sep_columns = tf.repeat(
+            sep_value[tf.newaxis, :], repeats=batch_size, axis=0
+        )
+        end_columns = tf.repeat(
+            end_value[tf.newaxis, :], repeats=batch_size, axis=0
+        )
+        ones_sep_columns = tf.ones_like(sep_columns, dtype="int32")
+        ones_end_columns = tf.ones_like(end_columns, dtype="int32")
+
+        segments_to_combine = []
+        segment_ids_to_combine = []
+        if add_start_value:
+            segments_to_combine.append(start_columns)
+            start_segment = tf.zeros_like(start_columns, dtype="int32")
+            segment_ids_to_combine.append(start_segment)
+
+        for i, seg in enumerate(segments):
+            # Combine all segments.
+            segments_to_combine.append(seg)
+
+            # Combine segment ids.
+            segment_ids_to_combine.append(tf.ones_like(seg, dtype="int32") * i)
+
+            # Account for the sep/end tokens here.
+            if i == len(segments) - 1:
+                if add_end_value:
+                    segments_to_combine.append(end_columns)
+                    segment_ids_to_combine.append(ones_end_columns * i)
+            else:
+                segments_to_combine.append(sep_columns)
+                segment_ids_to_combine.append(ones_sep_columns * i)
+
+        token_ids = tf.concat(segments_to_combine, 1)
+        segment_ids = tf.concat(segment_ids_to_combine, 1)
+        return token_ids, segment_ids
+
+    def call(
+        self,
+        inputs,
+        sequence_length=None,
+        add_start_value=True,
+        add_end_value=True,
+    ):
+        inputs, unbatched = self._sanitize_inputs(inputs)
+
+        segments = self._trim_inputs(inputs)
+        token_ids, segment_ids = self._combine_inputs(
+            segments,
+            add_start_value=add_start_value,
+            add_end_value=add_end_value,
+        )
+        # Pad to dense tensor output.
+        sequence_length = sequence_length or self.sequence_length
+        shape = tf.cast([-1, sequence_length], "int64")
+        token_ids = token_ids.to_tensor(
+            shape=shape, default_value=self.pad_value
+        )
+        segment_ids = segment_ids.to_tensor(shape=shape)
+        # Remove the batch dim if added.
+        if unbatched:
+            token_ids = tf.squeeze(token_ids, 0)
+            segment_ids = tf.squeeze(segment_ids, 0)
+
+        return (token_ids, segment_ids)
+
+    def compute_output_shape(self, inputs_shape):
+        if isinstance(inputs_shape[0], tuple):
+            inputs_shape = inputs_shape[0]
+        inputs_shape = list(inputs_shape)
+        inputs_shape[-1] = self.sequence_length
+        return tuple(inputs_shape)

keras_hub/src/layers/preprocessing/preprocessing_layer.py

@@ -0,0 +1,62 @@
+# Copyright 2024 The KerasHub Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import keras
+from keras import tree
+
+from keras_hub.src.utils.tensor_utils import assert_tf_libs_installed
+from keras_hub.src.utils.tensor_utils import (
+    convert_to_backend_tensor_or_python_list,
+)
+
+try:
+    import tensorflow as tf
+except ImportError:
+    tf = None
+
+
+class PreprocessingLayer(keras.layers.Layer):
+    """Preprocessing layer base class."""
+
+    def __init__(self, **kwargs):
+        assert_tf_libs_installed(self.__class__.__name__)
+
+        super().__init__(**kwargs)
+        self._convert_input_args = False
+        self._allow_non_tensor_positional_args = True
+        # Most pre-preprocessing has no build.
+        if not hasattr(self, "build"):
+            self.built = True
+
+    def get_build_config(self):
+        return None
+
+    def __call__(self, *args, **kwargs):
+        # Always place on CPU for preprocessing, to avoid expensive back and
+        # forth copies to GPU before the trainable model.
+        with tf.device("cpu"):
+            outputs = super().__call__(*args, **kwargs)
+
+            # Jax and Torch lack native string and ragged types.
+            # If we are running on those backends and not running with tf.data
+            # (we are outside a tf.function), we covert all ragged and string
+            # tensor to pythonic types.
+            is_tf_backend = keras.config.backend() == "tensorflow"
+            is_in_tf_graph = not tf.executing_eagerly()
+            if not is_tf_backend and not is_in_tf_graph:
+                outputs = tree.map_structure(
+                    convert_to_backend_tensor_or_python_list, outputs
+                )
+
+        return outputs

keras_hub/src/layers/preprocessing/random_deletion.py

@@ -0,0 +1,271 @@
+# Copyright 2024 The KerasHub Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import random
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.layers.preprocessing.preprocessing_layer import (
+    PreprocessingLayer,
+)
+from keras_hub.src.utils.tensor_utils import convert_to_ragged_batch
+from keras_hub.src.utils.tensor_utils import is_int_dtype
+from keras_hub.src.utils.tensor_utils import is_string_dtype
+
+try:
+    import tensorflow as tf
+except ImportError:
+    tf = None
+
+
+@keras_hub_export("keras_hub.layers.RandomDeletion")
+class RandomDeletion(PreprocessingLayer):
+    """Augments input by randomly deleting tokens.
+
+    This layer comes in handy when you need to generate new data using deletion
+    augmentation as described in the paper [EDA: Easy Data Augmentation
+    Techniques for Boosting Performance on Text Classification Tasks]
+    (https://arxiv.org/pdf/1901.11196.pdf). The layer expects the inputs to be
+    pre-split into token level inputs. This allows control over the level of
+    augmentation, you can split by character for character level swaps, or by
+    word for word level swaps.
+
+    Input data should be passed as tensors, `tf.RaggedTensor`s, or lists. For
+    batched input, inputs should be a list of lists or a rank two tensor. For
+    unbatched inputs, each element should be a list or a rank one tensor.
+
+    Args:
+        rate: The probability of a token being chosen for deletion.
+        max_deletions: The maximum number of tokens to delete.
+        skip_list: A list of token values that should not be considered
+            candidates for deletion.
+        skip_fn: A function that takes as input a scalar tensor token and
+            returns as output a scalar tensor True/False value. A value of
+            True indicates that the token should not be considered a
+            candidate for deletion. This function must be tracable--it
+            should consist of tensorflow operations.
+        skip_py_fn: A function that takes as input a python token value and
+            returns as output `True` or `False`. A value of True
+            indicates that should not be considered a candidate for deletion.
+            Unlike the `skip_fn` argument, this argument need not be
+            tracable--it can be any python function.
+        seed: A seed for the random number generator.
+
+    Examples:
+
+    Word level usage.
+    >>> keras.utils.set_random_seed(1337)
+    >>> inputs=tf.strings.split(["Hey I like", "Keras and Tensorflow"])
+    >>> augmenter=keras_hub.layers.RandomDeletion(rate=0.4, seed=42)
+    >>> augmented=augmenter(inputs)
+    >>> tf.strings.reduce_join(augmented, separator=" ", axis=-1)
+    <tf.Tensor: shape=(2,), dtype=string, numpy=array([b'I like', b'and'],
+    dtype=object)>
+
+    Character level usage.
+    >>> keras.utils.set_random_seed(1337)
+    >>> inputs=tf.strings.unicode_split(["Hey Dude", "Speed Up"], "UTF-8")
+    >>> augmenter=keras_hub.layers.RandomDeletion(rate=0.4, seed=42)
+    >>> augmented=augmenter(inputs)
+    >>> tf.strings.reduce_join(augmented, axis=-1)
+    <tf.Tensor: shape=(2,), dtype=string, numpy=array([b'H Dude', b'pedUp'],
+    dtype=object)>
+
+    Usage with skip_list.
+    >>> keras.utils.set_random_seed(1337)
+    >>> inputs=tf.strings.split(["Hey I like", "Keras and Tensorflow"])
+    >>> augmenter=keras_hub.layers.RandomDeletion(rate=0.4,
+    ...     skip_list=["Keras", "Tensorflow"], seed=42)
+    >>> augmented=augmenter(inputs)
+    >>> tf.strings.reduce_join(augmented, separator=" ", axis=-1)
+    <tf.Tensor: shape=(2,), dtype=string,
+    numpy=array([b'I like', b'Keras Tensorflow'], dtype=object)>
+
+    Usage with skip_fn.
+    >>> def skip_fn(word):
+    ...     return tf.strings.regex_full_match(word, r"\\pP")
+    >>> keras.utils.set_random_seed(1337)
+    >>> inputs=tf.strings.split(["Hey I like", "Keras and Tensorflow"])
+    >>> augmenter=keras_hub.layers.RandomDeletion(rate=0.4,
+    ...     skip_fn=skip_fn, seed=42)
+    >>> augmented=augmenter(inputs)
+    >>> tf.strings.reduce_join(augmented, separator=" ", axis=-1)
+    <tf.Tensor: shape=(2,), dtype=string, numpy=array([b'I like', b'and'],
+    dtype=object)>
+
+    Usage with skip_py_fn.
+    >>> def skip_py_fn(word):
+    ...     return len(word) < 4
+    >>> keras.utils.set_random_seed(1337)
+    >>> inputs=tf.strings.split(["Hey I like", "Keras and Tensorflow"])
+    >>> augmenter=RandomDeletion(rate=0.4,
+    ...     skip_py_fn=skip_py_fn, seed=42)
+    >>> augmented=augmenter(inputs)
+    >>> tf.strings.reduce_join(augmented, separator=" ", axis=-1)
+    <tf.Tensor: shape=(2,), dtype=string,
+    numpy=array([b'Hey I', b'and Tensorflow'], dtype=object)>
+    """
+
+    def __init__(
+        self,
+        rate,
+        max_deletions=None,
+        skip_list=None,
+        skip_fn=None,
+        skip_py_fn=None,
+        seed=None,
+        name=None,
+        dtype="int32",
+        **kwargs,
+    ):
+        if not is_int_dtype(dtype) and not is_string_dtype(dtype):
+            raise ValueError(
+                "Output dtype must be an integer type or a string. "
+                f"Received: dtype={dtype}"
+            )
+
+        super().__init__(dtype=dtype, name=name, **kwargs)
+
+        self.rate = rate
+        self.max_deletions = max_deletions
+        self.seed = random.randint(1, 1e9) if seed is None else seed
+        self._generator = tf.random.Generator.from_seed(self.seed)
+        self.skip_list = skip_list
+        self.skip_fn = skip_fn
+        self.skip_py_fn = skip_py_fn
+        if self.max_deletions is not None and self.max_deletions < 0:
+            raise ValueError(
+                "max_deletions must be non-negative."
+                f"Received max_deletions={max_deletions}."
+            )
+
+        if self.rate > 1 or self.rate < 0:
+            raise ValueError(
+                "Rate must be between 0 and 1 (both inclusive)."
+                f"Received: rate={rate}"
+            )
+
+        if [self.skip_list, self.skip_fn, self.skip_py_fn].count(None) < 2:
+            raise ValueError(
+                "Exactly one of `skip_list`, `skip_fn`, `skip_py_fn` must be "
+                "provided."
+            )
+
+        if self.skip_list:
+            self.StaticHashTable = tf.lookup.StaticHashTable(
+                tf.lookup.KeyValueTensorInitializer(
+                    tf.convert_to_tensor(self.skip_list),
+                    tf.convert_to_tensor([True] * len(self.skip_list)),
+                ),
+                default_value=False,
+            )
+
+    def call(self, inputs):
+        inputs, unbatched, _ = convert_to_ragged_batch(inputs)
+
+        skip_masks = None
+        if self.skip_list:
+            skip_masks = self.StaticHashTable.lookup(inputs.flat_values)
+        elif self.skip_fn:
+            skip_masks = tf.map_fn(
+                self.skip_fn, inputs.flat_values, fn_output_signature="bool"
+            )
+        elif self.skip_py_fn:
+
+            def string_fn(token):
+                return self.skip_py_fn(token.numpy().decode("utf-8"))
+
+            def int_fn(token):
+                return self.skip_py_fn(token.numpy())
+
+            py_fn = string_fn if inputs.dtype == tf.string else int_fn
+
+            skip_masks = tf.map_fn(
+                lambda x: tf.py_function(py_fn, [x], "bool"),
+                inputs.flat_values,
+                fn_output_signature="bool",
+            )
+
+        positions_flat = tf.range(tf.size(inputs.flat_values))
+        positions = inputs.with_flat_values(positions_flat)
+        if skip_masks is not None:
+            skip_masks = tf.logical_not(skip_masks)
+            skip_masks.set_shape([None])
+            positions = tf.ragged.boolean_mask(
+                positions, inputs.with_flat_values(skip_masks)
+            )
+
+        # Figure out how many we are going to select.
+        token_counts = tf.cast(positions.row_lengths(), "float32")
+        num_to_select = tf.random.stateless_binomial(
+            shape=tf.shape(token_counts),
+            seed=self._generator.make_seeds()[:, 0],
+            counts=token_counts,
+            probs=self.rate,
+        )
+        if self.max_deletions is not None:
+            num_to_select = tf.math.minimum(num_to_select, self.max_deletions)
+        num_to_select = tf.cast(num_to_select, "int64")
+
+        # Shuffle and trim to items that are going to be selected.
+        def _shuffle_and_trim(x):
+            positions, top_n = x
+            shuffled = tf.random.shuffle(positions, seed=self.seed)
+            return shuffled[:top_n]
+
+        selected_for_mask = tf.map_fn(
+            _shuffle_and_trim,
+            (positions, num_to_select),
+            fn_output_signature=tf.RaggedTensorSpec(
+                ragged_rank=positions.ragged_rank - 1, dtype=positions.dtype
+            ),
+        )
+        selected_for_mask.flat_values.set_shape([None])
+
+        # Construct the mask which is a boolean RT
+        # Scatter 0's to positions that have been selector for deletion.
+        update_values = tf.zeros_like(selected_for_mask.flat_values, "int32")
+        update_indices = selected_for_mask.flat_values
+        update_indices = tf.expand_dims(update_indices, -1)
+        update_indices = tf.cast(update_indices, "int32")
+        mask_flat = tf.ones_like(inputs.flat_values, dtype="int32")
+        mask_flat = tf.tensor_scatter_nd_update(
+            mask_flat, update_indices, update_values
+        )
+        mask = tf.cast(inputs.with_flat_values(mask_flat), "bool")
+
+        inputs = tf.ragged.boolean_mask(inputs, mask)
+
+        if unbatched:
+            inputs = tf.squeeze(inputs, axis=0)
+
+        return inputs
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "rate": self.rate,
+                "max_deletions": self.max_deletions,
+                "seed": self.seed,
+                "skip_list": self.skip_list,
+                "skip_fn": self.skip_fn,
+                "skip_py_fn": self.skip_py_fn,
+            }
+        )
+        return config
+
+    def compute_output_shape(self, inputs_shape):
+        inputs_shape = list(inputs_shape)
+        inputs_shape[-1] = None
+        return tuple(inputs_shape)