diff --git a/keras_nlp/layers/preprocessing/mlm_mask_generator.py b/keras_nlp/layers/preprocessing/mlm_mask_generator.py
index 19c6dfc..3c2bbd6 100644
--- a/keras_nlp/layers/preprocessing/mlm_mask_generator.py
+++ b/keras_nlp/layers/preprocessing/mlm_mask_generator.py
@@ -129,7 +129,7 @@ class MLMMaskGenerator(keras.layers.Layer):
 
     def call(self, inputs):
         input_is_ragged = isinstance(inputs, tf.RaggedTensor)
-        input_is_1d = tf.rank(inputs) == 1
+        input_is_1d = inputs.shape.rank == 1
         if input_is_1d:
             # If inputs is of rank 1, we manually add the batch axis.
             inputs = inputs[tf.newaxis, :]
@@ -137,6 +137,7 @@ class MLMMaskGenerator(keras.layers.Layer):
             # `tf_text.mask_language_model` requires a ragged tensor, so
             # convert dense to ragged.
             inputs = tf.RaggedTensor.from_tensor(inputs)
+
         (tokens, mask_positions, mask_ids,) = tf_text.mask_language_model(
             inputs,
             item_selector=self._random_selector,

diff --git a/keras_nlp/layers/preprocessing/mlm_mask_generator_test.py b/keras_nlp/layers/preprocessing/mlm_mask_generator_test.py
index 6654798..492a271 100644
--- a/keras_nlp/layers/preprocessing/mlm_mask_generator_test.py
+++ b/keras_nlp/layers/preprocessing/mlm_mask_generator_test.py
@@ -202,7 +202,7 @@ class MLMMaskGeneratorTest(tf.test.TestCase):
             mask_token_rate=1,
             random_token_rate=0,
         )
-        inputs = [unselectable_token_ids]
+        inputs = tf.convert_to_tensor([unselectable_token_ids])
         outputs = mlm_masker(inputs)
         # Verify that no token is masked out.
         self.assertEqual(tf.reduce_sum(outputs["mask_positions"]), 0)
@@ -232,3 +232,36 @@ class MLMMaskGeneratorTest(tf.test.TestCase):
             [[0, 1, 2], [0, 1, 2, 3, 4, 5], [0, 1, 2, 3, 4]]
         )
         cloned_mlm_masker(inputs)
+
+    def test_graph_mode_execution(self):
+        mlm_masker = MLMMaskGenerator(
+            vocabulary_size=self.vocabulary_size,
+            mask_selection_rate=0.5,
+            mask_token_id=self.mask_token_id,
+            mask_selection_length=5,
+        )
+
+        @tf.function
+        def masker(inputs):
+            return mlm_masker(inputs)
+
+        masker(tf.constant([1, 2, 3]))
+        masker(tf.constant([[1, 2, 3], [1, 2, 3]]))
+        masker(tf.ragged.constant([[3, 5, 7, 7], [4, 6, 7, 5]]))
+
+    def test_with_tf_data(self):
+        ds = tf.data.Dataset.from_tensor_slices(
+            tf.ones((100, 10), dtype="int32")
+        )
+        mlm_masker = MLMMaskGenerator(
+            vocabulary_size=self.vocabulary_size,
+            mask_selection_rate=0.5,
+            mask_token_id=self.mask_token_id,
+            mask_selection_length=5,
+        )
+        batch_first = ds.batch(8).map(mlm_masker)
+        batch_second = ds.map(mlm_masker).batch(8)
+        self.assertEqual(
+            batch_first.take(1).get_single_element()["tokens"].shape,
+            batch_second.take(1).get_single_element()["tokens"].shape,
+        )

Error with tf.data and MLMMaskGenerator for dense inputs
When using the MLMMaskGenerator for to map over dense, batched inputs in a tf.data.Dataset, we get the following error...

`TypeError: unsupported operand type(s) for +: 'int' and 'NoneType'`

This colab has a reproduction https://colab.research.google.com/gist/mattdangerw/4596df85105ff6e6731128fc79d16bf3/mlmmaskgenerator-bug.ipynb

tf.data + dense, batched inputs might be the most common use case for this layer, so this is an important one to fix.
keras_nlp/layers/preprocessing/mlm_mask_generator_test.py::MLMMaskGeneratorTest::test_graph_mode_execution keras_nlp/layers/preprocessing/mlm_mask_generator_test.py::MLMMaskGeneratorTest::test_with_tf_data
keras_nlp/layers/preprocessing/mlm_mask_generator_test.py::MLMMaskGeneratorTest::test_apply_random_token_not_mask keras_nlp/layers/preprocessing/mlm_mask_generator_test.py::MLMMaskGeneratorTest::test_config keras_nlp/layers/preprocessing/mlm_mask_generator_test.py::MLMMaskGeneratorTest::test_invalid_mask_token keras_nlp/layers/preprocessing/mlm_mask_generator_test.py::MLMMaskGeneratorTest::test_mask_1d_input keras_nlp/layers/preprocessing/mlm_mask_generator_test.py::MLMMaskGeneratorTest::test_mask_ragged_tensor keras_nlp/layers/preprocessing/mlm_mask_generator_test.py::MLMMaskGeneratorTest::test_mask_tensor keras_nlp/layers/preprocessing/mlm_mask_generator_test.py::MLMMaskGeneratorTest::test_number_of_masked_position_as_expected keras_nlp/layers/preprocessing/mlm_mask_generator_test.py::MLMMaskGeneratorTest::test_unselectable_tokens