def lookup_digits

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

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

69# Coordinate k of `Shape.digits(dims, x)` is what the pick of k computes.
70def lookup_digits(+dims: List<&2, Nat>, +x: Nat, +k: Nat) ->
71  {Axis.lookup(U32, Shape.words(Shape.digits(dims, x)), k) == coordinate_of(Build.pick_at(dims, k), x) : Maybe<&2, U32>}:
72  match dims:
73    case Nil{}:
74      {==}
75    case +head <> +tail:
76      match k:
77        case 0n:
78          {==}
79        case 1n+j:
80          Equal.trans(Maybe<&2, U32>, Axis.lookup(U32, Shape.words(Shape.digits(tail, Nat.mod(x, Shape.numel(tail)))), j),
81            coordinate_of(Build.pick_at(tail, j), Nat.mod(x, Shape.numel(tail))),
82            coordinate_of(Build.widen(Shape.numel(tail), Build.pick_at(tail, j)), x),
83            lookup_digits(tail, Nat.mod(x, Shape.numel(tail)), j),
84            Equal.sym(Maybe<&2, U32>, coordinate_of(Build.widen(Shape.numel(tail), Build.pick_at(tail, j)), x),
85              coordinate_of(Build.pick_at(tail, j), Nat.mod(x, Shape.numel(tail))),
86              widen_coordinate(Shape.numel(tail), Build.pick_at(tail, j), x)))