def BroadcastWith

bendtorch-core@0.1.0.0 · src/autograd/broadcast_pointwise_LAWS.bend:58

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/LAWS.bend as Contracts
import bendtorch-core@0.1.0.0/src/tensor/contracts/replayed.bend as Replayed
import bendtorch-core@0.1.0.0/src/autograd/broadcast.bend as Broadcast

58# What the transport form of broadcast_to_replayed asks of P: P holds for the
59# input, replayed with any history that keeps its words, next to an output of
60# the plan's shape, replayed over +0.0 with any history whose words are
61# `cells`.
62def BroadcastWith(+plan: Shape.BroadcastPlan, +shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>,
63  -P: Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic> -> Type) -> Type:
64  @+kept: List<&2, StorageLaws.Op> ->
65  @kept_words: {Contracts.words(shape, fill, kept, Shape.logical_numel(Shape.dimensions(shape))) ==
66    Contracts.words(shape, fill, ops, Shape.logical_numel(Shape.dimensions(shape))) : List<&2, F32>} ->
67  @-input_fits: {Lifecycle.logical_length(Replayed.storage(shape, fill, kept)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat} ->
68  @+output_ops: List<&2, StorageLaws.Op> ->
69  @values: {Contracts.words(Shape.output_shape(plan), 0.0, output_ops, Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) ==
70    cells(plan, shape, fill, ops, Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) : List<&2, F32>} ->
71  @-output_fits: {Lifecycle.logical_length(Replayed.storage(Shape.output_shape(plan), 0.0, output_ops)) ==
72    Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan))) : Nat} ->
73    P((Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, kept), input_fits}},
74      Done{Core.Packed{Shape.output_shape(plan), Core.TTensor{Replayed.storage(Shape.output_shape(plan), 0.0, output_ops), output_fits}}}))