Properly implement differential witness of intermediate context type. (#2699)

* Properly implement differential witness of intermediate context type.

* Modify test to include a loop.

---------

Co-authored-by: Yong He <yhe@nvidia.com>

1 files changed, 8 insertions, 2 deletions

diff --git a/tests/autodiff/high-order-backward-diff-3.slang b/tests/autodiff/high-order-backward-diff-3.slang
index eb3866b96..100a9a1e0 100644
--- a/tests/autodiff/high-order-backward-diff-3.slang
+++ b/tests/autodiff/high-order-backward-diff-3.slang
@@ -14,14 +14,20 @@ struct A : IDifferentiable
 [BackwardDifferentiable]
 float f(A x)
 {
-    return x.x * x.x;
+    A rs;
+    rs.x = 1.0;
+    for (int i = 0; i < 2; i++)
+        rs.x = rs.x * x.x;
+    return rs.x;
 }
 
 [BackwardDifferentiable]
 float outerF(A x)
 {
     A nx;
-    nx.x = x.x * x.x;
+    nx.x = 1.0;
+    for (int i = 0; i < 2; i++)
+        nx.x = nx.x * x.x;
     nx.nx = 2;//x.nx;
     return f(nx);
 }