import torch


def extract_into_tensor(arr, timesteps, broadcast_shape):
  """Extract values from a 1-D numpy array for a batch of indices.

  Args:
    arr: the 1-D numpy array.
    timesteps: a tensor of indices into the array to extract.
    broadcast_shape: a larger shape of K dimensions with the batch dimension
      equal to the length of timesteps.

  Returns:
    res: a tensor of shape [batch_size, 1, ...] where the shape has K dims.
  """
  device = timesteps.device
  assert arr.device == device
  res = arr[timesteps].float()
  new_dims = [1] * (len(broadcast_shape) - res.ndim)
  res = res.view(*res.shape, *new_dims)
  return torch.broadcast_to(res, broadcast_shape)