src/tensor/transpose_spec.bend

bendtorch-core@0.1.0.0 · 47 lines · 9 definitions

1import Base
2import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
3import bendtorch-core@0.1.0.0/src/tensor/axis_spec.bend as Axis
4
5def axis_limit(rank: Nat) -> Nat:
6  match rank:
7    case 0n:
8      1n
9    case 1n+p:
10      1n+p
11
12def axes_valid(+rank: Nat, +a: Nat, +b: Nat) -> Bool:
13  Bool.and(Nat.is_lt(a, axis_limit(rank)), Nat.is_lt(b, axis_limit(rank)))
14
15def valid(+shape: Shape.Shape, +a: U32, +b: U32) -> Bool:
16  axes_valid(Shape.rank_dimensions(Shape.dimensions(shape)), U32.to_nat(a), U32.to_nat(b))
17
18def replace(-A: Data, items: List<&2, A>, index: Nat, +value: A) -> List<&2, A>:
19  match items index:
20    case Nil{} index:
21      Nil{}
22    case head <> tail 0n:
23      value <> tail
24    case head <> tail 1n+p:
25      head <> replace(A, tail, p, value)
26
27def swap_selected(-A: Data, +items: List<&2, A>, +a: Nat, +b: Nat, left: Maybe<&2, A>, right: Maybe<&2, A>) -> List<&2, A>:
28  match left right:
29    case Some{x} Some{y}:
30      replace(A, replace(A, items, a, y), b, x)
31    case _ _:
32      items
33
34def swap(-A: Data, +items: List<&2, A>, +a: Nat, +b: Nat) -> List<&2, A>:
35  left right = Axis.lookup(A, items, a) Axis.lookup(A, items, b)
36  swap_selected(A, items, a, b, left, right)
37
38def natural_axes(+rank: Nat, +a: Nat, +b: Nat) -> List<&2, Nat>:
39  swap(Nat, Axis.indices(rank, 0n), a, b)
40
41def requested_axes(+rank: Nat, +a: Nat, +b: Nat) -> List<&2, U32>:
42  Shape.words(natural_axes(rank, a, b))
43
44# The permutation `transpose_copy(input, a, b)` runs: the identity of the
45# input's rank with axes a and b swapped.
46def swapped_axes(+input: Shape.Shape, +a: U32, +b: U32) -> List<&2, U32>:
47  requested_axes(Shape.rank_dimensions(Shape.dimensions(input)), U32.to_nat(a), U32.to_nat(b))