code: def expand_index_like(
      index: torch.Tensor, tokens: torch.Tensor
  ) -> torch.Tensor:
      """Expands the index along the last dimension of the input tokens.
  
      Args:
          index:
              Index tensor with shape (batch_size, idx_length) where each entry is
              an index in [0, sequence_length).
          tokens:
              Tokens tensor with shape (batch_size, sequence_length, dim).
  
      Returns:
          Index tensor with shape (batch_size, idx_length, dim) where the original
          indices are repeated dim times along the last dimension.
  
      """
      dim = tokens.shape[-1]
      index = index.unsqueeze(-1).expand(-1, -1, dim)
      return index
  
imports: import math
  import warnings
  from typing import Iterable, List, Optional, Tuple, Union
  import numpy
  import torch
  import torch.distributed
  import torch.nn
  from numpy.typing import NDArray
  from torch.nn import Module
  from torch.nn.modules.batchnorm import _BatchNorm
  from torch.nn.parameter import Parameter
label_desc: Use image node_id 12 for calling a global function with example usage: expand_index_like(index, tokens) and returns: index