src/tensor/contracts/binary_LAWS.bend

bendtorch-core@0.1.0.0 · 74 lines · 3 definitions

1import Base
2import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
3import bendtorch-core@0.1.0.0/src/storage/LAWS.bend as StorageLaws
4import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
5import bendtorch-core@0.1.0.0/src/tensor/lifecycle.bend as Lifecycle
6import bendtorch-core@0.1.0.0/src/tensor/elementwise.bend as Elementwise
7import bendtorch-core@0.1.0.0/src/tensor/contracts/replayed.bend as Replayed
8import bendtorch-core@0.1.0.0/src/tensor/contracts/pointwise.bend as Pointwise
9import bendtorch-core@0.1.0.0/src/tensor/contracts/broadcast.bend as Broadcast
10import bendtorch-core@0.1.0.0/src/tensor/contracts/LAWS.bend as Contracts
11
12# The words of a pointwise product in index order: word k is
13# scalar(operation, left[left_map(plan, k)], right[right_map(plan, k)]), the
14# value the pointwise family's broadcast instance writes at k.
15def cells(+operation: Elementwise.Operation, +plan: Shape.BroadcastPlan,
16  +left: Shape.Shape, +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>,
17  +right: Shape.Shape, +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>, +count: Nat) -> List<&2, F32>:
18  match count:
19    case 0n:
20      Nil{}
21    case 1n+previous:
22      List.append(&2, F32, cells(operation, plan, left, left_fill, left_ops, right, right_fill, right_ops, previous),
23        [Pointwise.cell(~Elementwise.Operation, ~Elementwise.scalar, ~Shape.BroadcastPlan, ~Broadcast.left_map,
24          ~Broadcast.right_map, operation, plan, left, left_fill, left_ops, right, right_fill, right_ops, previous)])
25
26# What the transport form of binary_replayed asks of P: P holds for both
27# inputs, replayed with any histories that keep their words, next to the
28# output replayed over +0.0 with any history whose words are `cells`.
29def BinaryWith(+operation: Elementwise.Operation, +plan: Shape.BroadcastPlan,
30  +left: Shape.Shape, +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>,
31  +right: Shape.Shape, +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>,
32  -P: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic> -> Type) -> Type:
33  @+left_kept: List<&2, StorageLaws.Op> ->
34  @left_words: {Contracts.words(left, left_fill, left_kept, Shape.logical_numel(Shape.dimensions(left))) ==
35    Contracts.words(left, left_fill, left_ops, Shape.logical_numel(Shape.dimensions(left))) : List<&2, F32>} ->
36  @-left_fits: {Lifecycle.logical_length(Replayed.storage(left, left_fill, left_kept)) == Shape.logical_numel(Shape.dimensions(left)) : Nat} ->
37  @+right_kept: List<&2, StorageLaws.Op> ->
38  @right_words: {Contracts.words(right, right_fill, right_kept, Shape.logical_numel(Shape.dimensions(right))) ==
39    Contracts.words(right, right_fill, right_ops, Shape.logical_numel(Shape.dimensions(right))) : List<&2, F32>} ->
40  @-right_fits: {Lifecycle.logical_length(Replayed.storage(right, right_fill, right_kept)) == Shape.logical_numel(Shape.dimensions(right)) : Nat} ->
41  @+output_ops: List<&2, StorageLaws.Op> ->
42  @values: {Contracts.words(Shape.output_shape(plan), 0.0, output_ops, Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) ==
43    cells(operation, plan, left, left_fill, left_ops, right, right_fill, right_ops,
44      Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) : List<&2, F32>} ->
45  @-output_fits: {Lifecycle.logical_length(Replayed.storage(Shape.output_shape(plan), 0.0, output_ops)) ==
46    Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan))) : Nat} ->
47    P(((Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, left_kept), left_fits}},
48      Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, right_kept), right_fits}}),
49      Done{Core.Packed{Shape.output_shape(plan), Core.TTensor{Replayed.storage(Shape.output_shape(plan), 0.0, output_ops), output_fits}}}))
50
51# The parallel elementwise build exports its replayed result. For two
52# replayed inputs whose shapes broadcast, for every P, if P holds for every
53# result that hands back both inputs with their words and gives an output of
54# the plan's shape whose words are `cells`, then P holds for
55# `Elementwise.binary`. The output word at k is the value the pointwise family
56# writes at k, whatever split of the output into parallel halves the build
57# takes, and the output is a replayed tensor, so a law about replayed tensors
58# applies to it.
59law binary_replayed:
60  for +operation: Elementwise.Operation
61  for +left: Shape.Shape
62  for +right: Shape.Shape
63  for +plan: Shape.BroadcastPlan
64  for planned: {Shape.broadcast_checked(left, right) == Done{plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}
65  for +left_fill: F32
66  for +left_ops: List<&2, StorageLaws.Op>
67  for -left_fits: {Lifecycle.logical_length(Replayed.storage(left, left_fill, left_ops)) == Shape.logical_numel(Shape.dimensions(left)) : Nat}
68  for +right_fill: F32
69  for +right_ops: List<&2, StorageLaws.Op>
70  for -right_fits: {Lifecycle.logical_length(Replayed.storage(right, right_fill, right_ops)) == Shape.logical_numel(Shape.dimensions(right)) : Nat}
71  for -P: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic> -> Type
72  for h: BinaryWith(operation, plan, left, left_fill, left_ops, right, right_fill, right_ops, P)
73  P(Elementwise.binary(operation, Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, left_ops), left_fits}},
74    Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, right_ops), right_fits}}))