def vjp_mse

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

import bendtorch-core@0.1.0.0/src/autograd/tensor_backward.bend as TensorBackward
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

481# d mse_loss(a, b)/da: (2 / N * d) * g, and d/db: (2 / N * e) * g, PyTorch's
482# mse_loss_backward, with the differences e = right - left and d = left -
483# right saved at record time over the broadcast shape of N cells; each is
484# reduced to its operand's shape. The steps are the Mul VJP's, with
485# MseGrad for Mul.
486def vjp_mse(+a: Nat, +b: Nat, +left: Shape.Shape, +right: Shape.Shape, saved: Saved, g: Core.Dynamic) -> Pass:
487  match saved:
488    case Unsaved{}:
489      Pass{Unsaved{}, g, Fail{TensorError{Core.StorageLengthMismatch{}}}}
490    case Operands{e, d}:
491      mse_shaped(a, b, left, right, e, g, Core.shape(d))

Uses (10)

Used by (1)