Internal change

PiperOrigin-RevId: 381540013

Author: tensorflower-gardener

official/nlp/keras_nlp/layers/position_embedding_test.py

@@ -48,12 +48,12 @@ def test_non_default_axis_static(self):
     test_layer = position_embedding.PositionEmbedding(
         max_length=sequence_length, seq_axis=2)
     width = 30
-    input_tensor = tf.keras.Input(shape=(sequence_length, width, width))
+    input_tensor = tf.keras.Input(shape=(width, sequence_length, width))
     output_tensor = test_layer(input_tensor)
 
     # When using static positional embedding shapes, the output is expected
     # to be the same as the input shape in all dimensions save batch.
-    expected_output_shape = [None, sequence_length, width, width]
+    expected_output_shape = [None, width, sequence_length, width]
     self.assertEqual(expected_output_shape, output_tensor.shape.as_list())
     # The default output dtype for this layer should be tf.float32.
     self.assertEqual(tf.float32, output_tensor.dtype)