search

iree-org / iree-llvm-sandbox

A sandbox for quick iteration and experimentation on projects related to IREE, MLIR, and LLVM
Apache License 2.0
54 stars 31 forks source link

[Indexing] Implement mapping of advanced indexing to `GatherOp` #714

Closed makslevental closed 1 year ago

makslevental commented 1 year ago

This PR implements advanced indexing (gather) with rectangular array-like objects (Python lists of lists, and index Tensors). Again, the approach closely follows the implementation of index_to_gather in jax.

@nicolasvasilache please take a close look at the lit tests to see whether the semantics are correct (at least the type level).

Currently stacked on top of

https://github.com/iree-org/iree-llvm-sandbox/pull/711

https://github.com/iree-org/iree-llvm-sandbox/pull/712

ingomueller-net commented 1 year ago

The link to JAX was super helpful! I still have trouble really understanding what these advanced indexing methods do, and I found the pointer to the great numpy documentation that seems to explain it. Maybe we should also link to it?

makslevental commented 1 year ago

The link to JAX was super helpful! I still have trouble really understanding what these advanced indexing methods do, and I found the pointer to the great numpy documentation that seems to explain it. Maybe we should also link to it?

That link is the first link here https://github.com/iree-org/iree-llvm-sandbox/blob/main/examples/advanced-indexing.ipynb :)