joelberkeley/spidr

NumPy-like indexing and slicing

Closed this issue · 0 comments

It would be really nice to have NumPy-like slicing syntax, where you can slice and index a number of dimensions at once. Obviously we can't use the same notation and it's not valid Idris, so how do we want to do this?

One options is to have an operator that acts on a Tensor and a list of Slices, like

xs : Tensor [3, 4, 5] S32

ys : Tensor [2, 3] S32
ys = xs // [..3, 2, 2..5]

where

  • a..b means slicing from a to b. If either is omitted, this means everything up to/from the other, so .. means everything in that dim
  • plain integers mean indexing and are possible via fromInteger

It would also be good to have the ... syntax of NumPy, which means "all remaining leading/trailing dims if at the start/end". We obviously also need to make sure the types work.

Note XLA supports multi-slicing, so with that and reshape this should be trivial on the XLA side.

Q: What do haskell libraries do?