def index_of

bendtorch-core@0.1.0.0 · src/shape/core.bend:828

import Base

828# Index rows are formed only below a nonzero output count. Every padded input
829# extent is then 1 or the output extent, so no product exceeds that count.
830def index_of(+left: Shape, +right: Shape, +output: Shape, empty: BoolCase<Nat.is_eq(elements(output), 0n)>) -> Index<extents(left), extents(right), extents(output)>:
831  match empty:
832    case BTrue{equal}:
833      Unindexed{Equal.trans(Nat, 0n, elements(output), numel(extents(output)), Equal.sym(Nat, elements(output), 0n, equal_from(elements(output), 0n, equal)), elements_exact(output))}
834    case BFalse{equal}:
835      Indexed{index_of.rows(left, right, output, rank_dimensions(dimensions(output))), index_of.exact(left, right, output)}