[add-acc]Add primitive ast nodes to compiled impl

themetaschemer · Sep 14, 2024 · 7d16ce3 · 7d16ce3
1 parent 8b1827c
commit 7d16ce3
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/accelerated-tensors/autodiff/B-prims.rkt b/accelerated-tensors/autodiff/B-prims.rkt
@@ -178,7 +178,7 @@
   (λ (ρ-fn ra rb shape-fn)
     (let* ((in-shape-a (flat-shape ra))
            (in-size-a (size-of in-shape-a))
-           (in-shape-b (flat-shape ra))
+           (in-shape-b (flat-shape rb))
            (in-size-b (size-of in-shape-b))
            (out-shape (shape-fn in-shape-a in-shape-b))
            (out-size (size-of out-shape)))
@@ -202,7 +202,7 @@
   (λ (∇-fn ra rb z shape-fn)
     (let* ((in-shape-a (flat-shape ra))
            (in-size-a (size-of in-shape-a))
-           (in-shape-b (flat-shape ra))
+           (in-shape-b (flat-shape rb))
            (in-size-b (size-of in-shape-b))
            (out-shape (shape-fn in-shape-a in-shape-b))
            (out-size (size-of out-shape)))