Repository: google/TensorNetwork

Issue:
new pytorch release breaks build
the test `tests/network_components_free_test.py` fails


Discussion:

Resolving patch:
diff --git a/tensornetwork/network_components.py b/tensornetwork/network_components.py
--- a/tensornetwork/network_components.py
+++ b/tensornetwork/network_components.py
@@ -557,8 +557,8 @@ def __init__(self,
         backend=backend_obj,
         shape=backend_obj.shape_tuple(self._tensor))
 
-  def op_protection(self, other: Union[int, float, "Node"]) -> "Node":
-    if not isinstance(other, (int, float, Node)):
+  def op_protection(self, other: Union[int, float, complex, "Node"]) -> Tensor:
+    if not isinstance(other, (int, float, complex, Node)):
       raise TypeError("Operand should be one of int, float, Node type")
     if not hasattr(self, '_tensor'):
       raise AttributeError("Please provide a valid tensor for this Node.")
@@ -569,61 +569,43 @@ def op_protection(self, other: Union[int, float, "Node"]) -> "Node":
                             self.backend.name, other.backend.name))
       if not hasattr(other, '_tensor'):
         raise AttributeError("Please provide a valid tensor for this Node.")
-    else:
-      other_tensor = self.backend.convert_to_tensor(other)
-      other = Node(tensor=other_tensor, backend=self.backend.name)
+      return other._tensor
     return other
 
   def __add__(self, other: Union[int, float, "Node"]) -> "Node":
-    other = self.op_protection(other)
-    new_tensor = self.backend.addition(self.tensor, other.tensor)
-    if len(self.axis_names) > len(other.axis_names):
-      axis_names = self.axis_names
-    else:
-      axis_names = other.axis_names
+    other_tensor = self.op_protection(other)
+    new_tensor = self.backend.addition(self.tensor, other_tensor)
     return Node(
         tensor=new_tensor,
         name=self.name,
-        axis_names=axis_names,
+        axis_names=None,
         backend=self.backend.name)
 
   def __sub__(self, other: Union[int, float, "Node"]) -> "Node":
-    other = self.op_protection(other)
-    new_tensor = self.backend.subtraction(self.tensor, other.tensor)
-    if len(self.axis_names) > len(other.axis_names):
-      axis_names = self.axis_names
-    else:
-      axis_names = other.axis_names
+    other_tensor = self.op_protection(other)
+    new_tensor = self.backend.subtraction(self.tensor, other_tensor)
     return Node(
         tensor=new_tensor,
         name=self.name,
-        axis_names=axis_names,
+        axis_names=None,
         backend=self.backend.name)
 
   def __mul__(self, other: Union[int, float, "Node"]) -> "Node":
-    other = self.op_protection(other)
-    new_tensor = self.backend.multiply(self.tensor, other.tensor)
-    if len(self.axis_names) > len(other.axis_names):
-      axis_names = self.axis_names
-    else:
-      axis_names = other.axis_names
+    other_tensor = self.op_protection(other)
+    new_tensor = self.backend.multiply(self.tensor, other_tensor)
     return Node(
         tensor=new_tensor,
         name=self.name,
-        axis_names=axis_names,
+        axis_names=None,
         backend=self.backend.name)
 
   def __truediv__(self, other: Union[int, float, "Node"]) -> "Node":
-    other = self.op_protection(other)
-    new_tensor = self.backend.divide(self.tensor, other.tensor)
-    if len(self.axis_names) > len(other.axis_names):
-      axis_names = self.axis_names
-    else:
-      axis_names = other.axis_names
+    other_tensor = self.op_protection(other)
+    new_tensor = self.backend.divide(self.tensor, other_tensor)
     return Node(
         tensor=new_tensor,
         name=self.name,
-        axis_names=axis_names,
+        axis_names=None,
         backend=self.backend.name)
 
   def get_tensor(self) -> Tensor:


Test patch:
diff --git a/tensornetwork/tests/backend_contextmanager_test.py b/tensornetwork/tests/backend_contextmanager_test.py
--- a/tensornetwork/tests/backend_contextmanager_test.py
+++ b/tensornetwork/tests/backend_contextmanager_test.py
@@ -40,8 +40,8 @@ def test_contextmanager_nested():
 def test_contextmanager_wrong_item():
   a = tn.Node(np.ones((10,)))
   with pytest.raises(ValueError):
-    with tn.DefaultBackend(a):  # pytype: disable=wrong-arg-types
-      pass
+    tn.DefaultBackend(a)  # pytype: disable=wrong-arg-types
+
 
 
 def test_contextmanager_BaseBackend():
diff --git a/tensornetwork/tests/network_components_free_test.py b/tensornetwork/tests/network_components_free_test.py
--- a/tensornetwork/tests/network_components_free_test.py
+++ b/tensornetwork/tests/network_components_free_test.py
@@ -6,8 +6,9 @@
 from collections import namedtuple
 import h5py
 import re
-#pylint: disable=line-too-long
-from tensornetwork.network_components import Node, CopyNode, Edge, NodeCollection, AbstractNode, _remove_trace_edge, _remove_edges
+from tensornetwork.network_components import (Node, CopyNode, Edge,
+                                              NodeCollection, AbstractNode,
+                                              _remove_trace_edge, _remove_edges)
 import tensornetwork as tn
 from tensornetwork.backends.abstract_backend import AbstractBackend
 
@@ -17,6 +18,21 @@
 DoubleNodeEdgeTensor = namedtuple('DoubleNodeEdgeTensor',
                                   'node1 node2 edge1 edge12 tensor')
 
+op_backend_dtype_values = [('numpy', np.float32, np.float32),
+                           ('numpy', np.float64, np.float64),
+                           ('numpy', np.complex64, np.complex64),
+                           ('numpy', np.complex128, np.complex128),
+                           ('pytorch', np.float32, torch.float32),
+                           ('pytorch', np.float64, torch.float64),
+                           ('tensorflow', np.float32, tf.float32),
+                           ('tensorflow', np.float64, tf.float64),
+                           ('tensorflow', np.complex64, tf.complex64),
+                           ('tensorflow', np.complex128, tf.complex128),
+                           ('jax', np.float32, np.float32),
+                           ('jax', np.float64, np.float64),
+                           ('jax', np.complex64, np.complex64),
+                           ('jax', np.complex128, np.complex128)]
+
 
 class TestNode(AbstractNode):
 
@@ -26,18 +42,6 @@ def get_tensor(self):  #pylint: disable=useless-super-delegation
   def set_tensor(self, tensor):  #pylint: disable=useless-super-delegation
     return super().set_tensor(tensor)
 
-  def __add__(self, other):  #pylint: disable=useless-super-delegation
-    return super().__add__(other)
-
-  def __sub__(self, other):  #pylint: disable=useless-super-delegation
-    return super().__sub__(other)
-
-  def __mul__(self, other):  #pylint: disable=useless-super-delegation
-    return super().__mul__(other)
-
-  def __truediv__(self, other):  #pylint: disable=useless-super-delegation
-    return super().__truediv__(other)
-
   @property
   def shape(self):
     return super().shape
@@ -429,20 +433,14 @@ def test_node_and_scalar_add_op(backend):
   expected = np.array([[3, 4], [5, 6]])
   result = (node + 2).tensor
   np.testing.assert_almost_equal(result, expected)
-  if backend == 'jax':
-    assert result.dtype == 'int64'
-  else:
-    assert node.tensor.dtype == result.dtype
+  assert node.tensor.dtype == result.dtype
 
   node = Node(
       tensor=np.array([[1, 2], [3, 4]], dtype=np.float32), backend=backend)
   expected = np.array([[3.5, 4.5], [5.5, 6.5]])
   result = (node + 2.5).tensor
   np.testing.assert_almost_equal(result, expected)
-  if backend == 'jax':
-    assert result.dtype == 'float64'
-  else:
-    assert node.tensor.dtype == result.dtype
+  assert node.tensor.dtype == result.dtype
 
 
 def test_between_node_sub_op(backend):
@@ -482,20 +480,13 @@ def test_node_and_scalar_sub_op(backend):
   expected = np.array([[-1, 0], [1, 2]])
   result = (node - 2).tensor
   np.testing.assert_almost_equal(result, expected)
-  if backend == 'jax':
-    assert result.dtype == 'int64'
-  else:
-    assert node.tensor.dtype == result.dtype
-
+  assert node.tensor.dtype == result.dtype
   node = Node(
       tensor=np.array([[1, 2], [3, 4]], dtype=np.float32), backend=backend)
   expected = np.array([[-1.5, -0.5], [0.5, 1.5]])
   result = (node - 2.5).tensor
   np.testing.assert_almost_equal(result, expected)
-  if backend == 'jax':
-    assert result.dtype == 'float64'
-  else:
-    assert node.tensor.dtype == result.dtype
+  assert node.tensor.dtype == result.dtype
 
 
 def test_between_node_mul_op(backend):
@@ -532,26 +523,36 @@ def test_node_and_scalar_mul_op(backend):
   expected = np.array([[2, 4], [6, 8]])
   result = (node * 2).tensor
   np.testing.assert_almost_equal(result, expected)
-  if backend == 'jax':
-    assert result.dtype == 'int64'
-  else:
-    assert node.tensor.dtype == result.dtype
+  assert node.tensor.dtype == result.dtype
 
   node = Node(
       tensor=np.array([[1, 2], [3, 4]], dtype=np.float32), backend=backend)
   expected = np.array([[2.5, 5], [7.5, 10]])
   result = (node * 2.5).tensor
   np.testing.assert_almost_equal(result, expected)
-  if backend == 'jax':
-    assert result.dtype == 'float64'
-  else:
-    assert node.tensor.dtype == result.dtype
+  assert node.tensor.dtype == result.dtype
+
+
+@pytest.mark.parametrize("backend, npdtype, dtype", op_backend_dtype_values)
+def test_between_node_truediv_op(backend, npdtype, dtype):
+  node1 = Node(
+      tensor=np.array([[1., 2.], [3., 4.]], dtype=npdtype), backend=backend)
+  node2 = Node(
+      tensor=np.array([[10., 10.], [10., 10.]], dtype=npdtype), backend=backend)
+  expected = np.array([[0.1, 0.2], [0.3, 0.4]], dtype=npdtype)
+  result = (node1 / node2).tensor
+  np.testing.assert_almost_equal(result, expected)
+  assert node1.tensor.dtype == node2.tensor.dtype == result.dtype == dtype
 
 
 def test_between_node_div_op(backend):
+  if backend == 'pytorch':
+    pytest.skip("pytorch integer division no longer supported")
+
   node1 = Node(tensor=np.array([[1., 2.], [3., 4.]]), backend=backend)
   node2 = Node(tensor=np.array([[10., 10.], [10., 10.]]), backend=backend)
-  node3 = Node(tensor=np.array([[1, 2], [3, 4]]), backend=backend)
+  node3 = Node(
+      tensor=np.array([[1, 2], [3, 4]], dtype=np.int64), backend=backend)
   int_node = Node(tensor=np.array(2, dtype=np.int64), backend=backend)
   float_node = Node(tensor=np.array(2.5, dtype=np.float64), backend=backend)
 
@@ -559,28 +560,17 @@ def test_between_node_div_op(backend):
   result = (node1 / node2).tensor
   np.testing.assert_almost_equal(result, expected)
   assert node1.tensor.dtype == node2.tensor.dtype == result.dtype
-
   expected = np.array([[0.5, 1.], [1.5, 2.]])
-  expected_pytorch = np.array([[0, 1], [1, 2]])
   result = (node3 / int_node).tensor
-  if backend == 'pytorch':
-    np.testing.assert_almost_equal(result, expected_pytorch)
-    assert node3.tensor.dtype == result.dtype == torch.int64
-  else:
-    np.testing.assert_almost_equal(result, expected)
-    assert node3.tensor.dtype == 'int64'
-    assert result.dtype == 'float64'
+  np.testing.assert_almost_equal(result, expected)
+  assert node3.tensor.dtype == 'int64'
+  assert result.dtype == 'float64'
 
   expected = np.array([[2., 1.], [2 / 3, 0.5]])
-  expected_pytorch = np.array([[2, 1], [0, 0]])
   result = (int_node / node3).tensor
-  if backend == 'pytorch':
-    np.testing.assert_almost_equal(result, expected_pytorch)
-    assert node3.tensor.dtype == result.dtype == torch.int64
-  else:
-    np.testing.assert_almost_equal(result, expected)
-    assert node3.tensor.dtype == 'int64'
-    assert result.dtype == 'float64'
+  np.testing.assert_almost_equal(result, expected)
+  assert node3.tensor.dtype == 'int64'
+  assert result.dtype == 'float64'
 
   expected = np.array([[4., 4.], [4., 4.]])
   result = (node2 / float_node).tensor
@@ -592,30 +582,15 @@ def test_between_node_div_op(backend):
   assert node2.dtype == float_node.dtype == result.dtype
 
 
-def test_node_and_scalar_div_op(backend):
+@pytest.mark.parametrize("backend, npdtype, dtype", op_backend_dtype_values)
+def test_node_and_scalar_div_op(backend, npdtype, dtype):
   node = Node(
-      tensor=np.array([[5, 10], [15, 20]], dtype=np.int32), backend=backend)
-  expected = np.array([[0.5, 1.], [1.5, 2.]])
-  expected_pytorch = np.array([[0, 1], [1, 2]])
+      tensor=np.array([[5, 10], [15, 20]], dtype=npdtype), backend=backend)
+  expected = np.array([[0.5, 1.0], [1.5, 2.0]], dtype=npdtype)
   result = (node / 10).tensor
-  if backend == 'pytorch':
-    np.testing.assert_almost_equal(result, expected_pytorch)
-    assert node.tensor.dtype == result.dtype == torch.int32
-  else:
-    np.testing.assert_almost_equal(result, expected)
-    assert result.dtype == 'float64'
-    assert node.tensor.dtype == 'int32'
-
-  node = Node(
-      tensor=np.array([[5., 10.], [15., 20.]], dtype=np.float32),
-      backend=backend)
-  expected = np.array([[2., 4.], [6., 8.]])
-  result = (node / 2.5).tensor
   np.testing.assert_almost_equal(result, expected)
-  if backend == 'jax':
-    assert result.dtype == 'float64'
-  else:
-    assert node.tensor.dtype == result.dtype
+  assert result.dtype == dtype
+  assert node.tensor.dtype == dtype
 
 
 def test_node_add_input_error():
@@ -627,6 +602,7 @@ def test_node_add_input_error():
   del node1._tensor
   with pytest.raises(AttributeError):
     result = node1 + node2
+  with pytest.raises(AttributeError):    
     result = node2 + node1
 
   node1.tensor = 1
@@ -634,6 +610,7 @@ def test_node_add_input_error():
   copynode = tn.CopyNode(rank=4, dimension=3)
   with pytest.raises(TypeError):
     result = node1 + node2
+  with pytest.raises(TypeError):    
     result = node1 + copynode
 
   node2 = Node(tensor=np.array([[10, 10], [10, 10]]), backend='pytorch')
@@ -658,6 +635,7 @@ def test_node_sub_input_error():
   copynode = tn.CopyNode(rank=4, dimension=3)
   with pytest.raises(TypeError):
     result = node1 - node2
+  with pytest.raises(TypeError):    
     result = node1 - copynode
 
   node2 = Node(tensor=np.array([[10, 10], [10, 10]]), backend='pytorch')
@@ -682,6 +660,7 @@ def test_node_mul_input_error():
   copynode = tn.CopyNode(rank=4, dimension=3)
   with pytest.raises(TypeError):
     result = node1 * node2
+  with pytest.raises(TypeError):    
     result = node1 * copynode
 
   node2 = Node(tensor=np.array([[10, 10], [10, 10]]), backend='pytorch')
@@ -700,6 +679,7 @@ def test_node_div_input_error():
   del node1._tensor
   with pytest.raises(AttributeError):
     result = node1 / node2
+  with pytest.raises(AttributeError):    
     result = node2 / node1
 
   node1.tensor = 1
@@ -707,6 +687,7 @@ def test_node_div_input_error():
   copynode = tn.CopyNode(rank=4, dimension=3)
   with pytest.raises(TypeError):
     result = node1 / node2
+  with pytest.raises(TypeError):    
     result = node1 / copynode
 
   node2 = Node(tensor=np.array([[10, 10], [10, 10]]), backend='pytorch')
@@ -1295,6 +1276,7 @@ def test_repr_for_Nodes_and_Edges(double_node_edge):
 
 def test_base_node_name_list_throws_error():
   with pytest.raises(TypeError,):
+    #pylint: disable=line-too-long
     TestNode(name=["A"], axis_names=['a', 'b'])  # pytype: disable=wrong-arg-types
 
 
diff --git a/tensornetwork/tests/tensornetwork_symmetric_test.py b/tensornetwork/tests/tensornetwork_symmetric_test.py
--- a/tensornetwork/tests/tensornetwork_symmetric_test.py
+++ b/tensornetwork/tests/tensornetwork_symmetric_test.py
@@ -44,7 +44,6 @@ def get_random_symmetric(shape, flows, num_charges, seed=10, dtype=np.float64):
   indices = [Index(charge, flows[n]) for n in range(R)]
   return BlockSparseTensor.random(indices=indices, dtype=dtype)
 
-
 def get_square_matrix(shape, dtype=np.float64):
   charge = U1Charge(np.random.randint(-5, 5, shape))
   flows = [True, False]
@@ -67,7 +66,6 @@ def get_ones(shape, dtype=np.float64):
   indices = [Index(charges[n], flows[n]) for n in range(R)]
   return BlockSparseTensor.ones(indices=indices, dtype=dtype)
 
-
 @pytest.mark.parametrize("dtype", [np.float64, np.complex128])
 @pytest.mark.parametrize("num_charges", [1, 2, 3])
 def test_network_copy_reordered(dtype, num_charges):