Update to always using ops.shape (#1231)

With recently released changes to keras-core, ops.shape will always
return a tuple.

keras_nlp/samplers/beam_sampler.py

@@ -124,12 +124,12 @@ def create_beams(x):
 
         def flatten_beams(x):
             """Combine the beam dim and batch dim."""
-            flat_shape = (batch_size * self.num_beams,) + tuple(x.shape)[2:]
+            flat_shape = (batch_size * self.num_beams,) + ops.shape(x)[2:]
             return ops.reshape(x, flat_shape)
 
         def unflatten_beams(x):
             """Separate the beam dim and batch dim."""
-            unflat_shape = (batch_size, self.num_beams) + tuple(x.shape)[1:]
+            unflat_shape = (batch_size, self.num_beams) + ops.shape(x)[1:]
             return ops.reshape(x, unflat_shape)
 
         if mask is None:

keras_nlp/samplers/contrastive_sampler.py

@@ -103,17 +103,17 @@ def __call__(
         def create_beams(x):
             """Add initial beam state."""
             x = ops.repeat(x, self.k, axis=0)
-            flat_shape = (batch_size * self.k,) + tuple(x.shape)[1:]
+            flat_shape = (batch_size * self.k,) + ops.shape(x)[1:]
             return ops.reshape(x, flat_shape)
 
         def flatten_beams(x):
             """Combine the beam dim and batch dim."""
-            flat_shape = (batch_size * self.k,) + tuple(x.shape)[2:]
+            flat_shape = (batch_size * self.k,) + ops.shape(x)[2:]
             return ops.reshape(x, flat_shape)
 
         def unflatten_beams(x):
             """Separate the beam dim and batch dim."""
-            unflat_shape = (batch_size, self.k) + tuple(x.shape)[1:]
+            unflat_shape = (batch_size, self.k) + ops.shape(x)[1:]
             return ops.reshape(x, unflat_shape)
 
         mask = ops.zeros_like(prompt, dtype="bool") if mask is None else mask