proofs/math_gelu_lanes_certificate.bend
stelliferous@0.0.3.0 · 44 lines · 0 definitions
1# Shape preservation reuses the traversal's gather/scatter witnesses.
2import Base
3import stelliferous@0.0.3.0/tensor_access.bend as Access
4import stelliferous@0.0.3.0/math_activation.bend as Act
5import stelliferous@0.0.3.0/math_gelu_lanes.bend as Gelu
6import stelliferous@0.0.3.0/kernel_tile.bend as Tile
7import stelliferous@0.0.3.0/traversal_lanes.bend as Lanes
8import stelliferous@0.0.3.0/traversal_cells.bend as Cells
9import stelliferous@0.0.3.0/proofs/traversal_lanes_witness.bend as LaneWitness
10import stelliferous@0.0.3.0/proofs/storage_witness.bend as Witness
11import stelliferous@0.0.3.0/proofs/storage_tree_model.bend as Tree
12import stelliferous@0.0.3.0/proofs/storage_certificate.bend as Certificate
13import stelliferous@0.0.3.0/proofs/traversal_witness.bend as TraversalWitness
14
15def gathered(-values: Array<F32>,+position: U32,+context: Unit,
16 evidence: LaneWitness.Gathered<values,Lanes.gather(~Unit,~Gelu.identity,values,position,context)>) ->
17 LaneWitness.Gathered<values,Gelu.gather(values,position,context)>:
18 LaneWitness.Gathered{cell,equation} = evidence
19 LaneWitness.Gathered{Gelu.second(Gelu.first(cell)),
20 Equal.cong(Array<F32> & Tile.Tile8,Array<F32> & Tile.Tile8,Gelu.gathered,
21 Lanes.gather(~Unit,~Gelu.identity,values,position,context),(values,cell),equation)}
22
23def gather(+depth: Nat,-values: Array<F32>,+position: U32,+context: Unit,+witness: Witness.Witness(F32,depth,values)) ->
24 LaneWitness.Gathered<values,Gelu.gather(values,position,context)>:
25 gathered(values,position,context,LaneWitness.gather(~Unit,~Gelu.identity,depth,values,position,context,witness))
26
27def mapped(+depth: Nat,count: Nat,+position: U32,-values: Array<F32>,witness: Witness.Witness(F32,depth,values)) ->
28 Witness.Witness(F32,depth,Gelu.map(count,position,values)):
29 LaneWitness.rounds_with(~Unit,~Access.apply_element(~F32,~Act.gelu_tanh_value),~Gelu.gather,~gather,depth,count,Cells.State{values,Lanes.scratch(),position,0,Unit{}},
30 LaneWitness.Evidence{values,Lanes.scratch(),position,0,Unit{},{==},witness,Witness.fresh(Tile.Tile8,6n,Tile.filled(0.0))})
31
32def shape(covered: Bool,+depth: Nat,count: Nat,+position: U32,-values: Array<F32>,witness: Witness.Witness(F32,depth,values)) ->
33 Witness.Witness(F32,depth,Gelu.map_when(covered,count,position,values)):
34 match covered:
35 case True{}: mapped(depth,count,position,values,witness)
36 case False{}: TraversalWitness.map_in_place(~F32,~Unit,~Access.apply_element(~F32,~Act.gelu_tanh_value),depth,count,position,Unit{},values,witness)
37
38def map_when(+covered: Bool,+count: Nat,+position: U32,-values: Array<F32>,certificate: Certificate.Certificate<F32,values>) ->
39 Certificate.Certificate<F32,Gelu.map_when(covered,count,position,values)>:
40 Certificate.Certificate{+depth,equation} = certificate
41 Certificate.Certificate{depth,
42 %Certificate.restored(F32,depth,Gelu.map_when(covered,count,position,values),shape(covered,depth,count,position,values,Certificate.witness_at(F32,depth,values,equation))) :
43 {_ == Tree.pack(F32,Certificate.normalized(F32,depth,Tree.reflect(F32,Gelu.map_when(covered,count,position,values)))) : Array<F32>}
44 {==}}