Gradients with State effect and slice indices

[srush]

In the CKY implementation in #586 when I try to take a gradient I get the error ->

Can't differentiate wrt type (Ref h' ((tmp77<..) => Float32))
CallStack (from HasCallStack):
error, called at src/lib/Autodiff.hs:402:23 in dex-0.1.0.0-1c1KJ1laoEjKvvidxV1QZ9:Autodiff

I think this is an okay thing to do. I've been able to use State and grad together before. But I think the slice indexing breaks it.

[chenzizhao]

Thanks @apaszke, for reviewing this issue today!
Normal indices give a similar error: Can't differentiate wrt type (Ref h' ((Fin 3) => Float32))

def myfunc2 (mat0: a=>a=>Float) (x:Float) : Float = 
    mat = yieldState mat0 \m.
        i = 0@a
        m!i!i := x
    sum for i. sum for j. mat.i.j

mat0 : (Fin 3)=>(Fin 3)=>Float = zero
myfunc2 mat0 7.0
-- (grad (myfunc2 mat0)) 7.0

[apaszke]

@chenzizhao Thanks for making the broken program smaller! I managed to fix that in #593, but it seems like it's still not enough to make the CKY example run. I'll keep digging...

[apaszke]

I think I drilled down to the core of the CKY issue and it turns out to be something that we knew is not implemented, but weren't checking for it. An improvement that should at least give a reasonable error message is coming in #594. The real fix is a little involved, but I'll think about it.

[axch]

Looks like this issue has since been fixed. The simple repro by @chenzizhao remains functional, and I also hack-ported the CKY example to current Dex syntax, and it runs to completion there too. (I have no idea whether the results are any good, but at least the program executes now.)