chickadee » nanograd » reduce-tensor

Generic reduction operation that maintains gradient flow. The reducer function is applied to each element in the forward pass. An optional compute-gradient function specifies how gradients are distributed in the backward pass.

tensor

input tensor to reduce

reducer

function (element accumulator) -> new-accumulator

compute-gradient

optional function (grad-out index value all-values) -> grad-in

If not provided, assumes uniform distribution (like 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 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)))))