law reduce_fold_fiber

bendtorch-core@0.1.0.0 · src/autograd/reduce_fold_LAWS.bend:22

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/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/LAWS.bend as Contracts
import bendtorch-core@0.1.0.0/src/autograd/reduce_LAWS.bend as ReduceLaws
import bendtorch-core@0.1.0.0/src/autograd/reduce_fold.bend as ReduceFold

22# Ordered-fold family, reduce instance: for a plan whose output has the
23# contribution's shape, every target index and every replayed contribution,
24# the reduce cell gives `ReduceLaws.fiber` over the contribution's words in
25# index order, and returns the contribution with every word kept.
26law reduce_fold_fiber:
27  for +plan: Shape.BroadcastPlan
28  for same: {Shape.extents(Shape.output_shape(plan)) == Shape.extents(Shape.right_shape(plan)) : List<&2, Nat>}
29  for +index: U32
30  for +fill: F32
31  for +ops: List<&2, StorageLaws.Op>
32  for -fits: {Lifecycle.logical_length(Replayed.storage(Shape.right_shape(plan), fill, ops)) ==
33    Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan))) : Nat}
34  FiberContract<Shape.right_shape(plan),
35    ReduceLaws.fiber(plan,
36      Contracts.words(Shape.right_shape(plan), fill, ops, Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan)))),
37      index, Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan)))),
38    other => Replayed.value(Shape.right_shape(plan), fill, ops, other),
39    ReduceFold.cell(plan, index,
40      Core.Packed{Shape.right_shape(plan), Core.TTensor{Replayed.storage(Shape.right_shape(plan), fill, ops), fits}})>