1	# Shape preservation reuses the traversal's gather/scatter witnesses.
2	import Base
3	import ../tensor_access.bend as Access
4	import ../math_activation.bend as Act
5	import ../math_gelu_lanes.bend as Gelu
6	import ../kernel_tile.bend as Tile
7	import ../traversal_lanes.bend as Lanes
8	import ../traversal_cells.bend as Cells
9	import ./traversal_lanes_witness.bend as LaneWitness
10	import ./storage_witness.bend as Witness
11	import ./storage_tree_model.bend as Tree
12	import ./storage_certificate.bend as Certificate
13	import ./traversal_witness.bend as TraversalWitness
14	
15	def 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	
23	def 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	
27	def 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	
32	def 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	
38	def 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	    {==}}