Fix TBRAnalyzer visiting the true branch twice. It ignored the false branch completely, causing silent wrong derivatives. Fixed the typo, added a regression test, and updated the assignments test.

Author: leetcodez

lib/Differentiator/TBRAnalyzer.cpp

@@ -247,7 +247,7 @@ bool TBRAnalyzer::TraverseConditionalOperator(clang::ConditionalOperator* CO) {
 
   auto thenBranch = std::move(m_BlockData[m_CurBlockID]);
   m_BlockData[m_CurBlockID] = std::move(elseBranch);
-  TraverseStmt(CO->getTrueExpr());
+  TraverseStmt(CO->getFalseExpr());
 
   merge(m_BlockData[m_CurBlockID].get(), thenBranch.get());
   return false;

test/Gradient/Assignments.C

@@ -416,26 +416,30 @@ double f12(double x, double y) {
 
 //CHECK:   void f12_grad(double x, double y, double *_d_x, double *_d_y) {
 //CHECK-NEXT:       double _t0;
+//CHECK-NEXT:       double _t1;
 //CHECK-NEXT:       double _d_t = 0.;
 //CHECK-NEXT:       double t;
 //CHECK-NEXT:       bool _cond0 = x > y;
 //CHECK-NEXT:       if (_cond0)
 //CHECK-NEXT:           _t0 = t;
-//CHECK-NEXT:       double *_t1 = &(_cond0 ? (t = x) : (t = y));
-//CHECK-NEXT:       double _t2 = *_t1;
-//CHECK-NEXT:       *_t1 *= y;
+//CHECK-NEXT:       else
+//CHECK-NEXT:           _t1 = t;
+//CHECK-NEXT:       double *_t2 = &(_cond0 ? (t = x) : (t = y));
+//CHECK-NEXT:       double _t3 = *_t2;
+//CHECK-NEXT:       *_t2 *= y;
 //CHECK-NEXT:       _d_t += 1;
 //CHECK-NEXT:       {
-//CHECK-NEXT:           *_t1 = _t2;
+//CHECK-NEXT:           *_t2 = _t3;
 //CHECK-NEXT:           double _r_d0 = (_cond0 ? _d_t : _d_t);
 //CHECK-NEXT:           (_cond0 ? _d_t : _d_t) = 0.;
 //CHECK-NEXT:           (_cond0 ? _d_t : _d_t) += _r_d0 * y;
-//CHECK-NEXT:           *_d_y += *_t1 * _r_d0;
+//CHECK-NEXT:           *_d_y += *_t2 * _r_d0;
 //CHECK-NEXT:           if (_cond0) {
 //CHECK-NEXT:               t = _t0;
 //CHECK-NEXT:               *_d_x += _d_t;
 //CHECK-NEXT:               _d_t = 0.;
 //CHECK-NEXT:           } else {
+//CHECK-NEXT:               t = _t1;
 //CHECK-NEXT:               *_d_y += _d_t;
 //CHECK-NEXT:               _d_t = 0.;
 //CHECK-NEXT:           }

test/Gradient/TBRTernary.C

@@ -0,0 +1,28 @@
+// RUN: %cladclang %s -I%S/../../include -oReverseMode.out 2>&1 | %filecheck %s
+// RUN: ./ReverseMode.out | %filecheck_exec %s
+
+#include "clad/Differentiator/Differentiator.h"
+#include <iostream>
+
+double f(bool cond, double a, double b) {
+    double x = cond ? (a * a) : (b * b);
+    a = 0;
+    b = 0;
+    return x;
+}
+
+int main() {
+    auto df = clad::gradient(f, "a,b");
+    double da = 0.0, db = 0.0;
+    
+    df.execute(false, 2.0, 3.0, &da, &db);
+    std::cout << "da: " << da << ", db: " << db << std::endl;
+    // CHECK-EXEC: da: 0, db: 6
+    
+    da = 0.0; db = 0.0;
+    df.execute(true, 2.0, 3.0, &da, &db);
+    std::cout << "da: " << da << ", db: " << db << std::endl;
+    // CHECK-EXEC: da: 4, db: 0
+    
+    return 0;
+}