New indexable TensorizedTensor

[JeanKossaifi]

This creates a new class, TensorizedTensor, which replaces and generalizes TensorizedMatrix. It can represent arbitrary tensorized tensors (including vectors, matrices, batched matrices, etc) and supports generalized indexing.

The proper generalized indexing for Tucker still remains to be done. I am thinking of merging these "tensorizedTensors" with regular factorized tensors and just giving the option of having either regular tensors or tensorized ones.

The idea is to have a tensor_shape.

For a regular tensor, that's just a tuple of ints, e.g. (2, 3, 4). To create a tensorized_tensor, one simply passes a nested tuple: (2, (2, 3, 4), (3, 4, 5)).

When indexing a tensorized_tensor, I iterate through this tensor shape and:

1. if the current element is an int, I index it (in the case of block decomposition it represents a batched dimension!)
2. if the current element is a shape, we're dealing with a tensorized dimension, I convert the index into indices for each of the tensorized dims.

I've generalized the concept of TTMatrix to any orders and consequently renamed it BlockTT, to differentiate it from TT. In a block TT, the above shape (2, (2, 3, 4), (3, 4, 5)) will result in a batch size of 2 and two dimensions, tensorized respectively to (2, 3, 4) and (3, 4, 5).

[JeanKossaifi]

Happy for feedback! @colehawkins, @merajhashemi

[colehawkins]

This is a really nice and powerful generalization of TTMatrix. Very cool, and very cool that it also supports block-term naturally.

A quibble, but I don't think that TTTensorized needs the separate naming convention. The list CPTensorized,BlockTT,TuckerTensorized seems to imply that CPTensorized and TuckerTensorized are not block decompositions. I think unifying these as *Tensorized or Block* would be more natural.

[JeanKossaifi]

CPTensorized is not a block decomposition though, it just splits a single mode into several factors. In fact we are working on adding a BlockCP decomposition as well, which represents parts of multiple modes in each factor.

I was hoping to unify TT and Block TT but I am not sure it's easily doable: in the case of BlockTT, any non-tensorized dimension becomes a batch dimension, while in regular TT, each mode is represented in a separate core. Merging the two would require an intricate representation and API.

[colehawkins]

That makes sense, thanks for clearing up my misunderstanding!

[JeanKossaifi]

Great, merging then!