# Indexing and Slicing

Haliax supports Numpy-style indexing, including so-called [Advanced Indexing](https://numpy.org/doc/stable/user/basics.indexing.html#advanced-indexing),
though the syntax is necessarily different. Most forms of indexing are supporting, except we don't support indexing with
booleans right now. (JAX doesn't support indexing with non-constant bool arrays anyway,
so I don't think it's worth the effort to implement it in Haliax.)

## Basic Indexing

Basic indexing works basically like you would expect: you can use integers or slices to index into an array.
Haliax supports two syntaxes for indexing: one accepts a dict of axis names and indices, and the other accepts
an alternating sequence of axis names and indices. The latter is useful for indexing with a small number of indices.

```python
import haliax as hax
import jax

X = hax.Axis("X", 10)
Y = hax.Axis("Y", 20)
Z = hax.Axis("Z", 30)

a = hax.random.uniform(jax.random.PRNGKey(0), (X, Y, Z))

a[{"X": 1, "Y": 2, "Z": 3}]  # returns a scalar jnp.ndarray
a[{"X": 1, "Y": 2, "Z": slice(3, 5)}]  # return a NamedArray with axes = Axis("Z", 2)
a[{"X": 1, "Y": slice(2, 4), "Z": slice(3, 5)}]  # return a NamedArray with axes = Axis("Y", 2), Axis("Z", 2)

a["X", 1, "Y", 2, "Z", 3]  # returns a scalar jnp.ndarray
a["X", 1, "Y", 2, "Z", 3:5]  # return a NamedArray with axes = Axis("Z", 2)
a["X", 1, "Y", 2:4, "Z", 3:5]  # return a NamedArray with axes = Axis("Y", 2), Axis("Z", 2)
```

Unfortunately, Python won't let us use `:` slice syntax inside of a dictionary, so we have to use `slice` instead.
This is why we have the second syntax, which is a bit less idiomatic in some ways, but it's more convenient.

Otherwise, the idea is pretty straightforward: any unspecified axes are treated as though indexed with `:` in NumPy,
slices are kept in reduced dimensions, and integers eliminate dimensions. If all dimensions are eliminated, a scalar
JAX ndarray is returned.


### Shapes in JAX

Before we continue note on shapes in JAX. Most JAX code will be used inside `jit`, which means that the sizes of all arrays
must be determined at compile time (i.e. when JAX interprets your functions abstractly). This is a hard requirement in
XLA.

A consequence of this restriction is that certain indexing patterns aren't allowed in `jit`-ed JAX code:

```python
import jax.numpy as jnp
import jax

@jax.jit
def f(x, slice_size: int):
    num_blocks = x.shape[0] // slice_size
    def body(i, m):
        return i + jnp.mean(x[i * slice_size : (i + 1) * slice_size])
    jax.lax.fori_loop(0, num_blocks, lambda i, m: m + body(i, m), 0.0)


f(jnp.arange(10), 2)
# IndexError: Array slice indices must have static start/stop/step to be used with NumPy indexing syntax.
# Found slice(Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=2/0)>,
# Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=2/0)>, None). To index a statically sized
# array at a dynamic position, try lax.dynamic_slice/dynamic_update_slice (JAX does not support dynamically sized
# arrays within JIT compiled functions).
```

This is a not-uncommon pattern in situations where you want to process a large array in chunks. In Haliax, we provide
two solutions: [haliax.slice][] as well as dynamic slices ([haliax.dslice][] a.k.a. [haliax.ds][]).

## Dynamic Slices

In light of the requirement that all array sizes be known at compile time, Haliax provides both a simple [haliax.slice][]
function, as well as [haliax.dslice][], which can be used with `[]`. The simple slice function is just a wrapper
around [jax.lax.dynamic_slice][]] and not worth discussing here.

`dslice` is a trick borrowed from the new experimental [jax.experimental.pallas][] module. It's essentially a slice,
except that instead of a start and an end (and maybe a stride), it takes a start and a size. The size must be
statically known, but the start can be dynamic. This allows us to write the above example as follows:

```python
import jax
import haliax as hax

N = hax.Axis("N", 10)
q = hax.arange(N)

@hax.named_jit
def f(x, slice_size: int):
    num_blocks = N.size // slice_size
    def body(i, m):
        return i + hax.mean(x["N", hax.dslice(i * slice_size, slice_size)])
    jax.lax.fori_loop(0, num_blocks, body, 0.0)

f(q, 2)
```

For convenience/brevity, `dslice` is aliased as `ds`. In addition, we also expose `dblock`, which is a convenience
function for computing the start and size of a slice given a block index and the size of the slice. Thus, the above
example can be written as follows:

```python
import jax
import haliax as hax

N = hax.Axis("N", 10)
q = hax.arange(N)

@hax.named_jit
def f(x, slice_size: int):
    num_blocks = N.size // slice_size
    def body(i, m):
        return i + hax.mean(x["N", hax.ds.block(i, slice_size)])
    jax.lax.fori_loop(0, num_blocks, body, 0.0)

f(q, 2)
```

It's not a huge improvement, but it's a bit more convenient.


## Advanced Indexing

NumPy's [Advanced Indexing](https://numpy.org/doc/stable/user/basics.indexing.html#advanced-indexing) is supported, though we use named arrays for the indices instead of normal arrays.
In NumPy, the indexed arrays much be broadcastable to the same shape.Advanced indexing in Haliax is similar, except that
they follow Haliax's broadcasting rules, meaning that the axis names determine broadcasting.Axes with the same name must
have the same size.

```python
import haliax as hax
import jax

X = hax.Axis("X", 10)
Y = hax.Axis("Y", 20)
Z = hax.Axis("Z", 30)

a = hax.random.uniform(jax.random.PRNGKey(0), (X, Y, Z))

I1 = hax.Axis("I1", 5)
I2 = hax.Axis("I2", 5)
I3 = hax.Axis("I3", 5)
ind1 = hax.random.randint(jax.random.PRNGKey(0), (I1,), 0, 10)
ind2 = hax.random.randint(jax.random.PRNGKey(0), (I2, I3), 0, 20)

a[{"X": ind1}]  # returns a NamedArray with axes = Axis("I1", 5), Axis("Y", 20), Axis("Z", 30)

a[{"X": ind1, "Y": ind2}]  # returns a NamedArray with axes = Axis("I1", 5), Axis("I2", 5), Axis("I3", 5), Axis("Z", 30)
a[{"X": ind1, "Y": ind2, "Z": 3}]  # returns a NamedArray with axes = Axis("I1", 5), Axis("I2", 5), Axis("I3", 5)
```

The order of the indices in the dictionary doesn't matter, and you can mix and match basic and advanced indexing.
The actual sequence of axes is a bit complex, both in Haliax and in NumPy. If you need a specific order, it's probably
best to use rearrange.

In keeping with the one-axis-per-name rule, you are allowed to index using axes with a name present in the array,
if it would be eliminated by the indexing operation. For example:

```python
import haliax as hax
import jax

X = hax.Axis("X", 10)
Y = hax.Axis("Y", 20)
Z = hax.Axis("Z", 30)

X2 = hax.Axis("X", 5)
Y2 = hax.Axis("Y", 5)

a = hax.random.uniform(jax.random.PRNGKey(0), (X, Y, Z))
ind1 = hax.random.randint(jax.random.PRNGKey(0), (X2,), 0, 10)
ind2 = hax.random.randint(jax.random.PRNGKey(0), (Y2,), 0, 10)

a[{"X": ind1, "Y": ind2}]  # returns a NamedArray with axes = Axis("X", 5), Axis("Y", 5), Axis("Z", 30)

a[{"Y": ind1}]  # error, "X" is not eliminated by the indexing operation

a[{"X": ind2, "Y": ind1}]  # ok, because X and Y are eliminated by the indexing operation
```
