code: def batched_index(
      values: Array, indices: Array, keepdims: bool = False
  ) -> Array:
      """Index into the last dimension of a tensor, preserving all others dims.
  
      Args:
        values: a tensor of shape [..., D],
        indices: indices of shape [...].
        keepdims: whether to keep the final dimension.
  
      Returns:
        a tensor of shape [...] or [..., 1].
      """
      indexed = jnp.take_along_axis(values, indices[..., None], axis=-1)
      if not keepdims:
          indexed = jnp.squeeze(indexed, axis=-1)
      return indexed
  
file_path: rlax/rlax/_src/base.py
imports: from typing import Optional, Sequence, Union
  import chex
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: batched_index(values, indices, keepdims) and returns: indexed