src/autograd/broadcast_LAWS.bend

bendtorch-core@0.1.0.0 · 102 lines · 15 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.bend as Elementwise
5import bendtorch-core@0.1.0.0/src/tensor/elementwise_totality.bend as Totality
6import bendtorch-core@0.1.0.0/src/tensor/projection_LAWS.bend as Projection
7import bendtorch-core@0.1.0.0/src/autograd/broadcast.bend as Broadcast
8
9def owner(pair: Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic>) -> Core.Dynamic:
10  (input, result) = pair
11  input
12
13def succeeded(pair: Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic>) -> Bool:
14  (input, result) = pair
15  Result.is_done(&2, &1, Broadcast.Error, Core.Dynamic, result)
16
17# The input value that the specification assigns to one output position: the
18# input offset of its right-aligned broadcast projection.
19def source_value(input: Core.Dynamic, offsets: Maybe<&1, U32 & U32>) -> Maybe<&2, F32>:
20  match offsets:
21    case None{}:
22      None{}
23    case Some{pair}:
24      (input_offset, target_offset) = pair
25      Core.checked_observe_dynamic(input, input_offset)
26
27def planned_source(input: Core.Dynamic, +index: U32, plan: Result<&2, &2, Shape.Error, Shape.BroadcastPlan>) -> Maybe<&2, F32>:
28  match plan:
29    case Fail{error}:
30      None{}
31    case Done{+plan}:
32      source_value(input, Projection.offset_spec(plan, U32.to_nat(index)))
33
34def MappingWith(-input: Core.Dynamic, +target: Shape.Shape, plan: Result<&2, &2, Shape.Error, Shape.BroadcastPlan>,
35  pair: Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic>) -> Type:
36  (retained, result) = pair
37  match result:
38    case Fail{error}:
39      Unit
40    case Done{output}:
41      {Totality.dimensions(output) == Shape.dimensions(target) : List<&2, U32>} &
42        (@+index: U32 -> @inside: Shape.LT(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(target))) ->
43          {Core.observe_dynamic(output, index) == planned_source(input, index, plan) : Maybe<&2, F32>})
44
45def Mapping(input: Core.Dynamic, +target: Shape.Shape, pair: Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic>) -> Type:
46  MappingWith(input, target, Shape.broadcast_checked(Totality.dynamic_shape(input), target), pair)
47
48type Retained<-input: Core.Dynamic, -pair: Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic>> is Data:
49  Retained{-kept: Totality.SameSource(input, owner(pair))}
50
51type MappingModel<-input: Core.Dynamic, -target: Shape.Shape, -pair: Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic>> is Data:
52  MappingModel{-mapping: Mapping(input, target, pair)}
53
54type Success<-pair: Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic>> is Data:
55  Success{-successful: {succeeded(pair) == True{} : Bool}}
56
57law keep_left_scalar:
58  for +left: F32
59  for +right: F32
60  {Elementwise.scalar(Elementwise.KeepLeft{}, left, right) == left : F32}
61
62# The input owner comes back with the same dimensions and observations.
63law broadcast_retains_input:
64  for input: Core.Dynamic
65  for +target: Shape.Shape
66  Retained<input, Broadcast.broadcast_to(input, target)>
67
68# On success the output has the target's dimensions, and each in-bounds output
69# position observes the input at the offset that the independent broadcast
70# projection specification assigns to it.
71law broadcast_mapping:
72  for input: Core.Dynamic
73  for +target: Shape.Shape
74  MappingModel<input, target, Broadcast.broadcast_to(input, target)>
75
76# A request whose right-aligned broadcast with the input shape is the target
77# itself always succeeds.
78law broadcast_valid_success:
79  for input: Core.Dynamic
80  for +target: Shape.Shape
81  for +alignment: Shape.Broadcast<Shape.extents(Totality.dynamic_shape(input)), Shape.extents(target), Shape.extents(target)>
82  Success<Broadcast.broadcast_to(input, target)>
83
84# When the broadcast planner rejects the input and target shapes, broadcast_to
85# fails with exactly that planner error and returns the input owner. For example [2] to
86# [0] fails with IncompatibleBroadcast instead of returning an empty tensor.
87law broadcast_shape_error:
88  for input: Core.Dynamic
89  for +target: Shape.Shape
90  for error: Shape.Error
91  for rejected: {Shape.broadcast_checked(Totality.dynamic_shape(input), target) == Fail{error} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}
92  {Broadcast.broadcast_to(input, target) == (input, Fail{Broadcast.TensorError{Core.ShapeError{error}}}) : Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic>}
93
94# When the planned broadcast is not the target itself (for example [3] to
95# [1]), broadcast_to fails with NotBroadcastable and returns the input owner.
96law broadcast_not_target:
97  for input: Core.Dynamic
98  for +target: Shape.Shape
99  for +plan: Shape.BroadcastPlan
100  for planned: {Shape.broadcast_checked(Totality.dynamic_shape(input), target) == Done{plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}
101  for differ: {Shape.dimensions(Shape.output_shape(plan)) != Shape.dimensions(target) : List<&2, U32>}
102  {Broadcast.broadcast_to(input, target) == (input, Fail{Broadcast.NotBroadcastable{}}) : Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic>}