import torch


def linspace(start, stop, num):
  """Linearly space values between multi-dimensional endpoints in PyTorch.

  Creates a tensor of shape [num, *start.shape] whose values are evenly spaced
  from start to end, inclusive. Replicates the multi-dimensional bahavior of
  numpy.linspace in PyTorch. Adds steps as a new, final dimension.
  Based on https://github.com/pytorch/pytorch/issues/61292.

  Args:
    start: Beginning values of range.
    stop: Ending values of range.
    num: Number of points in range, including endpoints.

  Returns:
    linspace
  """
  # create a tensor of "num" steps from 0 to 1.
  steps = torch.arange(num, dtype=start.dtype, device=start.device) / (num - 1.)
  # the output starts at "start" and increments until "stop" in each dimension,
  # adding a final dim.
  return start.unsqueeze(-1) + steps * (stop - start).unsqueeze(-1)