keras-nlp 0.6.3.dev0__tar.gz → 0.6.4.dev0__tar.gz

{keras-nlp-0.6.3.dev0 → keras-nlp-0.6.4.dev0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: keras-nlp
-Version: 0.6.3.dev0
+Version: 0.6.4.dev0
 Summary: Industry-strength Natural Language Processing extensions for Keras.
 Home-page: https://github.com/keras-team/keras-nlp
 Author: Keras team
@@ -39,11 +39,10 @@ Requires-Dist: sentencepiece; extra == "extras"
 ![Tensorflow](https://img.shields.io/badge/tensorflow-v2.5.0+-success.svg)
 [![contributions welcome](https://img.shields.io/badge/contributions-welcome-brightgreen.svg?style=flat)](https://github.com/keras-team/keras-nlp/issues)
 
-KerasNLP is a natural language processing library that works natively 
-with TensorFlow, JAX, or PyTorch. Built on [Keras Core](https://keras.io/keras_core/announcement/),
-these models, layers, metrics, callbacks, etc., can be trained and serialized
-in any framework and re-used in another without costly migrations. See "Using 
-KerasNLP with Keras Core" below for more details on multi-framework KerasNLP.
+KerasNLP is a natural language processing library that works natively
+with TensorFlow, JAX, or PyTorch. Built on Keras 3, these models, layers,
+metrics, and tokenizers can be trained and serialized in any framework and
+re-used in another without costly migrations.
 
 KerasNLP supports users through their entire development cycle. Our workflows 
 are built from modular components that have state-of-the-art preset weights and 
@@ -82,17 +81,28 @@ We are a new and growing project and welcome [contributions](CONTRIBUTING.md).
 
 ## Installation
 
-To install the latest official release:
+KerasNLP supports both Keras 2 and Keras 3. We recommend Keras 3 for all new
+users, as it enables using KerasNLP models and layers with JAX, TensorFlow and
+PyTorch.
+
+### Keras 2 Installation
+
+To install the latest KerasNLP release with Keras 2, simply run:
 
 ```
-pip install keras-nlp --upgrade
+pip install --upgrade keras-nlp
 ```
 
-To install the latest unreleased changes to the library, we recommend using
-pip to install directly from the master branch on github:
+### Keras 3 Installation
+
+There are currently two ways to install Keras 3 with KerasNLP. To install the
+stable versions of KerasNLP and Keras 3, you should install Keras 3 **after**
+installing KerasNLP. This is a temporary step while TensorFlow is pinned to
+Keras 2, and will no longer be necessary after TensorFlow 2.16.
 
 ```
-pip install git+https://github.com/keras-team/keras-nlp.git --upgrade
+pip install --upgrade keras-nlp
+pip install --upgrade keras>=3
 ```
 ## Using KerasNLP with Keras Core
 
@@ -123,12 +133,28 @@ Until Keras Core is officially released as Keras 3.0, KerasNLP will use
 `.keras/keras_nlp.json`. You will need to restart the Python runtime for changes 
 to take effect.
 
+To install the latest nightly changes for both KerasNLP and Keras, you can use
+our nightly package.
+
+```
+pip install --upgrade keras-nlp-nightly
+```
+
+> [!IMPORTANT]
+> Keras 3 will not function with TensorFlow 2.14 or earlier.
+
+Read [Getting started with Keras](https://keras.io/getting_started/) for more information
+on installing Keras 3 and compatibility with different frameworks.
+
 ## Quickstart
 
 Fine-tune BERT on a small sentiment analysis task using the 
 [`keras_nlp.models`](https://keras.io/api/keras_nlp/models/) API:
 
 ```python
+import os
+os.environ["KERAS_BACKEND"] = "tensorflow"  # Or "jax" or "torch"!
+
 import keras_nlp
 import tensorflow_datasets as tfds
 
@@ -151,6 +177,29 @@ classifier.predict(["What an amazing movie!", "A total waste of my time."])
 
 For more in depth guides and examples, visit https://keras.io/keras_nlp/.
 
+## Configuring your backend
+
+If you have Keras 3 installed in your environment (see installation above),
+you can use KerasNLP with any of JAX, TensorFlow and PyTorch. To do so, set the
+`KERAS_BACKEND` environment variable. For example:
+
+```shell
+export KERAS_BACKEND=jax
+```
+
+Or in Colab, with:
+
+```python
+import os
+os.environ["KERAS_BACKEND"] = "jax"
+
+import keras_nlp
+```
+
+> [!IMPORTANT]
+> Make sure to set the `KERAS_BACKEND` before import any Keras libraries, it
+> will be used to set up Keras when it is first imported.
+
 ## Compatibility
 
 We follow [Semantic Versioning](https://semver.org/), and plan to

{keras-nlp-0.6.3.dev0 → keras-nlp-0.6.4.dev0}/README.md

@@ -4,11 +4,10 @@
 ![Tensorflow](https://img.shields.io/badge/tensorflow-v2.5.0+-success.svg)
 [![contributions welcome](https://img.shields.io/badge/contributions-welcome-brightgreen.svg?style=flat)](https://github.com/keras-team/keras-nlp/issues)
 
-KerasNLP is a natural language processing library that works natively 
-with TensorFlow, JAX, or PyTorch. Built on [Keras Core](https://keras.io/keras_core/announcement/),
-these models, layers, metrics, callbacks, etc., can be trained and serialized
-in any framework and re-used in another without costly migrations. See "Using 
-KerasNLP with Keras Core" below for more details on multi-framework KerasNLP.
+KerasNLP is a natural language processing library that works natively
+with TensorFlow, JAX, or PyTorch. Built on Keras 3, these models, layers,
+metrics, and tokenizers can be trained and serialized in any framework and
+re-used in another without costly migrations.
 
 KerasNLP supports users through their entire development cycle. Our workflows 
 are built from modular components that have state-of-the-art preset weights and 
@@ -47,17 +46,28 @@ We are a new and growing project and welcome [contributions](CONTRIBUTING.md).
 
 ## Installation
 
-To install the latest official release:
+KerasNLP supports both Keras 2 and Keras 3. We recommend Keras 3 for all new
+users, as it enables using KerasNLP models and layers with JAX, TensorFlow and
+PyTorch.
+
+### Keras 2 Installation
+
+To install the latest KerasNLP release with Keras 2, simply run:
 
 ```
-pip install keras-nlp --upgrade
+pip install --upgrade keras-nlp
 ```
 
-To install the latest unreleased changes to the library, we recommend using
-pip to install directly from the master branch on github:
+### Keras 3 Installation
+
+There are currently two ways to install Keras 3 with KerasNLP. To install the
+stable versions of KerasNLP and Keras 3, you should install Keras 3 **after**
+installing KerasNLP. This is a temporary step while TensorFlow is pinned to
+Keras 2, and will no longer be necessary after TensorFlow 2.16.
 
 ```
-pip install git+https://github.com/keras-team/keras-nlp.git --upgrade
+pip install --upgrade keras-nlp
+pip install --upgrade keras>=3
 ```
 ## Using KerasNLP with Keras Core
 
@@ -88,12 +98,28 @@ Until Keras Core is officially released as Keras 3.0, KerasNLP will use
 `.keras/keras_nlp.json`. You will need to restart the Python runtime for changes 
 to take effect.
 
+To install the latest nightly changes for both KerasNLP and Keras, you can use
+our nightly package.
+
+```
+pip install --upgrade keras-nlp-nightly
+```
+
+> [!IMPORTANT]
+> Keras 3 will not function with TensorFlow 2.14 or earlier.
+
+Read [Getting started with Keras](https://keras.io/getting_started/) for more information
+on installing Keras 3 and compatibility with different frameworks.
+
 ## Quickstart
 
 Fine-tune BERT on a small sentiment analysis task using the 
 [`keras_nlp.models`](https://keras.io/api/keras_nlp/models/) API:
 
 ```python
+import os
+os.environ["KERAS_BACKEND"] = "tensorflow"  # Or "jax" or "torch"!
+
 import keras_nlp
 import tensorflow_datasets as tfds
 
@@ -116,6 +142,29 @@ classifier.predict(["What an amazing movie!", "A total waste of my time."])
 
 For more in depth guides and examples, visit https://keras.io/keras_nlp/.
 
+## Configuring your backend
+
+If you have Keras 3 installed in your environment (see installation above),
+you can use KerasNLP with any of JAX, TensorFlow and PyTorch. To do so, set the
+`KERAS_BACKEND` environment variable. For example:
+
+```shell
+export KERAS_BACKEND=jax
+```
+
+Or in Colab, with:
+
+```python
+import os
+os.environ["KERAS_BACKEND"] = "jax"
+
+import keras_nlp
+```
+
+> [!IMPORTANT]
+> Make sure to set the `KERAS_BACKEND` before import any Keras libraries, it
+> will be used to set up Keras when it is first imported.
+
 ## Compatibility
 
 We follow [Semantic Versioning](https://semver.org/), and plan to

{keras-nlp-0.6.3.dev0 → keras-nlp-0.6.4.dev0}/keras_nlp/__init__.py

@@ -12,4 +12,4 @@ from keras_nlp import samplers
 from keras_nlp import tokenizers
 
 
-__version__ = "0.6.3.dev0"
+__version__ = "0.6.4.dev0"

{keras-nlp-0.6.3.dev0 → keras-nlp-0.6.4.dev0}/keras_nlp/src/__init__.py

@@ -28,5 +28,5 @@ from keras_nlp.src import tokenizers
 from keras_nlp.src import utils
 
 # This is the global source of truth for the version number.
-__version__ = "0.6.3.dev0"
+__version__ = "0.6.4.dev0"
 

{keras-nlp-0.6.3.dev0 → keras-nlp-0.6.4.dev0}/keras_nlp/src/backend/__init__.py

@@ -14,14 +14,16 @@
 """
 Keras backend module.
 
-This module adds a temporarily Keras API surface that is fully under KerasNLP
-control. This allows us to switch between `keras_core` and `tf.keras`, as well
-as add shims to support older version of `tf.keras`.
+This module adds a temporary Keras API surface that is fully under KerasNLP
+control. The goal is to allow us to write Keras 3-like code everywhere, while
+still supporting Keras 2. We do this by using the `keras_core` package with
+Keras 2 to backport Keras 3 numerics APIs (`keras.ops` and `keras.random`) into
+Keras 2. The sub-modules exposed are as follows:
 
-- `config`: check which backend is being run.
-- `keras`: The full `keras` API (via `keras_core` or `tf.keras`).
-- `ops`: `keras_core.ops`, always tf backed if using `tf.keras`.
-- `random`: `keras_core.random`, always tf backed if using `tf.keras`.
+- `config`: check which version of Keras is being run.
+- `keras`: The full `keras` API with compat shims for older Keras versions.
+- `ops`: `keras.ops` for Keras 3 or `keras_core.ops` for Keras 2.
+- `random`: `keras.random` for Keras 3 or `keras_core.ops` for Keras 2.
 """
 
 from keras_nlp.src.backend import config

keras-nlp-0.6.4.dev0/keras_nlp/src/backend/config.py

@@ -0,0 +1,64 @@
+# Copyright 2023 The KerasNLP 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 os
+
+
+def detect_if_tensorflow_uses_keras_3():
+    # We follow the version of keras that tensorflow is configured to use.
+    try:
+        from tensorflow import keras
+
+        # Note that only recent versions of keras have a `version()` function.
+        if hasattr(keras, "version") and keras.version().startswith("3."):
+            return True
+    except:
+        raise ValueError(
+            "Unable to import `keras` with `tensorflow`.  Please check your "
+            "Keras and Tensorflow version are compatible; Keras 3 requires "
+            "TensorFlow 2.15 or later. See keras.io/getting_started for more "
+            "information on installing Keras."
+        )
+
+    # No `keras.version()` means we are on an old version of keras.
+    return False
+
+
+_USE_KERAS_3 = detect_if_tensorflow_uses_keras_3()
+
+if not _USE_KERAS_3:
+    backend = os.environ.get("KERAS_BACKEND")
+    if backend and backend != "tensorflow":
+        raise RuntimeError(
+            "When running Keras 2, the `KERAS_BACKEND` environment variable "
+            f"must either be unset or `'tensorflow'`. Received: `{backend}`. "
+            "To set another backend, please install Keras 3. See "
+            "https://github.com/keras-team/keras-nlp#installation"
+        )
+
+
+def keras_3():
+    """Check if Keras 3 is being used."""
+    return _USE_KERAS_3
+
+
+def backend():
+    """Check the backend framework."""
+    if not keras_3():
+        return "tensorflow"
+
+    import keras
+
+    return keras.config.backend()
+

{keras-nlp-0.6.3.dev0 → keras-nlp-0.6.4.dev0}/keras_nlp/src/backend/keras.py

@@ -20,8 +20,6 @@ from keras_nlp.src.backend import config
 
 if config.keras_3():
     from keras import *  # noqa: F403, F401
-elif config.multi_backend():
-    from keras_core import *  # noqa: F403, F401
 else:
     from tensorflow.keras import *  # noqa: F403, F401
 

{keras-nlp-0.6.3.dev0 → keras-nlp-0.6.4.dev0}/keras_nlp/src/conftest.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
-
 import pytest
 import tensorflow as tf
 
@@ -21,21 +19,6 @@ from keras_nlp.src.backend import config as backend_config
 from keras_nlp.src.backend import keras
 
 
-@pytest.fixture(scope="session")
-def tpu_strategy():
-    tpu_name = os.getenv("KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS")
-    resolver = tf.distribute.cluster_resolver.TPUClusterResolver.connect(
-        tpu=tpu_name,
-    )
-    return tf.distribute.TPUStrategy(resolver)
-
-
-@pytest.fixture(scope="class")
-def tpu_test_class(request, tpu_strategy):
-    # set a class attribute on the invoking test context
-    request.cls.tpu_strategy = tpu_strategy
-
-
 def pytest_addoption(parser):
     parser.addoption(
         "--run_large",
@@ -49,18 +32,6 @@ def pytest_addoption(parser):
         default=False,
         help="run extra_large tests",
     )
-    parser.addoption(
-        "--run_tpu",
-        action="store_true",
-        default=False,
-        help="run tpu tests",
-    )
-    parser.addoption(
-        "--mixed_precision",
-        action="store_true",
-        default=False,
-        help="run with mixed precision",
-    )
     parser.addoption(
         "--docstring_module",
         action="store",
@@ -70,18 +41,13 @@ def pytest_addoption(parser):
 
 
 def pytest_configure(config):
-    if config.getoption("--mixed_precision"):
-        keras.mixed_precision.set_global_policy("mixed_float16")
-    config.addinivalue_line(
-        "markers", "large: mark test as being slow or requiring a network"
-    )
     config.addinivalue_line(
         "markers",
-        "extra_large: mark test as being too large to run continuously",
+        "large: mark test as being slow or requiring a network",
     )
     config.addinivalue_line(
         "markers",
-        "tpu: mark test as tpu test",
+        "extra_large: mark test as being too large to run continuously",
     )
     config.addinivalue_line(
         "markers",
@@ -93,7 +59,6 @@ def pytest_collection_modifyitems(config, items):
     run_extra_large_tests = config.getoption("--run_extra_large")
     # Run large tests for --run_extra_large or --run_large.
     run_large_tests = config.getoption("--run_large") or run_extra_large_tests
-    run_tpu = config.getoption("--run_tpu")
 
     # Messages to annotate skipped tests with.
     skip_large = pytest.mark.skipif(
@@ -104,32 +69,27 @@ def pytest_collection_modifyitems(config, items):
         not run_extra_large_tests,
         reason="need --run_extra_large option to run",
     )
-    skip_tpu = pytest.mark.skipif(
-        not run_tpu,
-        reason="need --run_tpu option to run",
-    )
-    skip_tf_only = pytest.mark.skipif(
+    tf_only = pytest.mark.skipif(
         not backend_config.backend() == "tensorflow",
         reason="tests only run on tf backend",
     )
+    keras_3_only = pytest.mark.skipif(
+        not backend_config.keras_3(),
+        reason="tests only run on with multi-backend keras",
+    )
     for item in items:
         if "large" in item.keywords:
             item.add_marker(skip_large)
         if "extra_large" in item.keywords:
             item.add_marker(skip_extra_large)
-        if "tpu" in item.keywords:
-            item.add_marker(skip_tpu)
         if "tf_only" in item.keywords:
-            item.add_marker(skip_tf_only)
+            item.add_marker(tf_only)
+        if "keras_3_only" in item.keywords:
+            item.add_marker(keras_3_only)
 
 
 # Disable traceback filtering for quicker debugging of tests failures.
 tf.debugging.disable_traceback_filtering()
-if backend_config.multi_backend():
+if backend_config.keras_3():
     keras.config.disable_traceback_filtering()
 
-# One off setup for dtensor tests.
-if not backend_config.multi_backend():
-    keras.backend.experimental.enable_tf_random_generator()
-    keras.utils.set_random_seed(1337)
-

{keras-nlp-0.6.3.dev0 → keras-nlp-0.6.4.dev0}/keras_nlp/src/layers/__init__.py

@@ -18,6 +18,7 @@ from keras_nlp.src.layers.modeling.cached_multi_head_attention import (
 from keras_nlp.src.layers.modeling.f_net_encoder import FNetEncoder
 from keras_nlp.src.layers.modeling.masked_lm_head import MaskedLMHead
 from keras_nlp.src.layers.modeling.position_embedding import PositionEmbedding
+from keras_nlp.src.layers.modeling.reversible_embedding import ReversibleEmbedding
 from keras_nlp.src.layers.modeling.rotary_embedding import RotaryEmbedding
 from keras_nlp.src.layers.modeling.sine_position_encoding import (
     SinePositionEncoding,

{keras-nlp-0.6.3.dev0 → keras-nlp-0.6.4.dev0}/keras_nlp/src/layers/modeling/cached_multi_head_attention_test.py

@@ -36,9 +36,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
+            # Keras 2 does not handle mixed precision correctly when not set
             # globally.
-            run_mixed_precision_check=config.multi_backend(),
+            run_mixed_precision_check=config.keras_3(),
         )
 
     def test_cache_call_is_correct(self):

{keras-nlp-0.6.3.dev0 → keras-nlp-0.6.4.dev0}/keras_nlp/src/models/albert/albert_backbone_test.py

@@ -102,34 +102,3 @@ class AlbertBackboneTest(TestCase):
             model_output["pooled_output"], restored_output["pooled_output"]
         )
 
-
-@pytest.mark.tpu
-@pytest.mark.usefixtures("tpu_test_class")
-class AlbertBackboneTPUTest(TestCase):
-    def setUp(self):
-        with self.tpu_strategy.scope():
-            self.backbone = AlbertBackbone(
-                vocabulary_size=10,
-                num_layers=2,
-                num_heads=2,
-                num_groups=1,
-                num_inner_repetitions=1,
-                embedding_dim=16,
-                hidden_dim=2,
-                intermediate_dim=2,
-                max_sequence_length=4,
-            )
-
-        self.input_batch = {
-            "token_ids": np.ones((8, 128), dtype="int32"),
-            "segment_ids": np.ones((8, 128), dtype="int32"),
-            "padding_mask": np.ones((8, 128), dtype="int32"),
-        }
-        self.input_dataset = tf.data.Dataset.from_tensor_slices(
-            self.input_batch
-        ).batch(2)
-
-    def test_predict(self):
-        self.backbone.compile()
-        self.backbone.predict(self.input_dataset)
-

{keras-nlp-0.6.3.dev0 → keras-nlp-0.6.4.dev0}/keras_nlp/src/models/bart/bart_backbone_test.py

@@ -92,31 +92,3 @@ class BartBackboneTest(TestCase):
             restored_output["decoder_sequence_output"],
         )
 
-
-@pytest.mark.tpu
-@pytest.mark.usefixtures("tpu_test_class")
-class BartBackboneTPUTest(TestCase):
-    def setUp(self):
-        with self.tpu_strategy.scope():
-            self.backbone = BartBackbone(
-                vocabulary_size=1000,
-                num_layers=2,
-                num_heads=2,
-                hidden_dim=64,
-                intermediate_dim=128,
-                max_sequence_length=128,
-            )
-        self.input_batch = {
-            "encoder_token_ids": np.ones((8, 128), dtype="int32"),
-            "encoder_padding_mask": np.ones((8, 128), dtype="int32"),
-            "decoder_token_ids": np.ones((8, 128), dtype="int32"),
-            "decoder_padding_mask": np.ones((8, 128), dtype="int32"),
-        }
-        self.input_dataset = tf.data.Dataset.from_tensor_slices(
-            self.input_batch
-        ).batch(2)
-
-    def test_predict(self):
-        self.backbone.compile()
-        self.backbone.predict(self.input_dataset)
-

{keras-nlp-0.6.3.dev0 → keras-nlp-0.6.4.dev0}/keras_nlp/src/models/bert/bert_backbone_test.py

@@ -86,30 +86,3 @@ class BertBackboneTest(TestCase):
         restored_output = restored_model(self.input_batch)
         self.assertAllClose(model_output, restored_output)
 
-
-@pytest.mark.tpu
-@pytest.mark.usefixtures("tpu_test_class")
-class BertBackboneTPUTest(TestCase):
-    def setUp(self):
-        with self.tpu_strategy.scope():
-            self.backbone = BertBackbone(
-                vocabulary_size=1000,
-                num_layers=2,
-                num_heads=2,
-                hidden_dim=64,
-                intermediate_dim=128,
-                max_sequence_length=128,
-            )
-        self.input_batch = {
-            "token_ids": np.ones((8, 128), dtype="int32"),
-            "segment_ids": np.ones((8, 128), dtype="int32"),
-            "padding_mask": np.ones((8, 128), dtype="int32"),
-        }
-        self.input_dataset = tf.data.Dataset.from_tensor_slices(
-            self.input_batch
-        ).batch(2)
-
-    def test_predict(self):
-        self.backbone.compile()
-        self.backbone.predict(self.input_dataset)
-

{keras-nlp-0.6.3.dev0 → keras-nlp-0.6.4.dev0}/keras_nlp/src/models/deberta_v3/deberta_v3_backbone_test.py

@@ -91,30 +91,3 @@ class DebertaV3BackboneTest(TestCase):
         restored_output = restored_model(self.input_batch)
         self.assertAllClose(model_output, restored_output)
 
-
-@pytest.mark.tpu
-@pytest.mark.usefixtures("tpu_test_class")
-class DebertaV3BackboneTPUTest(TestCase):
-    def setUp(self):
-        with self.tpu_strategy.scope():
-            self.backbone = DebertaV3Backbone(
-                vocabulary_size=10,
-                num_layers=2,
-                num_heads=2,
-                hidden_dim=2,
-                intermediate_dim=4,
-                max_sequence_length=5,
-                bucket_size=2,
-            )
-        self.input_batch = {
-            "token_ids": np.ones((2, 5), dtype="int32"),
-            "padding_mask": np.ones((2, 5), dtype="int32"),
-        }
-        self.input_dataset = tf.data.Dataset.from_tensor_slices(
-            self.input_batch
-        ).batch(2)
-
-    def test_predict(self):
-        self.backbone.compile()
-        self.backbone.predict(self.input_dataset)
-

{keras-nlp-0.6.3.dev0 → keras-nlp-0.6.4.dev0}/keras_nlp/src/models/deberta_v3/disentangled_self_attention.py

@@ -232,12 +232,13 @@ class DisentangledSelfAttention(keras.layers.Layer):
             x1=rel_pos,
             x2=log_pos * sign,
         )
-        bucket_pos = ops.cast(bucket_pos, dtype="int64")
+        bucket_pos = ops.cast(bucket_pos, dtype="int")
 
         return bucket_pos
 
     def _get_rel_pos(self, num_positions):
-        ids = ops.arange(num_positions, dtype="int64")
+        ids = ops.arange(num_positions)
+        ids = ops.cast(ids, dtype="int")
         query_ids = ops.expand_dims(ids, axis=-1)
         key_ids = ops.expand_dims(ids, axis=0)
         key_ids = ops.repeat(key_ids, repeats=num_positions, axis=0)

{keras-nlp-0.6.3.dev0 → keras-nlp-0.6.4.dev0}/keras_nlp/src/models/distil_bert/distil_bert_backbone_test.py

@@ -84,29 +84,3 @@ class DistilBertTest(TestCase):
         restored_output = restored_model(self.input_batch)
         self.assertAllClose(model_output, restored_output)
 
-
-@pytest.mark.tpu
-@pytest.mark.usefixtures("tpu_test_class")
-class DistilBertTPUTest(TestCase):
-    def setUp(self):
-        with self.tpu_strategy.scope():
-            self.backbone = DistilBertBackbone(
-                vocabulary_size=1000,
-                num_layers=2,
-                num_heads=2,
-                hidden_dim=64,
-                intermediate_dim=128,
-                max_sequence_length=128,
-            )
-        self.input_batch = {
-            "token_ids": np.ones((8, 128), dtype="int32"),
-            "padding_mask": np.ones((8, 128), dtype="int32"),
-        }
-        self.input_dataset = tf.data.Dataset.from_tensor_slices(
-            self.input_batch
-        ).batch(2)
-
-    def test_predict(self):
-        self.backbone.compile()
-        self.backbone.predict(self.input_dataset)
-