def reorder

bendtorch-core@0.1.0.0 · src/tensor/axis_algebra.bend:33

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/tensor/axis_spec.bend as Axis

33def reorder(+axes: List<&2, U32>, +values: List<&2, U32>) ->
34  {Layout.reorder(values, axes) == Axis.gather(U32, values, Shape.naturals(axes)) : Maybe<&2, List<&2, U32>>}:
35  match axes:
36    case Nil{}:
37      {==}
38    case +axis <> +tail:
39      %Equal.sym(Maybe<&2, U32>, Layout.lookup(values, U32.to_nat(axis)), Axis.lookup(U32, values, U32.to_nat(axis)), lookup(values, U32.to_nat(axis))) :
40        {Layout.reorder_head(value => Layout.reorder_tail(value, Layout.reorder(values, tail)), _) == Axis.prepend(U32, Axis.lookup(U32, values, U32.to_nat(axis)), Axis.gather(U32, values, Shape.naturals(tail))) : Maybe<&2, List<&2, U32>>}
41      %Equal.sym(Maybe<&2, List<&2, U32>>, Layout.reorder(values, tail), Axis.gather(U32, values, Shape.naturals(tail)), reorder(tail, values)) :
42        {Layout.reorder_head(value => Layout.reorder_tail(value, _), Axis.lookup(U32, values, U32.to_nat(axis))) == Axis.prepend(U32, Axis.lookup(U32, values, U32.to_nat(axis)), Axis.gather(U32, values, Shape.naturals(tail))) : Maybe<&2, List<&2, U32>>}
43      reorder_step(Axis.lookup(U32, values, U32.to_nat(axis)), Axis.gather(U32, values, Shape.naturals(tail)))