flatten

dragon.vm.torch.flatten(
  input,
  start_dim=0,
  end_dim=- 1,
  out=None
)[source]

Return a tensor with dimensions flattened.

The argument start_dim and end_dim could be negative:

x = torch.tensor([[1, 2, 3], [4, 5, 6]])

# A negative axis is the last-k axis
print(torch.flatten(x, start_dim=0, end_dim=-1))
print(torch.flatten(x, start_dim=0, end_dim=1))  # Equivalent
Parameters:
  • input (torch.Tensor) – The input tensor.
  • start_dim (int, optional, default=0) – The start dimension to flatten.
  • end_dim (int, optional, default=-1) – The end dimension to flatten.
  • out (dragon.vm.torch.Tensor, optional) – The optional output tensor.
Returns:

dragon.vm.torch.Tensor – The output tensor.