def writes_value

bendtorch-core@0.1.0.0 · src/autograd/reduce_replayed_proof.bend:68

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/storage/LAWS.bend as StorageLaws
import bendtorch-core@0.1.0.0/src/storage/arith.bend as Arith
import bendtorch-core@0.1.0.0/src/tensor/word_conversion.bend as Conversion
import bendtorch-core@0.1.0.0/src/tensor/contracts/replayed.bend as Replayed
import bendtorch-core@0.1.0.0/src/tensor/contracts/pointwise_proof.bend as PointwiseProof

68def writes_value(+plan: Shape.BroadcastPlan, +map: PlanMap(plan), +shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +output: Shape.Shape, +count: Nat, +index: U32,
69  below: {Nat.is_lt(U32.to_nat(index), count) == True{} : Bool},
70  +room: {Nat.is_le(count, Shape.logical_numel(Shape.dimensions(output))) == True{} : Bool}) ->
71  {Replayed.value(output, 0.0, writes(plan, map, shape, fill, ops, count), index) ==
72    cell_word(plan, map, shape, fill, ops, U32.to_nat(index)) : F32}:
73  match count:
74    case 0n:
75      Empty.absurd({Replayed.value(output, 0.0, Nil{}, index) ==
76        cell_word(plan, map, shape, fill, ops, U32.to_nat(index)) : F32},
77        Arith.lt_zero(U32.to_nat(index), below))
78    case 1n+previous:
79      PointwiseProof.write_split(output, writes(plan, map, shape, fill, ops, previous),
80        previous, cell_word(plan, map, shape, fill, ops, previous), index,
81        cell_word(plan, map, shape, fill, ops, U32.to_nat(index)),
82        below, room,
83        at => Equal.cong(Nat, F32, position => cell_word(plan, map, shape, fill, ops, position),
84          previous, U32.to_nat(index), Equal.sym(Nat, U32.to_nat(index), previous, at)),
85        earlier => writes_value(plan, map, shape, fill, ops, output, previous, index, earlier,
86          Conversion.le_predecessor(previous, Shape.logical_numel(Shape.dimensions(output)), room)),
87        Shape.bool_case(Nat.is_eq(U32.to_nat(index), previous)))