def gather_naturals

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

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

283def gather_naturals(+axes: List<&2, Nat>, +items: List<&2, U32>) ->
284  {natural_values(Axis.gather(U32, items, axes)) == Axis.gather(Nat, Shape.naturals(items), axes) : Maybe<&2, List<&2, Nat>>}:
285  match axes:
286    case Nil{}:
287      {==}
288    case +head <> +tail:
289      %Equal.sym(Maybe<&2, List<&2, Nat>>, natural_values(Axis.prepend(U32, Axis.lookup(U32, items, head), Axis.gather(U32, items, tail))),
290        Axis.prepend(Nat, natural_value(Axis.lookup(U32, items, head)), natural_values(Axis.gather(U32, items, tail))), prepend_naturals(Axis.lookup(U32, items, head), Axis.gather(U32, items, tail))) :
291        {_ == Axis.prepend(Nat, Axis.lookup(Nat, Shape.naturals(items), head), Axis.gather(Nat, Shape.naturals(items), tail)) : Maybe<&2, List<&2, Nat>>}
292      %Equal.sym(Maybe<&2, Nat>, natural_value(Axis.lookup(U32, items, head)), Axis.lookup(Nat, Shape.naturals(items), head), lookup_naturals(items, head)) :
293        {Axis.prepend(Nat, _, natural_values(Axis.gather(U32, items, tail))) == Axis.gather(Nat, Shape.naturals(items), head <> tail) : Maybe<&2, List<&2, Nat>>}
294      %Equal.sym(Maybe<&2, List<&2, Nat>>, natural_values(Axis.gather(U32, items, tail)), Axis.gather(Nat, Shape.naturals(items), tail), gather_naturals(tail, items)) :
295        {Axis.prepend(Nat, Axis.lookup(Nat, Shape.naturals(items), head), _) == Axis.gather(Nat, Shape.naturals(items), head <> tail) : Maybe<&2, List<&2, Nat>>}
296      {==}