Fix dead code elimination for Expr::Slice

Author: samestep

crates/wasm/src/lib.rs

@@ -139,7 +139,10 @@ impl<'a, O: Eq + Hash, T: Refs<'a, Opaque = O>> Topsort<'a, O, T> {
                     vars.follow(index, instr.var);
                 }
                 &Expr::Member { tuple, .. } => vars.follow(tuple, instr.var),
-                &Expr::Slice { .. } => vars.live(instr.var),
+                &Expr::Slice { index, .. } => {
+                    vars.live(instr.var);
+                    vars.follow(index, instr.var);
+                }
                 &Expr::Field { .. } => vars.live(instr.var),
                 &Expr::Unary { arg, .. } => vars.follow(arg, instr.var),
                 &Expr::Binary { left, right, .. } => {

packages/core/src/index.test.ts

@@ -934,4 +934,12 @@ describe("valid", () => {
     const h = await compile(g);
     expect(h()).toEqual([[0]]);
   });
+
+  test("compile gradient with dynamic index", async () => {
+    const T = struct({ v: Vec(1, Real), i: 1 });
+    const f = fn([T], Real, ({ v, i }) => v[i]);
+    const g = fn([T], T, (x) => vjp(f)(x).grad(1));
+    const h = await compile(g);
+    expect(h({ v: [2], i: 0 })).toEqual({ v: [1], i: 0 });
+  });
 });