Support for tensor outputs in TensorNet

[RaulPPelaez]

This PR makes it possible to get tensor outputs from TensorNet.

Also, it implements some tests to ensure equivariance.

TODO:

• [ ] Write a somewhat general tensor output model