Repository: tensorly/tensorly

Issue:
weights in parafac decomposition
In parafac decomposition, the docs mention a "weights" parameter being returned as given below:

> weights : ndarray, optional
> Array of length rank of weights for each factor matrix. See the with_weights keyword attribute.

But they're not returned. Nor does the source have any mention of it. 

Discussion:

Patch:
diff --git a/tensorly/__init__.py b/tensorly/__init__.py
index 1f09dae..bdb0acb 100644
--- a/tensorly/__init__.py
+++ b/tensorly/__init__.py
@@ -6,7 +6,8 @@ from .base import tensor_to_vec, vec_to_tensor
 from .base import partial_unfold, partial_fold
 from .base import partial_tensor_to_vec, partial_vec_to_tensor
 
-from .kruskal_tensor import kruskal_to_tensor, kruskal_to_unfolded, kruskal_to_vec
+from .kruskal_tensor import (kruskal_to_tensor, kruskal_to_unfolded,
+                             kruskal_to_vec, unfolding_dot_khatri_rao)
 from .tucker_tensor import tucker_to_tensor, tucker_to_unfolded, tucker_to_vec
 from .mps_tensor import mps_to_tensor, mps_to_unfolded, mps_to_vec
 
diff --git a/tensorly/backend/core.py b/tensorly/backend/core.py
index 4e8247f..5b91c76 100644
--- a/tensorly/backend/core.py
+++ b/tensorly/backend/core.py
@@ -568,7 +568,7 @@ class Backend(object):
         b = self.reshape(b, (1, s3, 1, s4))
         return self.reshape(a * b, (s1 * s3, s2 * s4))
 
-    def kr(self, matrices, mask=None):
+    def kr(self, matrices, weights=None, mask=None):
         """Khatri-Rao product of a list of matrices
 
         This can be seen as a column-wise kronecker product.
@@ -603,7 +603,10 @@ class Backend(object):
         n_col = self.shape(matrices[0])[1]
         for i, e in enumerate(matrices[1:]):
             if not i:
-                res = matrices[0]
+                if weights is None:
+                    res = matrices[0]
+                else:
+                    res = matrices[0]*self.reshape(weights, (1, -1))
             s1, s2 = self.shape(res)
             s3, s4 = self.shape(e)
             if not s2 == s4 == n_col:
@@ -661,7 +664,7 @@ class Backend(object):
 
         if n_eigenvecs >= min_dim:
             if n_eigenvecs > max_dim:
-                warnings.warn(('Trying to compute SVD with n_eigenvecs={1}, which '
+                warnings.warn(('Trying to compute SVD with n_eigenvecs={0}, which '
                                'is larger than max(matrix.shape)={1}. Setting '
                                'n_eigenvecs to {1}').format(n_eigenvecs, max_dim))
                 n_eigenvecs = max_dim
diff --git a/tensorly/backend/numpy_backend.py b/tensorly/backend/numpy_backend.py
index d99bb63..1d7f49e 100644
--- a/tensorly/backend/numpy_backend.py
+++ b/tensorly/backend/numpy_backend.py
@@ -52,8 +52,7 @@ class NumpyBackend(Backend):
         else:
             return np.sum(np.abs(tensor)**order, axis=axis)**(1 / order)
 
-    @staticmethod
-    def kr(matrices, mask=None):
+    def kr(self, matrices, weights=None, mask=None):
         if mask is None: mask = 1
         n_columns = matrices[0].shape[1]
         n_factors = len(matrices)
@@ -63,6 +62,10 @@ class NumpyBackend(Backend):
         target = ''.join(chr(start + i) for i in range(n_factors))
         source = ','.join(i + common_dim for i in target)
         operation = source + '->' + target + common_dim
+
+        if weights is not None:
+            matrices = [m if i else m*self.reshape(weights, (1, -1)) for i, m in enumerate(matrices)]
+
         return np.einsum(operation, *matrices).reshape((-1, n_columns))*mask
 
     @property
diff --git a/tensorly/contrib/sparse/kruskal_tensor.py b/tensorly/contrib/sparse/kruskal_tensor.py
index 63a0f1c..7a85eb8 100644
--- a/tensorly/contrib/sparse/kruskal_tensor.py
+++ b/tensorly/contrib/sparse/kruskal_tensor.py
@@ -1,4 +1,5 @@
-from ...kruskal_tensor import kruskal_to_tensor
+from ...kruskal_tensor import kruskal_to_tensor,unfolding_dot_khatri_rao
 from .core import wrap
 
 kruskal_to_tensor = wrap(kruskal_to_tensor)
+unfolding_dot_khatri_rao = wrap(unfolding_dot_khatri_rao)
diff --git a/tensorly/contrib/sparse/tenalg.py b/tensorly/contrib/sparse/tenalg.py
index 32f43eb..83136f2 100644
--- a/tensorly/contrib/sparse/tenalg.py
+++ b/tensorly/contrib/sparse/tenalg.py
@@ -1,6 +1,5 @@
 from ...tenalg import (mode_dot, multi_mode_dot, 
-                       kronecker, khatri_rao, inner,
-                       unfolding_dot_khatri_rao)
+                       kronecker, khatri_rao, inner)
 from .core import wrap
 
 
@@ -9,4 +8,3 @@ multi_mode_dot = wrap(multi_mode_dot)
 kronecker = wrap(kronecker)
 khatri_rao = wrap(khatri_rao)
 inner = wrap(inner)
-unfolding_dot_khatri_rao = wrap(unfolding_dot_khatri_rao)
diff --git a/tensorly/decomposition/candecomp_parafac.py b/tensorly/decomposition/candecomp_parafac.py
index 6e3b2a4..4ee0a91 100755
--- a/tensorly/decomposition/candecomp_parafac.py
+++ b/tensorly/decomposition/candecomp_parafac.py
@@ -4,6 +4,9 @@ import warnings
 import tensorly as tl
 from ..random import check_random_state
 from ..base import unfold
+from ..kruskal_tensor import (kruskal_to_tensor, KruskalTensor,
+                              unfolding_dot_khatri_rao)
+from ..tenalg import khatri_rao
 
 # Authors: Jean Kossaifi <jean.kossaifi+tensors@gmail.com>
 #          Chris Swierczewski <csw@amazon.com>
@@ -12,48 +15,6 @@ from ..base import unfold
 
 # License: BSD 3 clause
 
-
-def normalize_factors(factors):
-    """Normalizes factors to unit length and returns factor magnitudes
-
-    Turns ``factors = [|U_1, ... U_n|]`` into ``[weights; |V_1, ... V_n|]``,
-    where the columns of each `V_k` are normalized to unit Euclidean length
-    from the columns of `U_k` with the normalizing constants absorbed into
-    `weights`. In the special case of a symmetric tensor, `weights` holds the
-    eigenvalues of the tensor.
-
-    Parameters
-    ----------
-    factors : ndarray list
-        list of matrices, all with the same number of columns
-        i.e.::
-            for u in U:
-                u[i].shape == (s_i, R)
-
-        where `R` is fixed while `s_i` can vary with `i`
-
-    Returns
-    -------
-    normalized_factors : list of ndarrays
-        list of matrices with the same shape as `factors`
-    weights : ndarray
-        vector of length `R` holding normalizing constants
-
-    """
-    # allocate variables for weights, and normalized factors
-    rank = factors[0].shape[1]
-    weights = tl.ones(rank, **tl.context(factors[0]))
-    normalized_factors = []
-
-    # normalize columns of factor matrices
-    for factor in factors:
-        scales = tl.norm(factor, axis=0)
-        weights *= scales
-        scales_non_zero = tl.where(scales==0, tl.ones(tl.shape(scales), **tl.context(factors[0])), scales)
-        normalized_factors.append(factor/scales_non_zero)
-    return normalized_factors, weights
-
-
 def initialize_factors(tensor, rank, init='svd', svd='numpy_svd', random_state=None, non_negative=False):
     r"""Initialize factors used in `parafac`.
 
@@ -115,15 +76,13 @@ def initialize_factors(tensor, rank, init='svd', svd='numpy_svd', random_state=N
     raise ValueError('Initialization method "{}" not recognized'.format(init))
 
 
-def parafac(tensor, rank, n_iter_max=100, init='svd', svd='numpy_svd', tol=1e-8,
-            orthogonalise=False, random_state=None, verbose=False,
-            return_errors=False, non_negative=False, mask=None):
-    """
-    CANDECOMP/PARAFAC decomposition via alternating least squares (ALS)
-
+def parafac(tensor, rank, n_iter_max=100, init='svd', svd='numpy_svd', normalize_factors=False,
+            tol=1e-8, orthogonalise=False, random_state=None, verbose=False, return_errors=False,
+            non_negative=False, mask=None):
+    """CANDECOMP/PARAFAC decomposition via alternating least squares (ALS)
     Computes a rank-`rank` decomposition of `tensor` [1]_ such that,
 
-        ``tensor = [| factors[0], ..., factors[-1] |]``.
+        ``tensor = [|weights; factors[0], ..., factors[-1] |]``.
 
     Parameters
     ----------
@@ -136,6 +95,8 @@ def parafac(tensor, rank, n_iter_max=100, init='svd', svd='numpy_svd', tol=1e-8,
         Type of factor matrix initialization. See `initialize_factors`.
     svd : str, default is 'numpy_svd'
         function to use to compute the SVD, acceptable values in tensorly.SVD_FUNS
+    normalize_factors : if True, aggregate the weights of each factor in a 1D-tensor
+        of shape (rank, ), which will contain the norms of the factors
     tol : float, optional
         (Default: 1e-6) Relative reconstruction error tolerance. The
         algorithm is considered to have found the global minimum when the
@@ -156,9 +117,13 @@ def parafac(tensor, rank, n_iter_max=100, init='svd', svd='numpy_svd', tol=1e-8,
 
     Returns
     -------
-    factors : ndarray list
-        List of factors of the CP decomposition element `i` is of shape
-        (tensor.shape[i], rank)
+    KruskalTensor : (weight, factors)
+        * weights : 1D array of shape (rank, )
+            all ones if normalize_factors is False (default), 
+            weights of the (normalized) factors otherwise
+        * factors : List of factors of the CP decomposition element `i` is of shape
+            (tensor.shape[i], rank)
+
     errors : list
         A list of reconstruction errors at each iteration of the algorithms.
 
@@ -182,6 +147,7 @@ def parafac(tensor, rank, n_iter_max=100, init='svd', svd='numpy_svd', tol=1e-8,
                                  non_negative=non_negative)
     rec_errors = []
     norm_tensor = tl.norm(tensor, 2)
+    weights = tl.ones(rank, **tl.context(tensor))
 
     for iteration in range(n_iter_max):
         if orthogonalise and iteration <= orthogonalise:
@@ -211,7 +177,7 @@ def parafac(tensor, rank, n_iter_max=100, init='svd', svd='numpy_svd', tol=1e-8,
             if mask is not None:
                 tensor = tensor*mask + tl.kruskal_to_tensor(factors, mask=1-mask)
 
-            mttkrp = tl.tenalg.unfolding_dot_khatri_rao(tensor, factors, mode)
+            mttkrp = unfolding_dot_khatri_rao(tensor, (weights, factors), mode)
 
             if non_negative:
                 numerator = tl.clip(mttkrp, a_min=epsilon, a_max=None)
@@ -220,6 +186,14 @@ def parafac(tensor, rank, n_iter_max=100, init='svd', svd='numpy_svd', tol=1e-8,
                 factor = factors[mode] * numerator / denominator
             else:
                 factor = tl.transpose(tl.solve(tl.conj(tl.transpose(pseudo_inverse)), tl.transpose(mttkrp)))
+                
+            if normalize_factors:
+                factor_norm = tl.norm(factor, axis=0)
+                weights *= factor_norm
+                positive_factor_norms = tl.where(factor_norm==0, 
+                                                 tl.ones(tl.shape(factor_norm), **tl.context(factors[0])),
+                                                 factor_norm)
+                factor = factor/positive_factor_norms
 
             factors[mode] = factor
 
@@ -241,15 +215,17 @@ def parafac(tensor, rank, n_iter_max=100, init='svd', svd='numpy_svd', tol=1e-8,
                 if tol and abs(rec_errors[-2] - rec_errors[-1]) < tol:
                     if verbose:
                         print('converged in {} iterations.'.format(iteration))
-                    break
+                    break       
             else:
                 if verbose:
                     print('reconstruction error={}'.format(rec_errors[-1]))
 
+    kruskal_tensor = KruskalTensor((weights, factors))
+
     if return_errors:
-        return factors, rec_errors
+        return kruskal_tensor, rec_errors
     else:
-        return factors
+        return kruskal_tensor
 
 
 def non_negative_parafac(tensor, rank, n_iter_max=100, init='svd', svd='numpy_svd',
@@ -411,6 +387,7 @@ def randomised_parafac(tensor, rank, n_samples, n_iter_max=100, init='random', s
     norm_tensor = tl.norm(tensor, 2)
     min_error = 0
 
+    weights = tl.ones(rank, **tl.context(tensor))
     for iteration in range(n_iter_max):
         for mode in range(n_dims):
             kr_prod, indices_list = sample_khatri_rao(factors, n_samples, skip_matrix=mode, random_state=rng)
@@ -430,7 +407,7 @@ def randomised_parafac(tensor, rank, n_samples, n_iter_max=100, init='random', s
             factors[mode] = factor
 
         if max_stagnation or tol:
-            rec_error = tl.norm(tensor - tl.kruskal_to_tensor(factors), 2) / norm_tensor
+            rec_error = tl.norm(tensor - kruskal_to_tensor((weights, factors)), 2) / norm_tensor
             if not min_error or rec_error < min_error:
                 min_error = rec_error
                 stagnation = -1
@@ -449,4 +426,4 @@ def randomised_parafac(tensor, rank, n_samples, n_iter_max=100, init='random', s
                         print('converged in {} iterations.'.format(iteration))
                     break
 
-    return factors
+    return KruskalTensor((weights, factors))
diff --git a/tensorly/kruskal_tensor.py b/tensorly/kruskal_tensor.py
index 7647116..cdd47af 100644
--- a/tensorly/kruskal_tensor.py
+++ b/tensorly/kruskal_tensor.py
@@ -4,16 +4,147 @@ Core operations on Kruskal tensors.
 
 from . import backend as T
 from .base import fold, tensor_to_vec
-from .tenalg import khatri_rao
+from .tenalg import khatri_rao, multi_mode_dot
+
+import warnings
+from collections.abc import Mapping
 
 # Author: Jean Kossaifi
 
 # License: BSD 3 clause
 
+class KruskalTensor(Mapping):
+    def __init__(self, kruskal_tensor):
+        super().__init__()
+
+        shape, rank = _validate_kruskal_tensor(kruskal_tensor)
+        weights, factors = kruskal_tensor
+
+        # Should we allow None weights?
+        if weights is None:
+            weights = T.ones(rank, **T.context(factors[0]))
+
+        self.shape = shape
+        self.rank = rank
+        self.factors = factors
+        self.weights = weights
 
-def kruskal_to_tensor(factors, weights=None, mask=None):
+    
+    def __getitem__(self, index):
+        if index == 0:
+            return self.weights
+        elif index == 1:
+            return self.factors
+        else: 
+            raise IndexError('You tried to access index {} of a Kruskal tensor.\n'
+                             'You can only access index 0 and 1 of a Kruskal tensor'
+                             '(corresponding respectively to the weights and factors)'.format(index))
+    
+    def __iter__(self):
+        yield self.weights
+        yield self.factors
+        
+    def __len__(self):
+        return 2
+    
+    def __repr__(self):
+        message = '(weights, factors) : rank-{} KruskalTensor of shape {} '.format(self.rank, self.shape)
+        return message
+
+
+def _validate_kruskal_tensor(kruskal_tensor):
+    """Validates a kruskal_tensor in the form (weights, factors)
+    
+        Returns the rank and shape of the validated tensor
+    
+    Parameters
+    ----------
+    kruskal_tensor : KruskalTensor or (weights, factors)
+    
+    Returns
+    -------
+    (shape, rank) : (int tuple, int)
+        size of the full tensor and rank of the Kruskal tensor
     """
-    Turns the Khatri-product of matrices into a full tensor
+    if isinstance(kruskal_tensor, KruskalTensor):
+        # it's already been validated at creation
+        return kruskal_tensor.shape, kruskal_tensor.rank
+
+    weights, factors = kruskal_tensor
+            
+    if len(factors) < 2:
+        raise ValueError('A Kruskal tensor should be composed of at least two factors.'
+                         'However, {} factor was given.'.format(len(factors)))
+
+    rank = int(T.shape(factors[0])[1])
+    shape = []
+    for i, factor in enumerate(factors):
+        current_mode_size, current_rank = T.shape(factor)
+        if current_rank != rank:
+            raise ValueError('All the factors of a Kruskal tensor should have the same number of column.'
+                             'However, factors[0].shape[1]={} but factors[{}].shape[1]={}.'.format(
+                                 rank, i, T.shape(factor)[1]))
+        shape.append(current_mode_size)
+
+    if weights is not None and len(weights) != rank:
+        raise ValueError('Given factors for a rank-{} Kruskal tensor but len(weights)={}.'.format(
+            rank, len(weights)))
+        
+    return tuple(shape), rank
+
+
+def kruskal_normalise(kruskal_tensor, copy=False):
+    """Returns kruskal_tensor with factors normalised to unit length
+
+    Turns ``factors = [|U_1, ... U_n|]`` into ``[weights; |V_1, ... V_n|]``,
+    where the columns of each `V_k` are normalized to unit Euclidean length
+    from the columns of `U_k` with the normalizing constants absorbed into
+    `weights`. In the special case of a symmetric tensor, `weights` holds the
+    eigenvalues of the tensor.
+
+    Parameters
+    ----------
+    kruskal_tensor : KruskalTensor = (weight, factors)
+        factors is list of matrices, all with the same number of columns
+        i.e.::
+            for u in U:
+                u[i].shape == (s_i, R)
+
+        where `R` is fixed while `s_i` can vary with `i`
+
+    Returns
+    -------
+    KruskalTensor = (normalisation_weights, normalised_factors)
+    """
+    # allocate variables for weights, and normalized factors
+    _, rank = _validate_kruskal_tensor(kruskal_tensor)
+    weights, factors = kruskal_tensor
+    
+    if (not copy) and (weights is None):
+        warnings.warn('Provided copy=False and weights=None: a new KruskalTensor'
+                      'with new weights and factors normalised inplace will be returned.')
+        weights = T.ones(rank, **T.context(factors[0]))
+    
+    if copy:
+        factors = [T.copy(f) for f in factors]
+        if weights is not None:
+            factors[0] *= weights
+        weights = T.ones(rank, **T.context(factors[0]))
+    else:
+        factors[0] *= weights
+        weights = 0
+        
+    for factor in factors:
+        scales = T.norm(factor, axis=0)
+        weights *= scales
+        scales_non_zero = T.where(scales==0, T.ones(T.shape(scales), **T.context(factors[0])), scales)
+        factor /= scales_non_zero
+        
+    return KruskalTensor((weights, factors))
+    
+
+def kruskal_to_tensor(kruskal_tensor):
+    """Turns the Khatri-product of matrices into a full tensor
 
         ``factor_matrices = [|U_1, ... U_n|]`` becomes
         a tensor shape ``(U[1].shape[0], U[2].shape[0], ... U[-1].shape[0])``
@@ -43,7 +174,9 @@ def kruskal_to_tensor(factors, weights=None, mask=None):
     summing over r and updating an outer product of vectors.
 
     """
-    shape = [T.shape(factor)[0] for factor in factors]
+    shape, _ = _validate_kruskal_tensor(kruskal_tensor)
+    weights, factors = kruskal_tensor
+
     if weights is not None:
         full_tensor = T.dot(factors[0]*weights,
                              T.transpose(khatri_rao(factors, skip_matrix=0)))
@@ -53,7 +186,7 @@ def kruskal_to_tensor(factors, weights=None, mask=None):
     return fold(full_tensor, 0, shape)
 
 
-def kruskal_to_unfolded(factors, mode):
+def kruskal_to_unfolded(kruskal_tensor, mode):
     """Turns the khatri-product of matrices into an unfolded tensor
 
         turns ``factors = [|U_1, ... U_n|]`` into a mode-`mode`
@@ -78,10 +211,16 @@ def kruskal_to_unfolded(factors, mode):
     Writing factors = [U_1, ..., U_n], we exploit the fact that
     ``U_k = U[k].dot(khatri_rao(U_1, ..., U_k-1, U_k+1, ..., U_n))``
     """
-    return T.dot(factors[mode], T.transpose(khatri_rao(factors, skip_matrix=mode)))
+    _validate_kruskal_tensor(kruskal_tensor)
+    weights, factors = kruskal_tensor
 
+    if weights is not None:
+        return T.dot(factors[mode]*weights, T.transpose(khatri_rao(factors, skip_matrix=mode)))
+    else:
+        return T.dot(factors[mode], T.transpose(khatri_rao(factors, skip_matrix=mode)))
 
-def kruskal_to_vec(factors):
+
+def kruskal_to_vec(kruskal_tensor):
     """Turns the khatri-product of matrices into a vector
 
         (the tensor ``factors = [|U_1, ... U_n|]``
@@ -103,4 +242,140 @@ def kruskal_to_vec(factors):
     ndarray
         vectorised tensor
     """
-    return tensor_to_vec(kruskal_to_tensor(factors))
+    return tensor_to_vec(kruskal_to_tensor(kruskal_tensor))
+
+
+def kruskal_mode_dot(kruskal_tensor, matrix_or_vector, mode, keep_dim=False, copy=False):
+        """n-mode product of a Kruskal tensor and a matrix or vector at the specified mode
+
+        Parameters
+        ----------
+        kruskal_tensor : tl.KruskalTensor or (core, factors)
+                        
+        matrix_or_vector : ndarray
+            1D or 2D array of shape ``(J, i_k)`` or ``(i_k, )``
+            matrix or vectors to which to n-mode multiply the tensor
+        mode : int
+
+        Returns
+        -------
+        KruskalTensor = (core, factors)
+            `mode`-mode product of `tensor` by `matrix_or_vector`
+            * of shape :math:`(i_1, ..., i_{k-1}, J, i_{k+1}, ..., i_N)` if matrix_or_vector is a matrix
+            * of shape :math:`(i_1, ..., i_{k-1}, i_{k+1}, ..., i_N)` if matrix_or_vector is a vector
+
+        See also
+        --------
+        kruskal_multi_mode_dot : chaining several mode_dot in one call
+        """
+        shape, _ = _validate_kruskal_tensor(kruskal_tensor)
+        weights, factors = kruskal_tensor
+        contract = False
+        
+        if T.ndim(matrix_or_vector) == 2:  # Tensor times matrix
+            # Test for the validity of the operation
+            if matrix_or_vector.shape[1] != shape[mode]:
+                raise ValueError(
+                    'shapes {0} and {1} not aligned in mode-{2} multiplication: {3} (mode {2}) != {4} (dim 1 of matrix)'.format(
+                        shape, matrix_or_vector.shape, mode, shape[mode], matrix_or_vector.shape[1]
+                    ))
+
+        elif T.ndim(matrix_or_vector) == 1:  # Tensor times vector
+            if matrix_or_vector.shape[0] != shape[mode]:
+                raise ValueError(
+                    'shapes {0} and {1} not aligned for mode-{2} multiplication: {3} (mode {2}) != {4} (vector size)'.format(
+                        shape, matrix_or_vector.shape, mode, shape[mode], matrix_or_vector.shape[0]
+                    ))
+            if not keep_dim:
+                contract = True # Contract over that mode
+        else:
+            raise ValueError('Can only take n_mode_product with a vector or a matrix.')
+                             
+        if copy:
+            factors = [T.copy(f) for f in factors]
+            weights = T.copy(weights)   
+            
+        if contract:
+            factor = factors.pop(mode)
+            factor = T.dot(matrix_or_vector, factor)
+            mode = max(mode - 1, 0)
+            factors[mode] *= factor
+        else:
+             factors[mode] = T.dot(matrix_or_vector, factors[mode])
+
+        return KruskalTensor((weights, factors))
+    
+
+def unfolding_dot_khatri_rao(tensor, kruskal_tensor, mode):
+    """mode-n unfolding times khatri-rao product of factors
+    
+    Parameters
+    ----------
+    tensor : tl.tensor
+        tensor to unfold
+    factors : tl.tensor list
+        list of matrices of which to the khatri-rao product
+    mode : int
+        mode on which to unfold `tensor`
+    
+    Returns
+    -------
+    mttkrp
+        dot(unfold(tensor, mode), khatri-rao(factors))
+
+    Notes
+    -----
+    This is a variant of::
+    
+        unfolded = unfold(tensor, mode)
+        kr_factors = khatri_rao(factors, skip_matrix=mode)
+        mttkrp2 = tl.dot(unfolded, kr_factors)
+
+    Multiplying with the Khatri-Rao product is equivalent to multiplying,
+    for each rank, with the kronecker product of each factor. 
+    In code::
+
+        mttkrp_parts = []
+        for r in range(rank):
+            component = tl.tenalg.multi_mode_dot(tensor, [f[:, r] for f in factors], skip=mode)
+            mttkrp_parts.append(component)
+        mttkrp = tl.stack(mttkrp_parts, axis=1)
+        return mttkrp 
+
+    This can be done by taking n-mode-product with the full factors 
+    (faster but more memory consumming)::
+
+        projected = multi_mode_dot(tensor, factors, skip=mode, transpose=True)
+        ndims = T.ndim(tensor)
+        res = []
+        for i in range(factors[0].shape[1]):
+            index = tuple([slice(None) if k == mode  else i for k in range(ndims)])
+            res.append(projected[index])
+        return T.stack(res, axis=-1)
+
+    
+    The same idea could be expressed using einsum::
+    
+        ndims = tl.ndim(tensor)
+        tensor_idx = ''.join(chr(ord('a') + i) for i in range(ndims))
+        rank = chr(ord('a') + ndims + 1)
+        op = tensor_idx
+        for i in range(ndims):
+            if i != mode:
+                op += ',' + ''.join([tensor_idx[i], rank])
+            else:
+                result = ''.join([tensor_idx[i], rank])
+        op += '->' + result
+        factors = [f for (i, f) in enumerate(factors) if i != mode]
+        return tl_einsum(op, tensor, *factors)
+    """
+    mttkrp_parts = []
+    _, rank = _validate_kruskal_tensor(kruskal_tensor)
+    weights, factors = kruskal_tensor
+    for r in range(rank):
+        component = multi_mode_dot(tensor, [f[:, r] for f in factors], skip=mode)
+        mttkrp_parts.append(component)
+    mttkrp = T.stack(mttkrp_parts, axis=1)*T.reshape(weights, (1, -1))
+    return mttkrp 
+
+
diff --git a/tensorly/mps_tensor.py b/tensorly/mps_tensor.py
index fb6ae11..3e982f9 100644
--- a/tensorly/mps_tensor.py
+++ b/tensorly/mps_tensor.py
@@ -5,6 +5,50 @@ Core operations on tensors in Matrix Product State (MPS) format, also known as T
 import tensorly as tl
 
 
+def _validate_mps_tensor(mps_tensor):
+    factors = mps_tensor
+    n_factors = len(factors)
+    
+    if n_factors < 2:
+        raise ValueError('A Matrix-Product-State (ttrain) tensor should be composed of at least two factors and a core.'
+                         'However, {} factor was given.'.format(n_factors))
+
+    rank = []
+    shape = []
+    for index, factor in enumerate(factors):
+        current_rank, current_shape, next_rank = tl.shape(factor)
+
+        # Check that factors are third order tensors
+        if not tl.ndim(factor)==3:
+            raise ValueError('MPS expresses a tensor as third order factors (tt-cores).\n'
+                             'However, tl.ndim(factors[{}]) = {}'.format(
+                                 index, tl.ndim(factor)))
+        # Consecutive factors should have matching ranks
+        if index and tl.shape(factors[index - 1])[2] != current_rank:
+            raise ValueError('Consecutive factors should have matchine ranks\n'
+                             ' -- e.g. tl.shape(factors[0])[2]) == tl.shape(factors[1])[0])\n'
+                             'However, tl.shape(factor[{}])[2] == {} but'
+                             ' tl.shape(factor[{}])[0] == {} '.format(
+                              index - 1, tl.shape(factors[index - 1])[2], index, current_rank))
+        # Check for boundary conditions
+        if (index == 0) and current_rank != 1:
+            raise ValueError('Boundary conditions dictate factor[0].shape[0] == 1.'
+                             'However, got factor[0].shape[0] = {}.'.format(
+                              current_rank))
+        if (index == n_factors - 1) and next_rank != 1:
+            raise ValueError('Boundary conditions dictate factor[-1].shape[2] == 1.'
+                             'However, got factor[{}].shape[2] = {}.'.format(
+                              n_factors, next_rank))
+    
+        shape.append(current_shape)
+        rank.append(current_rank)
+        
+    # Add last rank (boundary condition)
+    rank.append(next_rank)
+        
+    return tuple(shape), tuple(rank)
+
+
 def mps_to_tensor(factors):
     """Returns the full tensor whose MPS decomposition is given by 'factors'
 
diff --git a/tensorly/random/base.py b/tensorly/random/base.py
index 5dc8a24..70186c1 100644
--- a/tensorly/random/base.py
+++ b/tensorly/random/base.py
@@ -1,6 +1,7 @@
 import numpy as np
 from numpy.linalg import qr
-from ..kruskal_tensor import kruskal_to_tensor
+from ..kruskal_tensor import (kruskal_to_tensor, KruskalTensor,
+                              kruskal_normalise)
 from ..tucker_tensor import tucker_to_tensor
 from ..mps_tensor import mps_to_tensor
 from .. import backend as T
@@ -39,7 +40,8 @@ def check_random_state(seed):
 
     raise ValueError('Seed should be None, int or np.random.RandomState')
 
-def random_kruskal(shape, rank, full=False, orthogonal=False, random_state=None, **context):
+def random_kruskal(shape, rank, full=False, orthogonal=False, 
+                   random_state=None, normalise_factors=True, **context):
     """Generates a random CP tensor
 
     Parameters
@@ -69,13 +71,16 @@ def random_kruskal(shape, rank, full=False, orthogonal=False, random_state=None,
 
     rns = check_random_state(random_state)
     factors = [T.tensor(rns.random_sample((s, rank)), **context) for s in shape]
+    weights = T.ones(rank, **context)
     if orthogonal:
         factors = [T.qr(factor)[0] for factor in factors]
 
     if full:
-        return kruskal_to_tensor(factors)
+        return kruskal_to_tensor((weights, factors))
+    elif normalise_factors:
+        return kruskal_normalise((weights, factors))
     else:
-        return factors
+        return KruskalTensor((weights, factors))
 
 def random_tucker(shape, rank, full=False, orthogonal=False, random_state=None, **context):
     """Generates a random Tucker tensor
diff --git a/tensorly/regression/kruskal_regression.py b/tensorly/regression/kruskal_regression.py
index 4831175..2da4954 100644
--- a/tensorly/regression/kruskal_regression.py
+++ b/tensorly/regression/kruskal_regression.py
@@ -74,6 +74,7 @@ class KruskalRegressor():
 
         # Norm of the weight tensor at each iteration
         norm_W = []
+        weights = T.ones(self.weight_rank, **T.context(X))
 
         for iteration in range(self.n_iter_max):
 
@@ -86,7 +87,7 @@ class KruskalRegressor():
                 inv_term = T.dot(T.transpose(phi), phi) + self.reg_W*T.tensor(np.eye(phi.shape[1]), **T.context(X))
                 W[i] = T.reshape(T.solve(inv_term, T.dot(T.transpose(phi), y)), (X.shape[i + 1], self.weight_rank))
 
-            weight_tensor_ = kruskal_to_tensor(W)
+            weight_tensor_ = kruskal_to_tensor((weights, W))
             norm_W.append(T.norm(weight_tensor_, 2))
 
             # Convergence check
@@ -99,8 +100,9 @@ class KruskalRegressor():
                     break
 
         self.weight_tensor_ = weight_tensor_
-        self.kruskal_weight_ = W
-        self.vec_W_ = kruskal_to_vec(W)
+        self.kruskal_weight_ = (weights, W)
+
+        self.vec_W_ = kruskal_to_vec((weights, W))
         self.n_iterations_ = iteration + 1
         self.norm_W_ = norm_W
 
diff --git a/tensorly/tenalg/__init__.py b/tensorly/tenalg/__init__.py
index 0edfd1c..dfe3c84 100644
--- a/tensorly/tenalg/__init__.py
+++ b/tensorly/tenalg/__init__.py
@@ -3,8 +3,7 @@ The :mod:`tensorly.tenalg` module contains utilities for Tensor Algebra
 operations such as khatri-rao or kronecker product, n-mode product, etc.
 """ 
 
-from .n_mode_product import (mode_dot, multi_mode_dot,
-                             unfolding_dot_khatri_rao)
+from .n_mode_product import mode_dot, multi_mode_dot
 from ._kronecker import kronecker
 from ._khatri_rao import khatri_rao
 from .generalised_inner_product import inner
diff --git a/tensorly/tenalg/_khatri_rao.py b/tensorly/tenalg/_khatri_rao.py
index 5134b09..244f072 100644
--- a/tensorly/tenalg/_khatri_rao.py
+++ b/tensorly/tenalg/_khatri_rao.py
@@ -6,7 +6,7 @@ from .. import backend as T
 
 
 
-def khatri_rao(matrices, skip_matrix=None, reverse=False, mask=None):
+def khatri_rao(matrices, weights=None, skip_matrix=None, reverse=False, mask=None):
     """Khatri-Rao product of a list of matrices
 
         This can be seen as a column-wise kronecker product.
@@ -16,11 +16,15 @@ def khatri_rao(matrices, skip_matrix=None, reverse=False, mask=None):
 
     Parameters
     ----------
-    matrices : ndarray list
+    matrices : 2D-array list
         list of matrices with the same number of columns, i.e.::
 
             for i in len(matrices):
                 matrices[i].shape = (n_i, m)
+    
+    weights : 1D-array
+        array of weights for each rank, of length m, the number of column of the factors
+        (i.e. m == factor[i].shape[1] for any factor)
 
     skip_matrix : None or int, optional, default is None
         if not None, index of a matrix to skip
@@ -83,5 +87,5 @@ def khatri_rao(matrices, skip_matrix=None, reverse=False, mask=None):
     if reverse:
         matrices = matrices[::-1]
         # Note: we do NOT use .reverse() which would reverse matrices even outside this function
-
-    return T.kr(matrices, mask=mask)
+        
+    return T.kr(matrices, weights=weights, mask=mask)
diff --git a/tensorly/tenalg/n_mode_product.py b/tensorly/tenalg/n_mode_product.py
index fd85b16..f34ef4d 100644
--- a/tensorly/tenalg/n_mode_product.py
+++ b/tensorly/tenalg/n_mode_product.py
@@ -125,61 +125,3 @@ def multi_mode_dot(tensor, matrix_or_vec_list, modes=None, skip=None, transpose=
 
     return res
 
-def unfolding_dot_khatri_rao(tensor, factors, mode):
-    """mode-n unfolding times khatri-rao product of factors
-    
-    Parameters
-    ----------
-    tensor : tl.tensor
-        tensor to unfold
-    factors : tl.tensor list
-        list of matrices of which to the khatri-rao product
-    mode : int
-        mode on which to unfold `tensor`
-    
-    Returns
-    -------
-    mttkrp
-        dot(unfold(tensor, mode), khatri-rao(factors))
-
-    Notes
-    -----
-    This is a variant of::
-    
-        unfolded = unfold(tensor, mode)
-        kr_factors = khatri_rao(factors, skip_matrix=mode)
-        mttkrp2 = tl.dot(unfolded, kr_factors)
-
-    Multiplying with the Khatri-Rao product is equivalent to multiplying,
-    for each rank, with the kronecker product of each factor. 
-    In code::
-
-        mttkrp_parts = []
-        for r in range(rank):
-            component = tl.tenalg.multi_mode_dot(tensor, [f[:, r] for f in factors], skip=mode)
-            mttkrp_parts.append(component)
-        mttkrp = tl.stack(mttkrp_parts, axis=1)
-        return mttkrp 
-    
-    The same idea could be expressed using einsum::
-    
-        ndims = tl.ndim(tensor)
-        tensor_idx = ''.join(chr(ord('a') + i) for i in range(ndims))
-        rank = chr(ord('a') + ndims + 1)
-        op = tensor_idx
-        for i in range(ndims):
-            if i != mode:
-                op += ',' + ''.join([tensor_idx[i], rank])
-            else:
-                result = ''.join([tensor_idx[i], rank])
-        op += '->' + result
-        factors = [f for (i, f) in enumerate(factors) if i != mode]
-        return tl_einsum(op, tensor, *factors)
-    """
-    projected = multi_mode_dot(tensor, factors, skip=mode, transpose=True)
-    ndims = T.ndim(tensor)
-    res = []
-    for i in range(factors[0].shape[1]):
-        index = tuple([slice(None) if k == mode  else i for k in range(ndims)])
-        res.append(projected[index])
-    return T.stack(res, axis=-1)
diff --git a/tensorly/tucker_tensor.py b/tensorly/tucker_tensor.py
index 09f527a..4541825 100644
--- a/tensorly/tucker_tensor.py
+++ b/tensorly/tucker_tensor.py
@@ -3,13 +3,40 @@ Core operations on Tucker tensors.
 """
 
 from .base import unfold, tensor_to_vec
-from .tenalg import multi_mode_dot
+from .tenalg import multi_mode_dot, mode_dot
 from .tenalg import kronecker
+from . import backend as tl
+import warnings
 
 # Author: Jean Kossaifi <jean.kossaifi+tensors@gmail.com>
 
 # License: BSD 3 clause
 
+def _validate_tucker_tensor(tucker_tensor):
+    core, factors = tucker_tensor
+    
+    if len(factors) < 2:
+        raise ValueError('A Tucker tensor should be composed of at least two factors and a core.'
+                         'However, {} factor was given.'.format(len(factors)))
+
+    if len(factors) != tl.ndim(core):
+        raise ValueError('Tucker decompositions should have one factor per more of the core tensor.'
+                         'However, core has {} modes but {} factors have been provided'.format(
+                         tl.ndim(core), len(factors)))
+
+    shape = []
+    rank = []
+    for i, factor in enumerate(factors):
+        current_shape, current_rank = tl.shape(factor)
+        if current_rank != tl.shape(core)[i]:
+            raise ValueError('Factor `n` of Tucker decomposition should verify:\n'
+                             'factors[n].shape[1] = core.shape[n].'
+                             'However, factors[{0}].shape[1]={1} but core.shape[{0}]={2}.'.format(
+                                 i, tl.shape(factor)[1], tl.shape(core)[i]))
+        shape.append(current_shape)
+        rank.append(current_rank)
+
+    return tuple(shape), tuple(rank)
 
 def tucker_to_tensor(core, factors, skip_factor=None, transpose_factors=False):
     """Converts the Tucker tensor into a full tensor
@@ -88,3 +115,71 @@ def tucker_to_vec(core, factors, skip_factor=None, transpose_factors=False):
     """
     return tensor_to_vec(tucker_to_tensor(core, factors, skip_factor=skip_factor, transpose_factors=transpose_factors))
 
+
+import warnings
+
+def tucker_mode_dot(tucker_tensor, matrix_or_vector, mode, keep_dim=False, copy=False):
+        """n-mode product of a Tucker tensor and a matrix or vector at the specified mode
+
+        Parameters
+        ----------
+        tucker_tensor : tl.TuckerTensor or (core, factors)
+                        
+        matrix_or_vector : ndarray
+            1D or 2D array of shape ``(J, i_k)`` or ``(i_k, )``
+            matrix or vectors to which to n-mode multiply the tensor
+        mode : int
+
+        Returns
+        -------
+        TuckerTensor = (core, factors)
+            `mode`-mode product of `tensor` by `matrix_or_vector`
+            * of shape :math:`(i_1, ..., i_{k-1}, J, i_{k+1}, ..., i_N)` if matrix_or_vector is a matrix
+            * of shape :math:`(i_1, ..., i_{k-1}, i_{k+1}, ..., i_N)` if matrix_or_vector is a vector
+
+        See also
+        --------
+        tucker_multi_mode_dot : chaining several mode_dot in one call
+        """
+        shape, rank = _validate_tucker_tensor(tucker_tensor)
+        core, factors = tucker_tensor
+        contract = False
+        
+        if tl.ndim(matrix_or_vector) == 2:  # Tensor times matrix
+            # Test for the validity of the operation
+            if matrix_or_vector.shape[1] != shape[mode]:
+                raise ValueError(
+                    'shapes {0} and {1} not aligned in mode-{2} multiplication: {3} (mode {2}) != {4} (dim 1 of matrix)'.format(
+                        shape, matrix_or_vector.shape, mode, shape[mode], matrix_or_vector.shape[1]
+                    ))
+
+        elif tl.ndim(matrix_or_vector) == 1:  # Tensor times vector
+            if matrix_or_vector.shape[0] != shape[mode]:
+                raise ValueError(
+                    'shapes {0} and {1} not aligned for mode-{2} multiplication: {3} (mode {2}) != {4} (vector size)'.format(
+                        shape, matrix_or_vector.shape, mode, shape[mode], matrix_or_vector.shape[0]
+                    ))
+            if not keep_dim:
+                contract = True # Contract over that mode
+        else:
+            raise ValueError('Can only take n_mode_product with a vector or a matrix.')
+                             
+        if copy:
+            factors = [tl.copy(f) for f in factors]
+            core = tl.copy(core)
+            #if not contract:
+            #    core = tl.copy(core)
+            #else:
+            #    warnings.warn('copy=True and keepdim=False, while contracting with a vector'
+            #                 ' will result in a new core with one less mode.')
+
+        if contract:
+            print('contracting mode')
+            f = factors.pop(mode)
+            core = mode_dot(core, tl.dot(matrix_or_vector, f), mode=mode)
+        else:
+            factors[mode] = tl.dot(matrix_or_vector, factors[mode])            
+
+        return core, factors
+        #return TuckerTensor(core, factors)
+    
\ No newline at end of file


Test patch:
diff --git a/tensorly/contrib/sparse/tests/test_decomposition.py b/tensorly/contrib/sparse/tests/test_decomposition.py
index ba5df12..7fd7407 100644
--- a/tensorly/contrib/sparse/tests/test_decomposition.py
+++ b/tensorly/contrib/sparse/tests/test_decomposition.py
@@ -9,6 +9,7 @@ if not tl.get_backend() == "numpy":
 pytest.importorskip("sparse")
 
 import sparse
+import numpy as np
 
 def test_sparse_parafac():
     """Test for sparse parafac"""
@@ -18,7 +19,7 @@ def test_sparse_parafac():
     rank = 3
     factors = [sparse.random((2862, rank), random_state=random_state),
                sparse.random((14036, rank), random_state=random_state)]
-
-    tensor = kruskal_to_tensor(factors)
+    weights = np.ones(rank)
+    tensor = kruskal_to_tensor((weights, factors))
     _ = parafac(tensor, rank=rank, init='random', 
                 n_iter_max=1, random_state=random_state)
diff --git a/tensorly/contrib/sparse/tests/test_tenalg.py b/tensorly/contrib/sparse/tests/test_tenalg.py
index cb00ae7..eb16cb2 100644
--- a/tensorly/contrib/sparse/tests/test_tenalg.py
+++ b/tensorly/contrib/sparse/tests/test_tenalg.py
@@ -3,12 +3,13 @@
 from  .... import backend as tl
 
 import pytest
+import numpy as np
 if not tl.get_backend() == "numpy":
     pytest.skip("Tests for sparse only with numpy backend", allow_module_level=True)
 pytest.importorskip("sparse")
 
 import tensorly.contrib.sparse as stl
-from tensorly.contrib.sparse.tenalg import unfolding_dot_khatri_rao as sparse_unfolding_dot_khatri_rao
+from tensorly.contrib.sparse.kruskal_tensor import unfolding_dot_khatri_rao as sparse_unfolding_dot_khatri_rao
 
 def test_sparse_unfolding_times_kruskal():
     """Test for unfolding_times_kruskal with sparse tensors
@@ -22,8 +23,9 @@ def test_sparse_unfolding_times_kruskal():
     shape = (100, 101, 102, 100, 100, 100)
     rank = 5
     factors = [sparse.random((i, rank)) for i in shape]
-    tensor = stl.kruskal_to_tensor(factors)
+    weights = np.ones(rank)
+    tensor = stl.kruskal_to_tensor((weights, factors))
     
     for mode in range(tl.ndim(tensor)):
         # Will blow-up memory if not sparse-safe
-        _ = sparse_unfolding_dot_khatri_rao(tensor, factors, mode)
\ No newline at end of file
+        _ = sparse_unfolding_dot_khatri_rao(tensor, (weights, factors), mode)
\ No newline at end of file
diff --git a/tensorly/decomposition/tests/test_candecomp_parafac.py b/tensorly/decomposition/tests/test_candecomp_parafac.py
index 56f13c3..8586bf3 100644
--- a/tensorly/decomposition/tests/test_candecomp_parafac.py
+++ b/tensorly/decomposition/tests/test_candecomp_parafac.py
@@ -3,7 +3,7 @@ import pytest
 
 import tensorly as tl
 from ..candecomp_parafac import (
-    parafac, non_negative_parafac, normalize_factors, initialize_factors,
+    parafac, non_negative_parafac, initialize_factors,
     sample_khatri_rao, randomised_parafac)
 from ...kruskal_tensor import kruskal_to_tensor
 from ...random import check_random_state, random_kruskal
@@ -19,10 +19,10 @@ def test_parafac():
     tol_norm_2 = 10e-2
     tol_max_abs = 10e-2
     tensor = T.tensor(rng.random_sample((3, 4, 2)))
-    factors_svd = parafac(tensor, rank=4, n_iter_max=200, init='svd', tol=10e-5)
-    factors_random = parafac(tensor, rank=4, n_iter_max=200, init='random', tol=10e-5, random_state=1234, verbose=0)
-    rec_svd = kruskal_to_tensor(factors_svd)
-    rec_random = kruskal_to_tensor(factors_random)
+    rec_svd = parafac(tensor, rank=4, n_iter_max=200, init='svd', tol=10e-5)
+    rec_random = parafac(tensor, rank=4, n_iter_max=200, init='random', tol=10e-5, random_state=1234, verbose=0)
+    rec_svd = kruskal_to_tensor(rec_svd)
+    rec_random = kruskal_to_tensor(rec_random)
     error = T.norm(rec_svd - tensor, 2)
     error /= T.norm(tensor, 2)
     assert_(error < tol_norm_2,
@@ -54,15 +54,16 @@ def test_non_negative_parafac():
     tol_max_abs = 1
     rng = check_random_state(1234)
     tensor = T.tensor(rng.random_sample((3, 3, 3))+1)
-    factors = parafac(tensor, rank=3, n_iter_max=120)
-    nn_factors = non_negative_parafac(tensor, rank=3, n_iter_max=100, tol=10e-4, init='svd', verbose=0)
+    res = parafac(tensor, rank=3, n_iter_max=120)
+    nn_res = non_negative_parafac(tensor, rank=3, n_iter_max=100, tol=10e-4, init='svd', verbose=0)
 
     # Make sure all components are positive
+    _, nn_factors = nn_res
     for factor in nn_factors:
         assert_(T.all(factor >= 0))
 
-    reconstructed_tensor = kruskal_to_tensor(factors)
-    nn_reconstructed_tensor = kruskal_to_tensor(nn_factors)
+    reconstructed_tensor = kruskal_to_tensor(res)
+    nn_reconstructed_tensor = kruskal_to_tensor(nn_res)
     error = T.norm(reconstructed_tensor - nn_reconstructed_tensor, 2)
     error /= T.norm(reconstructed_tensor, 2)
     assert_(error < tol_norm_2,
@@ -72,12 +73,12 @@ def test_non_negative_parafac():
     assert_(T.max(T.abs(reconstructed_tensor - nn_reconstructed_tensor)) < tol_max_abs,
             'abs norm of reconstruction error higher than tol')
 
-    factors_svd = non_negative_parafac(tensor, rank=3, n_iter_max=100,
+    res_svd = non_negative_parafac(tensor, rank=3, n_iter_max=100,
                                        tol=10e-4, init='svd')
-    factors_random = non_negative_parafac(tensor, rank=3, n_iter_max=100, tol=10e-4,
+    res_random = non_negative_parafac(tensor, rank=3, n_iter_max=100, tol=10e-4,
                                           init='random', random_state=1234, verbose=0)
-    rec_svd = kruskal_to_tensor(factors_svd)
-    rec_random = kruskal_to_tensor(factors_random)
+    rec_svd = kruskal_to_tensor(res_svd)
+    rec_random = kruskal_to_tensor(res_random)
     error = T.norm(rec_svd - rec_random, 2)
     error /= T.norm(rec_svd, 2)
     assert_(error < tol_norm_2,
@@ -95,7 +96,7 @@ def test_sample_khatri_rao():
     t_shape = (8, 9, 10)
     rank = 3
     tensor = T.tensor(rng.random_sample(t_shape)+1)
-    factors = parafac(tensor, rank=rank, n_iter_max=120)
+    weights, factors = parafac(tensor, rank=rank, n_iter_max=120)
     num_samples = 4
     skip_matrix = 1
     sampled_kr, sampled_indices, sampled_rows = sample_khatri_rao(factors, num_samples, skip_matrix=skip_matrix,
@@ -121,7 +122,7 @@ def test_randomised_parafac():
     n_samples = 8
     tensor = T.tensor(rng.random_sample(t_shape))
     rank = 4
-    factors_svd = randomised_parafac(tensor, rank, n_samples, n_iter_max=1000,
+    _, factors_svd = randomised_parafac(tensor, rank, n_samples, n_iter_max=1000,
                                      init='svd', tol=10e-5, verbose=True)
     for i, f in enumerate(factors_svd):
         assert_(T.shape(f) == (t_shape[i], rank),
@@ -130,7 +131,7 @@ def test_randomised_parafac():
     # test tensor reconstructed properly
     tolerance = 0.05
     tensor = random_kruskal(shape=(10, 10, 10), rank=4, full=True)
-    factors = randomised_parafac(tensor, rank=5, n_samples=100, max_stagnation=20, n_iter_max=100, tol=0, verbose=0)
-    reconstruction = kruskal_to_tensor(factors)
+    kruskal_tensor = randomised_parafac(tensor, rank=5, n_samples=100, max_stagnation=20, n_iter_max=100, tol=0, verbose=0)
+    reconstruction = kruskal_to_tensor(kruskal_tensor)
     error = float(T.norm(reconstruction - tensor, 2)/T.norm(tensor, 2))
     assert_(error < tolerance, msg='reconstruction of {} (higher than tolerance of {})'.format(error, tolerance))
diff --git a/tensorly/random/tests/test_base.py b/tensorly/random/tests/test_base.py
index 83f0bbf..bd98a6a 100644
--- a/tensorly/random/tests/test_base.py
+++ b/tensorly/random/tests/test_base.py
@@ -37,14 +37,14 @@ def test_random_kruskal():
     for i in range(T.ndim(tensor)):
         assert_equal(matrix_rank(T.to_numpy(unfold(tensor, i))), rank)
 
-    factors = random_kruskal(shape, rank, full=False)
+    weights, factors = random_kruskal(shape, rank, full=False)
     for i, factor in enumerate(factors):
         assert_equal(factor.shape, (shape[i], rank),
                 err_msg=('{}-th factor has shape {}, expected {}'.format(
                      i, factor.shape, (shape[i], rank))))
 
     # tests that the columns of each factor matrix are indeed orthogonal
-    factors = random_kruskal(shape, rank, full=False, orthogonal=True)
+    weights, factors = random_kruskal(shape, rank, full=False, orthogonal=True)
     for i, factor in enumerate(factors):
         for j in range(rank):
             for k in range(j):
diff --git a/tensorly/tenalg/tests/test_n_mode_product.py b/tensorly/tenalg/tests/test_n_mode_product.py
index 1d25410..47455b1 100644
--- a/tensorly/tenalg/tests/test_n_mode_product.py
+++ b/tensorly/tenalg/tests/test_n_mode_product.py
@@ -7,7 +7,7 @@ from ...base import fold, unfold
 from .._kronecker import kronecker
 from .._khatri_rao import khatri_rao
 from ...random import random_kruskal
-from ..n_mode_product import mode_dot, multi_mode_dot, unfolding_dot_khatri_rao
+from ..n_mode_product import mode_dot, multi_mode_dot
 from ...testing import (assert_array_equal, assert_equal,
                         assert_array_almost_equal, assert_raises)
 
@@ -119,7 +119,7 @@ def test_multi_mode_dot():
     factors = [T.tensor(np.random.rand(dims[i], X.shape[i])) for i in range(T.ndim(X))]
     true_res = T.dot(T.dot(factors[0], unfold(X, 0)), T.transpose(kronecker(factors[1:])))
     n_mode_res = multi_mode_dot(X, factors)
-    assert_array_almost_equal(true_res, unfold(n_mode_res, 0))
+    assert_array_almost_equal(true_res, unfold(n_mode_res, 0), decimal=5)
     for i in range(T.ndim(X)):
         indices = [j for j in range(T.ndim(X)) if j != i]
         sub_factors = [factors[j] for j in indices]
@@ -128,8 +128,8 @@ def test_multi_mode_dot():
         temp = multi_mode_dot(X, sub_factors, modes=indices)
         res2 = T.dot(factors[i], unfold(temp, i))
         assert_equal(true_res.shape, res.shape, err_msg='shape should be {}, is {}'.format(true_res.shape, res.shape))
-        assert_array_almost_equal(true_res, res)
-        assert_array_almost_equal(true_res, res2)
+        assert_array_almost_equal(true_res, res, decimal=5)
+        assert_array_almost_equal(true_res, res2, decimal=5)
 
     # Test skipping a factor
     dims = [2, 3, 4, 5]
@@ -154,24 +154,3 @@ def test_multi_mode_dot():
         res = multi_mode_dot(X, [vecs[i] for i in modes], modes=modes)
         assert_equal(res.shape, (1,))
         assert_equal(res[0], 1)
-
-def test_unfolding_dot_khatri_rao():
-    """Test for unfolding_dot_khatri_rao
-    
-    Check against other version check sparse safe
-    """
-    shape = (10, 10, 10, 4)
-    rank = 5
-    tensor = T.tensor(np.random.random(shape))
-    factors = random_kruskal(shape=shape, rank=rank, full=False)
-    
-    for mode in range(T.ndim(tensor)):
-        # Version forming explicitely the khatri-rao product
-        unfolded = unfold(tensor, mode)
-        kr_factors = khatri_rao(factors, skip_matrix=mode)
-        true_res = T.dot(unfolded, kr_factors)
-
-        # Efficient sparse-safe version
-        res = unfolding_dot_khatri_rao(tensor, factors, mode)
-
-        assert_array_almost_equal(true_res, res, decimal=4)
\ No newline at end of file
diff --git a/tensorly/tests/test_kruskal_tensor.py b/tensorly/tests/test_kruskal_tensor.py
index 4a6d28a..159b830 100644
--- a/tensorly/tests/test_kruskal_tensor.py
+++ b/tensorly/tests/test_kruskal_tensor.py
@@ -1,14 +1,48 @@
 import numpy as np
 
 import tensorly as tl
-from ..tenalg import khatri_rao
-from ..kruskal_tensor import kruskal_to_tensor, kruskal_to_unfolded, kruskal_to_vec
+from ..tenalg import khatri_rao, mode_dot
+from ..kruskal_tensor import (kruskal_to_tensor, kruskal_to_unfolded, 
+                              kruskal_to_vec, _validate_kruskal_tensor,
+                              kruskal_normalise, KruskalTensor,
+                              kruskal_mode_dot, unfolding_dot_khatri_rao)
 from ..base import unfold, tensor_to_vec
-from ..testing import assert_array_equal, assert_array_almost_equal
-
+from tensorly.random import check_random_state, random_kruskal
+from tensorly.testing import (assert_equal, assert_raises,
+                              assert_array_equal, assert_array_almost_equal)
 
 # Author: Jean Kossaifi <jean.kossaifi+tensors@gmail.com>
 
+def test_validate_kruskal_tensor():
+    rng = check_random_state(12345)
+    true_shape = (3, 4, 5)
+    true_rank = 3
+    kruskal_tensor = random_kruskal(true_shape, true_rank)
+    (weights, factors) = kruskal_normalise(kruskal_tensor)
+    
+    # Check correct rank and shapes are returned
+    shape, rank = _validate_kruskal_tensor((weights, factors))
+    assert_equal(shape, true_shape,
+                    err_msg='Returned incorrect shape (got {}, expected {})'.format(
+                        shape, true_shape))
+    assert_equal(rank, true_rank,
+                    err_msg='Returned incorrect rank (got {}, expected {})'.format(
+                        rank, true_rank))
+    
+    # One of the factors has the wrong rank
+    factors[0], copy = tl.tensor(rng.random_sample((4, 4))), factors[0]
+    with assert_raises(ValueError):
+        _validate_kruskal_tensor((weights, factors))
+    
+    # Not the correct amount of weights
+    factors[0] = copy
+    wrong_weights = weights[1:]
+    with assert_raises(ValueError):
+        _validate_kruskal_tensor((wrong_weights, factors))
+
+    # Not enought factors
+    with assert_raises(ValueError):
+        _validate_kruskal_tensor((weights[:1], factors[:1]))
 
 def test_kruskal_to_tensor():
     """Test for kruskal_to_tensor."""
@@ -54,13 +88,13 @@ def test_kruskal_to_tensor():
 
                            [[ 324728.,  358154.],
                             [ 366830.,  404594.]]]])
-    res = kruskal_to_tensor(U)
+    res = kruskal_to_tensor((tl.ones(3), U))
     assert_array_equal(res, true_res, err_msg='Khatri-rao incorrectly transformed into full tensor.')
 
     columns = 4
     rows = [3, 4, 2]
     matrices = [tl.tensor(np.arange(k * columns).reshape((k, columns))) for k in rows]
-    tensor = kruskal_to_tensor(matrices)
+    tensor = kruskal_to_tensor((tl.ones(columns), matrices))
     for i in range(len(rows)):
         unfolded = unfold(tensor, mode=i)
         U_i = matrices.pop(i)
@@ -73,7 +107,7 @@ def test_kruskal_to_tensor_with_weights():
     B = tl.reshape(tl.arange(5,9), (2,2))
     weights = tl.tensor([2,-1])
 
-    out = kruskal_to_tensor([A,B], weights=weights)
+    out = kruskal_to_tensor((weights, [A,B]))
     expected = tl.tensor([[-2,-2], [6, 10]])  # computed by hand
     assert_array_equal(out, expected)
 
@@ -87,10 +121,12 @@ def test_kruskal_to_unfolded():
     U3 = np.reshape(np.arange(22, 28), (2, 3))
     U4 = np.reshape(np.arange(28, 34), (2, 3))
     U = [tl.tensor(t) for t in [U1, U2, U3, U4]]
-    full_tensor = kruskal_to_tensor(U)
+    kruskal_tensor = KruskalTensor((tl.ones(3), U))
+
+    full_tensor = kruskal_to_tensor(kruskal_tensor)
     for mode in range(4):
         true_res = unfold(full_tensor, mode)
-        res = kruskal_to_unfolded(U, mode)
+        res = kruskal_to_unfolded(kruskal_tensor, mode)
         assert_array_equal(true_res, res, err_msg='khatri_rao product unfolded incorrectly for mode {}.'.format(mode))
 
 
@@ -101,7 +137,68 @@ def test_kruskal_to_vec():
     U3 = np.reshape(np.arange(22, 28), (2, 3))
     U4 = np.reshape(np.arange(28, 34), (2, 3))
     U = [tl.tensor(t) for t in [U1, U2, U3, U4]]
-    full_tensor = kruskal_to_tensor(U)
+    kruskal_tensor = KruskalTensor((tl.ones(3), U))
+    full_tensor = kruskal_to_tensor(kruskal_tensor)
     true_res = tensor_to_vec(full_tensor)
-    res = kruskal_to_vec(U)
+    res = kruskal_to_vec(kruskal_tensor)
     assert_array_equal(true_res, res, err_msg='khatri_rao product converted incorrectly to vec.')
+
+def test_kruskal_mode_dot():
+    """Test for kruskal_mode_dot
+    
+        We will compare kruskal_mode_dot 
+        (which operates directly on decomposed tensors)
+        with mode_dot (which operates on full tensors)
+        and check that the results are the same.
+    """
+    rng = check_random_state(12345)
+    shape = (5, 4, 6)
+    rank = 3
+    kruskal_ten = random_kruskal(shape, rank=rank, orthogonal=True, full=False)
+    full_tensor = tl.kruskal_to_tensor(kruskal_ten)
+    # matrix for mode 1
+    matrix = tl.tensor(rng.random_sample((7, shape[1])))
+    # vec for mode 2
+    vec = tl.tensor(rng.random_sample(shape[2]))
+
+    # Test kruskal_mode_dot with matrix
+    res = kruskal_mode_dot(kruskal_ten, matrix, mode=1, copy=True)
+    # Note that if copy=True is not respected, factors will be changes
+    # And the next test will fail
+    res = tl.kruskal_to_tensor(res)
+    true_res = mode_dot(full_tensor, matrix, mode=1)
+    assert_array_almost_equal(true_res, res)
+    
+    # Check that the data was indeed copied
+    rec = tl.kruskal_to_tensor(kruskal_ten)
+    assert_array_almost_equal(full_tensor, rec)
+    
+    # Test kruskal_mode_dot with vec
+    res = kruskal_mode_dot(kruskal_ten, vec, mode=2, copy=True)
+    res = tl.kruskal_to_tensor(res)
+    true_res = mode_dot(full_tensor, vec, mode=2)
+    assert_equal(res.shape, true_res.shape)
+    assert_array_almost_equal(true_res, res)
+
+
+def test_unfolding_dot_khatri_rao():
+    """Test for unfolding_dot_khatri_rao
+    
+    Check against other version check sparse safe
+    """
+    shape = (10, 10, 10, 4)
+    rank = 5
+    tensor = tl.tensor(np.random.random(shape))
+    weights, factors = random_kruskal(shape=shape, rank=rank, 
+                                      full=False, normalise_factors=True)
+    
+    for mode in range(tl.ndim(tensor)):
+        # Version forming explicitely the khatri-rao product
+        unfolded = unfold(tensor, mode)
+        kr_factors = khatri_rao(factors, weights=weights, skip_matrix=mode)
+        true_res = tl.dot(unfolded, kr_factors)
+
+        # Efficient sparse-safe version
+        res = unfolding_dot_khatri_rao(tensor, (weights, factors), mode)
+
+        assert_array_almost_equal(true_res, res, decimal=4)
\ No newline at end of file
diff --git a/tensorly/tests/test_mps_tensor.py b/tensorly/tests/test_mps_tensor.py
index e5c25d8..bcd6caf 100644
--- a/tensorly/tests/test_mps_tensor.py
+++ b/tensorly/tests/test_mps_tensor.py
@@ -2,8 +2,45 @@ import numpy as np
 
 import tensorly as tl
 from ..decomposition import matrix_product_state
-from ..mps_tensor import mps_to_tensor
-from ..testing import assert_array_almost_equal
+from ..mps_tensor import mps_to_tensor, _validate_mps_tensor
+from ..testing import assert_array_almost_equal, assert_equal, assert_raises
+from ..random import check_random_state, random_mps
+
+
+def test_validate_mps_tensor():
+    rng = check_random_state(12345)
+    true_shape = (3, 4, 5)
+    true_rank = (1, 3, 2, 1)
+    factors = random_mps(true_shape, rank=true_rank)
+    
+    # Check that the correct shape/rank are returned
+    shape, rank = _validate_mps_tensor(factors)
+    assert_equal(shape, true_shape,
+                    err_msg='Returned incorrect shape (got {}, expected {})'.format(
+                        shape, true_shape))
+    assert_equal(rank, true_rank,
+                    err_msg='Returned incorrect rank (got {}, expected {})'.format(
+                        rank, true_rank))
+
+    
+    # One of the factors has the wrong ndim
+    factors[0] = tl.tensor(rng.random_sample((4, 4)))
+    with assert_raises(ValueError):
+        _validate_mps_tensor(factors)
+    
+    # Consecutive factors ranks don't match
+    factors[0] = tl.tensor(rng.random_sample((1, 3, 2)))
+    with assert_raises(ValueError):
+        _validate_mps_tensor(factors)
+        
+    # Boundary conditions not respected
+    factors[0] = tl.tensor(rng.random_sample((3, 3, 2)))
+    with assert_raises(ValueError):
+        _validate_mps_tensor(factors)
+
+    # Not enought factors
+    with assert_raises(ValueError):
+        _validate_mps_tensor(factors[:1])
 
 
 def test_mps_to_tensor():
diff --git a/tensorly/tests/test_tucker_tensor.py b/tensorly/tests/test_tucker_tensor.py
index 8b34b3b..79cda56 100644
--- a/tensorly/tests/test_tucker_tensor.py
+++ b/tensorly/tests/test_tucker_tensor.py
@@ -1,15 +1,49 @@
 import numpy as np
 
-from .. import backend as T
+from .. import backend as tl
 from ..base import unfold, tensor_to_vec
-from ..tucker_tensor import tucker_to_tensor, tucker_to_unfolded, tucker_to_vec
-from ..tenalg import kronecker
-from ..testing import assert_array_equal, assert_array_almost_equal
+from ..tucker_tensor import (tucker_to_tensor, tucker_to_unfolded,
+                             tucker_to_vec, _validate_tucker_tensor,
+                             tucker_mode_dot)
+from ..tenalg import kronecker, mode_dot
+from ..testing import (assert_array_equal, assert_array_almost_equal, 
+                       assert_equal, assert_raises)
+from ..random import check_random_state, random_tucker
+
+
+def test_validate_tucker_tensor():
+    rng = check_random_state(12345)
+    true_shape = (3, 4, 5)
+    true_rank = (3, 2, 4)
+    core, factors = random_tucker(true_shape, rank=true_rank)
+    
+    # Check shape and rank returned
+    shape, rank = _validate_tucker_tensor((core, factors))
+    assert_equal(shape, true_shape,
+                    err_msg='Returned incorrect shape (got {}, expected {})'.format(
+                        shape, true_shape))
+    assert_equal(rank, true_rank,
+                    err_msg='Returned incorrect rank (got {}, expected {})'.format(
+                        rank, true_rank))
+
+    # One of the factors has the wrong rank
+    factors[0], copy = tl.tensor(rng.random_sample((4, 4))), factors[0]
+    with assert_raises(ValueError):
+        _validate_tucker_tensor((core, factors))
+    
+    # Not enough factors to match core
+    factors[0] = copy
+    with assert_raises(ValueError):
+        _validate_tucker_tensor((core, factors[1:]))
+
+    # Not enought factors
+    with assert_raises(ValueError):
+        _validate_tucker_tensor((core, factors[:1]))
 
 
 def test_tucker_to_tensor():
     """Test for tucker_to_tensor"""
-    X = T.tensor([[[1, 13],
+    X = tl.tensor([[[1, 13],
                    [4, 16],
                    [7, 19],
                    [10, 22]],
@@ -24,7 +58,7 @@ def test_tucker_to_tensor():
                    [9, 21],
                    [12, 24]]])
     ranks = [2, 3, 4]
-    U = [T.tensor(np.arange(R * s).reshape((R, s))) for (R, s) in zip(ranks, T.shape(X))]
+    U = [tl.tensor(np.arange(R * s).reshape((R, s))) for (R, s) in zip(ranks, tl.shape(X))]
     true_res = np.array([[[390, 1518, 2646, 3774],
                          [1310, 4966, 8622, 12278],
                          [2230, 8414, 14598, 20782]],
@@ -42,14 +76,14 @@ def test_tucker_to_unfolded():
     -----
     Assumes that tucker_to_tensor is properly tested
     """
-    G = T.tensor(np.random.random((4, 3, 5, 2)))
+    G = tl.tensor(np.random.random((4, 3, 5, 2)))
     ranks = [2, 2, 3, 4]
-    U = [T.tensor(np.random.random((ranks[i], G.shape[i]))) for i in range(T.ndim(G))]
+    U = [tl.tensor(np.random.random((ranks[i], G.shape[i]))) for i in range(tl.ndim(G))]
     full_tensor = tucker_to_tensor(G, U)
-    for mode in range(T.ndim(G)):
+    for mode in range(tl.ndim(G)):
         assert_array_almost_equal(tucker_to_unfolded(G, U, mode), unfold(full_tensor, mode))
         assert_array_almost_equal(tucker_to_unfolded(G, U, mode),
-                                    T.dot(T.dot(U[mode], unfold(G, mode)), T.transpose(kronecker(U, skip_matrix=mode))),
+                                    tl.dot(tl.dot(U[mode], unfold(G, mode)), tl.transpose(kronecker(U, skip_matrix=mode))),
                                     decimal=5)
 
 
@@ -60,9 +94,47 @@ def test_tucker_to_vec():
     -----
     Assumes that tucker_to_tensor works correctly
     """
-    G = T.tensor(np.random.random((4, 3, 5, 2)))
+    G = tl.tensor(np.random.random((4, 3, 5, 2)))
     ranks = [2, 2, 3, 4]
-    U = [T.tensor(np.random.random((ranks[i], G.shape[i]))) for i in range(T.ndim(G))]
+    U = [tl.tensor(np.random.random((ranks[i], G.shape[i]))) for i in range(tl.ndim(G))]
     vec = tensor_to_vec(tucker_to_tensor(G, U))
     assert_array_almost_equal(tucker_to_vec(G, U), vec)
-    assert_array_almost_equal(tucker_to_vec(G, U), T.dot(kronecker(U), tensor_to_vec(G)), decimal=5)
+    assert_array_almost_equal(tucker_to_vec(G, U), tl.dot(kronecker(U), tensor_to_vec(G)), decimal=5)
+
+
+def test_tucker_mode_dot():
+    """Test for tucker_mode_dot
+    
+        We will compare tucker_mode_dot 
+        (which operates directly on decomposed tensors)
+        with mode_dot (which operates on full tensors)
+        and check that the results are the same.
+    """
+    rng = check_random_state(12345)
+    shape = (5, 4, 6)
+    rank = (3, 2, 4)
+    tucker_ten = random_tucker(shape, rank=rank, full=False, random_state=rng)
+    full_tensor = tucker_to_tensor(*tucker_ten)
+    # matrix for mode 1
+    matrix = tl.tensor(rng.random_sample((7, shape[1])))
+    # vec for mode 2
+    vec = tl.tensor(rng.random_sample(shape[2]))
+
+    # Test tucker_mode_dot with matrix
+    res = tucker_mode_dot(tucker_ten, matrix, mode=1, copy=True)
+    # Note that if copy=True is not respected, factors will be changes
+    # And the next test will fail
+    res = tucker_to_tensor(*res)
+    true_res = mode_dot(full_tensor, matrix, mode=1)
+    assert_array_almost_equal(true_res, res, decimal=5)
+    
+    # Check that the data was indeed copied
+    rec = tucker_to_tensor(*tucker_ten)
+    assert_array_almost_equal(full_tensor, rec, decimal=5)
+    
+    # Test tucker_mode_dot with vec
+    res = tucker_mode_dot(tucker_ten, vec, mode=2, copy=True)
+    res = tucker_to_tensor(*res)
+    true_res = mode_dot(full_tensor, vec, mode=2)
+    assert_equal(res.shape, true_res.shape)
+    assert_array_almost_equal(true_res, res, decimal=5)
\ No newline at end of file