AD in interpreter produces NaN values where it should not

[FluxusMagna]

This code results in incorrect output from the interpreter:

def identity_mat n = tabulate_2d n n (\i j -> f32.bool (i==j))

def Jacobi [n] (f: [n]f32 -> [n]f32) (x:[n]f32) : [n][n]f32
    = map (\i -> jvp f x i) (identity_mat n) |> transpose

def Hessian [n] (f: [n]f32 -> f32) (x:[n]f32) : [n][n]f32
    = Jacobi (\x -> vjp f x 1) x

def Hessian_test (x:[3]f32) = Hessian (\x -> x[1]**2+x[2]**2) x

To test this code we can run

> Hessian_test [0,0,0]

which should result in (and does in the compiled code)

[[0.0f32, 0.0f32, 0.0f32],
[0.0f32, 2.0f32, 0.0f32],
[0.0f32, 0.0f32, 2.0f32]]

But in the interpreter this produces.

[[0.0, f32.nan, f32.nan],
 [0.0, f32.nan, f32.nan],
 [0.0, f32.nan, f32.nan]]

I'm not sure why it fails like this, but I've observed similar issues with other nested AD, which eventually led me to this test.

edit: transpose jacobi matrix for correctness

[athas]

Looks like a bug in **; it works if the squares are written out as x[1]*x[1].

[athas]

This one is a little tricky actually. The problem is that the partial derivative of x**y wrt. y is (x**y)*log(x), which is not defined for x==0. The only reason this produces the "right result" with the compiler is due to a simplification rule that turns 0*x into 0 (which is not correct for floating point, and so I will remove this rule).

[FluxusMagna]

Thanks, although this didn't directly solve my original issue, it led me to the source of problem, which was then pretty easy to solve.

[athas]

Good! In retrospect it is pretty obvious that we need to do something special for the partial derivative of this operator. Maybe we need to define a conditional PrimExp.