debug for remat

Author: pass-lin

keras_hub/src/models/rwkv7/rwkv7_layer.py

@@ -112,7 +112,7 @@ def call(self, x, last_cache_x=None, train_mode=True):
             Mixed output tensor.
         """
         xx = self.time_shift(x, last_cache_x) - x
-        if last_cache_x is not None or not train_mode:
+        if not train_mode and last_cache_x is not None:
             last_cache_x = x[:, -1:]
         k = x + xx * self.x_k
         k = ops.relu(self.key(k)) ** 2
@@ -390,7 +390,7 @@ def call(
         B, T, C = ops.shape(x)
         H = self.n_head
         xx = self.time_shift(x, last_cache_x) - x
-        if last_cache_x is not None or not train_mode:
+        if not train_mode and last_cache_x is not None:
             last_cache_x = x[:, -1:]
         if padding_mask is not None:
             xx *= padding_mask
@@ -472,13 +472,14 @@ def reshape_and_cast(x, new_shape, dtype="float32"):
         x = x + ops.reshape(rwkv, (B, T, C))
         x = self.output_layer(x * g)
         if train_mode:
-            return x, v_first
+            return x, v_first, finnal_state
         return x, v_first, last_cache_x, finnal_state
 
     def compute_output_shape(self, input_shape):
         output_shapes = [
             [None, None, self.hidden_size],
             [None, None, self.hidden_size],
+            [None, self.n_head, self.head_size, self.head_size],
         ]
         return output_shapes
 
@@ -621,7 +622,7 @@ def call(
         if self.use_initial_norm:
             x = self.ln0(x)
         if train_mode:
-            xx, v_first = self.att(
+            xx, v_first, state = self.att(
                 self.ln1(x),
                 v_first=v_first,
                 padding_mask=padding_mask,