Multiple indexing vs Slicing

[kmg3821]

Hello, I am quite noob to JAX.

I am trying to implement Alphazero by myself. So I referred to the "search" part of MCTX code from Deepmind.
(link -> https://github.com/google-deepmind/mctx/blob/main/mctx/_src/search.py)

They use batch_range = jnp.arange(batch_size) for indexing batch dimension like tree.node_visits[batch_range, node_index] or tree.children_index[batch_range, parent_index, action].

I wonder why they don't use tree.node_visits[:, node_index] or tree.children_index[:, parent_index, action].
Is multiple indexing faster than slicing? Then, why?

In short, for x = jnp.ones((10,)),
what is the difference between arr[jnp.arange(10)] and arr[:]?

[jakevdp]

In the simplest cases, slicing will lower to a slice HLO, while multiple indexing will lower to a gather HLO. You can see this by using jax.make_jaxpr to view the jaxpr representation of the operation:

import jax
import jax.numpy as jnp

x = jnp.ones((10,))
i = jnp.arange(5)

def f1(x):
  return x[:5]
print(jax.make_jaxpr(f1)(x))

{ lambda ; a:f32[10]. let
    b:f32[5] = slice[limit_indices=(5,) start_indices=(0,) strides=None] a
  in (b,) }

def f2(x):
  return x[jnp.arange(5)]
print(jax.make_jaxpr(f2)(x))

{ lambda ; a:f32[10]. let
    b:i32[5] = iota[dimension=0 dtype=int32 shape=(5,) sharding=None] 
    c:bool[5] = lt b 0
    d:i32[5] = add b 10
    e:i32[5] = select_n c b d
    f:i32[5,1] = broadcast_in_dim[
      broadcast_dimensions=(0,)
      shape=(5, 1)
      sharding=None
    ] e
    g:f32[5] = gather[
      dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0,), start_index_map=(0,), operand_batching_dims=(), start_indices_batching_dims=())
      fill_value=None
      indices_are_sorted=False
      mode=GatherScatterMode.PROMISE_IN_BOUNDS
      slice_sizes=(1,)
      unique_indices=False
    ] a f
  in (g,) }

In general, I'd expect slicing to be faster than multiple indexing because it gives the compiler better constraints about the operation: i.e. a slice is always a contiguous chunk of the array, while a sequence of indices in general requires a random access pattern.