Remove einops dep (#2006)

We shouldn't add a new dependency for just to reshape calls. We can
do this the less readable way for now.

Author: mattdangerw

keras_hub/src/models/flux/flux_maths.py

@@ -1,5 +1,4 @@
 import keras
-from einops import rearrange
 from keras import ops
 
 
@@ -58,7 +57,7 @@ def call(self, pos, dim, theta):
         out = ops.stack(
             [ops.cos(out), -ops.sin(out), ops.sin(out), ops.cos(out)], axis=-1
         )
-        out = rearrange(out, "... n d (i j) -> ... n d i j", i=2, j=2)
+        out = ops.reshape(out, ops.shape(out)[:-1] + (2, 2))
         return ops.cast(out, dtype="float32")
 
 
@@ -122,9 +121,9 @@ def call(self, q, k, v, positional_encoding):
         x = scaled_dot_product_attention(
             q, k, v, dropout_p=self.dropout_p, is_causal=self.is_causal
         )
-
-        x = rearrange(x, "B H L D -> B L (H D)")
-        return x
+        x = ops.transpose(x, (0, 2, 1, 3))
+        b, l, h, d = ops.shape(x)
+        return ops.reshape(x, (b, l, h * d))
 
 
 # TODO: This is probably already implemented in several places, but is needed to ensure numeric equivalence to the original

requirements-common.txt

@@ -19,6 +19,3 @@ sentencepiece
 tensorflow-datasets
 safetensors
 pillow
-# Will be replaced once https://github.com/keras-team/keras/issues/20332
-# is resolved
-einops