def indexed_lift

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

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/storage/main.bend as Storage
import bendtorch-core@0.1.0.0/src/storage/LAWS.bend as StorageLaws
import bendtorch-core@0.1.0.0/src/storage/array.bend as ArrayCore
import bendtorch-core@0.1.0.0/src/storage/api.bend as Api
import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
import bendtorch-core@0.1.0.0/src/tensor/lifecycle.bend as Lifecycle
import bendtorch-core@0.1.0.0/src/tensor/contracts/replayed.bend as Replayed
import bendtorch-core@0.1.0.0/src/tensor/contracts/matmul_parallel_proof.bend as Par
import bendtorch-core@0.1.0.0/src/autograd/reduce.bend as Reduce

1043def indexed_lift(+plan: Shape.BroadcastPlan, +target: Shape.Shape, +shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>,
1044  -fits: {Lifecycle.logical_length(Replayed.storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat},
1045  +facts: Facts<plan, shape, U32.to_nat(Shape.count(target))>,
1046  +d: Nat, +fits_d: {Storage.bounds(d, U32.to_nat(Shape.count(target))) == True{} : Bool},
1047  -f: {Storage.bounds(d, U32.to_nat(Shape.count(target))) == True{} : Bool},
1048  -fo: {Lifecycle.logical_length(Storage.Owned{d, U32.to_nat(Shape.count(target)), Array.new(F32, d, 0.0), ArrayCore.new_balanced(F32, d, 0.0), f}) ==
1049    Shape.logical_numel(Shape.dimensions(target)) : Nat},
1050  +map: PlanMap(plan), -P: Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic> -> Type,
1051  final: Final(plan, target, shape, fill, ops, map,
1052    Storage.Owned{d, U32.to_nat(Shape.count(target)), Array.new(F32, d, 0.0), ArrayCore.new_balanced(F32, d, 0.0), f}, P)) ->
1053  P(Reduce.parallel_indexed(Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}}, target, plan,
1054    Core.TTensor{Storage.Owned{d, U32.to_nat(Shape.count(target)), Array.new(F32, d, 0.0), ArrayCore.new_balanced(F32, d, 0.0), f}, fo},
1055    d, U32.to_nat(Shape.count(target)), f, fo, map)):
1056  match map:
1057    case Shape.Unindexed{zero}:
1058      final(ops)(other => {==})({==})(fits)(Nil{})(index => inside => {==})(fo)
1059    case Shape.Indexed{+rows, +rows_exact}:
1060      +n = U32.to_nat(Shape.count(target))
1061      +map : PlanMap(plan) = Shape.Indexed{rows, rows_exact}
1062      Par.input_take(shape, fill, ops,
1063        taken => P(Reduce.parallel_rows(plan, target, d, n, f, fo, taken, rows)),
1064        kept => values => words => h => dc => fc => source => kept_fits =>
1065          Par.storage_shift(target,
1066            StorageLaws.replay(writes(plan, map, shape, fill, ops, n), Storage.Owned{d, n, Array.new(F32, d, 0.0), ArrayCore.new_balanced(F32, d, 0.0), f}),
1067            Storage.Owned{d, n,
1068              Reduce.build(d, parallel_splits(plan, d, n), rows, Shape.extents(Shape.output_shape(plan)), n, 0n, Api.arr(h, dc, fc)),
1069              Reduce.build_balanced(d, parallel_splits(plan, d, n), rows, Shape.extents(Shape.output_shape(plan)), n, 0n, Api.arr(h, dc, fc)), f},
1070            output_eq(plan, rows, rows_exact, shape, fill, ops, h, dc, fc, source, target, facts, d, fits_d, f),
1071            st => k => P((Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, kept), kept_fits}}, Done{Core.Packed{target, Core.TTensor{st, k}}})),
1072            final(kept)(values)(words)(kept_fits)(writes(plan, map, shape, fill, ops, n))(
1073              index => inside => writes_value(plan, map, shape, fill, ops, target, n, index, inside_count(target, index, inside), room_of(target))))(fo))(fits)