chickadee » nanograd » set-backward-fn!

Makes tensor the result of a differentiable operation. fn is a procedure of no arguments that reads (tensor-grad tensor) and adds the gradient of each input with add-to-grad!. inputs is the list of input tensors. This is how new operations are defined:

;; y = 2x
(define (double x)
  (let ((y (make-base-tensor (f32vector-map (lambda (v) (* 2.0 v)) (tensor-data x))
                             (tensor-shape x) 'f32 #t)))
    (set-backward-fn! y
                      (lambda ()
                        (add-to-grad! x (f32vector-map (lambda (g) (* 2.0 g))
                                                       (tensor-grad y))))
                      (list x))
    y))

f32vector-map in this example stands for any element-wise map over an f32vector; it is not part of nanograd.