def vjp

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

import Base
import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core

673# The vector-Jacobian product of one node with gradient g: the contribution
674# each operand receives, in operand order, already reduced to its shape. A
675# constant sends nothing, and Frozen{id, node} sends node's contributions
676# without the ones to id.
677def vjp(node: Node, saved: Saved, g: Core.Dynamic) -> Pass:
678  match node:
679    case Leaf{}:
680      Pass{saved, g, Done{Nil{}}}
681    case Add{+a, +b, +left, +right}:
682      vjp_add(a, b, left, right, g)
683    case Sub{+a, +b, +left, +right}:
684      vjp_sub(a, b, left, right, g)
685    case Mul{+a, +b, +left, +right}:
686      vjp_mul(a, b, left, right, saved, g)
687    case Sum{+a, +input}:
688      vjp_sum(a, input, g)
689    case SumDims{+a, +input, +kept}:
690      vjp_sum_dims(a, input, kept, g)
691    case Matmul{+a, +b, +left, +right}:
692      vjp_matmul(a, b, left, right, saved, g)
693    case MatmulNT{+a, +b, +left, +right}:
694      vjp_matmul_nt(a, b, left, right, saved, g)
695    case Relu{+a, +input}:
696      vjp_relu(a, input, saved, g)
697    case MseLoss{+a, +b, +left, +right}:
698      vjp_mse(a, b, left, right, saved, g)
699    case CrossEntropy{+a, +input}:
700      vjp_cross_entropy(a, input, saved, g)
701    case BatchedMatmul{+a, +b, +left, +right}:
702      vjp_batched(a, b, left, right, saved, g)
703    case Slice{+a, +input, +kept, +plan}:
704      vjp_slice(a, input, kept, plan, g)
705    case Gathered{+a, +input, +plan, +entries}:
706      vjp_gather(a, input, plan, entries, g)
707    case Constant{}:
708      Pass{saved, g, Done{Nil{}}}
709    case Frozen{+id, inner}:
710      match inner:
711        case MatmulNT{+a, +b, +left, +right}:
712          nt_frozen(id, a, b, left, right, saved, g, Bool.and(Nat.is_eq(id, a), Bool.not(Nat.is_eq(id, b))))
713        case other:
714          unsent(id, vjp(other, saved, g))