oneflow.flatten

oneflow.flatten()

Flattens a contiguous range of dims into a tensor.

Parameters
  • start_dim – first dim to flatten (default = 0).

  • end_dim – last dim to flatten (default = -1).

For example:

>>> import numpy as np
>>> import oneflow as flow
>>> input = flow.randn(32, 1, 5, 5)
>>> output = flow.flatten(input, start_dim=1)
>>> output.shape
oneflow.Size([32, 25])