def vjp_gather

bendtorch-core@0.1.0.0 · src/autograd/tensor_backward.bend:666

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
import bendtorch-core@0.1.0.0/src/tensor/gather.bend as Gather

666# d(an index_select or a gather of a)/da: g, of the output's shape, added
667# into +0.0 words of a's shape at the entries (PyTorch's index_add and
668# scatter_add): each input cell gets the sum of the gradient words of the
669# output cells that read it, in output order.
670def vjp_gather(+a: Nat, +input: Shape.Shape, +plan: Gather.Plan, +entries: List<&2, U32>, g: Core.Dynamic) -> Pass:
671  gather_sent(a, input, Gather.scatter(g, input, plan, entries))