DOC: improve fact linear + add fact. embedding

Author: JeanKossaifi

doc/user_guide/factorized_embeddings.rst

@@ -0,0 +1,20 @@
+Factorized embedding layers
+===========================
+
+In TensorLy-Torch, we also provide out-of-the-box tensorized embedding layers. 
+
+Just as for the case of factorized linear, you can either create a factorized embedding from scratch, here automatically determine the 
+input and output tensorized shapes, to have 3 dimensions each:
+
+.. code-block:: python
+
+    import tltorch
+    import torch
+
+    from_embedding = tltorch.FactorizedEmbedding(num_embeddings, embedding_dim, auto_reshape=True, d=3, rank=0.4)
+
+
+Or, you can create it by decomposing an existing embedding layer:
+
+    from_embedding = tltorch.FactorizedEmbedding.from_embedding(embedding_layer, auto_reshape=True,
+                factorization='blocktt', n_tensorized_modes=3, rank=0.4)

doc/user_guide/tensorized_linear.rst

@@ -24,22 +24,39 @@ Now, imagine you already have a linear layer:
 
     linear = torch.nn.Linear(in_features=16, 10)
 
-You can easily compress it into a tensorized linear layer:
+You can easily compress it into a tensorized linear layer: here we specify the shape to which to tensorize the weights,
+and use `rank=0.5`, which means automatically determine the rank so that the factorization uses approximately half the
+number of parameters.
 
 .. code-block:: python
 
-    fact_linear = tltorch.FactorizedLinear.from_linear(linear, (4, 4), (2, 5), rank=0.5)
+    fact_linear = tltorch.FactorizedLinear.from_linear(linear, auto_tensorize=False,
+                        in_tensorized_features=(4, 4), out_tensorized_features=(2, 5), rank=0.5)
 
 
+The tensorized weights will have the following shape: 
+
 .. parsed-literal::
 
     torch.Size([4, 4, 2, 5])
 
 
+Note that you can also let TensorLy-Torch automatically determine the tensorization shape. In this case we just instruct it to 
+find ``in_tensorized_features`` and ``out_tensorized_features`` to have length `2`:
+
+.. code-block:: python
+
+    fact_linear = tltorch.FactorizedLinear.from_linear(linear, auto_tensorize=True, n_tensorized_modes=2, rank=0.5)
+
+
 You can also create tensorized layers from scratch:
 
 .. code-block:: python
 
     fact_linear = tltorch.FactorizedLinear(in_tensorized_features=(4, 4), 
                                            out_tensorized_features=(2, 5), 
                                            factorization='tucker', rank=0.5)
+
+Finally, during the forward pass, you can reconstruct the full weights (``implementation='reconstructed'``) and perform a regular linear layer forward pass. 
+ALternatively, you can let TensorLy-Torch automatically direction contract the input tensor with the *factors of the decomposition*  (``implementation='factorized'``),
+ which can be faster, particularly if you have a very small rank, e.g. very small factorization factors.