Repository: taichi-dev/taichi (language: python)

Issue:
[ir] Support GEP for in autodiff for reading vector / matrix in args buffer
Currently vector / matrix are compiled as structs in the arguments and thus not supported by autodiff (because lacking implementation for GEP)



Discussion / hints:
@bobcao3 are you asking for differentiable vector/matrix support? Or can you give an example of the code that you'd like to work?  
We can revert https://github.com/taichi-dev/taichi/commit/5cd1407b0057019f2e48f5569fec30a8dc540213 if it is not easy to support.

Fix patch:
diff --git a/taichi/transforms/auto_diff.cpp b/taichi/transforms/auto_diff.cpp
index 7957553525613..67a4b08eb29e3 100644
--- a/taichi/transforms/auto_diff.cpp
+++ b/taichi/transforms/auto_diff.cpp
@@ -1046,6 +1046,10 @@ class ADTransform : public IRVisitor {
     // do nothing.
   }
 
+  void visit(GetElementStmt *stmt) override {
+    // do nothing
+  }
+
   void visit(LoopIndexStmt *stmt) override {
     // do nothing.
   }


Test patch:
diff --git a/tests/python/test_ad_ndarray.py b/tests/python/test_ad_ndarray.py
index e10fe49aea213..ed016ca7d793e 100644
--- a/tests/python/test_ad_ndarray.py
+++ b/tests/python/test_ad_ndarray.py
@@ -1435,3 +1435,31 @@ def test(x: ti.types.ndarray(needs_grad=False), y: ti.types.ndarray()):
     test.grad(x, y)
     for i in range(N):
         assert x.grad[i] == 0.0
+
+
+@test_utils.test(arch=archs_support_ndarray_ad)
+def test_ad_vector_arg():
+    N = 10
+
+    @ti.kernel
+    def compute_sum(a: ti.types.ndarray(), p: ti.types.ndarray(), z: ti.math.vec2):
+        for i in p:
+            p[i] = a[i] * z[0]
+
+    a = ti.ndarray(ti.math.vec2, shape=N, needs_grad=True)
+    p = ti.ndarray(ti.math.vec2, shape=N, needs_grad=True)
+    z = ti.math.vec2([2.0, 3.0])
+    for i in range(N):
+        a[i] = [3, 3]
+
+    compute_sum(a, p, z)
+
+    for i in range(N):
+        assert p[i] == [a[i] * 2, a[i] * 2]
+        p.grad[i] = [1, 1]
+
+    compute_sum.grad(a, p, z)
+
+    for i in range(N):
+        for j in range(2):
+            assert a.grad[i][j] == 2