[autodiff] Add ternary operators for forward mode (#5405)

* [autodiff] Add ternary operators for forward mode

* add ternary test case for forward mode

taichi/transforms/auto_diff.cpp

@@ -1200,6 +1200,15 @@ class MakeDual : public ADTransform {
     }
   }
 
+  void visit(TernaryOpStmt *stmt) override {
+    TI_ASSERT(stmt->op_type == TernaryOpType::select);
+    auto zero = insert<ConstStmt>(TypedConstant(stmt->ret_type));
+    accumulate(stmt, insert<TernaryOpStmt>(TernaryOpType::select, stmt->op1,
+                                           load(dual(stmt->op2)), zero));
+    accumulate(stmt, insert<TernaryOpStmt>(TernaryOpType::select, stmt->op1,
+                                           zero, load(dual(stmt->op3))));
+  }
+
   void visit(IfStmt *if_stmt) override {
     if (if_stmt->true_statements) {
       std::vector<Stmt *> true_statements;

tests/python/test_ad_basics.py

@@ -288,6 +288,44 @@ def func():
         assert y.grad[i] == (not i % 2) * 1.0
 
 
+@test_utils.test()
+def test_select_fwd():
+    N = 5
+    loss = ti.field(ti.f32, shape=N)
+    x = ti.field(ti.f32, shape=N)
+    y = ti.field(ti.f32, shape=N)
+    ti.root.lazy_dual()
+
+    for i in range(N):
+        x[i] = i
+        y[i] = -i
+
+    @ti.kernel
+    def func():
+        for i in range(N):
+            loss[i] = ti.select(i % 2, x[i], y[i])
+
+    with ti.ad.FwdMode(loss=loss, parameters=x, seed=[1.0 for _ in range(N)]):
+        func()
+
+    for i in range(N):
+        if i % 2:
+            assert loss[i] == i
+        else:
+            assert loss[i] == -i
+        assert loss.dual[i] == i % 2 * 1.0
+
+    with ti.ad.FwdMode(loss=loss, parameters=y, seed=[1.0 for _ in range(N)]):
+        func()
+
+    for i in range(N):
+        if i % 2:
+            assert loss[i] == i
+        else:
+            assert loss[i] == -i
+        assert loss.dual[i] == (not i % 2) * 1.0
+
+
 @test_utils.test()
 def test_obey_kernel_simplicity():
     x = ti.field(ti.f32)