src/autograd/sum_to_size_LAWS.bend

bendtorch-core@0.1.0.0 · 45 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/elementwise_totality.bend as Totality
5import bendtorch-core@0.1.0.0/src/autograd/reduce.bend as Reduce
6import bendtorch-core@0.1.0.0/src/autograd/reduce_LAWS.bend as ReduceLaws
7import bendtorch-core@0.1.0.0/src/autograd/sum_to_size.bend as SumToSize
8
9def Identity(-input: Core.Dynamic, +target: Shape.Shape, pair: Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>) -> Type:
10  (retained, result) = pair
11  match result:
12    case Fail{error}:
13      Empty
14    case Done{output}:
15      {Totality.dimensions(output) == Shape.dimensions(target) : List<&2, U32>} &
16        (@+index: U32 ->
17          {Core.checked_observe_dynamic(output, index) == Core.checked_observe_dynamic(input, index) : Maybe<&2, F32>} &
18            {Core.observe_dynamic(output, index) == Core.observe_dynamic(input, index) : Maybe<&2, F32>})
19
20type IdentityModel<-input: Core.Dynamic, -target: Shape.Shape,
21  -pair: Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>> is Data:
22  IdentityModel{-identity: Identity(input, target, pair)}
23
24# The contribution owner comes back with the same dimensions and observations.
25law sum_to_size_retains_input:
26  for input: Core.Dynamic
27  for +target: Shape.Shape
28  ReduceLaws.Retained<input, SumToSize.sum_to_size(input, target)>
29
30# A target with the contribution's dimensions always succeeds, and the output
31# has the target's dimensions and the contribution's exact words at every
32# position: -0.0 stays -0.0, as in PyTorch's autograd.
33law sum_to_size_identity:
34  for input: Core.Dynamic
35  for +target: Shape.Shape
36  for same: {Totality.dimensions(input) == Shape.dimensions(target) : List<&2, U32>}
37  IdentityModel<input, target, SumToSize.sum_to_size(input, target)>
38
39# A target with other dimensions gives exactly the ordered reduce_to result,
40# so every reduce_to law applies to it.
41law sum_to_size_reduces:
42  for input: Core.Dynamic
43  for +target: Shape.Shape
44  for differ: {Totality.dimensions(input) != Shape.dimensions(target) : List<&2, U32>}
45  {SumToSize.sum_to_size(input, target) == Reduce.reduce_to(input, target) : Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>}