Some people were unable to run the code below using Python3.9. Multiple people had this experience, probably with different variants of 3.9. One person said that he fixed the issue by reverting to an older version of numpy, though I didn't get which one.
import pyttb as ttb
import numpy as np
# 2D Example dense
# ---------------------------------------------------------------
print("--------------------------------------------------------")
print("DENSE EXAMPLE (high-level; we'll dig into MTTKRP later):")
print("--------------------------------------------------------")
weights = np.array([1,2])
fm0 = np.array([[3,3], [4,5]])
fm1 = np.array([[3,3], [4,5], [6,7]])
K = ttb.ktensor.from_data(weights, [fm0, fm1])
print(f'dense Kruskal tensor: {K.full()}')
print("--------------------------------------------------------")
dcmp, init, stats = ttb.cp_als(K, 2)
print("--------------------------------------------------------")
print(f'decomposition after cp_als: {dcmp}')
print(dcmp)
ValidateK = ttb.ktensor.from_data(dcmp.weights,dcmp.factor_matrices)
print("--------------------------------------------------------")
print(f'validation of solution: {ValidateK.full()}')
print("--------------------------------------------------------")
# same 2D Example sparse
# ----------------------------------------------------------------
print("--------------------------------------------------------")
print("SPARSE EXAMPLE (high-level; we'll dig into MTTKRP later):")
print("--------------------------------------------------------")
weights = np.array([1,2])
subs = np.array([[0,0], [0,1], [0,2], [1,0], [1,1], [1,2]])
vals = np.array([[27],[42],[60],[42],[66],[94]])
shape = (2,3)
spk = ttb.sptensor.from_data(subs, vals, shape)
print(spk.full())
sp_dcmp, sp_init, sp_stats = ttb.cp_als(spk, 2)
print(sp_dcmp)
sp_ValidateK = ttb.ktensor.from_data(sp_dcmp.weights,sp_dcmp.factor_matrices)
print(sp_ValidateK.full())
print("-----------------------------------")
print("BACK TO DENSE EXAMPLE : MTTKRP demo")
print("-----------------------------------")
A = np.array([[3,3], [4,5]])
B = np.array([[3,3], [4,5], [6,7]])
print(f'A: {A}')
print(f'B: {B}')
kr = ttb.khatrirao(A, B)
print(f'Khatri-Rao(A,B): {kr}')
fm0 = np.array([[1,1], [1,1]])
fm1 = np.array([[1,1], [1,1], [1,1]])
fm2 = np.array([[1,1], [1,1], [1,1], [1,1]])
weights = np.array([1,1])
Kd = ttb.ktensor.from_data(weights, [fm0, fm1,fm2])
print(f'Kd.full: {Kd.full()}')
print(f'Show that MTTKRP is the time bottleneck')
Kf = ttb.ktensor.from_function(np.random.random_sample, (200, 30, 40), 2)
kfdcmp, kfinit, kfstats = ttb.cp_als(Kf.full(), 3)
U = [np.ones((2, 2)), np.ones((3, 2)), np.ones(((4, 2)))]
print(f'U: {U}')
print("---------------------------------------------")
print(f'COMPUTE MTTKRP using the pyttb tenmat class:')
print("---------------------------------------------")
Kdmat0 = ttb.tenmat.from_tensor_type(Kd.full(), rdims = np.array([0]))
print(f'-------------------------------------------:')
print(f'Maticized Tensor along way 0 is this matrix:')
print(f'-------------------------------------------:')
print(f'{Kdmat0}')
Kdmat1 = ttb.tenmat.from_tensor_type(Kd.full(), rdims = np.array([1]))
print(f'-------------------------------------------:')
print(f'Maticized Tensor along way 1 is this matrix:')
print(f'-------------------------------------------:')
print(f'{Kdmat1}')
print(f'-------------------------------------------------:')
print(f'Khatri-Rao product of U[1], U[2] (we wouldnt actually compute this):')
print(f'-------------------------------------------------:')
Ukr0 = ttb.khatrirao([U[1], U[2]])
print(f'{Ukr0}')
print(f'-------------------------------------------:')
print(f'Matricized tensor times Khatri-Rao product:')
print(f'-------------------------------------------:')
print(f'{Kdmat0.data @ Ukr0}')
print(f'-------------------------------------------:')
Kdmat1 = ttb.tenmat.from_tensor_type(Kd.full(), rdims = np.array([1]))
print(f'-------------------------------------------:')
print(f'Maticized Tensor along way 1 is this matrix:')
print(f'-------------------------------------------:')
print(f'{Kdmat1}')
Ukr1 = ttb.khatrirao([U[0], U[2]])
print(f'-------------------------------------------:')
print(f'Khatri-Rao product of U[0], U[2] (we wouldnt actually compute this):')
print(f'-------------------------------------------:')
print(f'{Ukr1}')
print(f'-------------------------------------------:')
print(f'Matricized tensor times Khatri-Rao product:')
print(f'-------------------------------------------:')
print(f'{Kdmat1.data @ Ukr1}')
print(f'-------------------------------------------:')
print(f'Next, COMPUTE MTTKRP using the pyttb mttkrp function:')
print(f'this uses indexing to avoid computing the Khatri-Rao product')
print(f'after that, we will break down exactly what an efficient MTTKRP does')
print(f'-----------along way0:---------------------:')
print(f'Kd.mttkrp(U, 0):')
print(f'{Kd.mttkrp(U, 0)}')
print(f'-----------along way1:---------------------:')
print(f'Kd.mttkrp(U, 1):')
print(f'{Kd.mttkrp(U, 1)}')
print(f'-------------------------------------------:')
print(f'-dissection an efficient MTTKRP-------------:')
print(f'-------------------------------------------:')
print(f'The idea is to have a factoring of the tensor in mind,')
print(f'then iteratively multiply each factor rather than computing')
print(f'the Khatri-Rao product')
print(f'-------------------------------------------:')
print(f'recall that this is our tensor. We show the whole tensor first:')
print(f'{Kd.full()}')
print(f'-------------------------------------------:')
print(f'then show its factor matrices. These typically start out with some')
print(f'random init, but we happened to define this example with structure:')
print(f'-------------------------------------------:')
print(f'{Kd}')
print(f'-------------------------------------------:')
print(f'To compute CP_ALS, we could start with U matrices equal to the factor')
print(f'matrices:')
print(f'-------------------------------------------:')
U = Kd.factor_matrices
print(f'U: {U}')
print(f'-------------------------------------------:')
print(f'First, we find the right dimensions for the mttkrp internal computation')
print(f'this will be the number of tensor factors by the number of columns')
print(f'in a factor matrix (these are the same; thus square).')
print(f'There are TWO factors in this tensor even though there are THREE factor matrices')
print(f'The number of weights is equal to the number of factors (TWO) and')
print(f'also equal to the number of columns in a factor matrix.')
print(f'The number of factor matrices corresponds to the number of "toes" in')
print(f'a "chicken foot" (outer product) in the tensor decomposition')
print(f'We initialize a 2 x 2 matrix with the factor weights')
print(f'The ultimate MTTKRP result along way 0 will be the dimension of ')
print(f'factor_matrices[0] @ W = (2 x 2) x (2 x 2) = (2 x 2)')
print(f'-------------------------------------------:')
print(f'Start with the factor weights W (dim 2 x 2)')
print(f'-------------------------------------------:')
W = np.tile(Kd.weights[:, None], (1, 2))
print(W)
print(f'-------------------------------------------:')
print(f'we now compute factor_matrix[1].T @ U[1] (full matrix mult)')
print(f'-------------------------------------------:')
W1 = Kd.factor_matrices[1].T @ U[1]
print(f'{W1}')
W *= W1
print(f'-------------------------------------------:')
print(f'running total of W (element-wise multiplication): {W}')
print(f'-------------------------------------------:')
print(f'we next compute factor_matrix[2].T @ U[2] (full matrix mult)')
W2 = Kd.factor_matrices[2].T @ U[2]
W *= W2
print(f'{W2}')
print(f'-------------------------------------------:')
print(f'running total: {W}')
print(f'-------------------------------------------:')
print(f'the final mttkrp result is now factor_matrices[0] @ W (full matrix mult):')
print(f'{Kd.factor_matrices[0] @ W}')
print(f'-------------------------------------------:')
print(f'Repeat for MTTKRP along way 1:')
print(f'-------------------------------------------:')
print(f'U: {U}')
print(f'-------------------------------------------:')
print(f'Again, we find the right dimensions for the mttkrp internal computation')
print(f'this is still the number of tensor factors by the number of columns')
print(f'in a factor matrix. (these are the same, thus square)')
print(f'There are TWO factors in this tensor even though there are THREE factor matrices')
print(f'The number of weights is equal to the number of factors (TWO) and')
print(f'also equal to the number of columns in a factor matrix.')
print(f'The number of factor matrices corresponds to the number of "toes" in')
print(f'a "chicken foot" (outer product) in the tensor decomposition')
print(f'We initialize a 2 x 2 matrix with the factor weights')
print(f'The ultimate MTTKRP result along way 1 will be the dimension of ')
print(f'factor_matrices[0].T @ W = (3 x 2) x (2 x 2) = (3 x 2)')
print(f'-------------------------------------------:')
W = np.tile(Kd.weights[:, None], (1, 2))
print(W)
print(f'-------------------------------------------:')
print(f'we now compute factor_matrix[0].T @ U[0] (full matrix mult)')
print(f'-------------------------------------------:')
W1 = Kd.factor_matrices[0].T @ U[0]
print(f'{W1}')
W *= W1
print(f'-------------------------------------------:')
print(f'running total of W (element-wise multiplication): {W}')
print(f'-------------------------------------------:')
print(f'we next compute factor_matrix[2].T @ U[2] (full matrix mult)')
print(f'-------------------------------------------:')
W2 = Kd.factor_matrices[2].T @ U[2]
W *= W2
print(f'{W2}')
print(f'-------------------------------------------:')
print(f'running total: {W}')
print(f'-------------------------------------------:')
print(f'the final mttkrp result is now factor_matrices[1] @ W (full matrix mult):')
print(f'-------------------------------------------:')
print(f'{Kd.factor_matrices[1] @ W}')