chickadee » nanograd » reduce-tensor

Generic reduction that keeps gradient flow. The reducer procedure is applied to each element in the forward pass. The optional compute-gradient procedure says how the gradient is distributed in the backward pass.

tensor

input tensor to reduce

reducer

procedure (element accumulator) -> new-accumulator

compute-gradient

optional procedure (grad-out index value all-values) -> grad-in. Without it, the gradient is distributed uniformly, as for a sum.

Returns a scalar tensor with the reduced value.

;; Sum all elements (uniform gradient distribution)
(define total (reduce-tensor x +))

;; Product of all elements (gradient uses the product rule)
(define prod (reduce-tensor x *
  compute-gradient: (lambda (grad-out idx val all-values)
                     ;; d(prod)/dx_i = prod / x_i
                     (let ((prod (fold * 1.0 all-values)))
                       (if (> val 0.0)
                           (* grad-out (/ prod val))
                           0.0)))))