def contains

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

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

99def contains(+axes: List<&2, U32>, +axis: U32) -> {Layout.contains(axis, axes) == Axis.contains(U32.to_nat(axis), Shape.naturals(axes)) : Bool}:
100  match axes:
101    case Nil{}:
102      {==}
103    case +head <> +tail:
104      %Equal.sym(Bool, U32.is_eq(axis, head), Nat.is_eq(U32.to_nat(axis), U32.to_nat(head)), u32_eq(axis, head)) :
105        {Bool.or(_, Layout.contains(axis, tail)) == Axis.contains(U32.to_nat(axis), Shape.naturals(head <> tail)) : Bool}
106      %Equal.sym(Bool, Layout.contains(axis, tail), Axis.contains(U32.to_nat(axis), Shape.naturals(tail)), contains(tail, axis)) :
107        {Bool.or(Nat.is_eq(U32.to_nat(axis), U32.to_nat(head)), _) == Axis.contains(U32.to_nat(axis), Shape.naturals(head <> tail)) : Bool}
108      {==}