Repository: google/jax

Issue:
JIT fails to optimise batched array indexing: huge performance difference
Batched array indexing is the batched version of this kind operation:

```python
x = np.array([[0, 2], [1, 4], [2, 3]])
y = 2
print(x[y])  # [2 3]
```

---

I implemented batched array indexing in two ways. However, after JIT compilation, the performance is very different:

Method 1: `jax.vmap`

```python
@jax.jit
def f1(x, y):
    return jax.vmap(getitem)(x, y)
```

Method 2: `np.take_along_axis`

```python
@jax.jit
def f2(x, y):
    return np.take_along_axis(x, y[..., None, None], axis=1)[:, 0]
```

---

Test:

I am testing on the TPU platform.

```python
import jax
import jax.numpy as np
import jax.random as rand
from operator import getitem

key = rand.PRNGKey(42)

x = rand.randint(key, (32, 128, 1024), 0, 100)
y = rand.randint(key, (32,), 0, 128)

@jax.jit
def f1(x, y):
    return jax.vmap(getitem)(x, y)

@jax.jit
def f2(x, y):
    return np.take_along_axis(x, y[..., None, None], axis=1)[:, 0]

a = f1(x, y)  # shape: (32, 1024)
b = f2(x, y)
assert np.array_equal(a, b)

import timeit
print(timeit.timeit('f1(x, y).block_until_ready()', globals=globals(), number=5000))
print(timeit.timeit('f2(x, y).block_until_ready()', globals=globals(), number=5000))
```

Output:

```
0.7506967119406909
10.882259733974934
```

This means that the `jax.vmap` method is much faster than the `np.take_along_axis` method after JIT compilation.

---

And the HLO of the `np.take_along_axis` method is also much longer.

Method 1: `jax.vmap`

```llvm
HloModule xla_computation_f1.21, is_scheduled=true

%fused_computation (param_0.2: s32[32,128,1024], param_1.2: s32[64]) -> s32[32,1024] {
  %param_0.2 = s32[32,128,1024]{2,1,0:T(8,128)} parameter(0)
  %param_1.2 = s32[64]{0:T(256)} parameter(1)
  %reshape.6 = s32[32,2]{1,0:T(8,128)} reshape(s32[64]{0:T(256)} %param_1.2)
  %transpose.5 = s32[32,2]{1,0:T(8,128)} transpose(s32[32,2]{1,0:T(8,128)} %reshape.6), dimensions={0,1}
  %gather.1 = s32[32,1024]{1,0:T(8,128)} gather(s32[32,128,1024]{2,1,0:T(8,128)} %param_0.2, s32[32,2]{1,0:T(8,128)} %transpose.5), offset_dims={1}, collapsed_slice_dims={0,1}, start_index_map={0,1}, index_vector_dim=1, slice_sizes={1,1,1024}, indices_are_sorted=true, metadata={op_type="gather" op_name="xla_computation(f1)/jit(f1)/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(1,), collapsed_slice_dims=(0, 1), start_index_map=(0, 1)) slice_sizes=(1, 1, 1024) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/ayaka/main.py" source_line=18}
  %transpose.4 = s32[32,1024]{1,0:T(8,128)} transpose(s32[32,1024]{1,0:T(8,128)} %gather.1), dimensions={0,1}
  ROOT %reshape.5 = s32[32,1024]{1,0:T(8,128)} reshape(s32[32,1024]{1,0:T(8,128)} %transpose.4)
}

%fused_computation.2 (param_0.9: s32[32]) -> s32[32,1] {
  %param_0.9 = s32[32]{0:T(256)} parameter(0)
  %constant.9 = s32[]{:T(256)} constant(0), metadata={op_type="lt" op_name="xla_computation(f1)/jit(f1)/lt" source_file="/home/ayaka/main.py" source_line=18}
  %broadcast.4 = s32[32]{0:T(256)} broadcast(s32[]{:T(256)} %constant.9), dimensions={}, metadata={op_type="lt" op_name="xla_computation(f1)/jit(f1)/lt" source_file="/home/ayaka/main.py" source_line=18}
  %compare.3 = pred[32]{0:T(1024)(128)(4,1)} compare(s32[32]{0:T(256)} %param_0.9, s32[32]{0:T(256)} %broadcast.4), direction=LT, metadata={op_type="lt" op_name="xla_computation(f1)/jit(f1)/lt" source_file="/home/ayaka/main.py" source_line=18}
  %constant.6 = s32[]{:T(256)} constant(128), metadata={op_type="add" op_name="xla_computation(f1)/jit(f1)/add" source_file="/home/ayaka/main.py" source_line=18}
  %broadcast.3 = s32[32]{0:T(256)} broadcast(s32[]{:T(256)} %constant.6), dimensions={}, metadata={op_type="add" op_name="xla_computation(f1)/jit(f1)/add" source_file="/home/ayaka/main.py" source_line=18}
  %add.3 = s32[32]{0:T(256)} add(s32[32]{0:T(256)} %param_0.9, s32[32]{0:T(256)} %broadcast.3), metadata={op_type="add" op_name="xla_computation(f1)/jit(f1)/add" source_file="/home/ayaka/main.py" source_line=18}
  %select.3 = s32[32]{0:T(256)} select(pred[32]{0:T(1024)(128)(4,1)} %compare.3, s32[32]{0:T(256)} %add.3, s32[32]{0:T(256)} %param_0.9), metadata={op_type="select_n" op_name="xla_computation(f1)/jit(f1)/select_n" source_file="/home/ayaka/main.py" source_line=18}
  ROOT %bitcast.1 = s32[32,1]{0,1:T(2,128)} bitcast(s32[32]{0:T(256)} %select.3), metadata={op_type="broadcast_in_dim" op_name="xla_computation(f1)/jit(f1)/broadcast_in_dim[shape=(32, 1) broadcast_dimensions=(0,)]" source_file="/home/ayaka/main.py" source_line=18}
}

%fused_computation.3 (param_0.11: s32[32,1], param_1.15: s32[32,1]) -> s32[32,2] {
  %param_1.15 = s32[32,1]{0,1:T(2,128)} parameter(1)
  %constant.5 = s32[]{:T(256)} constant(-2147483648)
  %pad.3 = s32[32,2]{0,1:T(2,128)} pad(s32[32,1]{0,1:T(2,128)} %param_1.15, s32[]{:T(256)} %constant.5), padding=0_0x0_1, metadata={op_type="concatenate" op_name="xla_computation(f1)/jit(f1)/concatenate[dimension=1]" source_file="/home/ayaka/main.py" source_line=18}
  %param_0.11 = s32[32,1]{0,1:T(2,128)} parameter(0)
  %pad.2 = s32[32,2]{0,1:T(2,128)} pad(s32[32,1]{0,1:T(2,128)} %param_0.11, s32[]{:T(256)} %constant.5), padding=0_0x1_0, metadata={op_type="concatenate" op_name="xla_computation(f1)/jit(f1)/concatenate[dimension=1]" source_file="/home/ayaka/main.py" source_line=18}
  ROOT %maximum.1 = s32[32,2]{0,1:T(2,128)} maximum(s32[32,2]{0,1:T(2,128)} %pad.3, s32[32,2]{0,1:T(2,128)} %pad.2), metadata={op_type="concatenate" op_name="xla_computation(f1)/jit(f1)/concatenate[dimension=1]" source_file="/home/ayaka/main.py" source_line=18}
}

ENTRY %xla_computation_f1.21 (parameter.1: s32[32,128,1024], parameter.2: s32[32]) -> (s32[32,1024]) {
  %parameter.2 = s32[32]{0:T(256)} parameter(1)
  %parameter.1 = s32[32,128,1024]{2,1,0:T(8,128)} parameter(0)
  %fusion.2 = s32[32,1]{0,1:T(2,128)} fusion(s32[32]{0:T(256)} %parameter.2), kind=kLoop, calls=%fused_computation.2, metadata={op_type="broadcast_in_dim" op_name="xla_computation(f1)/jit(f1)/broadcast_in_dim[shape=(32, 1) broadcast_dimensions=(0,)]" source_file="/home/ayaka/main.py" source_line=18}
  %iota.0 = s32[32,1]{0,1:T(2,128)} iota(), iota_dimension=0, metadata={op_type="iota" op_name="xla_computation(f1)/jit(f1)/iota[dtype=int32 shape=(32, 1) dimension=0]" source_file="/home/ayaka/main.py" source_line=18}
  %fusion.3 = s32[32,2]{0,1:T(2,128)} fusion(s32[32,1]{0,1:T(2,128)} %fusion.2, s32[32,1]{0,1:T(2,128)} %iota.0), kind=kLoop, calls=%fused_computation.3, metadata={op_type="concatenate" op_name="xla_computation(f1)/jit(f1)/concatenate[dimension=1]" source_file="/home/ayaka/main.py" source_line=18}
  %copy = s32[32,2]{1,0:T(8,128)} copy(s32[32,2]{0,1:T(2,128)} %fusion.3), metadata={op_type="concatenate" op_name="xla_computation(f1)/jit(f1)/concatenate[dimension=1]" source_file="/home/ayaka/main.py" source_line=18}
  %reshape.1 = s32[64]{0:T(256)} reshape(s32[32,2]{1,0:T(8,128)} %copy)
  %fusion = s32[32,1024]{1,0:T(8,128)} fusion(s32[32,128,1024]{2,1,0:T(8,128)} %parameter.1, s32[64]{0:T(256)} %reshape.1), kind=kCustom, calls=%fused_computation, metadata={op_type="gather" op_name="xla_computation(f1)/jit(f1)/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(1,), collapsed_slice_dims=(0, 1), start_index_map=(0, 1)) slice_sizes=(1, 1, 1024) unique_indices=True indices_are_sorted=True mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/ayaka/main.py" source_line=18}
  ROOT %tuple.20 = (s32[32,1024]{1,0:T(8,128)}) tuple(s32[32,1024]{1,0:T(8,128)} %fusion)
}
```

Method 2: `np.take_along_axis`

```llvm
HloModule xla_computation_f2.43, is_scheduled=true

%fused_computation (param_0.2: s32[32,128,1024], param_1.3: s32[131072]) -> s32[32768] {
  %param_0.2 = s32[32,128,1024]{2,1,0:T(8,128)} parameter(0)
  %param_1.3 = s32[131072]{0:T(1024)} parameter(1)
  %reshape.17 = s32[32,1,1024,4]{3,2,1,0:T(8,128)} reshape(s32[131072]{0:T(1024)} %param_1.3)
  %slice.1 = s32[32,1,1024,3]{3,2,1,0:T(8,128)} slice(s32[32,1,1024,4]{3,2,1,0:T(8,128)} %reshape.17), slice={[0:32], [0:1], [0:1024], [0:3]}
  %transpose.6 = s32[32,1,1024,3]{3,2,1,0:T(8,128)} transpose(s32[32,1,1024,3]{3,2,1,0:T(8,128)} %slice.1), dimensions={0,1,2,3}
  %gather.3 = s32[32,1,1024]{2,1,0:T(2,128)} gather(s32[32,128,1024]{2,1,0:T(8,128)} %param_0.2, s32[32,1,1024,3]{3,2,1,0:T(8,128)} %transpose.6), offset_dims={}, collapsed_slice_dims={0,1,2}, start_index_map={0,1,2}, index_vector_dim=3, slice_sizes={1,1,1}, metadata={op_type="gather" op_name="xla_computation(f2)/jit(f2)/jit(take_along_axis)/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0, 1, 2), start_index_map=(0, 1, 2)) slice_sizes=(1, 1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/ayaka/main.py" source_line=22}
  %transpose.5 = s32[32,1,1024]{2,1,0:T(2,128)} transpose(s32[32,1,1024]{2,1,0:T(2,128)} %gather.3), dimensions={0,1,2}
  ROOT %reshape.16 = s32[32768]{0:T(1024)} reshape(s32[32,1,1024]{2,1,0:T(2,128)} %transpose.5)
}

%fused_computation.1 (param_0.4: s32[32,1,1024,1], param_1.14: s32[32,1,1024,1], param_2.7: s32[32,1,1024,1]) -> s32[32,1,1024,3] {
  %param_2.7 = s32[32,1,1024,1]{2,0,3,1:T(8,128)} parameter(2)
  %constant.14 = s32[]{:T(256)} constant(-2147483648)
  %pad.6 = s32[32,1,1024,3]{2,0,3,1:T(8,128)} pad(s32[32,1,1024,1]{2,0,3,1:T(8,128)} %param_2.7, s32[]{:T(256)} %constant.14), padding=0_0x0_0x0_0x0_2, metadata={op_type="concatenate" op_name="xla_computation(f2)/jit(f2)/jit(take_along_axis)/concatenate[dimension=3]" source_file="/home/ayaka/main.py" source_line=22}
  %param_1.14 = s32[32,1,1024,1]{2,0,3,1:T(8,128)} parameter(1)
  %pad.5 = s32[32,1,1024,3]{2,0,3,1:T(8,128)} pad(s32[32,1,1024,1]{2,0,3,1:T(8,128)} %param_1.14, s32[]{:T(256)} %constant.14), padding=0_0x0_0x0_0x1_1, metadata={op_type="concatenate" op_name="xla_computation(f2)/jit(f2)/jit(take_along_axis)/concatenate[dimension=3]" source_file="/home/ayaka/main.py" source_line=22}
  %maximum.3 = s32[32,1,1024,3]{2,0,3,1:T(8,128)} maximum(s32[32,1,1024,3]{2,0,3,1:T(8,128)} %pad.6, s32[32,1,1024,3]{2,0,3,1:T(8,128)} %pad.5), metadata={op_type="concatenate" op_name="xla_computation(f2)/jit(f2)/jit(take_along_axis)/concatenate[dimension=3]" source_file="/home/ayaka/main.py" source_line=22}
  %param_0.4 = s32[32,1,1024,1]{2,0,3,1:T(8,128)} parameter(0)
  %pad.4 = s32[32,1,1024,3]{2,0,3,1:T(8,128)} pad(s32[32,1,1024,1]{2,0,3,1:T(8,128)} %param_0.4, s32[]{:T(256)} %constant.14), padding=0_0x0_0x0_0x2_0, metadata={op_type="concatenate" op_name="xla_computation(f2)/jit(f2)/jit(take_along_axis)/concatenate[dimension=3]" source_file="/home/ayaka/main.py" source_line=22}
  ROOT %maximum.2 = s32[32,1,1024,3]{2,0,3,1:T(8,128)} maximum(s32[32,1,1024,3]{2,0,3,1:T(8,128)} %maximum.3, s32[32,1,1024,3]{2,0,3,1:T(8,128)} %pad.4), metadata={op_type="concatenate" op_name="xla_computation(f2)/jit(f2)/jit(take_along_axis)/concatenate[dimension=3]" source_file="/home/ayaka/main.py" source_line=22}
}

%fused_computation.2 (param_0.6: s32[32], param_1.9: s32[32], param_2.6: pred[32]) -> s32[32,1,1024,1] {
  %param_2.6 = pred[32]{0:T(1024)(128)(4,1)} parameter(2)
  %broadcast.29 = pred[32,1,1024,1]{2,0,3,1:T(8,128)(4,1)} broadcast(pred[32]{0:T(1024)(128)(4,1)} %param_2.6), dimensions={0}, metadata={op_type="broadcast_in_dim" op_name="xla_computation(f2)/jit(f2)/jit(take_along_axis)/broadcast_in_dim[shape=(32, 1, 1024) broadcast_dimensions=(0, 1)]" source_file="/home/ayaka/main.py" source_line=22}
  %param_1.9 = s32[32]{0:T(256)} parameter(1)
  %broadcast.27 = s32[32,1,1024,1]{2,0,3,1:T(8,128)} broadcast(s32[32]{0:T(256)} %param_1.9), dimensions={0}, metadata={op_type="broadcast_in_dim" op_name="xla_computation(f2)/jit(f2)/jit(take_along_axis)/broadcast_in_dim[shape=(32, 1, 1024) broadcast_dimensions=(0, 1)]" source_file="/home/ayaka/main.py" source_line=22}
  %param_0.6 = s32[32]{0:T(256)} parameter(0)
  %broadcast.26 = s32[32,1,1024,1]{2,0,3,1:T(8,128)} broadcast(s32[32]{0:T(256)} %param_0.6), dimensions={0}, metadata={op_type="broadcast_in_dim" op_name="xla_computation(f2)/jit(f2)/jit(take_along_axis)/broadcast_in_dim[shape=(32, 1, 1024) broadcast_dimensions=(0, 1)]" source_file="/home/ayaka/main.py" source_line=22}
  ROOT %select.4 = s32[32,1,1024,1]{2,0,3,1:T(8,128)} select(pred[32,1,1024,1]{2,0,3,1:T(8,128)(4,1)} %broadcast.29, s32[32,1,1024,1]{2,0,3,1:T(8,128)} %broadcast.27, s32[32,1,1024,1]{2,0,3,1:T(8,128)} %broadcast.26), metadata={op_type="select_n" op_name="xla_computation(f2)/jit(f2)/jit(take_along_axis)/select_n" source_file="/home/ayaka/main.py" source_line=22}
}

%fused_computation.3 (param_0.7: s32[32]) -> s32[32] {
  %param_0.7 = s32[32]{0:T(256)} parameter(0)
  %constant.16 = s32[]{:T(256)} constant(128), metadata={op_type="add" op_name="xla_computation(f2)/jit(f2)/jit(take_along_axis)/add" source_file="/home/ayaka/main.py" source_line=22}
  %broadcast.30 = s32[32]{0:T(256)} broadcast(s32[]{:T(256)} %constant.16), dimensions={}, metadata={op_type="add" op_name="xla_computation(f2)/jit(f2)/jit(take_along_axis)/add" source_file="/home/ayaka/main.py" source_line=22}
  ROOT %add.4 = s32[32]{0:T(256)} add(s32[32]{0:T(256)} %param_0.7, s32[32]{0:T(256)} %broadcast.30), metadata={op_type="add" op_name="xla_computation(f2)/jit(f2)/jit(take_along_axis)/add" source_file="/home/ayaka/main.py" source_line=22}
}

%fused_computation.4 (param_0.8: s32[32]) -> pred[32] {
  %param_0.8 = s32[32]{0:T(256)} parameter(0)
  %constant.17 = s32[]{:T(256)} constant(0), metadata={op_type="lt" op_name="xla_computation(f2)/jit(f2)/jit(take_along_axis)/lt" source_file="/home/ayaka/main.py" source_line=22}
  %broadcast.31 = s32[32]{0:T(256)} broadcast(s32[]{:T(256)} %constant.17), dimensions={}, metadata={op_type="lt" op_name="xla_computation(f2)/jit(f2)/jit(take_along_axis)/lt" source_file="/home/ayaka/main.py" source_line=22}
  ROOT %compare.4 = pred[32]{0:T(1024)(128)(4,1)} compare(s32[32]{0:T(256)} %param_0.8, s32[32]{0:T(256)} %broadcast.31), direction=LT, metadata={op_type="lt" op_name="xla_computation(f2)/jit(f2)/jit(take_along_axis)/lt" source_file="/home/ayaka/main.py" source_line=22}
}

ENTRY %xla_computation_f2.43 (parameter.1: s32[32,128,1024], parameter.2: s32[32]) -> (s32[32,1024]) {
  %constant.2 = s32[]{:T(256)} constant(0), metadata={op_type="lt" op_name="xla_computation(f2)/jit(f2)/jit(take_along_axis)/lt" source_file="/home/ayaka/main.py" source_line=22}
  %parameter.2 = s32[32]{0:T(256)} parameter(1), metadata={op_type="broadcast_in_dim" op_name="xla_computation(f2)/jit(f2)/jit(take_along_axis)/broadcast_in_dim[shape=(32, 1, 1024) broadcast_dimensions=(0, 1)]" source_file="/home/ayaka/main.py" source_line=22}
  %parameter.1 = s32[32,128,1024]{2,1,0:T(8,128)} parameter(0)
  %fusion.3 = s32[32]{0:T(256)} fusion(s32[32]{0:T(256)} %parameter.2), kind=kLoop, calls=%fused_computation.3, metadata={op_type="add" op_name="xla_computation(f2)/jit(f2)/jit(take_along_axis)/add" source_file="/home/ayaka/main.py" source_line=22}
  %fusion.4 = pred[32]{0:T(1024)(128)(4,1)} fusion(s32[32]{0:T(256)} %parameter.2), kind=kLoop, calls=%fused_computation.4, metadata={op_type="lt" op_name="xla_computation(f2)/jit(f2)/jit(take_along_axis)/lt" source_file="/home/ayaka/main.py" source_line=22}
  %fusion.2 = s32[32,1,1024,1]{2,0,3,1:T(8,128)} fusion(s32[32]{0:T(256)} %parameter.2, s32[32]{0:T(256)} %fusion.3, pred[32]{0:T(1024)(128)(4,1)} %fusion.4), kind=kLoop, calls=%fused_computation.2, metadata={op_type="select_n" op_name="xla_computation(f2)/jit(f2)/jit(take_along_axis)/select_n" source_file="/home/ayaka/main.py" source_line=22}
  %iota.5 = s32[32,1,1024,1]{2,0,3,1:T(8,128)} iota(), iota_dimension=2, metadata={op_type="broadcast_in_dim" op_name="xla_computation(f2)/jit(f2)/jit(take_along_axis)/broadcast_in_dim[shape=(32, 1, 1024, 1) broadcast_dimensions=(2,)]" source_file="/home/ayaka/main.py" source_line=22}
  %iota.4 = s32[32,1,1024,1]{2,0,3,1:T(8,128)} iota(), iota_dimension=0, metadata={op_type="broadcast_in_dim" op_name="xla_computation(f2)/jit(f2)/jit(take_along_axis)/broadcast_in_dim[shape=(32, 1, 1024, 1) broadcast_dimensions=(0,)]" source_file="/home/ayaka/main.py" source_line=22}
  %fusion.1 = s32[32,1,1024,3]{2,0,3,1:T(8,128)} fusion(s32[32,1,1024,1]{2,0,3,1:T(8,128)} %iota.5, s32[32,1,1024,1]{2,0,3,1:T(8,128)} %fusion.2, s32[32,1,1024,1]{2,0,3,1:T(8,128)} %iota.4), kind=kLoop, calls=%fused_computation.1, metadata={op_type="concatenate" op_name="xla_computation(f2)/jit(f2)/jit(take_along_axis)/concatenate[dimension=3]" source_file="/home/ayaka/main.py" source_line=22}
  %copy.1 = s32[32,1,1024,3]{3,2,0,1:T(8,128)} copy(s32[32,1,1024,3]{2,0,3,1:T(8,128)} %fusion.1)
  %pad.clone = s32[32,1,1024,4]{3,2,0,1:T(8,128)} pad(s32[32,1,1024,3]{3,2,0,1:T(8,128)} %copy.1, s32[]{:T(256)} %constant.2), padding=0_0x0_0x0_0x0_1
  %reshape.12 = s32[131072]{0:T(1024)} reshape(s32[32,1,1024,4]{3,2,0,1:T(8,128)} %pad.clone)
  %fusion = s32[32768]{0:T(1024)} fusion(s32[32,128,1024]{2,1,0:T(8,128)} %parameter.1, s32[131072]{0:T(1024)} %reshape.12), kind=kCustom, calls=%fused_computation, metadata={op_type="gather" op_name="xla_computation(f2)/jit(f2)/jit(take_along_axis)/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0, 1, 2), start_index_map=(0, 1, 2)) slice_sizes=(1, 1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/ayaka/main.py" source_line=22}
  %reshape.7 = s32[32,1024]{1,0:T(8,128)} reshape(s32[32768]{0:T(1024)} %fusion), metadata={op_type="gather" op_name="xla_computation(f2)/jit(f2)/jit(take_along_axis)/gather[dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0, 1, 2), start_index_map=(0, 1, 2)) slice_sizes=(1, 1, 1) unique_indices=False indices_are_sorted=False mode=GatherScatterMode.PROMISE_IN_BOUNDS fill_value=None]" source_file="/home/ayaka/main.py" source_line=22}
  ROOT %tuple.42 = (s32[32,1024]{1,0:T(8,128)}) tuple(s32[32,1024]{1,0:T(8,128)} %reshape.7)
}
```


Discussion:
Thanks for the report. I agree with your conclusions. I filed a bug with the TPU compiler team (Google bug b/229234664).

The plot thickens.

On reflection this is sort of an XLA bug and sort of not an XLA bug. It's an XLA bug in the sense that XLA perhaps should be able to optimize the HLO we are generating, but JAX perhaps could generate more optimal HLO. In particular, we broadcast the input shapes in the implementation of `jnp.take_along_axis` and that does bad things to performance. Maybe we can avoid that broadcasting on the JAX side.

However, we added that broadcasting to fix #[1521](https://github.com/google/jax/issues/1521). So to remove the broadcasting I'll need to remember what went wrong for that issue.

Resolving patch:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -3425,6 +3425,10 @@ def _normalize_index(index, axis_size):
 @partial(jit, static_argnames=('axis',))
 def take_along_axis(arr, indices, axis: Optional[int]):
   _check_arraylike("take_along_axis", arr, indices)
+  index_dtype = dtypes.dtype(indices)
+  if not dtypes.issubdtype(index_dtype, integer):
+    raise TypeError("take_along_axis indices must be of integer type, got "
+                    f"{str(index_dtype)}")
   if axis is None:
     if ndim(indices) != 1:
       msg = "take_along_axis indices must be 1D if axis=None, got shape {}"
@@ -3442,13 +3446,9 @@ def replace(tup, val):
     return tuple(lst)
 
   use_64bit_index = _any([not core.is_constant_dim(d) or d >= (1 << 31) for d in arr.shape])
-  index_dtype = int64 if use_64bit_index else int32
+  index_dtype = dtype(int64 if use_64bit_index else int32)
   indices = lax.convert_element_type(indices, index_dtype)
 
-  bcast_shape = lax.broadcast_shapes(replace(arr.shape, 1), replace(indices.shape, 1))
-  indices = broadcast_to(indices, replace(bcast_shape, indices.shape[axis]))
-  arr     = broadcast_to(arr,     replace(bcast_shape, arr.shape[axis]))
-
   axis_size = arr.shape[axis]
   arr_shape = replace(arr.shape, 1)
   idx_shape = indices.shape
@@ -3471,26 +3471,38 @@ def replace(tup, val):
       start_index_map.append(i)
       collapsed_slice_dims.append(i)
       j += 1
-    elif not core.symbolic_equal_dim(idx_shape[i], 1):
+    elif core.symbolic_equal_dim(idx_shape[i], 1):
+      # If idx_shape[i] == 1, we can just take the entirety of the arr's axis
+      # and avoid forming an iota index.
+      offset_dims.append(i)
+      slice_sizes.append(arr_shape[i])
+    elif core.symbolic_equal_dim(arr_shape[i], 1):
+      # If the array dimension is 1 but the index dimension is not, we
+      # broadcast the array dimension to the index dimension by repeatedly
+      # gathering the first element.
+      gather_indices.append(zeros(gather_index_shape, dtype=index_dtype))
+      slice_sizes.append(1)
+      start_index_map.append(i)
+      collapsed_slice_dims.append(i)
+      j += 1
+    else:
+      # Otherwise, idx_shape[i] == arr_shape[i]. Use an iota index so
+      # corresponding elements of array and index are gathered.
       # TODO(mattjj): next line needs updating for dynamic shapes
-      iota = lax.iota(_dtype(indices), out_shape[i])  # type: ignore
-      iota = lax.broadcast_in_dim(iota, gather_index_shape, (j,))
+      iota = lax.broadcasted_iota(index_dtype, gather_index_shape, j)
       gather_indices.append(iota)
       slice_sizes.append(1)
       start_index_map.append(i)
       collapsed_slice_dims.append(i)
       j += 1
-    else:
-      # If idx_shape[i] == 1, we can just take the entirety of the arr's axis
-      # and avoid forming an iota index.
-      offset_dims.append(i)
-      slice_sizes.append(arr_shape[i])
+
 
   gather_indices = lax.concatenate(gather_indices, dimension=j)
   dnums = lax.GatherDimensionNumbers(
     offset_dims=tuple(offset_dims),
     collapsed_slice_dims=tuple(collapsed_slice_dims),
     start_index_map=tuple(start_index_map))
+  # TODO(phawkins): change the mode to "fill".
   return lax.gather(arr, gather_indices, dnums, tuple(slice_sizes))
 
 ### Indexing


Test patch: