~/jendhub

src/autograd/sum_to_size.bend

bendtorch-core@0.1.0.0 · 40 lines · 5 definitions

1import Base
2import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
3import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
4import bendtorch-core@0.1.0.0/src/tensor/layout.bend as Layout
5import bendtorch-core@0.1.0.0/src/autograd/reduce.bend as Reduce
6
7# The PyTorch-facing reducer for a broadcast contribution: the gradient
8# PyTorch's autograd engine gives an operand of a broadcast operation. When the
9# target has the contribution's dimensions nothing is reduced and the output is
10# a clone of the contribution, so every word, including -0.0 and NaN payloads,
11# is kept. Otherwise the result is exactly Reduce.reduce_to's ordered fold.
12
13def cloned(+shape: Shape.Shape, pair: Core.Tensor<shape> & Core.Tensor<shape>) ->
14  Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>:
15  (kept, copy) = pair
16  (Core.Packed{shape, kept}, Done{Core.Packed{shape, copy}})
17
18def decided(+shape: Shape.Shape, tensor: Core.Tensor<shape>, +target: Shape.Shape, same: Bool) ->
19  Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>:
20  match same:
21    case False{}:
22      Reduce.reduce_to(Core.Packed{shape, tensor}, target)
23    case True{}:
24      cloned(shape, Core.clone(shape, tensor))
25
26# The contribution owner comes back next to the result, as with reduce_to.
27def sum_to_size(input: Core.Dynamic, +target: Shape.Shape) -> Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>:
28  Core.Packed{+shape, tensor} = input
29  decided(shape, tensor, target, Layout.words_same(Shape.dimensions(shape), Shape.dimensions(target)))
30
31def sum_to_size_dimensions_result(input: Core.Dynamic, -dimensions: List<&2, U32>, result: Shape.ConstructResult<dimensions>) ->
32  Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>:
33  match result:
34    case Shape.SError{error}:
35      (input, Fail{Reduce.TensorError{Core.ShapeError{error}}})
36    case Shape.SBuilt{+target, preserved}:
37      sum_to_size(input, target)
38
39def sum_to_size_dimensions(input: Core.Dynamic, +dimensions: List<&2, U32>) -> Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>:
40  sum_to_size_dimensions_result(input, dimensions, Shape.construct(dimensions))