Confusing Note on dense and flatten?

I was a bit confused whether nn.Dense support multi-dimensional inputs through flatten=False?

The doc says:

Note: the input must be a tensor with rank 2. Use flatten to convert it to rank 2 manually if necessary.


flatten (bool) – Whether the input tensor should be flattened. If true, all but the first axis of input data are collapsed together. If false, all but the last axis of input data are kept the same, and the transformation applies on the last axis.

In my experience:

I always use dense with flatten just to eliminate potential issues. I have never tested the flatten=False option to my satisfaction, but I may just lack trying.


‘flatten=False’ works exactly like keras.layers.TimeDistributed(). It can be useful you need to generate sequence results keeping the independence of each sequence.

Another common use case is to create positionwise fully connected layer. Here is an example. Check this paper to find out how it is used.

Thanks @indu @gogamza for the convincing answers!

I guess the “Note” on the doc can be removed then?