A static polyhedra shape index helper

[thisiscam]

Hi JAX community!

I implemented a small prototype library called HedraX to help with static shape indexing in JAX.

TLDR; it allows you to express static, but non-regtangularly shaped loop domains easily, using a set-like expression (Presburger arithmetic, islpy)

Example

With the library, you can express a loop like this one:

for i in range(N):
  for j in range(i):
    A[i] += f(i, j)

Using the following:

def triangular_sum_hdx(N: int):
    x_ref = jax.array_ref(jnp.zeros((N,)))
    addresses, unravel = hdx.compile_table_indexer(
        "[N] -> {[i, j]: 0 <= j < i < N}", N=N
    )

    @foreach(addresses)
    def _(addr):
        i, j = unravel(addr)
        x_ref[i] += f(i, j)

    return x_ref[...]

Note that I'm using the jax.array_ref here that I just found out today --- I naturally included it in this example because array_ref seems like would go well with hedrax!

---

Here's a blog post (part 1) https://people.csail.mit.edu/camyang/blog/2025/hedrax/, where I wrote a little bit more on the libraries' usage. There is also a compile_closed_form_indexer which also works for the above example like so:

 _range, unravel = hdx.compile_closed_form_indexer(
      "[N] -> {[i, j]: 0 <= j < i < N}", N=N
  )

  @foreach(jnp.arange(_range.stop))
  def _(addr):
      i, j = unravel(addr)
      x_ref[i] += f(i, j)

I'm sharing the idea here to get some feedbacks on improving the presentation and the library.