Initialization of orthogonal tensors with respect to a pivot

tensornetwork/backends/abstract_backend.py

@@ -1044,3 +1044,6 @@ def eps(self, dtype: Type[np.number]) -> float:
 
     raise NotImplementedError(
         f"Backend {self.name} has not implemented eps.")
+
+  def initialize_orthogonal_tensor_wrt_pivot(self,shape=Sequence[int],dtype:Optional[Type[np.number]]=None,pivot_axis:int=-1,seed=Optional[int]=None,backend: Optional[Union[Text, AbstractBackend]] = None,non_negative_diagonal: bool = False):->Tensor
+    raise NotImplementedError("Backend '{}' has not implemented initialize_orthogonal_tensor_wrt_pivot.".format(self.name))

tensornetwork/backends/numpy/numpy_backend.py

@@ -795,3 +795,13 @@ def eps(self, dtype: Type[np.number]) -> float:
       float: Machine epsilon.
     """
     return np.finfo(dtype).eps
+  def initialize_orthogonal_tensor_wrt_pivot(self,shape=Sequence[int],dtype:Optional[Type[np.number]]=None,pivot_axis:int=-1,seed=Optional[int]=None,backend: Optional[Union[Text, AbstractBackend]] = None,non_negative_diagonal: bool = False):->Tensor
+    if seed:
+      np.random.seed(seed)
+    dtype = dtype if dtype is not None else np.float64
+    if ((np.dtype(dtype) is np.dtype(np.complex128)) or
+        (np.dtype(dtype) is np.dtype(np.complex64))):
+      q,r= decompositions.qr(np,np.random.randn(
+          *shape).astype(dtype) + 1j * np.random.randn(*shape).astype(dtype),pivot_axis,non_negative_diagonal)
+    q,r= decompositions.qr(np,np.random.randn(*shape).astype(dtype),pivot_axis,non_negative_diagonal)
+    return q

tensornetwork/backends/numpy/numpy_backend_test.py

@@ -974,3 +974,15 @@ def test_item(dtype):
 def test_eps(dtype):
   backend = numpy_backend.NumPyBackend()
   assert backend.eps(dtype) == np.finfo(dtype).eps
+
+@pytest.mark.parametrize("dtype", np_dtypes)
+def test_initialize_orthogonal_tensor_wrt_pivot_dtype(dtype):
+  backend = numpy_backend.NumPyBackend()
+  a = backend.initialize_orthogonal_tensor_wrt_pivot((4,4), dtype=dtype,pivot_axis=-1,seed=10,non_negative_diagonal=False)
+  assert a.dtype == dtype
+
+@pytest.mark.parametrize("dtype", np_dtypes)
+def test_initialize_orthogonal_tensor_wrt_pivot_shape(dtype):
+  backend = numpy_backend.NumPyBackend()
+  a = backend.initialize_orthogonal_tensor_wrt_pivot((4,4), dtype=dtype,pivot_axis=-1,seed=10,non_negative_diagonal=False)
+  assert a.shape[0] == 4

tensornetwork/linalg/initialization.py

@@ -21,6 +21,7 @@
 from tensornetwork import backend_contextmanager
 from tensornetwork import backends
 from tensornetwork.tensor import Tensor
+from tensornetwork.linalg import linalg
 
 AbstractBackend = abstract_backend.AbstractBackend
 
@@ -200,3 +201,7 @@ def random_uniform(shape: Sequence[int],
   the_tensor = initialize_tensor("random_uniform", shape, backend=backend,
                                  seed=seed, boundaries=boundaries, dtype=dtype)
   return the_tensor
+def initialize_orthogonal_tensor_wrt_pivot(shape=Sequence[int],dtype:Optional[Type[np.number]]=None,pivot_axis:int=-1,seed=Optional[int]=None,backend: Optional[Union[Text, AbstractBackend]] = None,non_negative_diagonal:bool=False) ->Tensor:
+    the_tensor=initialize_tensor("randn",shape,backend=backend,seed=seed,dtype=dtype)
+    q,r=linalg.qr(the_tensor,pivot_axis,non_negative_diagonal)
+    return q

tensornetwork/linalg/tests/initialization_test.py

@@ -177,3 +177,18 @@ def inner_zero_test(dtype):
         numpyCheck = backend_obj.zeros(n.shape, dtype=dtype)
         np.testing.assert_allclose(tensor.array, tensorCheck)
         np.testing.assert_allclose(numpyT.array, numpyCheck)
+
+def test_initialize_orthogonal_tensor_wrt_pivot(backend):
+  shape=(5, 10, 3, 2)
+  pivot_axis=1
+  seed = int(time.time())
+  np.random.seed(seed=seed)
+  backend_obj = backends.backend_factory.get_backend(backend)
+  for dtype in dtypes[backend]["rand"]:
+    tnI = tensornetwork.initialize_orthogonal_tensor_wrt_pivot(
+        shape,
+        dtype=dtype,pivot_axis,
+        seed=seed,
+        backend=backend,non_negative_diagonal)
+    npI = backend_obj.initialize_orthogonal_tensor_wrt_pivot(shape, dtype=dtype, pivot_axis, seed=seed,non_negative_diagonal)
+    np.testing.assert_allclose(tnI.array, npI)