Repository: keras-team/keras-nlp
Pull request #1251 (2023-09-21 00:53:03)

Quick fix for 0.1.7 keras-core release

mattdangerw: /gcbrun

Patch:
diff --git a/keras_nlp/layers/modeling/f_net_encoder.py b/keras_nlp/layers/modeling/f_net_encoder.py
index f8de353e..a5370d96 100644
--- a/keras_nlp/layers/modeling/f_net_encoder.py
+++ b/keras_nlp/layers/modeling/f_net_encoder.py
@@ -97,11 +97,13 @@ class FNetEncoder(keras.layers.Layer):
         # Layer Norm layers.
         self._mixing_layer_norm = keras.layers.LayerNormalization(
             epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
             name="mixing_layer_norm",
         )
         self._mixing_layer_norm.build(inputs_shape)
         self._output_layer_norm = keras.layers.LayerNormalization(
             epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
             name="output_layer_norm",
         )
         self._output_layer_norm.build(inputs_shape)
@@ -112,6 +114,7 @@ class FNetEncoder(keras.layers.Layer):
             activation=self.activation,
             kernel_initializer=clone_initializer(self.kernel_initializer),
             bias_initializer=clone_initializer(self.bias_initializer),
+            dtype=self.dtype_policy,
             name="intermediate_dense",
         )
         self._intermediate_dense.build(inputs_shape)
@@ -119,12 +122,17 @@ class FNetEncoder(keras.layers.Layer):
             feature_size,
             kernel_initializer=clone_initializer(self.kernel_initializer),
             bias_initializer=clone_initializer(self.bias_initializer),
+            dtype=self.dtype_policy,
             name="output_dense",
         )
         self._output_dense.build(
             self._intermediate_dense.compute_output_shape(inputs_shape)
         )
-        self._output_dropout = keras.layers.Dropout(rate=self.dropout)
+        self._output_dropout = keras.layers.Dropout(
+            rate=self.dropout,
+            dtype=self.dtype_policy,
+            name="output_dropout",
+        )
         self.built = True
 
     def call(self, inputs):
@@ -140,9 +148,12 @@ class FNetEncoder(keras.layers.Layer):
 
         def fourier_transform(input):
             # Apply FFT on the input and take the real part.
+            input_dtype = input.dtype
+            # FFT transforms do not support float16.
+            input = ops.cast(input, "float32")
             real_in, imaginary_in = (input, ops.zeros_like(input))
             real_out, _ = ops.fft2((real_in, imaginary_in))
-            return real_out
+            return ops.cast(real_out, input_dtype)
 
         def add_and_norm(input1, input2, norm_layer):
             return norm_layer(input1 + input2)
diff --git a/keras_nlp/layers/modeling/masked_lm_head.py b/keras_nlp/layers/modeling/masked_lm_head.py
index 5bbcb491..eacee7e8 100644
--- a/keras_nlp/layers/modeling/masked_lm_head.py
+++ b/keras_nlp/layers/modeling/masked_lm_head.py
@@ -141,10 +141,12 @@ class MaskedLMHead(keras.layers.Layer):
             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
@@ -185,6 +187,7 @@ class MaskedLMHead(keras.layers.Layer):
             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.
diff --git a/keras_nlp/layers/modeling/reversible_embedding.py b/keras_nlp/layers/modeling/reversible_embedding.py
index 51de8dfd..e01d6698 100644
--- a/keras_nlp/layers/modeling/reversible_embedding.py
+++ b/keras_nlp/layers/modeling/reversible_embedding.py
@@ -49,6 +49,9 @@ class ReversibleEmbedding(keras.layers.Embedding):
             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.
+            For stability, it is usually best to use full precision even when
+            working with half or mixed precision training.
 
     Call args:
         inputs: The tensor inputs to the layer.
@@ -87,6 +90,7 @@ class ReversibleEmbedding(keras.layers.Embedding):
         embeddings_regularizer=None,
         embeddings_constraint=None,
         mask_zero=False,
+        reverse_dtype="float32",
         **kwargs,
     ):
         super().__init__(
@@ -99,6 +103,7 @@ class ReversibleEmbedding(keras.layers.Embedding):
             **kwargs,
         )
         self.tie_weights = tie_weights
+        self.reverse_dtype = reverse_dtype
 
     def build(self, inputs_shape=None):
         super().build(inputs_shape)
@@ -114,12 +119,12 @@ class ReversibleEmbedding(keras.layers.Embedding):
     def call(self, inputs, reverse=False):
         if reverse:
             if self.tie_weights:
-                reverse_embeddings = ops.transpose(
-                    ops.convert_to_tensor(self.embeddings)
-                )
+                kernel = ops.transpose(ops.convert_to_tensor(self.embeddings))
             else:
-                reverse_embeddings = self.reverse_embeddings
-            return ops.matmul(inputs, reverse_embeddings)
+                kernel = self.reverse_embeddings
+            inputs = ops.cast(inputs, self.reverse_dtype)
+            kernel = ops.cast(kernel, self.reverse_dtype)
+            return ops.matmul(inputs, kernel)
 
         return super().call(inputs)
 
@@ -128,6 +133,7 @@ class ReversibleEmbedding(keras.layers.Embedding):
         config.update(
             {
                 "tie_weights": self.tie_weights,
+                "reverse_dtype": self.reverse_dtype,
             }
         )
         return config
diff --git a/keras_nlp/layers/modeling/token_and_position_embedding.py b/keras_nlp/layers/modeling/token_and_position_embedding.py
index ad89af05..f3dffe34 100644
--- a/keras_nlp/layers/modeling/token_and_position_embedding.py
+++ b/keras_nlp/layers/modeling/token_and_position_embedding.py
@@ -93,11 +93,13 @@ class TokenAndPositionEmbedding(keras.layers.Layer):
                 self.embeddings_initializer
             ),
             mask_zero=mask_zero,
+            dtype=self.dtype_policy,
             name="token_embedding",
         )
         self.position_embedding = PositionEmbedding(
             sequence_length=sequence_length,
             initializer=clone_initializer(self.embeddings_initializer),
+            dtype=self.dtype_policy,
             name="position_embedding",
         )
         self.supports_masking = self.token_embedding.supports_masking
diff --git a/keras_nlp/layers/modeling/transformer_decoder.py b/keras_nlp/layers/modeling/transformer_decoder.py
index 4117a82e..92734fca 100644
--- a/keras_nlp/layers/modeling/transformer_decoder.py
+++ b/keras_nlp/layers/modeling/transformer_decoder.py
@@ -111,7 +111,6 @@ class TransformerDecoder(keras.layers.Layer):
         kernel_initializer="glorot_uniform",
         bias_initializer="zeros",
         normalize_first=False,
-        name=None,
         **kwargs,
     ):
         # Work around for model saving, we need to ensure our model is built
@@ -119,7 +118,7 @@ class TransformerDecoder(keras.layers.Layer):
         decoder_sequence_shape = kwargs.pop("decoder_sequence_shape", None)
         encoder_sequence_shape = kwargs.pop("encoder_sequence_shape", None)
 
-        super().__init__(name=name, **kwargs)
+        super().__init__(**kwargs)
         self.intermediate_dim = intermediate_dim
         self.num_heads = num_heads
         self.dropout = dropout
@@ -160,6 +159,7 @@ class TransformerDecoder(keras.layers.Layer):
             dropout=self.dropout,
             kernel_initializer=clone_initializer(self.kernel_initializer),
             bias_initializer=clone_initializer(self.bias_initializer),
+            dtype=self.dtype_policy,
             name="self_attention",
         )
         if hasattr(self._self_attention_layer, "_build_from_signature"):
@@ -174,11 +174,14 @@ class TransformerDecoder(keras.layers.Layer):
             )
         self._self_attention_layer_norm = keras.layers.LayerNormalization(
             epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
             name="self_attention_layer_norm",
         )
         self._self_attention_layer_norm.build(decoder_sequence_shape)
         self._self_attention_dropout = keras.layers.Dropout(
             rate=self.dropout,
+            dtype=self.dtype_policy,
+            name="self_attention_dropout",
         )
 
         # Cross attention layers are optional.
@@ -191,6 +194,7 @@ class TransformerDecoder(keras.layers.Layer):
                 dropout=self.dropout,
                 kernel_initializer=clone_initializer(self.kernel_initializer),
                 bias_initializer=clone_initializer(self.bias_initializer),
+                dtype=self.dtype_policy,
                 name="cross_attention",
             )
             if hasattr(self._cross_attention_layer, "_build_from_signature"):
@@ -205,11 +209,14 @@ class TransformerDecoder(keras.layers.Layer):
                 )
             self._cross_attention_layer_norm = keras.layers.LayerNormalization(
                 epsilon=self.layer_norm_epsilon,
+                dtype=self.dtype_policy,
                 name="cross_attention_layer_norm",
             )
             self._cross_attention_layer_norm.build(encoder_sequence_shape)
             self._cross_attention_dropout = keras.layers.Dropout(
                 rate=self.dropout,
+                dtype=self.dtype_policy,
+                name="cross_attention_dropout",
             )
 
         # Feedforward layers.
@@ -218,25 +225,30 @@ class TransformerDecoder(keras.layers.Layer):
             activation=self.activation,
             kernel_initializer=clone_initializer(self.kernel_initializer),
             bias_initializer=clone_initializer(self.bias_initializer),
-            name="intermediate_dense",
+            dtype=self.dtype_policy,
+            name="feedforward_intermediate_dense",
         )
         self._feedforward_intermediate_dense.build(decoder_sequence_shape)
         self._feedforward_output_dense = keras.layers.Dense(
             hidden_dim,
             kernel_initializer=clone_initializer(self.kernel_initializer),
             bias_initializer=clone_initializer(self.bias_initializer),
-            name="output_dense",
+            dtype=self.dtype_policy,
+            name="feedforward_output_dense",
         )
         intermediate_shape = list(decoder_sequence_shape)
         intermediate_shape[-1] = self.intermediate_dim
         self._feedforward_output_dense.build(tuple(intermediate_shape))
         self._feedforward_layer_norm = keras.layers.LayerNormalization(
             epsilon=self.layer_norm_epsilon,
-            name="output_layer_norm",
+            dtype=self.dtype_policy,
+            name="feedforward_layer_norm",
         )
         self._feedforward_layer_norm.build(decoder_sequence_shape)
         self._feedforward_dropout = keras.layers.Dropout(
             rate=self.dropout,
+            dtype=self.dtype_policy,
+            name="feedforward_dropout",
         )
         # Create layers based on input shape.
         self.built = True
diff --git a/keras_nlp/layers/modeling/transformer_encoder.py b/keras_nlp/layers/modeling/transformer_encoder.py
index 6f8f2ad9..32cdd355 100644
--- a/keras_nlp/layers/modeling/transformer_encoder.py
+++ b/keras_nlp/layers/modeling/transformer_encoder.py
@@ -92,10 +92,9 @@ class TransformerEncoder(keras.layers.Layer):
         kernel_initializer="glorot_uniform",
         bias_initializer="zeros",
         normalize_first=False,
-        name=None,
         **kwargs,
     ):
-        super().__init__(name=name, **kwargs)
+        super().__init__(**kwargs)
         self.intermediate_dim = intermediate_dim
         self.num_heads = num_heads
         self.dropout = dropout
@@ -125,6 +124,7 @@ class TransformerEncoder(keras.layers.Layer):
             dropout=self.dropout,
             kernel_initializer=clone_initializer(self.kernel_initializer),
             bias_initializer=clone_initializer(self.bias_initializer),
+            dtype=self.dtype_policy,
             name="self_attention_layer",
         )
         if hasattr(self._self_attention_layer, "_build_from_signature"):
@@ -139,17 +139,21 @@ class TransformerEncoder(keras.layers.Layer):
             )
         self._self_attention_layer_norm = keras.layers.LayerNormalization(
             epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
             name="self_attention_layer_norm",
         )
         self._self_attention_layer_norm.build(inputs_shape)
         self._self_attention_dropout = keras.layers.Dropout(
             rate=self.dropout,
+            dtype=self.dtype_policy,
             name="self_attention_dropout",
         )
 
         # Feedforward layers.
         self._feedforward_layer_norm = keras.layers.LayerNormalization(
             epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
+            name="feedforward_layer_norm",
         )
         self._feedforward_layer_norm.build(inputs_shape)
         self._feedforward_intermediate_dense = keras.layers.Dense(
@@ -157,6 +161,7 @@ class TransformerEncoder(keras.layers.Layer):
             activation=self.activation,
             kernel_initializer=clone_initializer(self.kernel_initializer),
             bias_initializer=clone_initializer(self.bias_initializer),
+            dtype=self.dtype_policy,
             name="feedforward_intermediate_dense",
         )
         self._feedforward_intermediate_dense.build(inputs_shape)
@@ -164,6 +169,7 @@ class TransformerEncoder(keras.layers.Layer):
             hidden_dim,
             kernel_initializer=clone_initializer(self.kernel_initializer),
             bias_initializer=clone_initializer(self.bias_initializer),
+            dtype=self.dtype_policy,
             name="feedforward_output_dense",
         )
         intermediate_shape = list(inputs_shape)
@@ -171,6 +177,8 @@ class TransformerEncoder(keras.layers.Layer):
         self._feedforward_output_dense.build(tuple(intermediate_shape))
         self._feedforward_dropout = keras.layers.Dropout(
             rate=self.dropout,
+            dtype=self.dtype_policy,
+            name="feedforward_dropout",
         )
         self.built = True
 
diff --git a/keras_nlp/models/xlnet/xlnet_content_and_query_embedding.py b/keras_nlp/models/xlnet/xlnet_content_and_query_embedding.py
index 1e7a63ed..a2bc2d0d 100644
--- a/keras_nlp/models/xlnet/xlnet_content_and_query_embedding.py
+++ b/keras_nlp/models/xlnet/xlnet_content_and_query_embedding.py
@@ -67,12 +67,12 @@ class ContentAndQueryEmbedding(keras.layers.Layer):
 
     def relative_positional_encoding(self, qlen, klen, bsz=None, clamp_len=-1):
         """create relative positional encoding."""
-        freq_seq = ops.arange(0, self.hidden_dim, 2.0)
+        freq_seq = ops.arange(0, self.hidden_dim, 2.0, dtype=self.compute_dtype)
         inv_freq = 1 / (10000 ** (freq_seq / self.hidden_dim))
 
         beg, end = klen, -qlen
 
-        fwd_pos_seq = ops.arange(beg, end, -1.0)
+        fwd_pos_seq = ops.arange(beg, end, -1.0, dtype=self.compute_dtype)
         if clamp_len > 0:
             fwd_pos_seq = ops.clip(
                 fwd_pos_seq, x_min=-clamp_len, x_max=clamp_len
diff --git a/keras_nlp/samplers/beam_sampler.py b/keras_nlp/samplers/beam_sampler.py
index f6f4c8e4..9562f95d 100644
--- a/keras_nlp/samplers/beam_sampler.py
+++ b/keras_nlp/samplers/beam_sampler.py
@@ -16,7 +16,6 @@ import tensorflow as tf
 import tree
 
 from keras_nlp.api_export import keras_nlp_export
-from keras_nlp.backend import keras
 from keras_nlp.backend import ops
 from keras_nlp.samplers.sampler import Sampler
 from keras_nlp.samplers.sampler import call_args_docstring
@@ -161,7 +160,7 @@ class BeamSampler(Sampler):
             # Compute the softmax distribution for the next token.
             logits, _, cache = next(prompt, cache, index)
             vocab_size = ops.shape(logits)[-1]
-            probs = keras.activations.softmax(logits / self.temperature)
+            probs = self.compute_probabilities(logits)
 
             # Compute the running log-likelihood of each new candidate.
             next_log_probs = ops.log(probs) + log_probs[..., None]
diff --git a/keras_nlp/samplers/contrastive_sampler.py b/keras_nlp/samplers/contrastive_sampler.py
index a5a91a79..bac65bcf 100644
--- a/keras_nlp/samplers/contrastive_sampler.py
+++ b/keras_nlp/samplers/contrastive_sampler.py
@@ -15,7 +15,6 @@
 import tree
 
 from keras_nlp.api_export import keras_nlp_export
-from keras_nlp.backend import keras
 from keras_nlp.backend import ops
 from keras_nlp.samplers.sampler import Sampler
 from keras_nlp.samplers.sampler import call_args_docstring
@@ -131,7 +130,7 @@ class ContrastiveSampler(Sampler):
 
         def body(prompt, cache, index, logits, hidden_states):
             # Compute the softmax distribution for the next token.
-            probabilities = keras.activations.softmax(logits / self.temperature)
+            probabilities = self.compute_probabilities(logits)
 
             # Replicate for `self.k` times to find the best token in top-k
             # candidates.
diff --git a/keras_nlp/samplers/sampler.py b/keras_nlp/samplers/sampler.py
index f7f08322..e28fbe9d 100644
--- a/keras_nlp/samplers/sampler.py
+++ b/keras_nlp/samplers/sampler.py
@@ -143,7 +143,7 @@ class Sampler:
         def body(prompt, cache, index):
             # Compute the softmax distribution for the next token.
             logits, _, cache = next(prompt, cache, index)
-            probabilities = keras.activations.softmax(logits / self.temperature)
+            probabilities = self.compute_probabilities(logits)
             # Compute the next token.
             next_token = self.get_next_token(probabilities)
             # Don't overwrite anywhere mask is True.
@@ -164,6 +164,17 @@ class Sampler:
         )
         return prompt
 
+    def compute_probabilities(self, logits):
+        """Compute token probabilities from logits.
+
+        This will always be done in full precision, regardless of dtype, and
+        scale by `temperature`.
+        """
+        logits_dtype = logits.dtype
+        logits = ops.cast(logits, "float32")
+        probs = keras.activations.softmax(logits / self.temperature)
+        return ops.cast(probs, logits_dtype)
+
     def run_loop(self, cond, body, loop_vars=None, maximum_iterations=None):
         """Run ops.while_loops with a `StatelessScope` if necessary."""
         if config.backend() == "jax":


Test patch:
diff --git a/keras_nlp/layers/modeling/cached_multi_head_attention_test.py b/keras_nlp/layers/modeling/cached_multi_head_attention_test.py
index 3f6ecbc7..fdaab606 100644
--- a/keras_nlp/layers/modeling/cached_multi_head_attention_test.py
+++ b/keras_nlp/layers/modeling/cached_multi_head_attention_test.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from keras_nlp.backend import config
 from keras_nlp.backend import ops
 from keras_nlp.layers.modeling.cached_multi_head_attention import (
     CachedMultiHeadAttention,
@@ -34,6 +35,9 @@ class CachedMultiHeadAttentionTest(TestCase):
             expected_output_shape=(2, 4, 6),
             expected_num_trainable_weights=8,
             expected_num_non_trainable_variables=1,
+            # tf.keras does not handle mixed precision correctly when not set
+            # globally.
+            run_mixed_precision_check=config.multi_backend(),
         )
 
     def test_cache_call_is_correct(self):
diff --git a/keras_nlp/layers/modeling/reversible_embedding_test.py b/keras_nlp/layers/modeling/reversible_embedding_test.py
index 16cc3f8f..ceb04578 100644
--- a/keras_nlp/layers/modeling/reversible_embedding_test.py
+++ b/keras_nlp/layers/modeling/reversible_embedding_test.py
@@ -17,6 +17,7 @@ import os
 import numpy as np
 from absl.testing import parameterized
 
+from keras_nlp.backend import config
 from keras_nlp.backend import keras
 from keras_nlp.backend import ops
 from keras_nlp.layers.modeling.reversible_embedding import ReversibleEmbedding
@@ -73,3 +74,22 @@ class ReversibleEmbeddingTest(TestCase):
 
         input_data = ops.ones(shape=(4, 10), dtype="int32")
         self.assertAllClose(untied_model(input_data), tied_model(input_data))
+
+    def test_reverse_dtype(self):
+        embedding = ReversibleEmbedding(100, 16, reverse_dtype="float32")
+        input_data = ops.ones(shape=(4, 10, 16))
+        output_data = embedding(input_data, reverse=True)
+        self.assertEqual(output_data.shape, (4, 10, 100))
+        self.assertDTypeEqual(output_data, "float32")
+
+        if config.backend() == "torch":
+            import torch
+
+            if not torch.cuda.is_available():
+                self.skipTest("Torch CPU does not support float16")
+
+        embedding = ReversibleEmbedding(100, 16, reverse_dtype="float16")
+        input_data = ops.ones(shape=(4, 10, 16))
+        output_data = embedding(input_data, reverse=True)
+        self.assertEqual(output_data.shape, (4, 10, 100))
+        self.assertDTypeEqual(output_data, "float16")
diff --git a/keras_nlp/tests/test_case.py b/keras_nlp/tests/test_case.py
index a53c29a6..d0af5a47 100644
--- a/keras_nlp/tests/test_case.py
+++ b/keras_nlp/tests/test_case.py
@@ -17,7 +17,8 @@ import json
 import tensorflow as tf
 import tree
 from absl.testing import parameterized
-from keras_core.src.backend import standardize_dtype
+from keras_core.backend import is_float_dtype
+from keras_core.backend import standardize_dtype
 
 from keras_nlp.backend import config
 from keras_nlp.backend import keras
@@ -67,6 +68,10 @@ class TestCase(tf.test.TestCase, parameterized.TestCase):
         x2 = tree.map_structure(convert_to_comparible_type, x2)
         super().assertAllEqual(x1, x2, msg=msg)
 
+    def assertDTypeEqual(self, x, expected_dtype, msg=None):
+        input_dtype = standardize_dtype(x.dtype)
+        super().assertEqual(input_dtype, expected_dtype, msg=msg)
+
     def run_layer_test(
         self,
         layer_cls,
@@ -78,6 +83,7 @@ class TestCase(tf.test.TestCase, parameterized.TestCase):
         expected_num_non_trainable_weights=0,
         expected_num_non_trainable_variables=0,
         run_training_check=True,
+        run_mixed_precision_check=True,
     ):
         # Serialization test.
         layer = layer_cls(**init_kwargs)
@@ -168,6 +174,25 @@ class TestCase(tf.test.TestCase, parameterized.TestCase):
         if run_training_check:
             run_training_step(layer, input_data, output_data)
 
+        # Never test mixed precision on torch CPU. Torch lacks support.
+        if run_mixed_precision_check and config.backend() == "torch":
+            import torch
+
+            run_mixed_precision_check = torch.cuda.is_available()
+
+        if run_mixed_precision_check:
+            layer = layer_cls(**{**init_kwargs, "dtype": "mixed_float16"})
+            if isinstance(input_data, dict):
+                output_data = layer(**input_data)
+            else:
+                output_data = layer(input_data)
+            for tensor in tree.flatten(output_data):
+                if is_float_dtype(tensor.dtype):
+                    self.assertDTypeEqual(tensor, "float16")
+            for weight in layer.weights:
+                if is_float_dtype(weight.dtype):
+                    self.assertDTypeEqual(weight, "float32")
+
     def run_class_serialization_test(self, instance):
         # get_config roundtrip
         cls = instance.__class__