Repository: keras-team/keras-nlp
Pull request #1268 (2023-10-09 21:35:12)

Remove dtensor
We will replace this with the work on
https://github.com/keras-team/keras-nlp/pull/1267
But we have no coverage for that PR till we run tests against Keras 3, which will probably still be about a week.

For now, let's just remove this usage, which is no longer needed and will break a Keras 3 install.

mattdangerw: /gcbrun

Patch:
diff --git a/.github/workflows/actions.yml b/.github/workflows/actions.yml
index 926d6795..7d4a7d4e 100644
--- a/.github/workflows/actions.yml
+++ b/.github/workflows/actions.yml
@@ -30,7 +30,6 @@ jobs:
       - name: Install dependencies
         run: |
           pip install -r requirements.txt --progress-bar off
-          pip install jax[cpu] --progress-bar off
           pip install --no-deps -e "." --progress-bar off
       - name: Test with pytest
         run: |
@@ -66,8 +65,6 @@ jobs:
     - name: Install dependencies
       run: |
           pip install -r requirements.txt --progress-bar off
-          pip install torch>=2.0.1+cpu --progress-bar off
-          pip install jax[cpu] --progress-bar off
           pip install --no-deps -e "." --progress-bar off
     - name: Test with pytest
       env:
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index d5af724c..688c6311 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -81,7 +81,7 @@ Once the pull request is approved, a team member will take care of merging.
 
 ## Setting up an Environment
 
-Python 3.7 or later is required.
+Python 3.9 or later is required.
 
 Setting up your KerasNLP development environment requires you to fork the
 KerasNLP repository and clone it locally. With the
@@ -93,72 +93,57 @@ cd keras-nlp
 ```
 
 Next we must setup a python environment with the correct dependencies. We
-recommend using `conda` to install tensorflow dependencies (such as CUDA), and
-`pip` to install python packages from PyPI. The exact method will depend on your
-OS.
+recommend using `conda` to set up a base environment, and `pip` to install
+python packages from PyPI. The exact method will depend on your OS.
 
-**Note**: Please be careful not to use the `tensorflow` pre-packaged with conda,
-which is incompatible with `tensorflow-text` on PyPi, and follow the
-instructions below.
+**Note**: Be careful not to use mix pre-packaged tensorflow and jax libraries in
+`conda` with PyPI packages from `pip`. We recommend pulling *all* KerasNLP
+dependencies via `pip` as described below.
 
 ### Linux (recommended)
 
-To setup a complete environment with TensorFlow, a local install of keras-nlp,
-and all development tools, run the following or adapt it to suit your needs.
+For developing and unit testing the library, a CPU-only environment is often
+sufficient. For any training or inference with the library, you will quickly
+want accelerator support. The easiest way to get GPU support across all of our
+backends is to set up a few different python environements and pull in all cuda
+dependencies via `pip`.
+
+The shell snippet below will install four conda environments: `keras-nlp-cpu`,
+`keras-nlp-jax`, `keras-nlp-torch`, and `keras-nlp-tensorflow`. The cpu
+environement supports all backends without cuda, and each backend environement
+has cuda support.
 
 ```shell
-# Create and activate conda environment.
-conda create -n keras-nlp python=3.9
-conda activate keras-nlp
-
-# The following can be omitted if GPU support is not required.
-conda install -c conda-forge cudatoolkit-dev=11.2 cudnn=8.1.0
-mkdir -p $CONDA_PREFIX/etc/conda/activate.d/
-echo 'export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$CONDA_PREFIX/lib/' >> $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh
-echo 'export XLA_FLAGS=--xla_gpu_cuda_data_dir=$CONDA_PREFIX/' >> $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh
-source $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh
-
-# Install dependencies.
-python -m pip install --upgrade pip
-python -m pip install -r requirements.txt
-python -m pip install -e "."
+conda create -y -n keras-nlp-cpu python=3.10
+conda activate keras-nlp-cpu
+pip install -r requirements.txt  # install deps
+python pip_build.py --install  # install keras-nlp
+
+for backend in "jax" "torch" "tensorflow"; do
+    conda create -y -n keras-nlp-${backend} python=3.10
+    conda activate keras-nlp-${backend}
+    pip install -r requirements-${backend}-cuda.txt  # install deps
+    python pip_build.py --install  # install keras-nlp
+done
 ```
 
-### MacOS
-
-⚠️⚠️⚠️ MacOS binaries are for the M1 architecture are not currently available from
-official sources. You can try experimental development workflow leveraging the
-[tensorflow metal plugin](https://developer.apple.com/metal/tensorflow-plugin/)
-and a [community maintained build](https://github.com/sun1638650145/Libraries-and-Extensions-for-TensorFlow-for-Apple-Silicon)
-of `tensorflow-text`. These binaries are not provided by Google, so proceed at
-your own risk.
-
-#### Experimental instructions for Arm (M1)
+To activate the jax environment and set keras to use jax, run:
 
 ```shell
-# Create and activate conda environment.
-conda create -n keras-nlp python=3.9
-conda activate keras-nlp
-
-# Install dependencies.
-conda install -c apple tensorflow-deps=2.9
-python -m pip install --upgrade pip
-python -m pip install -r requirements-macos-m1.txt
-python -m pip install -e "."
+conda activate keras-nlp-jax && export KERAS_BACKEND=jax
 ```
 
-#### Instructions for x86 (Intel)
+### MacOS
 
-```shell
-# Create and activate conda environment.
-conda create -n keras-nlp python=3.9
-conda activate keras-nlp
-
-# Install dependencies.
-python -m pip install --upgrade pip
-python -m pip install -r requirements.txt
-python -m pip install -e "."
-```
+`tensorflow-text` does not release precompiled binaries for MacOS M-series
+chips, though the library does support building from source on MacOS.
+
+We strongly recommend a Linux development environment for an easy contribution
+experience. To build a dev environement from scratch on MacOS, see the following
+guides:
+
+https://developer.apple.com/metal/tensorflow-plugin/
+https://github.com/tensorflow/text
 
 ### Windows
 
diff --git a/keras_nlp/backend/config.py b/keras_nlp/backend/config.py
index 578e1746..9b85907d 100644
--- a/keras_nlp/backend/config.py
+++ b/keras_nlp/backend/config.py
@@ -15,9 +15,11 @@
 import json
 import os
 
-import keras_core
+import keras
+from packaging import version
 
 _MULTI_BACKEND = False
+_IS_KERAS_3 = False
 
 # Set Keras base dir path given KERAS_HOME env variable, if applicable.
 # Otherwise either ~/.keras or /tmp.
@@ -59,16 +61,32 @@ if not os.path.exists(_config_path):
         # Except permission denied.
         pass
 
-# Use keras-core if KERAS_BACKEND is set in the environment.
+# If KERAS_BACKEND is set in the environment use multi-backend keras.
 if "KERAS_BACKEND" in os.environ and os.environ["KERAS_BACKEND"]:
     _MULTI_BACKEND = True
 
+# If keras is version 3, use multi-backend keras (our only option).
+_IS_KERAS_3 = version.parse(keras.__version__) >= version.parse("3.0.0")
+if _IS_KERAS_3:
+    _MULTI_BACKEND = True
+
+
+def keras_3():
+    """Check if Keras 3 is installed."""
+    return _IS_KERAS_3
+
 
 def multi_backend():
-    """Check if keras_core is enabled."""
+    """Check if multi-backend Keras is enabled."""
     return _MULTI_BACKEND
 
 
 def backend():
     """Check the backend framework."""
-    return "tensorflow" if not multi_backend() else keras_core.config.backend()
+    if not multi_backend():
+        return "tensorflow"
+    if not keras_3():
+        import keras_core
+
+        return keras_core.config.backend()
+    return keras.config.backend()
diff --git a/keras_nlp/backend/keras.py b/keras_nlp/backend/keras.py
index 85f0ebcb..865f62d1 100644
--- a/keras_nlp/backend/keras.py
+++ b/keras_nlp/backend/keras.py
@@ -16,9 +16,11 @@ import types
 
 import tensorflow as tf
 
-from keras_nlp.backend.config import multi_backend
+from keras_nlp.backend import config
 
-if multi_backend():
+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
diff --git a/keras_nlp/backend/ops.py b/keras_nlp/backend/ops.py
index 516834b6..f36a2d2d 100644
--- a/keras_nlp/backend/ops.py
+++ b/keras_nlp/backend/ops.py
@@ -12,22 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import keras_core
-import tensorflow as tf
+from keras_nlp.backend import config
 
-from keras_nlp.backend.config import multi_backend
-
-if multi_backend():
-    from keras_core.src.ops import *  # noqa: F403, F401
+if config.keras_3():
+    from keras.ops import *  # noqa: F403, F401
 else:
-    from keras_core.src.backend.tensorflow import *  # noqa: F403, F401
-    from keras_core.src.backend.tensorflow.core import *  # noqa: F403, F401
-    from keras_core.src.backend.tensorflow.math import *  # noqa: F403, F401
-    from keras_core.src.backend.tensorflow.nn import *  # noqa: F403, F401
-    from keras_core.src.backend.tensorflow.numpy import *  # noqa: F403, F401
-
+    from keras_core.ops import *  # noqa: F403, F401
 
-if keras_core.config.backend() == "tensorflow" or not multi_backend():
+if config.backend() == "tensorflow":
+    import tensorflow as tf
+    from tensorflow.experimental import numpy as tfnp
 
     def take_along_axis(x, indices, axis=None):
         # TODO: move this workaround for dynamic shapes into keras-core.
@@ -46,6 +40,4 @@ if keras_core.config.backend() == "tensorflow" or not multi_backend():
                 indices = tf.squeeze(indices, leftover_axes)
             return tf.gather(x, indices, batch_dims=axis)
         # Otherwise, fall back to the tfnp call.
-        return keras_core.src.backend.tensorflow.numpy.take_along_axis(
-            x, indices, axis=axis
-        )
+        return tfnp.take_along_axis(x, indices, axis=axis)
diff --git a/keras_nlp/backend/random.py b/keras_nlp/backend/random.py
index 70be5910..c4f4e6f4 100644
--- a/keras_nlp/backend/random.py
+++ b/keras_nlp/backend/random.py
@@ -12,9 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from keras_nlp.backend.config import multi_backend
+from keras_nlp.backend import config
 
-if multi_backend():
-    from keras_core.random import *  # noqa: F403, F401
+if config.keras_3():
+    from keras.random import *  # noqa: F403, F401
 else:
-    from keras_core.src.backend.tensorflow.random import *  # noqa: F403, F401
+    from keras_core.random import *  # noqa: F403, F401
diff --git a/keras_nlp/layers/preprocessing/random_deletion.py b/keras_nlp/layers/preprocessing/random_deletion.py
index 41289c00..061290ba 100644
--- a/keras_nlp/layers/preprocessing/random_deletion.py
+++ b/keras_nlp/layers/preprocessing/random_deletion.py
@@ -21,7 +21,7 @@ from keras_nlp.layers.preprocessing.preprocessing_layer import (
     PreprocessingLayer,
 )
 from keras_nlp.utils.tensor_utils import convert_to_ragged_batch
-from keras_nlp.utils.tensor_utils import is_integer_dtype
+from keras_nlp.utils.tensor_utils import is_int_dtype
 from keras_nlp.utils.tensor_utils import is_string_dtype
 
 
@@ -125,7 +125,7 @@ class RandomDeletion(PreprocessingLayer):
         dtype="int32",
         **kwargs,
     ):
-        if not is_integer_dtype(dtype) and not is_string_dtype(dtype):
+        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}"
diff --git a/keras_nlp/layers/preprocessing/random_swap.py b/keras_nlp/layers/preprocessing/random_swap.py
index ad6f1c09..27873f0f 100644
--- a/keras_nlp/layers/preprocessing/random_swap.py
+++ b/keras_nlp/layers/preprocessing/random_swap.py
@@ -21,7 +21,7 @@ from keras_nlp.layers.preprocessing.preprocessing_layer import (
     PreprocessingLayer,
 )
 from keras_nlp.utils.tensor_utils import convert_to_ragged_batch
-from keras_nlp.utils.tensor_utils import is_integer_dtype
+from keras_nlp.utils.tensor_utils import is_int_dtype
 from keras_nlp.utils.tensor_utils import is_string_dtype
 
 
@@ -127,7 +127,7 @@ class RandomSwap(PreprocessingLayer):
         dtype="int32",
         **kwargs,
     ):
-        if not is_integer_dtype(dtype) and not is_string_dtype(dtype):
+        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}"
diff --git a/keras_nlp/metrics/bleu.py b/keras_nlp/metrics/bleu.py
index 51d224c6..750a1b70 100644
--- a/keras_nlp/metrics/bleu.py
+++ b/keras_nlp/metrics/bleu.py
@@ -20,7 +20,7 @@ import tensorflow as tf
 from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.backend import keras
 from keras_nlp.backend import ops
-from keras_nlp.utils.tensor_utils import is_floating_dtype
+from keras_nlp.utils.tensor_utils import is_float_dtype
 from keras_nlp.utils.tensor_utils import tensor_to_list
 
 REPLACE_SUBSTRINGS = [
@@ -112,7 +112,7 @@ class Bleu(keras.metrics.Metric):
     ):
         super().__init__(name=name, dtype=dtype, **kwargs)
 
-        if not is_floating_dtype(dtype):
+        if not is_float_dtype(dtype):
             raise ValueError(
                 "`dtype` must be a floating point type. "
                 f"Received: dtype={dtype}"
diff --git a/keras_nlp/metrics/edit_distance.py b/keras_nlp/metrics/edit_distance.py
index 899f7f1f..263ff829 100644
--- a/keras_nlp/metrics/edit_distance.py
+++ b/keras_nlp/metrics/edit_distance.py
@@ -16,7 +16,7 @@ import tensorflow as tf
 
 from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.backend import keras
-from keras_nlp.utils.tensor_utils import is_floating_dtype
+from keras_nlp.utils.tensor_utils import is_float_dtype
 
 
 @keras_nlp_export("keras_nlp.metrics.EditDistance")
@@ -87,7 +87,7 @@ class EditDistance(keras.metrics.Metric):
     ):
         super().__init__(name=name, dtype=dtype, **kwargs)
 
-        if not is_floating_dtype(dtype):
+        if not is_float_dtype(dtype):
             raise ValueError(
                 "`dtype` must be a floating point type. "
                 f"Received: dtype={dtype}"
diff --git a/keras_nlp/metrics/perplexity.py b/keras_nlp/metrics/perplexity.py
index eb742fc3..4a7e626b 100644
--- a/keras_nlp/metrics/perplexity.py
+++ b/keras_nlp/metrics/perplexity.py
@@ -15,7 +15,7 @@
 from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.backend import keras
 from keras_nlp.backend import ops
-from keras_nlp.utils.tensor_utils import is_floating_dtype
+from keras_nlp.utils.tensor_utils import is_float_dtype
 
 
 @keras_nlp_export("keras_nlp.metrics.Perplexity")
@@ -88,7 +88,7 @@ class Perplexity(keras.metrics.Metric):
         name="perplexity",
         **kwargs,
     ):
-        if not is_floating_dtype(dtype):
+        if not is_float_dtype(dtype):
             raise ValueError(
                 "`dtype` must be a floating point type. "
                 f"Received: dtype={dtype}"
diff --git a/keras_nlp/metrics/rouge_base.py b/keras_nlp/metrics/rouge_base.py
index f84e7180..824a6c4b 100644
--- a/keras_nlp/metrics/rouge_base.py
+++ b/keras_nlp/metrics/rouge_base.py
@@ -16,7 +16,7 @@ import tensorflow as tf
 
 from keras_nlp.backend import keras
 from keras_nlp.backend import ops
-from keras_nlp.utils.tensor_utils import is_floating_dtype
+from keras_nlp.utils.tensor_utils import is_float_dtype
 from keras_nlp.utils.tensor_utils import tensor_to_list
 
 try:
@@ -65,7 +65,7 @@ class RougeBase(keras.metrics.Metric):
                 "package. Please install it with `pip install rouge-score`."
             )
 
-        if not is_floating_dtype(dtype):
+        if not is_float_dtype(dtype):
             raise ValueError(
                 "`dtype` must be a floating point type. "
                 f"Received: dtype={dtype}"
diff --git a/keras_nlp/models/gpt2/gpt2_backbone.py b/keras_nlp/models/gpt2/gpt2_backbone.py
index caf73c06..3f357d64 100644
--- a/keras_nlp/models/gpt2/gpt2_backbone.py
+++ b/keras_nlp/models/gpt2/gpt2_backbone.py
@@ -14,10 +14,6 @@
 
 import copy
 
-from tensorflow.experimental import dtensor
-from tensorflow.experimental.dtensor import Layout
-from tensorflow.keras.dtensor.experimental import LayoutMap
-
 from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.backend import keras
 from keras_nlp.layers.modeling.position_embedding import PositionEmbedding
@@ -190,71 +186,3 @@ class GPT2Backbone(Backbone):
     @classproperty
     def presets(cls):
         return copy.deepcopy(backbone_presets)
-
-    @classmethod
-    def create_layout_map(cls, mesh):
-        """Create a DTensor layout map for a GPT2Backbone.
-
-        Given a DTensor mesh describing a list of devices, this method returns a
-        DTensor layout map for creating a `keras_nlp.models.GPT2Backbone`
-        instance. This mapping describes how to distribute all model weights
-        across multiple devices. For an overview of DTensor concepts, see
-        [this guide](https://www.tensorflow.org/guide/dtensor_overview).
-
-        Args:
-            mesh: A 2D `tf.experimental.dtensor.Mesh` describing the arrangement
-                of devices for running distributed computation. The
-                first dimension in the mesh is expected to be for data parallel
-                distribution, and the second for model parallel distribution.
-
-        Returns:
-            A `tf.keras.dtensor.experimental.LayoutMap` which contains the
-            proper layout to weights mapping for the model parallel setting.
-
-        Examples:
-        ```python
-        keras.backend.experimental.enable_tf_random_generator()
-        keras.utils.set_random_seed(1337)
-
-        # Update both dimensions below for a multi-device setting.
-        mesh = dtensor.create_mesh([("batch", 1), ("model", 1)])
-        layout_map = keras_nlp.models.GPT2Backbone.create_layout_map(mesh)
-
-        with layout_map.scope():
-            model = keras_nlp.models.GPT2Backbone.from_preset("gpt2_base_en")
-        ```
-        """
-        # We assert the mesh is 2D, and assume the first mesh dim is for data
-        # parallel and the second dim is for model parallel.
-        mesh_shape = mesh.shape()
-        if len(mesh_shape) != 2:
-            raise ValueError(
-                f"Expect to create layout based on 2D mesh, received {mesh}"
-            )
-        _, model_dim = mesh.dim_names
-        unshard_dim = dtensor.UNSHARDED
-
-        layout_map = LayoutMap(mesh=mesh)
-        # Embedding sharding
-        layout_map[r".*embeddings"] = Layout([unshard_dim, model_dim], mesh)
-
-        # Transformer block sharding
-        layout_map[r".*_(query|key|value)_dense.kernel"] = Layout(
-            [unshard_dim, unshard_dim, model_dim], mesh
-        )
-        layout_map[r".*_(query|key|value)_dense.bias"] = Layout(
-            [model_dim, unshard_dim], mesh
-        )
-        layout_map[r".*_feedforward_intermediate_dense.kernel"] = Layout(
-            [unshard_dim, model_dim], mesh
-        )
-        layout_map[r".*_feedforward_intermediate_dense.bias"] = Layout(
-            [model_dim], mesh
-        )
-        layout_map[r".*_feedforward_output_dense.kernel"] = Layout(
-            [model_dim, unshard_dim], mesh
-        )
-        layout_map[r".*_feedforward_output_dense.bias"] = Layout(
-            [unshard_dim], mesh
-        )
-        return layout_map
diff --git a/keras_nlp/models/gpt2/gpt2_causal_lm.py b/keras_nlp/models/gpt2/gpt2_causal_lm.py
index 23dcc416..44eebd0a 100644
--- a/keras_nlp/models/gpt2/gpt2_causal_lm.py
+++ b/keras_nlp/models/gpt2/gpt2_causal_lm.py
@@ -325,39 +325,3 @@ class GPT2CausalLM(GenerativeTask):
             "token_ids": token_ids,
             "padding_mask": padding_mask,
         }
-
-    @classmethod
-    def create_layout_map(cls, mesh):
-        """Create a DTensor layout map for a GPT2CausalLM.
-
-        Given a DTensor mesh describing a list of devices, this method returns a
-        DTensor layout map for creating a `keras_nlp.models.GPT2CausalLM`
-        instance. This mapping describes how to distribute all model weights
-        across multiple devices. For an overview of DTensor concepts, see
-        [this guide](https://www.tensorflow.org/guide/dtensor_overview).
-
-        Args:
-            mesh: A 2D `tf.experimental.dtensor.Mesh` describing the arrangement
-                of devices for running distributed computation. The
-                first dimension in the mesh is expected to be for data parallel
-                distribution, and the second for model parallel distribution.
-
-        Returns:
-            A `keras.dtensor.experimental.LayoutMap` which contains the
-            proper layout to weights mapping for the model parallel setting.
-
-        Examples:
-        ```python
-        keras.backend.experimental.enable_tf_random_generator()
-        keras.utils.set_random_seed(1337)
-
-        # Update both dimensions below for a multi-device setting.
-        mesh = tf.experimental.dtensor.create_mesh([("batch", 1), ("model", 1)])
-        layout_map = keras_nlp.models.GPT2CausalLM.create_layout_map(mesh)
-
-        with layout_map.scope():
-            gpt2_lm = keras_nlp.models.GPT2CausalLM.from_preset("gpt2_base_en")
-        ```
-        """
-        # As this task has no new variables, we just re-use the backbone method.
-        return cls.backbone_cls.create_layout_map(mesh)
diff --git a/keras_nlp/models/opt/opt_backbone.py b/keras_nlp/models/opt/opt_backbone.py
index 8fe37472..ff1495ba 100644
--- a/keras_nlp/models/opt/opt_backbone.py
+++ b/keras_nlp/models/opt/opt_backbone.py
@@ -14,10 +14,6 @@
 
 import copy
 
-from tensorflow.experimental import dtensor
-from tensorflow.experimental.dtensor import Layout
-from tensorflow.keras.dtensor.experimental import LayoutMap
-
 from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.backend import keras
 from keras_nlp.layers.modeling.token_and_position_embedding import (
@@ -168,71 +164,3 @@ class OPTBackbone(Backbone):
     @classproperty
     def presets(cls):
         return copy.deepcopy(backbone_presets)
-
-    @classmethod
-    def create_layout_map(cls, mesh):
-        """Create a DTensor layout map for an OPTBackbone.
-
-        Given a DTensor mesh describing a list of devices, this method returns a
-        DTensor layout map for creating a `keras_nlp.models.OPTBackbone`
-        instance. This mapping describes how to distribute all model weights
-        across multiple devices. For an overview of DTensor concepts, see
-        [this guide](https://www.tensorflow.org/guide/dtensor_overview).
-
-        Args:
-            mesh: A 2D `tf.experimental.dtensor.Mesh` describing the arrangement
-                of devices for running distributed computation. The
-                first dimension in the mesh is expected to be for data parallel
-                distribution, and the second for model parallel distribution.
-
-        Returns:
-            A `tf.keras.dtensor.experimental.LayoutMap` which contains the
-            proper layout to weights mapping for the model parallel setting.
-
-        Examples:
-        ```python
-        keras.backend.experimental.enable_tf_random_generator()
-        keras.utils.set_random_seed(1337)
-
-        # Update both dimensions below for a multi-device setting.
-        mesh = dtensor.create_mesh([("batch", 1), ("model", 1)])
-        layout_map = keras_nlp.models.OPTBackbone.create_layout_map(mesh)
-
-        with layout_map.scope():
-            model = keras_nlp.models.OPTBackbone.from_preset("opt_125m_en")
-        ```
-        """
-        # We assert the mesh is 2D, and assume the first mesh dim is for data
-        # parallel and the second dim is for model parallel.
-        mesh_shape = mesh.shape()
-        if len(mesh_shape) != 2:
-            raise ValueError(
-                f"Expect to create layout based on 2D mesh, received {mesh}"
-            )
-        _, model_dim = mesh.dim_names
-        unshard_dim = dtensor.UNSHARDED
-
-        layout_map = LayoutMap(mesh=mesh)
-        # Embedding sharding
-        layout_map[r".*embeddings"] = Layout([unshard_dim, model_dim], mesh)
-
-        # Transformer block sharding
-        layout_map[r".*_(query|key|value)_dense.kernel"] = Layout(
-            [unshard_dim, unshard_dim, model_dim], mesh
-        )
-        layout_map[r".*_(query|key|value)_dense.bias"] = Layout(
-            [model_dim, unshard_dim], mesh
-        )
-        layout_map[r".*_feedforward_intermediate_dense.kernel"] = Layout(
-            [unshard_dim, model_dim], mesh
-        )
-        layout_map[r".*_feedforward_intermediate_dense.bias"] = Layout(
-            [model_dim], mesh
-        )
-        layout_map[r".*_feedforward_output_dense.kernel"] = Layout(
-            [model_dim, unshard_dim], mesh
-        )
-        layout_map[r".*_feedforward_output_dense.bias"] = Layout(
-            [unshard_dim], mesh
-        )
-        return layout_map
diff --git a/keras_nlp/models/opt/opt_causal_lm.py b/keras_nlp/models/opt/opt_causal_lm.py
index f0b06827..6197a87f 100644
--- a/keras_nlp/models/opt/opt_causal_lm.py
+++ b/keras_nlp/models/opt/opt_causal_lm.py
@@ -321,39 +321,3 @@ class OPTCausalLM(GenerativeTask):
             "token_ids": token_ids,
             "padding_mask": padding_mask,
         }
-
-    @classmethod
-    def create_layout_map(cls, mesh):
-        """Create a DTensor layout map for an OPTCausalLM.
-
-        Given a DTensor mesh describing a list of devices, this method returns a
-        DTensor layout map for creating a `keras_nlp.models.OPTCausalLM`
-        instance. This mapping describes how to distribute all model weights
-        across multiple devices. For an overview of DTensor concepts, see
-        [this guide](https://www.tensorflow.org/guide/dtensor_overview).
-
-        Args:
-            mesh: A 2D `tf.experimental.dtensor.Mesh` describing the arrangement
-                of devices for running distributed computation. The
-                first dimension in the mesh is expected to be for data parallel
-                distribution, and the second for model parallel distribution.
-
-        Returns:
-            A `tf.keras.dtensor.experimental.LayoutMap` which contains the
-            proper layout to weights mapping for the model parallel setting.
-
-        Examples:
-        ```python
-        keras.backend.experimental.enable_tf_random_generator()
-        keras.utils.set_random_seed(1337)
-
-        # Update both dimensions below for a multi-device setting.
-        mesh = tf.experimental.dtensor.create_mesh([("batch", 1), ("model", 1)])
-        layout_map = keras_nlp.models.OPTCausalLM.create_layout_map(mesh)
-
-        with layout_map.scope():
-            opt_lm = keras_nlp.models.OPTCausalLM.from_preset("opt_125m_en")
-        ```
-        """
-        # As this task has no new variables, we just re-use the backbone method.
-        return cls.backbone_cls.create_layout_map(mesh)
diff --git a/keras_nlp/models/task.py b/keras_nlp/models/task.py
index 5d9a6054..d4c61804 100644
--- a/keras_nlp/models/task.py
+++ b/keras_nlp/models/task.py
@@ -14,11 +14,11 @@
 
 import os
 
-import keras_core
 from rich import console as rich_console
 from rich import markup
 from rich import table as rich_table
 
+from keras_nlp.backend import config
 from keras_nlp.backend import keras
 from keras_nlp.utils.keras_utils import print_msg
 from keras_nlp.utils.pipeline_model import PipelineModel
@@ -315,11 +315,21 @@ class Task(PipelineModel):
             if print_fn:
                 print_fn(console.end_capture(), line_break=False)
 
-        # Hardcode summary from keras_core for now.
-        keras_core.Model.summary(
-            self,
-            line_length=line_length,
-            positions=positions,
-            print_fn=print_fn,
-            **kwargs,
-        )
+        # Avoid `tf.keras.Model.summary()`, so the above output matches.
+        if config.multi_backend():
+            super().summary(
+                line_length=line_length,
+                positions=positions,
+                print_fn=print_fn,
+                **kwargs,
+            )
+        else:
+            import keras_core
+
+            keras_core.Model.summary(
+                self,
+                line_length=line_length,
+                positions=positions,
+                print_fn=print_fn,
+                **kwargs,
+            )
diff --git a/keras_nlp/models/whisper/whisper_audio_feature_extractor.py b/keras_nlp/models/whisper/whisper_audio_feature_extractor.py
index 73b15a3a..e41519bb 100644
--- a/keras_nlp/models/whisper/whisper_audio_feature_extractor.py
+++ b/keras_nlp/models/whisper/whisper_audio_feature_extractor.py
@@ -18,14 +18,16 @@ import numpy as np
 import tensorflow as tf
 
 from keras_nlp.api_export import keras_nlp_export
-from keras_nlp.backend import keras
+from keras_nlp.layers.preprocessing.preprocessing_layer import (
+    PreprocessingLayer,
+)
 from keras_nlp.models.whisper.whisper_presets import backbone_presets
 from keras_nlp.utils.python_utils import classproperty
 from keras_nlp.utils.python_utils import format_docstring
 
 
 @keras_nlp_export("keras_nlp.models.WhisperAudioFeatureExtractor")
-class WhisperAudioFeatureExtractor(keras.layers.Layer):
+class WhisperAudioFeatureExtractor(PreprocessingLayer):
     """
     Whisper audio feature extractor layer.
 
@@ -163,9 +165,10 @@ class WhisperAudioFeatureExtractor(keras.layers.Layer):
         weights *= enorm[:, np.newaxis]
 
         weights = np.transpose(weights)
-        return tf.constant(weights, dtype=self.dtype)
+        return tf.constant(weights, dtype=self.compute_dtype)
 
     def _extract_audio_features(self, audio):
+        audio = tf.cast(audio, self.compute_dtype)
         # Use "reflection" padding - `tf.signal.stft` uses symmetric padding
         # internally.
         audio = tf.pad(
@@ -246,6 +249,8 @@ class WhisperAudioFeatureExtractor(keras.layers.Layer):
 
         # Find the log mel spectrogram.
         log_spec = self._extract_audio_features(audio)
+        if rank_1_input:
+            log_spec = tf.squeeze(log_spec, 0)
         return log_spec
 
     def get_config(self):
diff --git a/keras_nlp/models/whisper/whisper_preprocessor.py b/keras_nlp/models/whisper/whisper_preprocessor.py
index 88fd60ca..8545890c 100644
--- a/keras_nlp/models/whisper/whisper_preprocessor.py
+++ b/keras_nlp/models/whisper/whisper_preprocessor.py
@@ -14,7 +14,6 @@
 
 import copy
 
-import tensorflow as tf
 from absl import logging
 
 from keras_nlp.api_export import keras_nlp_export
@@ -278,9 +277,6 @@ class WhisperPreprocessor(Preprocessor):
             )
 
         encoder_features = self.audio_feature_extractor(encoder_audio[0])
-        if encoder_audio[0].shape.rank < 2:
-            encoder_features = tf.squeeze(encoder_features, axis=0)
-
         decoder_sequence_length = (
             decoder_sequence_length or self.decoder_sequence_length
         )
diff --git a/keras_nlp/tokenizers/byte_pair_tokenizer.py b/keras_nlp/tokenizers/byte_pair_tokenizer.py
index 6ec140a1..f92d9e6a 100644
--- a/keras_nlp/tokenizers/byte_pair_tokenizer.py
+++ b/keras_nlp/tokenizers/byte_pair_tokenizer.py
@@ -34,7 +34,7 @@ from keras_nlp.utils.python_utils import classproperty
 from keras_nlp.utils.python_utils import format_docstring
 from keras_nlp.utils.tensor_utils import assert_tf_text_installed
 from keras_nlp.utils.tensor_utils import convert_to_ragged_batch
-from keras_nlp.utils.tensor_utils import is_integer_dtype
+from keras_nlp.utils.tensor_utils import is_int_dtype
 from keras_nlp.utils.tensor_utils import is_string_dtype
 
 try:
@@ -283,7 +283,7 @@ class BytePairTokenizer(tokenizer.Tokenizer):
     ) -> None:
         assert_tf_text_installed(self.__class__.__name__)
 
-        if not is_integer_dtype(dtype) and not is_string_dtype(dtype):
+        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}"
diff --git a/keras_nlp/tokenizers/byte_tokenizer.py b/keras_nlp/tokenizers/byte_tokenizer.py
index f6b60bfe..3aefc4a0 100644
--- a/keras_nlp/tokenizers/byte_tokenizer.py
+++ b/keras_nlp/tokenizers/byte_tokenizer.py
@@ -19,7 +19,7 @@ from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.tokenizers import tokenizer
 from keras_nlp.utils.tensor_utils import assert_tf_text_installed
 from keras_nlp.utils.tensor_utils import convert_to_ragged_batch
-from keras_nlp.utils.tensor_utils import is_integer_dtype
+from keras_nlp.utils.tensor_utils import is_int_dtype
 
 try:
     import tensorflow_text as tf_text
@@ -165,7 +165,7 @@ class ByteTokenizer(tokenizer.Tokenizer):
     ):
         assert_tf_text_installed(self.__class__.__name__)
 
-        if not is_integer_dtype(dtype):
+        if not is_int_dtype(dtype):
             raise ValueError(
                 "Output dtype must be an integer type. "
                 f"Received: dtype={dtype}"
diff --git a/keras_nlp/tokenizers/sentence_piece_tokenizer.py b/keras_nlp/tokenizers/sentence_piece_tokenizer.py
index 2308c335..3a53c12a 100644
--- a/keras_nlp/tokenizers/sentence_piece_tokenizer.py
+++ b/keras_nlp/tokenizers/sentence_piece_tokenizer.py
@@ -26,7 +26,7 @@ from keras_nlp.utils.python_utils import classproperty
 from keras_nlp.utils.python_utils import format_docstring
 from keras_nlp.utils.tensor_utils import assert_tf_text_installed
 from keras_nlp.utils.tensor_utils import convert_to_ragged_batch
-from keras_nlp.utils.tensor_utils import is_integer_dtype
+from keras_nlp.utils.tensor_utils import is_int_dtype
 from keras_nlp.utils.tensor_utils import is_string_dtype
 from keras_nlp.utils.tensor_utils import tensor_to_list
 
@@ -113,7 +113,7 @@ class SentencePieceTokenizer(tokenizer.Tokenizer):
     ) -> None:
         assert_tf_text_installed(self.__class__.__name__)
 
-        if not is_integer_dtype(dtype) and not is_string_dtype(dtype):
+        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}"
diff --git a/keras_nlp/tokenizers/unicode_codepoint_tokenizer.py b/keras_nlp/tokenizers/unicode_codepoint_tokenizer.py
index 5a16a76f..5fe8f014 100644
--- a/keras_nlp/tokenizers/unicode_codepoint_tokenizer.py
+++ b/keras_nlp/tokenizers/unicode_codepoint_tokenizer.py
@@ -18,7 +18,7 @@ from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.tokenizers import tokenizer
 from keras_nlp.utils.tensor_utils import assert_tf_text_installed
 from keras_nlp.utils.tensor_utils import convert_to_ragged_batch
-from keras_nlp.utils.tensor_utils import is_integer_dtype
+from keras_nlp.utils.tensor_utils import is_int_dtype
 
 try:
     import tensorflow_text as tf_text
@@ -219,7 +219,7 @@ class UnicodeCodepointTokenizer(tokenizer.Tokenizer):
     ) -> None:
         assert_tf_text_installed(self.__class__.__name__)
 
-        if not is_integer_dtype(dtype):
+        if not is_int_dtype(dtype):
             raise ValueError(
                 "Output dtype must be an integer type. "
                 f"Received: dtype={dtype}"
diff --git a/keras_nlp/tokenizers/word_piece_tokenizer.py b/keras_nlp/tokenizers/word_piece_tokenizer.py
index dc9ce494..fe37bebf 100644
--- a/keras_nlp/tokenizers/word_piece_tokenizer.py
+++ b/keras_nlp/tokenizers/word_piece_tokenizer.py
@@ -25,7 +25,7 @@ from keras_nlp.utils.python_utils import classproperty
 from keras_nlp.utils.python_utils import format_docstring
 from keras_nlp.utils.tensor_utils import assert_tf_text_installed
 from keras_nlp.utils.tensor_utils import convert_to_ragged_batch
-from keras_nlp.utils.tensor_utils import is_integer_dtype
+from keras_nlp.utils.tensor_utils import is_int_dtype
 from keras_nlp.utils.tensor_utils import is_string_dtype
 
 try:
@@ -305,7 +305,7 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
     ) -> None:
         assert_tf_text_installed(self.__class__.__name__)
 
-        if not is_integer_dtype(dtype) and not is_string_dtype(dtype):
+        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}"
diff --git a/keras_nlp/utils/tensor_utils.py b/keras_nlp/utils/tensor_utils.py
index 97df75e7..1c9ad1e3 100644
--- a/keras_nlp/utils/tensor_utils.py
+++ b/keras_nlp/utils/tensor_utils.py
@@ -15,6 +15,7 @@
 import tensorflow as tf
 
 from keras_nlp.backend import config
+from keras_nlp.backend import keras
 from keras_nlp.backend import ops
 
 try:
@@ -151,19 +152,21 @@ def is_tensor_type(x):
     return hasattr(x, "__array__")
 
 
-def is_floating_dtype(dtype):
+def standardize_dtype(dtype):
+    if config.multi_backend():
+        return keras.backend.standardize_dtype(dtype)
     if hasattr(dtype, "name"):
-        dtype = dtype.name
-    return "float" in dtype
+        return dtype.name
+    return dtype
 
 
-def is_integer_dtype(dtype):
-    if hasattr(dtype, "name"):
-        dtype = dtype.name
-    return "int" in dtype
+def is_float_dtype(dtype):
+    return "float" in standardize_dtype(dtype)
+
+
+def is_int_dtype(dtype):
+    return "int" in standardize_dtype(dtype)
 
 
 def is_string_dtype(dtype):
-    if hasattr(dtype, "name"):
-        dtype = dtype.name
-    return "string" in dtype
+    return "string" in standardize_dtype(dtype)
diff --git a/requirements-common.txt b/requirements-common.txt
index 21334a40..44661e31 100644
--- a/requirements-common.txt
+++ b/requirements-common.txt
@@ -16,3 +16,6 @@ namex
 # Optional deps.
 rouge-score
 sentencepiece
+tensorflow-datasets
+# Breakage fix.
+ml-dtypes==0.2.0
diff --git a/requirements-jax-cuda.txt b/requirements-jax-cuda.txt
new file mode 100644
index 00000000..bb115b14
--- /dev/null
+++ b/requirements-jax-cuda.txt
@@ -0,0 +1,14 @@
+# Tensorflow cpu-only version.
+tensorflow>=2.14.0
+tensorflow-text>=2.14.0
+
+# Torch cpu-only version.
+--extra-index-url https://download.pytorch.org/whl/cpu
+torch>=2.1.0
+torchvision>=0.16.0
+
+# Jax with cuda support.
+--find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+jax[cuda12_pip]
+
+-r requirements-common.txt
diff --git a/requirements-macos-m1.txt b/requirements-macos-m1.txt
deleted file mode 100644
index 05dd07e6..00000000
--- a/requirements-macos-m1.txt
+++ /dev/null
@@ -1,16 +0,0 @@
-# WARNING: KerasNLP has no official support for MacOS M1 at this time. The
-# following will pull required depenencies from the following external sources.
-# - https://developer.apple.com/metal/tensorflow-plugin/
-# - https://github.com/sun1638650145/Libraries-and-Extensions-for-TensorFlow-for-Apple-Silicon/
-# These are not provided by Google, please review both of these dependencies
-# before proceeding.
-
-# Core deps.
-tensorflow-macos~=2.9
-https://github.com/sun1638650145/Libraries-and-Extensions-for-TensorFlow-for-Apple-Silicon/releases/download/v2.9/tensorflow_text-2.9.0-cp39-cp39-macosx_11_0_arm64.whl
-tensorflow-datasets
-# The metal plugin breaks many tests, so is not enabled by default.
-# tensorflow-metal~=0.5
-
-# Common deps.
--r requirements-common.txt
diff --git a/requirements-nightly.txt b/requirements-nightly.txt
deleted file mode 100644
index 22a0c6e5..00000000
--- a/requirements-nightly.txt
+++ /dev/null
@@ -1,7 +0,0 @@
-# Core deps.
-tf-nightly
-tensorflow-text-nightly
-tfds-nightly
-
-# Common deps.
--r requirements-common.txt
diff --git a/requirements-tensorflow-cuda.txt b/requirements-tensorflow-cuda.txt
new file mode 100644
index 00000000..4b2cf167
--- /dev/null
+++ b/requirements-tensorflow-cuda.txt
@@ -0,0 +1,13 @@
+# Tensorflow with cuda support.
+tensorflow[and-cuda]>=2.14.0
+tensorflow-text>=2.14.0
+
+# Torch cpu-only version.
+--extra-index-url https://download.pytorch.org/whl/cpu
+torch>=2.1.0
+torchvision>=0.16.0
+
+# Jax cpu-only version.
+jax[cpu]
+
+-r requirements-common.txt
diff --git a/requirements-torch-cuda.txt b/requirements-torch-cuda.txt
new file mode 100644
index 00000000..14e94dd8
--- /dev/null
+++ b/requirements-torch-cuda.txt
@@ -0,0 +1,13 @@
+# Tensorflow cpu-only version.
+tensorflow>=2.14.0
+tensorflow-text>=2.14.0
+
+# Torch with cuda support.
+--extra-index-url https://download.pytorch.org/whl/cu118
+torch==2.1.0
+torchvision==0.16.0
+
+# Jax cpu-only version.
+jax[cpu]
+
+-r requirements-common.txt
diff --git a/requirements.txt b/requirements.txt
index 99adbbd6..aa289402 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,7 +1,13 @@
-# Core deps.
-tensorflow~=2.13.0
-tensorflow-text~=2.13.0
-tensorflow-datasets
+# Tensorflow.
+tensorflow>=2.14.0
+tensorflow-text>=2.14.0
+
+# Torch.
+--extra-index-url https://download.pytorch.org/whl/cpu
+torch>=2.1.0
+torchvision>=0.16.0
+
+# Jax.
+jax[cpu]
 
-# Common deps.
 -r requirements-common.txt


Test patch:
diff --git a/keras_nlp/conftest.py b/keras_nlp/conftest.py
index 04daf5cd..3a23ee47 100644
--- a/keras_nlp/conftest.py
+++ b/keras_nlp/conftest.py
@@ -86,8 +86,3 @@ def pytest_collection_modifyitems(config, items):
 tf.debugging.disable_traceback_filtering()
 if backend_config.multi_backend():
     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)
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 fdaab606..8e233f10 100644
--- a/keras_nlp/layers/modeling/cached_multi_head_attention_test.py
+++ b/keras_nlp/layers/modeling/cached_multi_head_attention_test.py
@@ -14,6 +14,7 @@
 
 from keras_nlp.backend import config
 from keras_nlp.backend import ops
+from keras_nlp.backend import random
 from keras_nlp.layers.modeling.cached_multi_head_attention import (
     CachedMultiHeadAttention,
 )
@@ -29,8 +30,8 @@ class CachedMultiHeadAttentionTest(TestCase):
                 "key_dim": 4,
             },
             input_data={
-                "query": ops.random.uniform(shape=(2, 4, 6)),
-                "value": ops.random.uniform(shape=(2, 4, 6)),
+                "query": random.uniform(shape=(2, 4, 6)),
+                "value": random.uniform(shape=(2, 4, 6)),
             },
             expected_output_shape=(2, 4, 6),
             expected_num_trainable_weights=8,
@@ -48,7 +49,7 @@ class CachedMultiHeadAttentionTest(TestCase):
         hidden_dim = num_heads * key_dim
 
         input_shape = (batch_size, seq_len, hidden_dim)
-        x = ops.random.uniform(shape=input_shape)
+        x = random.uniform(shape=input_shape)
         input_cache = ops.zeros((batch_size, 2, seq_len, num_heads, key_dim))
         # Use a causal mask.
         mask = ops.tril(ops.ones((seq_len, seq_len)))
diff --git a/keras_nlp/layers/modeling/f_net_encoder_test.py b/keras_nlp/layers/modeling/f_net_encoder_test.py
index ffafc8a7..06d759ad 100644
--- a/keras_nlp/layers/modeling/f_net_encoder_test.py
+++ b/keras_nlp/layers/modeling/f_net_encoder_test.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from keras_nlp.backend import ops
+from keras_nlp.backend import random
 from keras_nlp.layers.modeling.f_net_encoder import FNetEncoder
 from keras_nlp.tests.test_case import TestCase
 
@@ -29,7 +29,7 @@ class FNetEncoderTest(TestCase):
                 "kernel_initializer": "HeNormal",
                 "bias_initializer": "Zeros",
             },
-            input_data=ops.random.uniform(shape=(2, 4, 6)),
+            input_data=random.uniform(shape=(2, 4, 6)),
             expected_output_shape=(2, 4, 6),
             expected_num_trainable_weights=8,
             expected_num_non_trainable_variables=1,
diff --git a/keras_nlp/layers/modeling/masked_lm_head_test.py b/keras_nlp/layers/modeling/masked_lm_head_test.py
index 9ccdcae1..703c5652 100644
--- a/keras_nlp/layers/modeling/masked_lm_head_test.py
+++ b/keras_nlp/layers/modeling/masked_lm_head_test.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from keras_nlp.backend import ops
+from keras_nlp.backend import random
 from keras_nlp.layers.modeling.masked_lm_head import MaskedLMHead
 from keras_nlp.layers.modeling.reversible_embedding import ReversibleEmbedding
 from keras_nlp.tests.test_case import TestCase
@@ -29,8 +29,8 @@ class MaskedLMHeadTest(TestCase):
                 "bias_initializer": "Zeros",
             },
             input_data={
-                "inputs": ops.random.uniform(shape=(4, 10, 16)),
-                "mask_positions": ops.random.randint(
+                "inputs": random.uniform(shape=(4, 10, 16)),
+                "mask_positions": random.randint(
                     minval=0, maxval=10, shape=(4, 5)
                 ),
             },
@@ -51,8 +51,8 @@ class MaskedLMHeadTest(TestCase):
                 "token_embedding": embedding,
             },
             input_data={
-                "inputs": ops.random.uniform(shape=(4, 10, 16)),
-                "mask_positions": ops.random.randint(
+                "inputs": random.uniform(shape=(4, 10, 16)),
+                "mask_positions": random.randint(
                     minval=0, maxval=10, shape=(4, 5)
                 ),
             },
diff --git a/keras_nlp/layers/modeling/position_embedding_test.py b/keras_nlp/layers/modeling/position_embedding_test.py
index ec16f541..e80cef8c 100644
--- a/keras_nlp/layers/modeling/position_embedding_test.py
+++ b/keras_nlp/layers/modeling/position_embedding_test.py
@@ -16,6 +16,7 @@ import numpy as np
 
 from keras_nlp.backend import keras
 from keras_nlp.backend import ops
+from keras_nlp.backend import random
 from keras_nlp.layers.modeling.position_embedding import PositionEmbedding
 from keras_nlp.tests.test_case import TestCase
 
@@ -34,7 +35,7 @@ class PositionEmbeddingTest(TestCase):
             init_kwargs={
                 "sequence_length": 21,
             },
-            input_data=ops.random.uniform(shape=(4, 21, 30)),
+            input_data=random.uniform(shape=(4, 21, 30)),
             expected_output_shape=(4, 21, 30),
             expected_num_trainable_weights=1,
         )
@@ -45,7 +46,7 @@ class PositionEmbeddingTest(TestCase):
             init_kwargs={
                 "sequence_length": 21,
             },
-            input_data=ops.random.uniform(shape=(4, 5, 21, 30)),
+            input_data=random.uniform(shape=(4, 5, 21, 30)),
             expected_output_shape=(4, 5, 21, 30),
             expected_num_trainable_weights=1,
         )
@@ -145,7 +146,7 @@ class PositionEmbeddingTest(TestCase):
     def test_start_index(self):
         batch_size, seq_length, feature_size = 2, 3, 4
         layer = PositionEmbedding(seq_length)
-        data = ops.random.uniform(shape=(batch_size, seq_length, feature_size))
+        data = random.uniform(shape=(batch_size, seq_length, feature_size))
         full_output = layer(data)
         sequential_output = ops.zeros((batch_size, seq_length, feature_size))
         for i in range(seq_length):
diff --git a/keras_nlp/layers/modeling/reversible_embedding_test.py b/keras_nlp/layers/modeling/reversible_embedding_test.py
index ceb04578..dcd4599c 100644
--- a/keras_nlp/layers/modeling/reversible_embedding_test.py
+++ b/keras_nlp/layers/modeling/reversible_embedding_test.py
@@ -20,6 +20,7 @@ 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.backend import random
 from keras_nlp.layers.modeling.reversible_embedding import ReversibleEmbedding
 from keras_nlp.tests.test_case import TestCase
 
@@ -38,7 +39,7 @@ class ReversibleEmbeddingTest(TestCase):
                 "tie_weights": tie_weights,
                 "embeddings_initializer": "HeNormal",
             },
-            input_data=ops.random.randint(minval=0, maxval=100, shape=(4, 10)),
+            input_data=random.randint(minval=0, maxval=100, shape=(4, 10)),
             expected_output_shape=(4, 10, 32),
             expected_num_trainable_weights=1 if tie_weights else 2,
         )
diff --git a/keras_nlp/layers/modeling/rotary_embedding_test.py b/keras_nlp/layers/modeling/rotary_embedding_test.py
index 3fdac028..cb502f45 100644
--- a/keras_nlp/layers/modeling/rotary_embedding_test.py
+++ b/keras_nlp/layers/modeling/rotary_embedding_test.py
@@ -14,6 +14,7 @@
 
 from keras_nlp.backend import keras
 from keras_nlp.backend import ops
+from keras_nlp.backend import random
 from keras_nlp.layers.modeling.rotary_embedding import RotaryEmbedding
 from keras_nlp.tests.test_case import TestCase
 
@@ -28,7 +29,7 @@ class RotaryEmbeddingTest(TestCase):
                 "sequence_axis": 1,
                 "feature_axis": -1,
             },
-            input_data=ops.random.uniform(shape=(2, 4, 6)),
+            input_data=random.uniform(shape=(2, 4, 6)),
             expected_output_shape=(2, 4, 6),
         )
 
@@ -38,7 +39,7 @@ class RotaryEmbeddingTest(TestCase):
             init_kwargs={
                 "max_wavelength": 1000,
             },
-            input_data=ops.random.uniform(shape=(2, 8, 4, 6)),
+            input_data=random.uniform(shape=(2, 8, 4, 6)),
             expected_output_shape=(2, 8, 4, 6),
         )
 
@@ -86,7 +87,7 @@ class RotaryEmbeddingTest(TestCase):
     def test_start_index(self):
         batch_size, seq_length, feature_size = 2, 3, 4
         layer = RotaryEmbedding(seq_length)
-        data = ops.random.uniform(shape=(batch_size, seq_length, feature_size))
+        data = random.uniform(shape=(batch_size, seq_length, feature_size))
         full_output = layer(data)
         sequential_output = ops.zeros((batch_size, seq_length, feature_size))
         for i in range(seq_length):
diff --git a/keras_nlp/layers/modeling/sine_position_encoding_test.py b/keras_nlp/layers/modeling/sine_position_encoding_test.py
index 22d1d9d3..2163d4ee 100644
--- a/keras_nlp/layers/modeling/sine_position_encoding_test.py
+++ b/keras_nlp/layers/modeling/sine_position_encoding_test.py
@@ -14,6 +14,7 @@
 
 from keras_nlp.backend import keras
 from keras_nlp.backend import ops
+from keras_nlp.backend import random
 from keras_nlp.layers.modeling.sine_position_encoding import (
     SinePositionEncoding,
 )
@@ -27,7 +28,7 @@ class SinePositionEncodingTest(TestCase):
             init_kwargs={
                 "max_wavelength": 10000,
             },
-            input_data=ops.random.uniform(shape=(2, 4, 6)),
+            input_data=random.uniform(shape=(2, 4, 6)),
             expected_output_shape=(2, 4, 6),
         )
 
@@ -37,7 +38,7 @@ class SinePositionEncodingTest(TestCase):
             init_kwargs={
                 "max_wavelength": 10000,
             },
-            input_data=ops.random.uniform(shape=(1, 2, 4, 6)),
+            input_data=random.uniform(shape=(1, 2, 4, 6)),
             expected_output_shape=(1, 2, 4, 6),
         )
 
@@ -85,7 +86,7 @@ class SinePositionEncodingTest(TestCase):
                 pos_encoding,
             ]
         )
-        input = ops.random.uniform(shape=[1, 4, 6])
+        input = random.uniform(shape=[1, 4, 6])
         output = model(input)
 
         # comapre position encoding values for position 0 and 3
@@ -97,7 +98,7 @@ class SinePositionEncodingTest(TestCase):
     def test_start_index(self):
         batch_size, seq_length, feature_size = 2, 3, 4
         layer = SinePositionEncoding()
-        data = ops.random.uniform(shape=(batch_size, seq_length, feature_size))
+        data = random.uniform(shape=(batch_size, seq_length, feature_size))
         full_output = layer(data)
         sequential_output = ops.zeros((batch_size, seq_length, feature_size))
         for i in range(seq_length):
diff --git a/keras_nlp/layers/modeling/token_and_position_embedding_test.py b/keras_nlp/layers/modeling/token_and_position_embedding_test.py
index ceb96b11..b0c5949a 100644
--- a/keras_nlp/layers/modeling/token_and_position_embedding_test.py
+++ b/keras_nlp/layers/modeling/token_and_position_embedding_test.py
@@ -16,6 +16,7 @@ import numpy as np
 
 from keras_nlp.backend import keras
 from keras_nlp.backend import ops
+from keras_nlp.backend import random
 from keras_nlp.layers.modeling.token_and_position_embedding import (
     TokenAndPositionEmbedding,
 )
@@ -32,7 +33,7 @@ class TokenAndPositionEmbeddingTest(TestCase):
                 "embedding_dim": 3,
                 "embeddings_initializer": keras.initializers.Constant(1.0),
             },
-            input_data=ops.random.randint(minval=0, maxval=5, shape=(2, 4)),
+            input_data=random.randint(minval=0, maxval=5, shape=(2, 4)),
             expected_output_shape=(2, 4, 3),
             expected_output_data=ops.ones((2, 4, 3)) * 2,
             expected_num_trainable_weights=2,
diff --git a/keras_nlp/layers/modeling/transformer_decoder_test.py b/keras_nlp/layers/modeling/transformer_decoder_test.py
index 9a7bc9c1..f904d925 100644
--- a/keras_nlp/layers/modeling/transformer_decoder_test.py
+++ b/keras_nlp/layers/modeling/transformer_decoder_test.py
@@ -15,6 +15,7 @@
 from absl.testing import parameterized
 
 from keras_nlp.backend import ops
+from keras_nlp.backend import random
 from keras_nlp.layers.modeling.transformer_decoder import TransformerDecoder
 from keras_nlp.tests.test_case import TestCase
 
@@ -36,7 +37,7 @@ class TransformerDecoderTest(TestCase):
                 "kernel_initializer": "HeNormal",
                 "bias_initializer": "Zeros",
             },
-            input_data=ops.random.uniform(shape=(2, 4, 6)),
+            input_data=random.uniform(shape=(2, 4, 6)),
             expected_output_shape=(2, 4, 6),
             expected_num_trainable_weights=16,
             expected_num_non_trainable_variables=3,  # dropout rng seeds
@@ -60,8 +61,8 @@ class TransformerDecoderTest(TestCase):
                 "bias_initializer": "Zeros",
             },
             input_data={
-                "decoder_sequence": ops.random.uniform(shape=(2, 4, 6)),
-                "encoder_sequence": ops.random.uniform(shape=(2, 4, 6)),
+                "decoder_sequence": random.uniform(shape=(2, 4, 6)),
+                "encoder_sequence": random.uniform(shape=(2, 4, 6)),
             },
             expected_output_shape=(2, 4, 6),
             expected_num_trainable_weights=26,
@@ -106,8 +107,8 @@ class TransformerDecoderTest(TestCase):
             intermediate_dim=4,
             num_heads=2,
         )
-        decoder_sequence = ops.random.uniform(shape=[1, 4, 6])
-        encoder_sequence = ops.random.uniform(shape=[1, 4, 6])
+        decoder_sequence = random.uniform(shape=[1, 4, 6])
+        encoder_sequence = random.uniform(shape=[1, 4, 6])
         mask = ops.array([[True, True, False, False]])
         decoder_sequence._keras_mask = mask
         outputs = decoder(decoder_sequence, encoder_sequence)
@@ -118,7 +119,7 @@ class TransformerDecoderTest(TestCase):
             intermediate_dim=4,
             num_heads=2,
         )
-        decoder_sequence = ops.random.uniform(shape=[1, 4, 6])
+        decoder_sequence = random.uniform(shape=[1, 4, 6])
         mask = ops.array([[True, True, False, False]])
         decoder_sequence._keras_mask = mask
         outputs = decoder(decoder_sequence)
@@ -132,7 +133,7 @@ class TransformerDecoderTest(TestCase):
         hidden_dim = num_heads * key_dim
 
         input_shape = (batch_size, seq_len, hidden_dim)
-        x = ops.random.uniform(shape=input_shape)
+        x = random.uniform(shape=input_shape)
         input_cache = ops.zeros((batch_size, 2, seq_len, num_heads, key_dim))
         outputs = ops.zeros_like(x)
 
@@ -174,6 +175,6 @@ class TransformerDecoderTest(TestCase):
             intermediate_dim=4,
             num_heads=2,
         )
-        decoder_sequence = ops.random.uniform(shape=[1, 4, 6])
-        encoder_sequence = ops.random.uniform(shape=[1, 4, 5])
+        decoder_sequence = random.uniform(shape=[1, 4, 6])
+        encoder_sequence = random.uniform(shape=[1, 4, 5])
         decoder(decoder_sequence, encoder_sequence)
diff --git a/keras_nlp/layers/modeling/transformer_encoder_test.py b/keras_nlp/layers/modeling/transformer_encoder_test.py
index 9fe7c3ea..3882a0e7 100644
--- a/keras_nlp/layers/modeling/transformer_encoder_test.py
+++ b/keras_nlp/layers/modeling/transformer_encoder_test.py
@@ -16,6 +16,7 @@ from absl.testing import parameterized
 
 from keras_nlp.backend import keras
 from keras_nlp.backend import ops
+from keras_nlp.backend import random
 from keras_nlp.layers.modeling.transformer_encoder import TransformerEncoder
 from keras_nlp.tests.test_case import TestCase
 
@@ -37,7 +38,7 @@ class TransformerEncoderTest(TestCase):
                 "kernel_initializer": "HeNormal",
                 "bias_initializer": "Zeros",
             },
-            input_data=ops.random.uniform(shape=(2, 4, 6)),
+            input_data=random.uniform(shape=(2, 4, 6)),
             expected_output_shape=(2, 4, 6),
             expected_num_trainable_weights=16,
             expected_num_non_trainable_variables=3,  # dropout rng seeds
@@ -59,7 +60,7 @@ class TransformerEncoderTest(TestCase):
                 encoder,
             ]
         )
-        input = ops.random.uniform(shape=[2, 4, 6])
+        input = random.uniform(shape=[2, 4, 6])
         model(input)
 
     def test_valid_call_with_mask(self):
@@ -68,7 +69,7 @@ class TransformerEncoderTest(TestCase):
             num_heads=2,
         )
         encoder.build([2, 4, 6])
-        input = ops.random.uniform(shape=[2, 4, 6])
+        input = random.uniform(shape=[2, 4, 6])
         mask = input[:, :, 0] < 0.5
         encoder(input, mask)
 
@@ -86,7 +87,7 @@ class TransformerEncoderTest(TestCase):
             intermediate_dim=4,
             num_heads=2,
         )
-        inputs = ops.random.uniform(shape=[1, 4, 6])
+        inputs = random.uniform(shape=[1, 4, 6])
         mask = ops.array([[True, True, False, False]])
         inputs._keras_mask = mask
         outputs = encoder(inputs)
diff --git a/keras_nlp/layers/modeling/transformer_layer_utils_test.py b/keras_nlp/layers/modeling/transformer_layer_utils_test.py
index 7fc2013a..57df677e 100644
--- a/keras_nlp/layers/modeling/transformer_layer_utils_test.py
+++ b/keras_nlp/layers/modeling/transformer_layer_utils_test.py
@@ -14,6 +14,7 @@
 
 import keras_nlp.layers.modeling.transformer_layer_utils as utils
 from keras_nlp.backend import ops
+from keras_nlp.backend import random
 from keras_nlp.tests.test_case import TestCase
 
 
@@ -25,7 +26,7 @@ class TransformerLayerUtilsTest(TestCase):
     def test_merge_padding_and_attention_mask(self):
         padding_mask = ops.array([[1, 1, 0]])
         attention_mask = ops.array([[[0, 0, 1], [0, 1, 0], [1, 0, 0]]])
-        inputs = ops.random.uniform(shape=[1, 3, 2])
+        inputs = random.uniform(shape=[1, 3, 2])
         merged_mask = utils.merge_padding_and_attention_mask(
             inputs,
             padding_mask,
@@ -37,7 +38,7 @@ class TransformerLayerUtilsTest(TestCase):
         with self.assertRaises(ValueError):
             padding_mask = ops.array([[[1, 1, 0], [1, 0, 0]]])
             attention_mask = ops.array([[0, 0, 1], [0, 1, 0], [1, 0, 0]])
-            inputs = ops.random.uniform(shape=[1, 3, 2])
+            inputs = random.uniform(shape=[1, 3, 2])
             utils.merge_padding_and_attention_mask(
                 inputs,
                 padding_mask,
@@ -47,7 +48,7 @@ class TransformerLayerUtilsTest(TestCase):
         with self.assertRaises(ValueError):
             padding_mask = ops.array([[1, 1, 0]])
             attention_mask = ops.array([[0, 0, 1], [1, 0, 0]])
-            inputs = ops.random.uniform(shape=[1, 3, 2])
+            inputs = random.uniform(shape=[1, 3, 2])
             utils.merge_padding_and_attention_mask(
                 inputs,
                 padding_mask,
diff --git a/keras_nlp/models/albert/albert_presets_test.py b/keras_nlp/models/albert/albert_presets_test.py
index 86265a1a..f7576c67 100644
--- a/keras_nlp/models/albert/albert_presets_test.py
+++ b/keras_nlp/models/albert/albert_presets_test.py
@@ -16,6 +16,7 @@ import pytest
 from absl.testing import parameterized
 
 from keras_nlp.backend import ops
+from keras_nlp.backend import random
 from keras_nlp.models.albert.albert_backbone import AlbertBackbone
 from keras_nlp.models.albert.albert_classifier import AlbertClassifier
 from keras_nlp.models.albert.albert_preprocessor import AlbertPreprocessor
@@ -138,7 +139,7 @@ class AlbertPresetFullTest(TestCase):
                 preset, load_weights=load_weights
             )
             input_data = {
-                "token_ids": ops.random.uniform(
+                "token_ids": random.uniform(
                     shape=(1, 512), dtype="int64", maxval=model.vocabulary_size
                 ),
                 "segment_ids": ops.array([0] * 200 + [1] * 312, shape=(1, 512)),
@@ -171,7 +172,7 @@ class AlbertPresetFullTest(TestCase):
                 load_weights=load_weights,
             )
             input_data = {
-                "token_ids": ops.random.uniform(
+                "token_ids": random.uniform(
                     shape=(1, 512),
                     dtype="int64",
                     maxval=classifier.backbone.vocabulary_size,
diff --git a/keras_nlp/models/bart/bart_presets_test.py b/keras_nlp/models/bart/bart_presets_test.py
index 89518d84..59bf3b55 100644
--- a/keras_nlp/models/bart/bart_presets_test.py
+++ b/keras_nlp/models/bart/bart_presets_test.py
@@ -15,6 +15,7 @@
 #
 
 from keras_nlp.backend import ops
+from keras_nlp.backend import random
 from keras_nlp.tests.test_case import TestCase
 
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -117,7 +118,7 @@ class BartPresetFullTest(TestCase):
         for preset in BartBackbone.presets:
             model = BartBackbone.from_preset(preset, load_weights=load_weights)
             input_data = {
-                "encoder_token_ids": ops.random.uniform(
+                "encoder_token_ids": random.uniform(
                     shape=(1, 1024),
                     dtype="int64",
                     maxval=model.vocabulary_size,
@@ -125,7 +126,7 @@ class BartPresetFullTest(TestCase):
                 "encoder_padding_mask": ops.array(
                     [1] * 768 + [0] * 256, shape=(1, 1024)
                 ),
-                "decoder_token_ids": ops.random.uniform(
+                "decoder_token_ids": random.uniform(
                     shape=(1, 1024),
                     dtype="int64",
                     maxval=model.vocabulary_size,
diff --git a/keras_nlp/models/bert/bert_presets_test.py b/keras_nlp/models/bert/bert_presets_test.py
index 71e739bb..a84286c0 100644
--- a/keras_nlp/models/bert/bert_presets_test.py
+++ b/keras_nlp/models/bert/bert_presets_test.py
@@ -16,6 +16,7 @@ import pytest
 from absl.testing import parameterized
 
 from keras_nlp.backend import ops
+from keras_nlp.backend import random
 from keras_nlp.models.bert.bert_backbone import BertBackbone
 from keras_nlp.models.bert.bert_classifier import BertClassifier
 from keras_nlp.models.bert.bert_preprocessor import BertPreprocessor
@@ -186,7 +187,7 @@ class BertPresetFullTest(TestCase):
         for preset in BertBackbone.presets:
             model = BertBackbone.from_preset(preset, load_weights=load_weights)
             input_data = {
-                "token_ids": ops.random.uniform(
+                "token_ids": random.uniform(
                     shape=(1, 512), dtype="int64", maxval=model.vocabulary_size
                 ),
                 "segment_ids": ops.array([0] * 200 + [1] * 312, shape=(1, 512)),
@@ -219,7 +220,7 @@ class BertPresetFullTest(TestCase):
                 load_weights=load_weights,
             )
             input_data = {
-                "token_ids": ops.random.uniform(
+                "token_ids": random.uniform(
                     shape=(1, 512),
                     dtype="int64",
                     maxval=classifier.backbone.vocabulary_size,
diff --git a/keras_nlp/models/deberta_v3/deberta_v3_presets_test.py b/keras_nlp/models/deberta_v3/deberta_v3_presets_test.py
index a033825d..7b023a06 100644
--- a/keras_nlp/models/deberta_v3/deberta_v3_presets_test.py
+++ b/keras_nlp/models/deberta_v3/deberta_v3_presets_test.py
@@ -16,6 +16,7 @@ import pytest
 from absl.testing import parameterized
 
 from keras_nlp.backend import ops
+from keras_nlp.backend import random
 from keras_nlp.models.deberta_v3.deberta_v3_backbone import DebertaV3Backbone
 from keras_nlp.models.deberta_v3.deberta_v3_classifier import (
     DebertaV3Classifier,
@@ -150,7 +151,7 @@ class DebertaV3PresetFullTest(TestCase):
                 preset, load_weights=load_weights
             )
             input_data = {
-                "token_ids": ops.random.uniform(
+                "token_ids": random.uniform(
                     shape=(1, 512), dtype="int64", maxval=model.vocabulary_size
                 ),
                 "padding_mask": ops.array([1] * 512, shape=(1, 512)),
@@ -182,7 +183,7 @@ class DebertaV3PresetFullTest(TestCase):
                 preprocessor=None,
             )
             input_data = {
-                "token_ids": ops.random.uniform(
+                "token_ids": random.uniform(
                     shape=(1, 512),
                     dtype="int64",
                     maxval=classifier.backbone.vocabulary_size,
diff --git a/keras_nlp/models/distil_bert/distil_bert_presets_test.py b/keras_nlp/models/distil_bert/distil_bert_presets_test.py
index a974f54a..691f3dbb 100644
--- a/keras_nlp/models/distil_bert/distil_bert_presets_test.py
+++ b/keras_nlp/models/distil_bert/distil_bert_presets_test.py
@@ -16,6 +16,7 @@ import pytest
 from absl.testing import parameterized
 
 from keras_nlp.backend import ops
+from keras_nlp.backend import random
 from keras_nlp.models.distil_bert.distil_bert_backbone import DistilBertBackbone
 from keras_nlp.models.distil_bert.distil_bert_classifier import (
     DistilBertClassifier,
@@ -141,7 +142,7 @@ class DistilBertPresetFullTest(TestCase):
                 preset, load_weights=load_weights
             )
             input_data = {
-                "token_ids": ops.random.uniform(
+                "token_ids": random.uniform(
                     shape=(1, 512), dtype="int64", maxval=model.vocabulary_size
                 ),
                 "padding_mask": ops.array([1] * 512, shape=(1, 512)),
@@ -173,7 +174,7 @@ class DistilBertPresetFullTest(TestCase):
                 preprocessor=None,
             )
             input_data = {
-                "token_ids": ops.random.uniform(
+                "token_ids": random.uniform(
                     shape=(1, 512),
                     dtype="int64",
                     maxval=classifier.backbone.vocabulary_size,
diff --git a/keras_nlp/models/f_net/f_net_presets_test.py b/keras_nlp/models/f_net/f_net_presets_test.py
index 1f8f79c4..9f656b51 100644
--- a/keras_nlp/models/f_net/f_net_presets_test.py
+++ b/keras_nlp/models/f_net/f_net_presets_test.py
@@ -16,6 +16,7 @@ import pytest
 from absl.testing import parameterized
 
 from keras_nlp.backend import ops
+from keras_nlp.backend import random
 from keras_nlp.models.f_net.f_net_backbone import FNetBackbone
 from keras_nlp.models.f_net.f_net_classifier import FNetClassifier
 from keras_nlp.models.f_net.f_net_preprocessor import FNetPreprocessor
@@ -126,7 +127,7 @@ class FNetPresetFullTest(TestCase):
         for preset in FNetBackbone.presets:
             model = FNetBackbone.from_preset(preset, load_weights=load_weights)
             input_data = {
-                "token_ids": ops.random.uniform(
+                "token_ids": random.uniform(
                     shape=(1, 512), dtype="int64", maxval=model.vocabulary_size
                 ),
                 "segment_ids": ops.array([0] * 200 + [1] * 312, shape=(1, 512)),
@@ -158,7 +159,7 @@ class FNetPresetFullTest(TestCase):
                 load_weights=load_weights,
             )
             input_data = {
-                "token_ids": ops.random.uniform(
+                "token_ids": random.uniform(
                     shape=(1, 512),
                     dtype="int64",
                     maxval=classifier.backbone.vocabulary_size,
diff --git a/keras_nlp/models/gpt2/gpt2_backbone_test.py b/keras_nlp/models/gpt2/gpt2_backbone_test.py
index 8fc779c6..c82ec1b0 100644
--- a/keras_nlp/models/gpt2/gpt2_backbone_test.py
+++ b/keras_nlp/models/gpt2/gpt2_backbone_test.py
@@ -84,20 +84,3 @@ class GPT2Test(TestCase):
         # Check that output matches.
         restored_output = restored_model(self.input_batch)
         self.assertAllClose(model_output, restored_output)
-
-    def test_create_layout_map(self):
-        mesh = tf.experimental.dtensor.create_mesh([("batch", 1), ("model", 1)])
-        with GPT2Backbone.create_layout_map(mesh).scope():
-            GPT2Backbone(
-                vocabulary_size=10,
-                num_layers=2,
-                num_heads=2,
-                hidden_dim=2,
-                intermediate_dim=4,
-                max_sequence_length=5,
-            )
-        # Using DTensor enables the mlir bridge as a side effect. Eventually
-        # this will be default, but for now we have compile errors with the
-        # bridge elsewhere and must disable. See
-        # https://github.com/keras-team/keras-nlp/issues/1001
-        tf.config.experimental.disable_mlir_bridge()
diff --git a/keras_nlp/models/gpt2/gpt2_causal_lm_test.py b/keras_nlp/models/gpt2/gpt2_causal_lm_test.py
index c50b6c5c..412083b2 100644
--- a/keras_nlp/models/gpt2/gpt2_causal_lm_test.py
+++ b/keras_nlp/models/gpt2/gpt2_causal_lm_test.py
@@ -165,13 +165,3 @@ class GPT2CausalLMTest(TestCase):
         keras.utils.set_random_seed(42)
         restored_output = restored_model.predict(self.raw_batch)
         self.assertAllClose(model_output, restored_output)
-
-    def test_create_layout_map(self):
-        mesh = tf.experimental.dtensor.create_mesh([("batch", 1), ("model", 1)])
-        with GPT2CausalLM.create_layout_map(mesh).scope():
-            GPT2CausalLM(backbone=self.backbone)
-        # Using DTensor enables the mlir bridge as a side effect. Eventually
-        # this will be default, but for now we have compile errors with the
-        # bridge elsewhere and must disable. See
-        # https://github.com/keras-team/keras-nlp/issues/1001
-        tf.config.experimental.disable_mlir_bridge()
diff --git a/keras_nlp/models/gpt2/gpt2_presets_test.py b/keras_nlp/models/gpt2/gpt2_presets_test.py
index 37e7e53e..a1b64555 100644
--- a/keras_nlp/models/gpt2/gpt2_presets_test.py
+++ b/keras_nlp/models/gpt2/gpt2_presets_test.py
@@ -16,6 +16,7 @@ import pytest
 from absl.testing import parameterized
 
 from keras_nlp.backend import ops
+from keras_nlp.backend import random
 from keras_nlp.models.gpt2.gpt2_backbone import GPT2Backbone
 from keras_nlp.models.gpt2.gpt2_tokenizer import GPT2Tokenizer
 from keras_nlp.tests.test_case import TestCase
@@ -94,7 +95,7 @@ class GPT2PresetFullTest(TestCase):
         for preset in GPT2Backbone.presets:
             model = GPT2Backbone.from_preset(preset, load_weights=load_weights)
             input_data = {
-                "token_ids": ops.random.uniform(
+                "token_ids": random.uniform(
                     shape=(1, 1024),
                     dtype="int64",
                     maxval=model.vocabulary_size,
diff --git a/keras_nlp/models/opt/opt_backbone_test.py b/keras_nlp/models/opt/opt_backbone_test.py
index 1d7e5488..c8870010 100644
--- a/keras_nlp/models/opt/opt_backbone_test.py
+++ b/keras_nlp/models/opt/opt_backbone_test.py
@@ -84,20 +84,3 @@ class OPTBackboneTest(TestCase):
         # Check that output matches.
         restored_output = restored_model(self.input_batch)
         self.assertAllClose(model_output, restored_output)
-
-    def test_create_layout_map(self):
-        mesh = tf.experimental.dtensor.create_mesh([("batch", 1), ("model", 1)])
-        with OPTBackbone.create_layout_map(mesh).scope():
-            OPTBackbone(
-                vocabulary_size=10,
-                num_layers=2,
-                num_heads=2,
-                hidden_dim=2,
-                intermediate_dim=4,
-                max_sequence_length=5,
-            )
-        # Using DTensor enables the mlir bridge as a side effect. Eventually
-        # this will be default, but for now we have compile errors with the
-        # bridge elsewhere and must disable. See
-        # https://github.com/keras-team/keras-nlp/issues/1001
-        tf.config.experimental.disable_mlir_bridge()
diff --git a/keras_nlp/models/opt/opt_causal_lm_test.py b/keras_nlp/models/opt/opt_causal_lm_test.py
index 19325c7b..1e8fcb87 100644
--- a/keras_nlp/models/opt/opt_causal_lm_test.py
+++ b/keras_nlp/models/opt/opt_causal_lm_test.py
@@ -171,13 +171,3 @@ class OPTCausalLMTest(TestCase):
         keras.utils.set_random_seed(42)
         restored_output = restored_model.predict(self.raw_batch)
         self.assertAllClose(model_output, restored_output)
-
-    def test_create_layout_map(self):
-        mesh = tf.experimental.dtensor.create_mesh([("batch", 1), ("model", 1)])
-        with OPTCausalLM.create_layout_map(mesh).scope():
-            OPTCausalLM(backbone=self.backbone)
-        # Using DTensor enables the mlir bridge as a side effect. Eventually
-        # this will be default, but for now we have compile errors with the
-        # bridge elsewhere and must disable. See
-        # https://github.com/keras-team/keras-nlp/issues/1001
-        tf.config.experimental.disable_mlir_bridge()
diff --git a/keras_nlp/models/opt/opt_presets_test.py b/keras_nlp/models/opt/opt_presets_test.py
index a9426d29..2484e5f3 100644
--- a/keras_nlp/models/opt/opt_presets_test.py
+++ b/keras_nlp/models/opt/opt_presets_test.py
@@ -16,6 +16,7 @@ import pytest
 from absl.testing import parameterized
 
 from keras_nlp.backend import ops
+from keras_nlp.backend import random
 from keras_nlp.models.opt.opt_backbone import OPTBackbone
 from keras_nlp.models.opt.opt_tokenizer import OPTTokenizer
 from keras_nlp.tests.test_case import TestCase
@@ -94,7 +95,7 @@ class OPTPresetFullTest(TestCase):
         for preset in OPTBackbone.presets:
             model = OPTBackbone.from_preset(preset, load_weights=load_weights)
             input_data = {
-                "token_ids": ops.random.uniform(
+                "token_ids": random.uniform(
                     shape=(1, 1024),
                     dtype="int64",
                     maxval=model.vocabulary_size,
diff --git a/keras_nlp/models/roberta/roberta_presets_test.py b/keras_nlp/models/roberta/roberta_presets_test.py
index 22c1f2e8..657c4350 100644
--- a/keras_nlp/models/roberta/roberta_presets_test.py
+++ b/keras_nlp/models/roberta/roberta_presets_test.py
@@ -16,6 +16,7 @@ import pytest
 from absl.testing import parameterized
 
 from keras_nlp.backend import ops
+from keras_nlp.backend import random
 from keras_nlp.models.roberta.roberta_backbone import RobertaBackbone
 from keras_nlp.models.roberta.roberta_classifier import RobertaClassifier
 from keras_nlp.models.roberta.roberta_masked_lm import RobertaMaskedLM
@@ -167,7 +168,7 @@ class RobertaPresetFullTest(TestCase):
                 preset, load_weights=load_weights
             )
             input_data = {
-                "token_ids": ops.random.uniform(
+                "token_ids": random.uniform(
                     shape=(1, 512), dtype="int64", maxval=model.vocabulary_size
                 ),
                 "padding_mask": ops.array([1] * 512, shape=(1, 512)),
@@ -197,7 +198,7 @@ class RobertaPresetFullTest(TestCase):
                 load_weights=load_weights,
             )
             input_data = {
-                "token_ids": ops.random.uniform(
+                "token_ids": random.uniform(
                     shape=(1, 512),
                     dtype="int64",
                     maxval=classifier.backbone.vocabulary_size,
@@ -228,7 +229,7 @@ class RobertaPresetFullTest(TestCase):
                 load_weights=load_weights,
             )
             input_data = {
-                "token_ids": ops.random.uniform(
+                "token_ids": random.uniform(
                     shape=(1, 512),
                     dtype="int64",
                     maxval=classifier.backbone.vocabulary_size,
diff --git a/keras_nlp/models/whisper/whisper_audio_feature_extractor_test.py b/keras_nlp/models/whisper/whisper_audio_feature_extractor_test.py
index 8eca8f9c..1d282f76 100644
--- a/keras_nlp/models/whisper/whisper_audio_feature_extractor_test.py
+++ b/keras_nlp/models/whisper/whisper_audio_feature_extractor_test.py
@@ -42,10 +42,10 @@ class WhisperAudioFeatureExtractorTest(TestCase):
         outputs = self.audio_feature_extractor(audio_tensor)
 
         # Verify shape.
-        self.assertEqual(outputs.shape, (1, 5, self.num_mels))
+        self.assertEqual(outputs.shape, (5, self.num_mels))
         # Verify output.
         expected = [1.1656, 1.0151, -0.8343, -0.8343, -0.8343]
-        self.assertAllClose(outputs[0, :, 0], expected, atol=0.01, rtol=0.01)
+        self.assertAllClose(outputs[:, 0], expected, atol=0.01, rtol=0.01)
 
     def test_batched_inputs(self):
         audio_tensor_1 = tf.ones((2,), dtype="float32")
diff --git a/keras_nlp/models/whisper/whisper_preprocessor_test.py b/keras_nlp/models/whisper/whisper_preprocessor_test.py
index 3f07ef61..9f8f54e4 100644
--- a/keras_nlp/models/whisper/whisper_preprocessor_test.py
+++ b/keras_nlp/models/whisper/whisper_preprocessor_test.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import numpy as np
 import tensorflow as tf
 
 from keras_nlp.backend import keras
@@ -89,8 +90,8 @@ class WhisperPreprocessorTest(TestCase):
 
     def test_unbatched_preprocess(self):
         input_data = {
-            "encoder_audio": tf.ones((200,)),
-            "decoder_text": tf.constant(" airplane at airport"),
+            "encoder_audio": np.ones((200,)),
+            "decoder_text": " airplane at airport",
         }
 
         x = self.preprocessor(input_data)
@@ -106,8 +107,8 @@ class WhisperPreprocessorTest(TestCase):
 
     def test_preprocess_batch(self):
         input_data = {
-            "encoder_audio": tf.ones((4, 200)),
-            "decoder_text": tf.constant([" airplane at airport"] * 4),
+            "encoder_audio": np.ones((4, 200)),
+            "decoder_text": [" airplane at airport"] * 4,
         }
 
         x = self.preprocessor(input_data)
@@ -125,11 +126,11 @@ class WhisperPreprocessorTest(TestCase):
 
     def test_preprocess_labeled_batch(self):
         x = {
-            "encoder_audio": tf.ones((4, 200)),
-            "decoder_text": tf.constant([" airplane at airport"] * 4),
+            "encoder_audio": np.ones((4, 200)),
+            "decoder_text": [" airplane at airport"] * 4,
         }
-        y_in = tf.constant([1] * 4)
-        sw_in = tf.constant([1.0] * 4)
+        y_in = np.ones((4,))
+        sw_in = np.ones((4,))
         x, y, sw = self.preprocessor(x, y_in, sw_in)
         self.assertAllEqual(
             x["encoder_features"].shape, [4, self.output_length, self.num_mels]
@@ -147,8 +148,8 @@ class WhisperPreprocessorTest(TestCase):
 
     def test_preprocess_dataset(self):
         x = {
-            "encoder_audio": tf.ones((4, 200)),
-            "decoder_text": tf.constant([" airplane at airport"] * 4),
+            "encoder_audio": np.ones((4, 200)),
+            "decoder_text": [" airplane at airport"] * 4,
         }
         ds = tf.data.Dataset.from_tensor_slices(x)
         ds = ds.map(self.preprocessor)
@@ -167,8 +168,8 @@ class WhisperPreprocessorTest(TestCase):
 
     def test_sequence_length_override(self):
         input_data = {
-            "encoder_audio": tf.ones((200,)),
-            "decoder_text": tf.constant(" airplane at airport"),
+            "encoder_audio": np.ones((200,)),
+            "decoder_text": " airplane at airport",
         }
         x = self.preprocessor(input_data, decoder_sequence_length=6)
         self.assertAllEqual(x["decoder_token_ids"], [9, 14, 13, 11, 0, 10])
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_presets_test.py b/keras_nlp/models/xlm_roberta/xlm_roberta_presets_test.py
index b74de26b..a1824c86 100644
--- a/keras_nlp/models/xlm_roberta/xlm_roberta_presets_test.py
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_presets_test.py
@@ -16,6 +16,7 @@ import pytest
 from absl.testing import parameterized
 
 from keras_nlp.backend import ops
+from keras_nlp.backend import random
 from keras_nlp.models.xlm_roberta.xlm_roberta_backbone import XLMRobertaBackbone
 from keras_nlp.models.xlm_roberta.xlm_roberta_classifier import (
     XLMRobertaClassifier,
@@ -143,7 +144,7 @@ class XLMRobertaPresetFullTest(TestCase):
                 preset, load_weights=load_weights
             )
             input_data = {
-                "token_ids": ops.random.uniform(
+                "token_ids": random.uniform(
                     shape=(1, 512), dtype="int64", maxval=model.vocabulary_size
                 ),
                 "padding_mask": ops.array([1] * 512, shape=(1, 512)),
@@ -177,7 +178,7 @@ class XLMRobertaPresetFullTest(TestCase):
                 preprocessor=None,
             )
             input_data = {
-                "token_ids": ops.random.uniform(
+                "token_ids": random.uniform(
                     shape=(1, 512),
                     dtype="int64",
                     maxval=classifier.backbone.vocabulary_size,
diff --git a/keras_nlp/tests/test_case.py b/keras_nlp/tests/test_case.py
index d0af5a47..976f4eb3 100644
--- a/keras_nlp/tests/test_case.py
+++ b/keras_nlp/tests/test_case.py
@@ -17,12 +17,12 @@ import json
 import tensorflow as tf
 import tree
 from absl.testing import parameterized
-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
 from keras_nlp.backend import ops
+from keras_nlp.utils.tensor_utils import is_float_dtype
+from keras_nlp.utils.tensor_utils import standardize_dtype
 
 
 def convert_to_comparible_type(x):
@@ -39,7 +39,7 @@ def convert_to_comparible_type(x):
         return tree.map_structure(lambda x: x.decode("utf-8"), x)
     if isinstance(x, (tf.Tensor, tf.RaggedTensor)):
         return x
-    if ops.is_tensor(x):
+    if hasattr(x, "__array__"):
         return ops.convert_to_numpy(x)
     return x