code: import numpy as np
  import jax
  import jax.numpy as jnp
  
  from ttax.base_class import TT
  from ttax.base_class import TTMatrix
  
  
  def tensor(rng, shape, tt_rank=2, batch_shape=None, dtype=jnp.float32):
    """Generate a random `TT-Tensor` of the given shape and `TT-rank`.
    
    :param rng: JAX PRNG key
    :type rng: random state is described by two unsigned 32-bit integers
    :param shape: desired tensor shape
    :type shape: array
    :param tt_rank: desired `TT-ranks` of `TT-Tensor`
    :type tt_rank: single number for equal `TT-ranks` or array specifying all `TT-ranks`
    :param batch_shape: desired batch shape of `TT-Tensor`
    :type batch_shape: array
    :param dtype: type of elements in `TT-Tensor`
    :type dtype: `dtype`
    :return: generated `TT-Tensor`
    :rtype: TT
    """
    shape = np.array(shape)
    tt_rank = np.array(tt_rank)
    batch_shape = list(batch_shape) if batch_shape else []
  
    num_dims = shape.size
    if tt_rank.size == 1:
      tt_rank = tt_rank * np.ones(num_dims - 1)
      tt_rank = np.insert(tt_rank, 0, 1)
      tt_rank = np.append(tt_rank, 1)
  
    tt_rank = tt_rank.astype(int)
  
    tt_cores = []
    rng_arr = jax.random.split(rng, num_dims)
    for i in range(num_dims):
      curr_core_shape = [tt_rank[i], shape[i], tt_rank[i + 1]]
      curr_core_shape = batch_shape + curr_core_shape
      tt_cores.append(jax.random.normal(rng_arr[i], curr_core_shape, dtype=dtype))
  
    return TT(tt_cores)
  
  
  def matrix(rng, shape, tt_rank=2, batch_shape=None, dtype=jnp.float32):
    """Generate a random `TT-Matrix` of the given shape and `TT-rank`.
      
    :param rng: JAX PRNG key
    :type rng: random state is described by two unsigned 32-bit integers
    :param shape: desired tensor shape
    :type shape: array
    :param tt_rank: desired `TT-ranks` of `TT-Matrix`
    :type tt_rank: single number for equal `TT-ranks` or array specifying all `TT-ranks`
    :param batch_shape: desired batch shape of `TT-Matrix`
    :type batch_shape: array
    :param dtype: type of elements in `TT-Matrix`
    :type dtype: `dtype`
    :return: generated `TT-Matrix`
    :rtype: TTMatrix
    """
    shape = [np.array(shape[0]), np.array(shape[1])]
    tt_rank = np.array(tt_rank)
    batch_shape = list(batch_shape) if batch_shape else []
  
    num_dims = shape[0].size
    if tt_rank.size == 1:
      tt_rank = tt_rank * np.ones(num_dims - 1)
      tt_rank = np.insert(tt_rank, 0, 1)
      tt_rank = np.append(tt_rank, 1)
  
    tt_rank = tt_rank.astype(int)
  
    tt_cores = []
    rng_arr = jax.random.split(rng, num_dims)
    for i in range(num_dims):
      curr_core_shape = [tt_rank[i], shape[0][i], shape[1][i], tt_rank[i + 1]]
      curr_core_shape = batch_shape + curr_core_shape
      tt_cores.append(jax.random.normal(rng_arr[i], curr_core_shape, dtype=dtype))
  
    return TTMatrix(tt_cores)
   from absl.testing import absltest
  from absl.testing import parameterized
  
  import numpy as np
  import jax
  import jax.numpy as jnp
  import jax.test_util as jtu
  from jax.config import config
  
  from ttax.base_class import TT
  from ttax.base_class import TTMatrix
  from ttax import random_
  from ttax import ops
  
  config.parse_flags_with_absl()
  
  
  class TTTensorTest(jtu.JaxTestCase):
  
    def testFullTensor2d(self):
      np.random.seed(1)
      for rank in [1, 2]:
        a = np.random.rand(10, rank)
        b = np.random.rand(rank, 9)
        tt_cores = (a.reshape(1, 10, rank), b.reshape(rank, 9, 1))
        desired = np.dot(a, b)
        tt_tens = TT(tt_cores)
        actual = ops.full(tt_tens)
        self.assertAllClose(desired, actual)
  
    def testFullTensor2dBatch(self):
      np.random.seed(1)
      for rank in [1, 2]:
        a = np.random.rand(3, 10, rank)
        b = np.random.rand(3, rank, 9)
        tt_cores = (a.reshape(3, 1, 10, rank), b.reshape(3, rank, 9, 1))
        desired = np.einsum('bij,bjk->bik', a, b)
        tt_tens = TT(tt_cores)
        actual = ops.full(tt_tens)
        self.assertAllClose(desired, actual)
  
    def testMultiply(self):
      # Multiply two TT-tensors.
      rng1, rng2 = jax.random.split(jax.random.PRNGKey(0))
      dtype = jnp.float32
      tt_a = random_.tensor(rng1, (1, 2, 3, 4), tt_rank=2, dtype=dtype)
      tt_b = random_.tensor(rng2, (1, 2, 3, 4), tt_rank=[1, 1, 4, 3, 1],
                            dtype=dtype)
  
      res_actual1 = ops.full(ops.multiply(tt_a, tt_b))
      res_actual2 = ops.full(tt_a * tt_b)
      res_desired = ops.full(tt_a) * ops.full(tt_b)
      self.assertAllClose(res_actual1, res_desired)
      self.assertAllClose(res_actual2, res_desired)
  
    def testMultiplyBatch(self):
      # Multiply two batches of TT-tensors.
      rng1, rng2 = jax.random.split(jax.random.PRNGKey(0))
      dtype = jnp.float32
      tt_a = random_.tensor(rng1, (1, 2, 3, 4), tt_rank=2, batch_shape=(3,),
                            dtype=dtype)
      tt_b = random_.tensor(rng2, (1, 2, 3, 4), tt_rank=[1, 1, 4, 3, 1],
                            batch_shape=(3,), dtype=dtype)
  
      res_actual1 = ops.full(ops.multiply(tt_a, tt_b))
      res_actual2 = ops.full(tt_a * tt_b)
      res_desired = ops.full(tt_a) * ops.full(tt_b)
      self.assertAllClose(res_actual1, res_desired)
      self.assertAllClose(res_actual2, res_desired)
  
    def testFlatInner(self):
      # Multiply two TT-tensors.
      rng1, rng2 = jax.random.split(jax.random.PRNGKey(0))
      dtype = jnp.float32
      tt_a = random_.tensor(rng1, (1, 2, 3, 4), tt_rank=2, dtype=dtype)
      tt_b = random_.tensor(rng2, (1, 2, 3, 4), tt_rank=[1, 1, 4, 3, 1], dtype=dtype)
      res_actual = ops.flat_inner(tt_a, tt_b)
      res_desired = jnp.sum(ops.full(tt_a) * ops.full(tt_b))
      self.assertAllClose(res_actual, res_desired)
      
    def testAdd(self):
      # Add two TT-tensors.
      rng1, rng2 = jax.random.split(jax.random.PRNGKey(0))
      dtype = jnp.float32
      tt_a = random_.tensor(rng1, (2, 1, 3, 4), tt_rank=2, dtype=dtype)
      tt_b = random_.tensor(rng2, (2, 1, 3, 4), tt_rank=[1, 2, 4, 3, 1],
                            dtype=dtype)
  
      res_actual1 = ops.full(ops.add(tt_a, tt_b))
      res_actual2 = ops.full(tt_a + tt_b)
      res_desired = ops.full(tt_a) + ops.full(tt_b)
      self.assertAllClose(res_actual1, res_desired)
      self.assertAllClose(res_actual2, res_desired)
  
    def testAddSameBatchSize(self):
      # Add two batches of TT-tensors.
      rng1, rng2 = jax.random.split(jax.random.PRNGKey(0))
      dtype = jnp.float32
      tt_a = random_.tensor(rng1, (2, 1, 3, 4), tt_rank=2, batch_shape=(3,),
                            dtype=dtype)
      tt_b = random_.tensor(rng2, (2, 1, 3, 4), tt_rank=[1, 2, 4, 3, 1],
                            batch_shape=(3,), dtype=dtype)
  
      res_actual1 = ops.full(ops.add(tt_a, tt_b))
      res_actual2 = ops.full(tt_a + tt_b)
      res_desired = ops.full(tt_a) + ops.full(tt_b)
      self.assertAllClose(res_actual1, res_desired)
      self.assertAllClose(res_actual2, res_desired)
      
    def testAddBroadcasting(self):
      # Sum two TT-tensors with broadcasting.
      rng1, rng2 = jax.random.split(jax.random.PRNGKey(0))
      dtype = jnp.float32
      tt_a = random_.tensor(rng1, (2, 1, 4), tt_rank=2, 
                            batch_shape=(1,), dtype=dtype)
      tt_b = random_.tensor(rng2, (2, 1, 4), tt_rank=[1, 2, 4, 1],
                            batch_shape=(3,), dtype=dtype)
      
      res_actual1 = ops.full(ops.add(tt_a, tt_b))
      res_actual2 = ops.full(tt_b + tt_a)
      res_desired = ops.full(tt_a) + ops.full(tt_b)
      self.assertAllClose(res_actual1, res_desired)
      self.assertAllClose(res_actual2, res_desired)
  
    def testMultiplyByScalar(self):
      # Multiply batch of TT-tensor by scalar.
      c = 4.5
      rng = jax.random.PRNGKey(0)
      dtype = jnp.float32
      tt = random_.tensor(rng, (2, 1, 3, 4), tt_rank=[1, 2, 4, 3, 1],
                          dtype=dtype)
      res_actual1 = ops.full(ops.multiply(tt, c))
      res_actual2 = ops.full(tt * c)
      res_actual3 = ops.full(c * tt)
      res_desired = c * ops.full(tt)
      self.assertAllClose(res_actual1, res_desired, rtol=1e-4)
      self.assertAllClose(res_actual2, res_desired, rtol=1e-4)
      self.assertAllClose(res_actual3, res_desired, rtol=1e-4)
  
    def testMultiplyBatchByScalar(self):
      # Multiply batch of TT-tensor by scalar.
      c = 4.5
      rng = jax.random.PRNGKey(0)
      dtype = jnp.float32
      tt = random_.tensor(rng, (2, 1, 3, 4), tt_rank=[1, 2, 4, 3, 1],
                          batch_shape=(3,), dtype=dtype)
      res_actual1 = ops.full(ops.multiply(tt, c))
      res_actual2 = ops.full(tt * c)
      res_actual3 = ops.full(c * tt)
      res_desired = c * ops.full(tt)
      self.assertAllClose(res_actual1, res_desired, rtol=1e-4)
      self.assertAllClose(res_actual2, res_desired, rtol=1e-4)
      self.assertAllClose(res_actual3, res_desired, rtol=1e-4)
  
  
  class TTMatrixTest(jtu.JaxTestCase):
  
    def testFull2d(self):
      np.random.seed(1)
      for rank in [1, 2]:
        a = np.random.rand(9, rank)
        b = np.random.rand(rank, 10)
        tt_cores = (a.reshape(1, 3, 3, rank), b.reshape(rank, 2, 5, 1))
        desired = np.einsum('aijb,bpqc->ipjq', *tt_cores)
        desired = desired.reshape(6, 15)
        tt_tens = TTMatrix(tt_cores)
        actual = ops.full(tt_tens)
        self.assertAllClose(desired, actual)
  
    def testFull2dBatch(self):
      np.random.seed(1)
      for rank in [1, 2]:
        a = np.random.rand(7, 9, rank)
        b = np.random.rand(7, rank, 10)
        tt_cores = (a.reshape(7, 1, 3, 3, rank), b.reshape(7, rank, 2, 5, 1))
        desired = np.einsum('taijb,tbpqc->tipjq', *tt_cores)
        desired = desired.reshape(7, 6, 15)
        tt_tens = TTMatrix(tt_cores)
        actual = ops.full(tt_tens)
        self.assertAllClose(desired, actual)
  
    def testMatmul(self):
      # Multiply two TT-matrices.
      rng1, rng2 = jax.random.split(jax.random.PRNGKey(0))
      dtype = jnp.float32
      left_shape = (2, 3, 4)
      sum_shape = (4, 3, 5)
      right_shape = (4, 4, 4)
      tt_a = random_.matrix(rng1, (left_shape, sum_shape), tt_rank=3, dtype=dtype)
      tt_b = random_.matrix(rng2, (sum_shape, right_shape), tt_rank=[1, 4, 3, 1],
                            dtype=dtype)
  
      res_actual = ops.full(ops.matmul(tt_a, tt_b))
      res_desired = ops.full(tt_a) @ ops.full(tt_b)
      # TODO: why such low precision?
      self.assertAllClose(res_actual, res_desired, rtol=1e-3)
  
    def testMultiply(self):
      # Elementwise multiply two TT-matrices.
      rng1, rng2 = jax.random.split(jax.random.PRNGKey(0))
      dtype = jnp.float32
      left_shape = (2, 3, 4)
      right_shape = (4, 4, 4)
      tt_a = random_.matrix(rng1, (left_shape, right_shape), tt_rank=3,
                            dtype=dtype)
      tt_b = random_.matrix(rng2, (left_shape, right_shape), tt_rank=[1, 4, 3, 1],
                            dtype=dtype)
  
      res_actual1 = ops.full(ops.multiply(tt_a, tt_b))
      res_actual2 = ops.full(tt_a * tt_b)
      res_desired = ops.full(tt_a) * ops.full(tt_b)
      self.assertAllClose(res_actual1, res_desired, rtol=1e-4)
      self.assertAllClose(res_actual2, res_desired, rtol=1e-4)
  
    def testMultiplyBatch(self):
      # Elementwise multiply two batches of TT-matrices.
      rng1, rng2 = jax.random.split(jax.random.PRNGKey(0))
      dtype = jnp.float32
      left_shape = (2, 3, 4)
      right_shape = (4, 4, 4)
      tt_a = random_.matrix(rng1, (left_shape, right_shape), tt_rank=3,
                            batch_shape=(3,), dtype=dtype)
      tt_b = random_.matrix(rng2, (left_shape, right_shape), tt_rank=[1, 4, 3, 1],
                            batch_shape=(3,), dtype=dtype)
  
      res_actual1 = ops.full(ops.multiply(tt_a, tt_b))
      res_actual2 = ops.full(tt_a * tt_b)
      res_desired = ops.full(tt_a) * ops.full(tt_b)
      # TODO: why such low precision?
      self.assertAllClose(res_actual1, res_desired, rtol=1e-3)
      self.assertAllClose(res_actual2, res_desired, rtol=1e-3)
  
    def testFlatInner(self):
      # Multiply two TT-matrices.
      rng1, rng2 = jax.random.split(jax.random.PRNGKey(0))
      dtype = jnp.float32
      left_shape = (2, 3, 4)
      right_shape = (4, 4, 4)
      tt_a = random_.matrix(rng1, (left_shape, right_shape), tt_rank=3,
                            dtype=dtype)
      tt_b = random_.matrix(rng2, (left_shape, right_shape), tt_rank=[1, 4, 3, 1],
                            dtype=dtype)
      res_actual = ops.flat_inner(tt_a, tt_b)
      res_desired = jnp.sum(ops.full(tt_a) * ops.full(tt_b))
      self.assertAllClose(res_actual, res_desired)
      
    def testAdd(self):
      # Add two TT-matrices.
      rng1, rng2 = jax.random.split(jax.random.PRNGKey(0))
      dtype = jnp.float32
      left_shape = (2, 3, 4)
      right_shape = (4, 4, 4)
      tt_a = random_.matrix(rng1, (left_shape, right_shape), tt_rank=3,
                            dtype=dtype)
      tt_b = random_.matrix(rng2, (left_shape, right_shape), tt_rank=[1, 4, 3, 1],
                            dtype=dtype)
  
      res_actual1 = ops.full(ops.add(tt_a, tt_b))
      res_actual2 = ops.full(tt_a + tt_b)
      res_desired = ops.full(tt_a) + ops.full(tt_b)
      self.assertAllClose(res_actual1, res_desired, rtol=1e-5)
      self.assertAllClose(res_actual2, res_desired, rtol=1e-5)
  
    def testAddSameBatchSize(self):
      # Add two batches of TT-matrices.
      rng1, rng2 = jax.random.split(jax.random.PRNGKey(0))
      dtype = jnp.float32
      left_shape = (2, 3, 4)
      right_shape = (4, 4, 4)
      tt_a = random_.matrix(rng1, (left_shape, right_shape), tt_rank=3,
                            batch_shape=(3,), dtype=dtype)
      tt_b = random_.matrix(rng2, (left_shape, right_shape), tt_rank=[1, 4, 3, 1],
                            batch_shape=(3,), dtype=dtype)
  
      res_actual1 = ops.full(ops.add(tt_a, tt_b))
      res_actual2 = ops.full(tt_a + tt_b)
      res_desired = ops.full(tt_a) + ops.full(tt_b)
      self.assertAllClose(res_actual1, res_desired, rtol=1e-3)
      self.assertAllClose(res_actual2, res_desired, rtol=1e-3)  
      
    def testAddBroadcasting(self):
      # Sum two TT-Matrices with broadcasting.
      rng1, rng2 = jax.random.split(jax.random.PRNGKey(0))
      dtype = jnp.float32
      left_shape = (2, 3, 4)
      right_shape = (4, 4, 4)
      tt_a = random_.matrix(rng1, (left_shape, right_shape), tt_rank=3,
                            batch_shape=(3, 1, 3,), dtype=dtype)
      tt_b = random_.matrix(rng2, (left_shape, right_shape), tt_rank=[1, 4, 3, 1],
                            batch_shape=(3, 3, 3), dtype=dtype)
      
      res_actual1 = ops.full(ops.add(tt_a, tt_b))
      res_actual2 = ops.full(tt_b + tt_a)
      res_desired = ops.full(tt_a) + ops.full(tt_b)
      self.assertAllClose(res_actual1, res_desired, rtol=1e-4)
      self.assertAllClose(res_actual2, res_desired, rtol=1e-4)
  
    def testMultiplyByScalar(self):
      # Multiply TT-matrix by scalar.
      c = 4.5
      rng = jax.random.PRNGKey(0)
      dtype = jnp.float32
      left_shape = (2, 3, 4)
      right_shape = (4, 4, 4)
      tt = random_.matrix(rng, (left_shape, right_shape), tt_rank=3,
                          dtype=dtype)
      res_actual1 = ops.full(ops.multiply(tt, c))
      res_actual2 = ops.full(tt * c)
      res_actual3 = ops.full(c * tt)
      res_desired = c * ops.full(tt)
      self.assertAllClose(res_actual1, res_desired, rtol=1e-4)  
      self.assertAllClose(res_actual2, res_desired, rtol=1e-4) 
      self.assertAllClose(res_actual3, res_desired, rtol=1e-4)
  
    def testMultiplyBatchByScalar(self):
      # Multiply batch of TT-matrix by scalar.
      c = 4.5
      rng = jax.random.PRNGKey(0)
      dtype = jnp.float32
      left_shape = (2, 3, 4)
      right_shape = (4, 4, 4)
      tt = random_.matrix(rng, (left_shape, right_shape), tt_rank=3,
                            batch_shape=(3, 1, 3,), dtype=dtype)
      res_actual1 = ops.full(ops.multiply(tt, c))
      res_actual2 = ops.full(tt * c)
      res_actual3 = ops.full(c * tt)
      res_desired = c * ops.full(tt)
      self.assertAllClose(res_actual1, res_desired, rtol=1e-3) 
      self.assertAllClose(res_actual2, res_desired, rtol=1e-3)  
      self.assertAllClose(res_actual3, res_desired, rtol=1e-3)  
  
    def testNorm(self):
      rng = jax.random.PRNGKey(0)
      shape = (6, 6, 6, 6)
      tt = random_.tensor(rng, shape)
      self.assertAllClose(ops.norm(tt), ops.norm(tt, True))
      self.assertAllClose(ops.norm(tt), np.linalg.norm(ops.full(tt)))
  
  if __name__ == '__main__':
    absltest.main(testLoader=jtu.JaxTestLoader())
  
apis: numpy.array numpy.ones numpy.append numpy.insert
apis: numpy.einsum numpy.dot numpy.random.seed numpy.random.rand