def permuted_value

bendtorch-core@0.1.0.0 · src/tensor/permute_values.bend:347

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
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/layout.bend as Layout
import bendtorch-core@0.1.0.0/src/tensor/from_values_contract.bend as FromValues
import bendtorch-core@0.1.0.0/src/tensor/contracts/replayed.bend as Replayed
import bendtorch-core@0.1.0.0/src/tensor/index_map.bend as IndexMap
import bendtorch-core@0.1.0.0/src/tensor/observation.bend as Values
import bendtorch-core@0.1.0.0/src/tensor/transform_observation.bend as Observation
import bendtorch-core@0.1.0.0/src/tensor/permute_build_proof.bend as PermuteProof

347def permuted_value(+input: Shape.Shape, +values: List<&2, F32>,
348  +counted: {List.length(&2, F32, values) == Shape.logical_numel(Shape.dimensions(input)) : Nat},
349  +axes: List<&2, U32>, +valid: {True{} == Layout.bijective(input, axes) : Bool}, +index: U32,
350  +inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(input))) == True{} : Bool},
351  -fits: {Lifecycle.logical_length(Replayed.storage(input, 0.0, FromValues.writes(values, 0n))) == Shape.logical_numel(Shape.dimensions(input)) : Nat}) ->
352  Permuted(input, values, axes, fits, pair => {Observation.applied_value_pair(input, index, pair) ==
353    Values.listed(values, Values.permute_source(input, axes, index)) : Maybe<&2, F32>}):
354  permute_lift(input, values, axes, valid, pair => {Observation.applied_value_pair(input, index, pair) ==
355      Values.listed(values, Values.permute_source(input, axes, index)) : Maybe<&2, F32>},
356    plan => source_fits => target_fits =>
357      +output = Layout.output(input, axes, plan)
358      +total = Shape.logical_numel(Shape.dimensions(output))
359      +history = FromValues.writes(values, 0n)
360      +below = output_inside(input, axes, plan, index, inside)
361      Equal.trans(Maybe<&2, F32>,
362        Core.dynamic_read_value(Core.read_dynamic(Core.Packed{output, Core.TTensor{written(input, axes, plan, 0.0, FromValues.writes(values, 0n), total), target_fits}}, index)),
363        Some{Replayed.value(input, 0.0, history, U32.from_nat(source_of(Request{input, axes}, U32.to_nat(index))))},
364        Values.listed(values, Values.permute_source(input, axes, index)),
365        IndexMap.copied_point(~Request, ~source_of, Request{input, axes}, input, 0.0, history, output, target_fits, index, below),
366        found_mapped(input, axes, plan, values, counted, index, Layout.map_index(input, axes, plan, index),
367          PermuteProof.mapping_done(input, axes, plan, index, below))), fits)