Fix

Author: Perseus14

src/maxdiffusion/models/lora_nnx.py

@@ -50,7 +50,9 @@ def __init__(
 
     # Infer dimensions from the base layer
     # nnx.Linear stores weights in 'kernel' usually shaped (in_features, out_features)
-    in_features, out_features = base_layer.kernel.shape
+    k_shape = base_layer.kernel.shape
+    in_features = k_shape[0]
+    out_features = int(jnp.prod(jnp.array(k_shape[1:])))
 
     # 1. Down Projection (A): Random Initialization
     # Projects inputs down to rank 'r'
@@ -73,6 +75,11 @@ def __call__(self, x):
     # Equation: (x @ A @ B) * scaling
     lora_out = (x @ self.A.value @ self.B.value) * self.scaling()
 
+    # If base layer kernel was >2D, its output might be >2D on feature axis.
+    # We need to reshape lora_out to match base_out's trailing dimensions if they don't match.
+    if len(self.base_layer.kernel.shape) > 2:
+      lora_out = lora_out.reshape(x.shape[:-1] + self.base_layer.kernel.shape[1:])
+
     # 3. Sum
     return base_out + lora_out