def plan_ranked_success

bendtorch-core@0.1.0.0 · src/tensor/permutation_mapping.bend:304

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/tensor/layout.bend as Layout
import bendtorch-core@0.1.0.0/src/storage/arith.bend as Arith

304def plan_ranked_success(+input: Shape.Shape, +axes: List<&2, U32>,
305  decision: Shape.BoolCase<Nat.is_eq(Shape.rank_dimensions(axes), Shape.rank_dimensions(Shape.dimensions(input)))>,
306  +valid: {True{} == Layout.bijective(input, axes) : Bool}) ->
307  {Result.is_done(&2, &2, Layout.Error, Layout.Plan<input, axes>, Layout.plan_ranked(input, axes, decision)) == True{} : Bool}:
308  match decision:
309    case Shape.BFalse{equal}:
310      Empty.absurd({False{} == True{} : Bool},
311        Arith.clash(Equal.trans(Bool, False{}, Nat.is_eq(Shape.rank_dimensions(axes), Shape.rank_dimensions(Shape.dimensions(input))), True{}, equal,
312          Equal.sym(Bool, True{}, Nat.is_eq(Shape.rank_dimensions(axes), Shape.rank_dimensions(Shape.dimensions(input))),
313            and_left(Nat.is_eq(Shape.rank_dimensions(axes), Shape.rank_dimensions(Shape.dimensions(input))),
314              Bool.and(Layout.bounded(axes, Shape.rank(input)), Layout.unique(axes)), valid)))))
315    case Shape.BTrue{equal}:
316      plan_bounded_success(input, axes, Shape.bool_case(Layout.bounded(axes, Shape.rank(input))), valid)