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

keras_hub/src/layers/modeling/masked_lm_head.py

@@ -0,0 +1,239 @@
+# 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 ops
+
+from keras_hub.src.api_export import keras_hub_export
+
+
+@keras_hub_export("keras_hub.layers.MaskedLMHead")
+class MaskedLMHead(keras.layers.Layer):
+    """Masked Language Model (MaskedLM) head.
+
+    This layer takes two inputs:
+
+     - `inputs`: which should be a tensor of encoded tokens with shape
+       `(batch_size, sequence_length, hidden_dim)`.
+     - `mask_positions`: which should be a tensor of integer positions to
+       predict with shape `(batch_size, masks_per_sequence)`.
+
+    The token encodings should usually be the last output of an encoder model,
+    and mask positions should be the integer positions you would like to
+    predict for the MaskedLM task.
+
+    The layer will first gather the token encodings at the mask positions. These
+    gathered tokens will be passed through a dense layer the same size as
+    encoding dimension, then transformed to predictions the same size as the
+    input vocabulary. This layer will produce a single output with shape
+    `(batch_size, masks_per_sequence, vocabulary_size)`, which can be used to
+    compute an MaskedLM loss function.
+
+    This layer is often be paired with `keras_hub.layers.MaskedLMMaskGenerator`,
+    which will help prepare inputs for the MaskedLM task.
+
+    Args:
+        vocabulary_size: The total size of the vocabulary for predictions.
+        token_embedding: Optional. A `keras_hub.layers.ReversibleEmbedding`
+            instance. If passed, the layer will be used to project from the
+            `hidden_dim` of the model to the output `vocabulary_size`.
+        intermediate_activation: The activation function of intermediate dense layer.
+        activation: The activation function for the outputs of the layer.
+            Usually either `None` (return logits), or `"softmax"`
+            (return probabilities).
+        layer_norm_epsilon: float. The epsilon value in layer
+            normalization components. Defaults to `1e-5`.
+        kernel_initializer: string or `keras.initializers` initializer.
+            The kernel initializer for the dense and multiheaded
+            attention layers. Defaults to `"glorot_uniform"`.
+        bias_initializer: string or `keras.initializers` initializer.
+            The bias initializer for the dense and multiheaded
+            attention layers. Defaults to `"zeros"`.
+        **kwargs: other keyword arguments passed to `keras.layers.Layer`,
+            including `name`, `trainable`, `dtype` etc.
+
+    Example:
+
+    ```python
+    batch_size = 16
+    vocab_size = 100
+    hidden_dim = 32
+    seq_length = 50
+
+    # Generate random inputs.
+    token_ids = np.random.randint(vocab_size, size=(batch_size, seq_length))
+    # Choose random positions as the masked inputs.
+    mask_positions = np.random.randint(seq_length, size=(batch_size, 5))
+
+    # Embed tokens in a `hidden_dim` feature space.
+    token_embedding = keras_hub.layers.ReversibleEmbedding(
+        vocab_size,
+        hidden_dim,
+    )
+    hidden_states = token_embedding(token_ids)
+
+    preds = keras_hub.layers.MaskedLMHead(
+        vocabulary_size=vocab_size,
+        token_embedding=token_embedding,
+        activation="softmax",
+    )(hidden_states, mask_positions)
+    ```
+
+    References:
+     - [Press and Wolf, 2016](https://arxiv.org/abs/1608.05859)
+    """
+
+    def __init__(
+        self,
+        vocabulary_size=None,
+        token_embedding=None,
+        intermediate_activation="relu",
+        activation=None,
+        layer_norm_epsilon=1e-05,
+        kernel_initializer="glorot_uniform",
+        bias_initializer="zeros",
+        **kwargs,
+    ):
+        super().__init__(**kwargs, autocast=False)
+
+        self.vocabulary_size = vocabulary_size
+        self.token_embedding = token_embedding
+        self.intermediate_activation = keras.activations.get(
+            intermediate_activation
+        )
+        self.activation = keras.activations.get(activation)
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.kernel_initializer = keras.initializers.get(kernel_initializer)
+        self.bias_initializer = keras.initializers.get(bias_initializer)
+
+        if vocabulary_size is None and token_embedding is None:
+            raise ValueError(
+                "One of `vocabulary_size` or `token_embedding` must be set. "
+                "Received: `vocabulary_size=None`, `token_embedding=None`"
+            )
+
+        if token_embedding:
+            if vocabulary_size and vocabulary_size != token_embedding.input_dim:
+                raise ValueError(
+                    "`vocabulary_size` should match the input dimension of the "
+                    "of `token_embedding`. Received: "
+                    f"`vocabulary_size={vocabulary_size}`, "
+                    f"`token_embedding.input_dim={token_embedding.input_dim}`"
+                )
+            self.vocabulary_size = token_embedding.input_dim
+
+    def build(self, inputs_shape, mask_positions_shape=None):
+        if self.token_embedding is not None:
+            feature_size = self.token_embedding.output_dim
+        else:
+            feature_size = inputs_shape[-1]
+
+        self._intermediate_dense = keras.layers.Dense(
+            feature_size,
+            activation=self.intermediate_activation,
+            kernel_initializer=self.kernel_initializer,
+            bias_initializer=self.bias_initializer,
+            dtype=self.dtype_policy,
+            name="intermediate_dense",
+        )
+        self._intermediate_layer_norm = keras.layers.LayerNormalization(
+            epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
+            name="intermediate_layer_norm",
+        )
+        # The gather length does not affect any of our built variables, so
+        # we can pass any value here.
+        gather_length = None
+        shape = (inputs_shape[0], gather_length, inputs_shape[-1])
+        self._intermediate_dense.build(shape)
+        shape = (inputs_shape[0], gather_length, feature_size)
+        self._intermediate_layer_norm.build(shape)
+        if self.token_embedding is None:
+            self._kernel = self.add_weight(
+                name="output_kernel",
+                shape=[feature_size, self.vocabulary_size],
+                initializer=self.kernel_initializer,
+                dtype=self.dtype,
+            )
+        self._bias = self.add_weight(
+            name="output_bias",
+            shape=[self.vocabulary_size],
+            initializer=self.bias_initializer,
+            dtype=self.dtype,
+        )
+        self.built = True
+
+    def call(self, inputs, mask_positions):
+        if keras.config.backend() == "tensorflow":
+            import tensorflow as tf
+
+            # On the tf backend, we need to work around an issue with dynamic
+            # shape broadcasting in take_along_axis.
+            x = tf.gather(inputs, mask_positions, batch_dims=1)
+        else:
+            # Gather the encoded tokens at the masked indices.
+            mask_positions = ops.expand_dims(mask_positions, axis=-1)
+            x = ops.take_along_axis(inputs, mask_positions, axis=1)
+
+        # Apply a trainable linear transformation and a layer norm.
+        x = self._intermediate_dense(x)
+        x = self._intermediate_layer_norm(x)
+
+        # Transform encodings to vocabulary_size predictions.
+        if self.token_embedding:
+            outputs = self.token_embedding(x, reverse=True)
+        else:
+            outputs = ops.matmul(x, self._kernel)
+        outputs = ops.cast(outputs, self.compute_dtype)
+        outputs = outputs + self._bias
+
+        # Apply a final activation.
+        if self.activation is not None:
+            outputs = self.activation(outputs)
+
+        return outputs
+
+    @classmethod
+    def from_config(cls, config):
+        embedding = config.get("token_embedding")
+        if embedding:
+            config["token_embedding"] = keras.layers.deserialize(embedding)
+        return super().from_config(config)
+
+    def get_config(self):
+        config = super().get_config()
+        embedding_config = None
+        if self.token_embedding:
+            embedding_config = keras.layers.serialize(self.token_embedding)
+        config.update(
+            {
+                "vocabulary_size": self.vocabulary_size,
+                "token_embedding": embedding_config,
+                "intermediate_activation": keras.activations.serialize(
+                    self.intermediate_activation
+                ),
+                "activation": keras.activations.serialize(self.activation),
+                "layer_norm_epsilon": self.layer_norm_epsilon,
+                "kernel_initializer": keras.initializers.serialize(
+                    self.kernel_initializer
+                ),
+                "bias_initializer": keras.initializers.serialize(
+                    self.bias_initializer
+                ),
+            }
+        )
+        return config
+
+    def compute_output_shape(self, inputs_shape, mask_positions_shape):
+        return mask_positions_shape + (self.vocabulary_size,)

keras_hub/src/layers/modeling/position_embedding.py

@@ -0,0 +1,123 @@
+# 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 ops
+
+from keras_hub.src.api_export import keras_hub_export
+
+
+@keras_hub_export("keras_hub.layers.PositionEmbedding")
+class PositionEmbedding(keras.layers.Layer):
+    """A layer which learns a position embedding for inputs sequences.
+
+    This class assumes that in the input tensor, the last dimension corresponds
+    to the features, and the dimension before the last corresponds to the
+    sequence.
+
+    This layer does not supporting masking, but can be combined with a
+    `keras.layers.Embedding` for padding mask support.
+
+    Args:
+        sequence_length: The maximum length of the dynamic sequence.
+        initializer: The initializer to use for the embedding weights. Defaults
+            to `"glorot_uniform"`.
+        seq_axis: The axis of the input tensor where we add the embeddings.
+        **kwargs: other keyword arguments passed to `keras.layers.Layer`,
+            including `name`, `trainable`, `dtype` etc.
+
+    Call arguments:
+        inputs: The tensor inputs to compute an embedding for, with shape
+            `(batch_size, sequence_length, hidden_dim)`. Only the input shape
+            will be used, as the position embedding does not depend on the
+            input sequence content.
+        start_index: An integer or integer tensor. The starting position to
+            compute the position embedding from. This is useful during cached
+            decoding, where each position is predicted separately in a loop.
+
+    Example:
+
+    Called directly on input.
+    >>> layer = keras_hub.layers.PositionEmbedding(sequence_length=10)
+    >>> layer(np.zeros((8, 10, 16)))
+
+    Combine with a token embedding.
+    ```python
+    seq_length = 50
+    vocab_size = 5000
+    embed_dim = 128
+    inputs = keras.Input(shape=(seq_length,))
+    token_embeddings = keras.layers.Embedding(
+        input_dim=vocab_size, output_dim=embed_dim
+    )(inputs)
+    position_embeddings = keras_hub.layers.PositionEmbedding(
+        sequence_length=seq_length
+    )(token_embeddings)
+    outputs = token_embeddings + position_embeddings
+    ```
+
+    Reference:
+     - [Devlin et al., 2019](https://arxiv.org/abs/1810.04805)
+    """
+
+    def __init__(
+        self,
+        sequence_length,
+        initializer="glorot_uniform",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        if sequence_length is None:
+            raise ValueError(
+                "`sequence_length` must be an Integer, received `None`."
+            )
+        self.sequence_length = int(sequence_length)
+        self.initializer = keras.initializers.get(initializer)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "sequence_length": self.sequence_length,
+                "initializer": keras.initializers.serialize(self.initializer),
+            }
+        )
+        return config
+
+    def build(self, inputs_shape):
+        feature_size = inputs_shape[-1]
+        self.position_embeddings = self.add_weight(
+            name="embeddings",
+            shape=[self.sequence_length, feature_size],
+            initializer=self.initializer,
+            trainable=True,
+        )
+        self.built = True
+
+    def call(self, inputs, start_index=0):
+        shape = ops.shape(inputs)
+        feature_length = shape[-1]
+        sequence_length = shape[-2]
+        # trim to match the length of the input sequence, which might be less
+        # than the sequence_length of the layer.
+        position_embeddings = ops.convert_to_tensor(self.position_embeddings)
+        position_embeddings = ops.slice(
+            position_embeddings,
+            (start_index, 0),
+            (sequence_length, feature_length),
+        )
+        return ops.broadcast_to(position_embeddings, shape)
+
+    def compute_output_shape(self, input_shape):
+        return input_shape

keras_hub/src/layers/modeling/reversible_embedding.py

@@ -0,0 +1,311 @@
+# 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 ops
+from packaging.version import parse
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.utils.keras_utils import assert_quantization_support
+
+
+@keras_hub_export("keras_hub.layers.ReversibleEmbedding")
+class ReversibleEmbedding(keras.layers.Embedding):
+    """An embedding layer which can project backwards to the input dim.
+
+    This layer is an extension of `keras.layers.Embedding` for language models.
+    This layer can be called "in reverse" with `reverse=True`, in which case the
+    layer will linearly project from `output_dim` back to `input_dim`.
+
+    By default, the reverse projection will use the transpose of the
+    `embeddings` weights to project to `input_dim` (weights are "tied"). If
+    `tie_weights=False`, the model will use a separate, trainable variable for
+    reverse projection.
+
+    This layer has no bias terms.
+
+    Args:
+        input_dim: Integer. Size of the vocabulary,
+            i.e. maximum integer index + 1.
+        output_dim: Integer. Dimension of the dense embedding.
+        tie_weights: Boolean, whether or not the matrix for embedding and
+            the matrix for the `reverse` projection should share the same
+            weights.
+        embeddings_initializer: Initializer for the `embeddings`
+            matrix (see `keras.initializers`).
+        embeddings_regularizer: Regularizer function applied to
+            the `embeddings` matrix (see `keras.regularizers`).
+        embeddings_constraint: Constraint function applied to
+            the `embeddings` matrix (see `keras.constraints`).
+        mask_zero: Boolean, whether or not the input value 0 is a special
+            "padding" value that should be masked out.
+        reverse_dtype: The dtype for the reverse projection computation.
+            Defaults to the `compute_dtype` of the layer.
+        logit_soft_cap: If `logit_soft_cap` is set and `reverse=True`, the
+            output logits will be scaled by
+            `tanh(logits / logit_soft_cap) * logit_soft_cap`. This narrows the
+            range of output logits and can improve training.
+        **kwargs: other keyword arguments passed to `keras.layers.Embedding`,
+            including `name`, `trainable`, `dtype` etc.
+
+    Call arguments:
+        inputs: The tensor inputs to the layer.
+        reverse: Boolean. If `True` the layer will perform a linear projection
+            from `output_dim` to `input_dim`, instead of a normal embedding
+            call. Default to `False`.
+
+    Example:
+    ```python
+    batch_size = 16
+    vocab_size = 100
+    hidden_dim = 32
+    seq_length = 50
+
+    # Generate random inputs.
+    token_ids = np.random.randint(vocab_size, size=(batch_size, seq_length))
+
+    embedding = keras_hub.layers.ReversibleEmbedding(vocab_size, hidden_dim)
+    # Embed tokens to shape `(batch_size, seq_length, hidden_dim)`.
+    hidden_states = embedding(token_ids)
+    # Project hidden states to shape `(batch_size, seq_length, vocab_size)`.
+    logits = embedding(hidden_states, reverse=True)
+    ```
+
+    References:
+    - [Vaswani et al., 2017](https://arxiv.org/abs/1706.03762)
+    - [Press and Wolf, 2016](https://arxiv.org/abs/1608.05859)
+    """
+
+    def __init__(
+        self,
+        input_dim,
+        output_dim,
+        tie_weights=True,
+        embeddings_initializer="uniform",
+        embeddings_regularizer=None,
+        embeddings_constraint=None,
+        mask_zero=False,
+        reverse_dtype=None,
+        logit_soft_cap=None,
+        **kwargs,
+    ):
+        super().__init__(
+            input_dim,
+            output_dim,
+            embeddings_initializer=embeddings_initializer,
+            embeddings_regularizer=embeddings_regularizer,
+            embeddings_constraint=embeddings_constraint,
+            mask_zero=mask_zero,
+            **kwargs,
+        )
+        self.tie_weights = tie_weights
+        self.reverse_dtype = reverse_dtype
+        self.logit_soft_cap = logit_soft_cap
+
+    def build(self, inputs_shape=None):
+        super().build(inputs_shape)
+        if (
+            not self.tie_weights
+            and getattr(self, "quantization_mode", None) != "int8"
+        ):
+            self.reverse_embeddings = self.add_weight(
+                name="reverse_embeddings",
+                shape=(self.output_dim, self.input_dim),
+                initializer=self.embeddings_initializer,
+                dtype=self.dtype,
+            )
+
+    def call(self, inputs, reverse=False):
+        if reverse:
+            if self.tie_weights:
+                kernel = ops.transpose(ops.convert_to_tensor(self.embeddings))
+            else:
+                kernel = self.reverse_embeddings
+            if self.reverse_dtype is not None:
+                inputs = ops.cast(inputs, self.reverse_dtype)
+                kernel = ops.cast(kernel, self.reverse_dtype)
+            logits = ops.matmul(inputs, kernel)
+            # Optionally soft-cap logits.
+            if self.logit_soft_cap is not None:
+                soft_cap = self.logit_soft_cap
+                logits = ops.tanh(logits / soft_cap) * soft_cap
+            return logits
+
+        return super().call(inputs)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "tie_weights": self.tie_weights,
+                "reverse_dtype": self.reverse_dtype,
+                "logit_soft_cap": self.logit_soft_cap,
+            }
+        )
+        return config
+
+    def save_own_variables(self, store):
+        if not self.built:
+            return
+        super().save_own_variables(store)
+        # Before Keras 3.2, the reverse weight is saved in the super() call.
+        # After Keras 3.2, the reverse weight must be saved manually.
+        if parse(keras.version()) < parse("3.2.0"):
+            return
+        target_variables = []
+        if not self.tie_weights:
+            # Store the reverse embedding weights as the last weights.
+            target_variables.append(self.reverse_embeddings)
+            if getattr(self, "quantization_mode", None) == "int8":
+                target_variables.append(self.reverse_embeddings_scale)
+            for i, variable in enumerate(target_variables, start=len(store)):
+                store[str(i)] = variable
+
+    def load_own_variables(self, store):
+        if not self.built:
+            self.build()
+        super().load_own_variables(store)
+        if not self.tie_weights:
+            # Last weights in the stores are the reverse embedding weights.
+            target_variables = [self.reverse_embeddings]
+            if getattr(self, "quantization_mode", None) == "int8":
+                target_variables.append(self.reverse_embeddings_scale)
+            for i, variable in enumerate(
+                target_variables, start=len(store) - len(target_variables)
+            ):
+                variable.assign(store[str(i)])
+
+    def compute_output_spec(self, inputs, reverse=False):
+        output_shape = list(inputs.shape)
+        if reverse:
+            output_shape[-1] = self.input_dim
+        else:
+            output_shape += [self.output_dim]
+        return keras.KerasTensor(output_shape, dtype=self.compute_dtype)
+
+    # Quantization-related (int8) methods
+
+    def quantized_call(self, inputs, reverse=False):
+        # TODO (hongyu): This function could be removed once we add `*args` and
+        # `**kwargs` for `Embedding.quantized_call`
+        if self.quantization_mode == "int8":
+            return self._int8_call(inputs, reverse=reverse)
+        else:
+            self._quantization_mode_error(self.quantization_mode)
+
+    def _int8_build(
+        self,
+        embeddings_initializer="zeros",
+        embeddings_scale_initializer="ones",
+        reverse_embeddings_initializer="zeros",
+        reverse_embeddings_scale_initializer="ones",
+    ):
+        super()._int8_build(
+            embeddings_initializer, embeddings_scale_initializer
+        )
+        self.inputs_quantizer = keras.quantizers.AbsMaxQuantizer(axis=-1)
+        if not self.tie_weights:
+            self.reverse_embeddings = self.add_weight(
+                name="reverse_embeddings",
+                shape=(self.output_dim, self.input_dim),
+                initializer=reverse_embeddings_initializer,
+                dtype="int8",
+                trainable=False,
+            )
+            self.reverse_embeddings_scale = self.add_weight(
+                name="reverse_embeddings_scale",
+                shape=(self.input_dim,),
+                initializer=reverse_embeddings_scale_initializer,
+                trainable=False,
+            )
+
+    def _int8_call(self, inputs, reverse=False):
+        if reverse:
+            if self.tie_weights:
+                kernel = ops.transpose(self._embeddings)
+                scale = ops.transpose(self.embeddings_scale)
+            else:
+                kernel = self.reverse_embeddings
+                scale = self.reverse_embeddings_scale
+            inputs, inputs_scale = self.inputs_quantizer(inputs)
+            logits = ops.matmul(inputs, kernel)
+            # De-scale outputs
+            logits = ops.cast(logits, self.compute_dtype)
+            logits = ops.divide(logits, ops.multiply(inputs_scale, scale))
+            # Optionally soft-cap logits.
+            if self.logit_soft_cap is not None:
+                soft_cap = self.logit_soft_cap
+                logits = ops.tanh(logits / soft_cap) * soft_cap
+            return logits
+
+        return super()._int8_call(inputs)
+
+    def quantize(self, mode, type_check=True):
+        import gc
+        import inspect
+
+        assert_quantization_support()
+        if type_check and type(self) is not ReversibleEmbedding:
+            raise NotImplementedError(
+                f"Layer {self.__class__.__name__} does not have a `quantize()` "
+                "method implemented."
+            )
+        self._check_quantize_args(mode, self.compute_dtype)
+
+        def abs_max_quantize(inputs, axis):
+            sig = inspect.signature(keras.quantizers.abs_max_quantize)
+            if "to_numpy" in sig.parameters:
+                return keras.quantizers.abs_max_quantize(
+                    inputs, axis=axis, to_numpy=True
+                )
+            else:
+                # `keras<=3.4.1` doesn't support `to_numpy`
+                return keras.quantizers.abs_max_quantize(inputs, axis=axis)
+
+        self._tracker.unlock()
+        if mode == "int8":
+            embeddings, embeddings_scale = abs_max_quantize(
+                self._embeddings, axis=-1
+            )
+            embeddings_scale = ops.squeeze(embeddings_scale, axis=-1)
+            self._untrack_variable(self._embeddings)
+            del self._embeddings
+            if not self.tie_weights:
+                reverse_embeddings, reverse_embeddings_scale = abs_max_quantize(
+                    self.reverse_embeddings, axis=0
+                )
+                reverse_embeddings_scale = ops.squeeze(
+                    reverse_embeddings_scale, axis=0
+                )
+                self._untrack_variable(self.reverse_embeddings)
+                del self.reverse_embeddings
+            else:
+                reverse_embeddings = None
+                reverse_embeddings_scale = None
+            self._int8_build(
+                lambda shape, dtype: embeddings,
+                lambda shape, dtype: embeddings_scale,
+                lambda shape, dtype: reverse_embeddings,
+                lambda shape, dtype: reverse_embeddings_scale,
+            )
+        else:
+            raise self._quantization_mode_error(mode)
+        self._tracker.lock()
+
+        if self.dtype_policy.quantization_mode is None:
+            policy = keras.dtype_policies.get(
+                f"{mode}_from_{self.dtype_policy.name}"
+            )
+            self.dtype_policy = policy
+        gc.collect()