def matches

bendtorch-core@0.1.0.0 · src/tensor/gather.bend:107

import Base

107# Whether the index's dimensions are the input's at every axis but `dim`.
108def matches(dimensions: List<&2, U32>, index: List<&2, U32>, +dim: Nat) -> Bool:
109  match dimensions index dim:
110    case extent <> rest other <> more 0n:
111      equal(rest, more)
112    case extent <> rest other <> more 1n+next:
113      Bool.and(Nat.is_eq(U32.to_nat(extent), U32.to_nat(other)), matches(rest, more, next))
114    case _ _ _:
115      False{}