def scatter_digits

bendtorch-core@0.1.0.0 · src/tensor/permute_build_proof.bend:109

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/tensor/axis_spec.bend as Axis
import bendtorch-core@0.1.0.0/src/tensor/permute_build.bend as Build

109# The list `Axis.scatter` builds from the output coordinates of x is the
110# picks' coordinates.
111def scatter_digits(+rank: Nat, +axes: List<&2, Nat>, +out: List<&2, Nat>, +x: Nat, +cursor: Nat) ->
112  {Axis.scatter(U32, rank, axes, Shape.words(Shape.digits(out, x)), cursor) == coordinates_of(Build.picks(rank, axes, out, cursor), x) :
113    Maybe<&2, List<&2, U32>>}:
114  match rank:
115    case 0n:
116      {==}
117    case 1n+p:
118      +position = Axis.position(axes, cursor)
119      %Equal.sym(Maybe<&2, U32>, Axis.select(U32, Shape.words(Shape.digits(out, x)), position), coordinate_of(Build.located(out, position), x),
120        select_digits(out, x, position)) :
121        {Axis.prepend(U32, _, Axis.scatter(U32, p, axes, Shape.words(Shape.digits(out, x)), 1n+cursor)) ==
122          coordinates_of(Axis.prepend(Build.Pick, Build.located(out, position), Build.picks(p, axes, out, 1n+cursor)), x) : Maybe<&2, List<&2, U32>>}
123      %Equal.sym(Maybe<&2, List<&2, U32>>, Axis.scatter(U32, p, axes, Shape.words(Shape.digits(out, x)), 1n+cursor),
124        coordinates_of(Build.picks(p, axes, out, 1n+cursor), x), scatter_digits(p, axes, out, x, 1n+cursor)) :
125        {Axis.prepend(U32, coordinate_of(Build.located(out, position), x), _) ==
126          coordinates_of(Axis.prepend(Build.Pick, Build.located(out, position), Build.picks(p, axes, out, 1n+cursor)), x) : Maybe<&2, List<&2, U32>>}
127      prepend_coordinates(Build.located(out, position), Build.picks(p, axes, out, 1n+cursor), x)