fixes

Author: MarioSieg

magnetron/magnetron.c

@@ -2873,15 +2873,16 @@ static void mag_collect_topo_iterative(mag_tensor_t* root, mag_tensor_array_t* o
     (*mag_alloc)(visited, 0);
 }
 
-static void mag_tensor_patch_grad(mag_tensor_t* dst, mag_tensor_t* new_grad, const char* stage) {
+static void mag_tensor_patch_grad(const mag_tensor_t* root, mag_tensor_t* dst, mag_tensor_t* new_grad, const char* stage) {
     if (!mag_tensor_is_shape_eq(dst, new_grad)) {
         char shape_dst[MAG_FMT_DIM_BUF_SIZE];
         char shape_grad[MAG_FMT_DIM_BUF_SIZE];
         mag_fmt_dims(&shape_dst, &dst->shape, dst->rank);
         mag_fmt_dims(&shape_grad, &new_grad->shape, new_grad->rank);
         const char* dst_op = mag_op_meta_of(dst->op)->mnemonic;
         const char* grad_op = mag_op_meta_of(new_grad->op)->mnemonic;
-        mag_panic("Shape mismatch: %s (%s) != %s (%s) !%s", shape_dst, dst_op, shape_grad, grad_op, stage);
+        const char* root_op = mag_op_meta_of(root->op)->mnemonic;
+        mag_panic("Shape mismatch: %s (%s) != %s (%s) Stage: %s, Root Op: %s\n", shape_dst, dst_op, shape_grad, grad_op, stage, root_op);
     }
     mag_tensor_fmt_name(new_grad, "%s (grad)", dst->name);
     new_grad->flags = (new_grad->flags | MAG_TFLAG_IS_GRAD) & ~MAG_TFLAG_REQUIRES_GRAD;
@@ -2906,7 +2907,7 @@ void mag_tensor_backward(mag_tensor_t* root) {
         if (!child->grad) {
             mag_tensor_t* grad = mag_tensor_create(child->ctx, child->dtype, child->shape, child->rank, NULL, 0);
             mag_tensor_fill(grad, 1.0f);
-            mag_tensor_patch_grad(child, grad, "init");
+            mag_tensor_patch_grad(child, child, grad, "init");
         }
         if (mag_unlikely(child->op == MAG_OP_NOP)) continue;
         mag_tensor_t* grads[MAG_MAX_OP_PARAMS] = {0};
@@ -2917,15 +2918,15 @@ void mag_tensor_backward(mag_tensor_t* root) {
         for (uint32_t i = 0; i < numin; ++i) {
             mag_tensor_t* input = child->op_inputs[i];
             mag_assert2(input);
-            if (!(input->flags & MAG_TFLAG_REQUIRES_GRAD) || input->op == MAG_OP_NOP) continue;
+            if (!(input->flags & MAG_TFLAG_REQUIRES_GRAD)) continue;
             mag_tensor_t* gri = grads[i];
             mag_assert(gri, "Gradient for op %s, input #%d is not computed", meta->mnemonic, i);
             if (!input->grad) {
-                mag_tensor_patch_grad(input, gri, "patch");
+                mag_tensor_patch_grad(child, input, gri, "patch");
             } else {
                 mag_tensor_t* acc = mag_add_(gri, input->grad);
                 mag_tensor_decref(input->grad);
-                mag_tensor_patch_grad(input, acc, "acc");
+                mag_tensor_patch_grad(child, input, acc, "acc");
                 mag_tensor_decref(gri);
             }
         }