differentiable.u

-- Autodiff using algebraic effects
-- Simple reverse mode with a gradient tape

use .base

ability Differentiable where
    cos : Float -> Float
    sin : Float -> Float
    exp : Float -> Float

Tape : List (Float -> Float)
Tape = []

pullback : List (Float -> Float) -> Request Differentiable a -> List (Float -> Float)
pullback l = cases
    { Differentiable.cos x -> k} -> handle
                                        k (Float.cos x) 
                                    with pullback (l List.++ [y -> -1.0 * y * Float.sin x])

    { Differentiable.exp x -> k} -> handle
                                        k (Float.exp x) 
                                    with pullback (l List.++ [y -> y * Float.exp x])

    { Differentiable.sin x -> k} -> handle 
                                        k (Float.sin x) 
                                    with pullback (l List.++ [y -> y * Float.cos x])
                                        
    { x } -> l List.++ [y -> y * 1.0]

gradient_tape : (Float -> Float) -> Float -> List (Float -> Float)
gradient_tape v x =
    handle v x with pullback []

grad : (Float -> Float) -> Float -> Float
grad f x = (foldl (x -> y -> y x) 1.0) (gradient_tape f x)


-- Test!
p : Float -> Float
p x = Differentiable.exp (Differentiable.cos x)

> grad p 5.0