Repository: google/TensorNetwork

Issue:
Add split_edge method.
Add an option that will split an edge into multiple edges. This plus `flatten_edges` are needed for certain distributed hardware optimization.


Discussion:
The API for this should look like `net.split_edge(edge, shape)`
Could you explain how this will help distributed computations? I would have thought "slice_edge" (direct-sum decomposition of index space) and "concat_edges()" would be more useful for splitting contractions across devices.
It's a TPU specific optimization, and I'm not sure I'm allowed to explain the reasoning here. :/
Note that this is different from https://github.com/google/TensorNetwork/issues/260 in that `split_edge` would turn 1 edge into multiple edges. `slice_edge` would take 1 edge of dimension X and cut it so that it only has dimension Y. The underlying tensors connected to the edge in `slice_edge` would be sliced, whereas with `split_edge` they would only be reshaped.
I'll have a go at this.
@Thenerdstation What's the convention for handling shapes in the codebase? I've found `Tuple[int]` being used in the constructor of `BaseNode` but the `flatten_edges` code uses backend tensors that then get concatenated to construct the new shape for reshaping. Writing tests for my `split_edge` implementation (that accepts `shape: Tuple[int]`) I'm seeing that the `reshape` wrappers of backends `pytorch` and `tensorflow` accept tuples whereas `numpy` and `jax` do not (due to the `numpy` `.astype` trying to cast the tuple immediately). Should this be made consistent? Or do the backends need helper functions to convert tuples to and from whatever backend tensor is preferred for shapes (like `LongTensor` for PyTorch)?
Aw yes, the dreaded tensor as a shape issue.

This weirdness is explictly just a qwerk of tensorflow. The reason why we do this is sometimes the shape of a tensor is actually unknown at construction time, so the only way to know the actual shape at run time is to store the shape as a tensor. It seems super janky but trust me we put it in for a reason.

I'm surprised numpy and jax are the ones giving you the issue. What's your current implementation and how is it breaking for them?
Implementation is pretty similar to flatten but you reshape with the user-specified shape instead of the product of the to-be-flattened dimensions. Like I said, numpy and jax complain because you cannot feed tuples to `reshape` for numpy and jax since the numpy backed expects a numpy array https://github.com/google/TensorNetwork/blob/acf77a14d08d8c1e355d02d81c1dc14da7c78d60/tensornetwork/backends/numpy/numpy_backend.py#L35 and will throw `AttributeError: 'tuple' object has no attribute 'astype'`.

I think casting to a numpy array in `network_components.py` is a bit ugly for what should be backend-agnostic code. Maybe we should add a numpy `asarray` before the typecasting to integer? That's what the PyTorch backend does in https://github.com/google/TensorNetwork/blob/acf77a14d08d8c1e355d02d81c1dc14da7c78d60/tensornetwork/backends/pytorch/pytorch_backend.py#L44
Oh I see what you mean now. Yes that is totally an oversight.

Go ahead and change that to `np.asarry(shape).astype(...)`.

Good catch btw! I'm surprised we didn't hit that earlier.

Resolving patch:
diff --git a/tensornetwork/__init__.py b/tensornetwork/__init__.py
--- a/tensornetwork/__init__.py
+++ b/tensornetwork/__init__.py
@@ -5,7 +5,7 @@
 
 from tensornetwork.network_operations import conj, copy, transpose, split_node, split_node_qr, split_node_rq, split_node_full_svd, reachable, check_connected, check_correct, get_all_nodes, get_all_edges, remove_node, contract_trace_edges, get_subgraph_dangling
 
-from tensornetwork.network_components import contract, contract_copy_node, contract_between, outer_product, outer_product_final_nodes, contract_parallel, flatten_edges, get_all_nondangling, get_all_dangling, flatten_all_edges, flatten_edges_between, get_parallel_edges, get_shared_edges
+from tensornetwork.network_components import contract, contract_copy_node, contract_between, outer_product, outer_product_final_nodes, contract_parallel, flatten_edges, split_edge, get_all_nondangling, get_all_dangling, flatten_all_edges, flatten_edges_between, get_parallel_edges, get_shared_edges
 
 from tensornetwork.network_components import connect, disconnect
 from tensornetwork.ncon_interface import ncon, ncon_network
diff --git a/tensornetwork/backends/numpy/numpy_backend.py b/tensornetwork/backends/numpy/numpy_backend.py
--- a/tensornetwork/backends/numpy/numpy_backend.py
+++ b/tensornetwork/backends/numpy/numpy_backend.py
@@ -32,7 +32,7 @@ def tensordot(self, a: Tensor, b: Tensor, axes: Sequence[Sequence[int]]):
     return self.np.tensordot(a, b, axes)
 
   def reshape(self, tensor: Tensor, shape: Tensor):
-    return self.np.reshape(tensor, shape.astype(self.np.int32))
+    return self.np.reshape(tensor, self.np.asarray(shape).astype(self.np.int32))
 
   def transpose(self, tensor, perm):
     return self.np.transpose(tensor, perm)
diff --git a/tensornetwork/network_components.py b/tensornetwork/network_components.py
--- a/tensornetwork/network_components.py
+++ b/tensornetwork/network_components.py
@@ -267,7 +267,13 @@ def reorder_axes(self, perm: List[int]) -> "BaseNode":
       edge = self.edges[position]
       edge.update_axis(position, self, i, self)
       tmp_edges.append(edge)
-    self.edges = tmp_edges
+    self.edges = tmp_edges    
+    if self.axis_names is not None:      
+      # Permute axis names accordingly.
+      tmp_axis_names = []
+      for i in perm:
+        tmp_axis_names.append(self.axis_names[i])
+      self.axis_names = tmp_axis_names
     return self
 
   def get_axis_number(self, axis: Union[Text, int]) -> int:
@@ -798,13 +804,13 @@ class Edge:
   Dangling Edge:
     A dangling edge is an edge that only connects to a single node and only one
     part of the edge connects to the node. The other end is left "dangling".
-    These types of edges can not be contrated and represent additional
+    These types of edges can not be contracted and represent additional
     dimensions on the underlying tensor. After all other edges are contracted,
     the final result will have the same rank as the number of dangling edges. If
     there are no dangling edges, then the final value will be a scalar.
 
   Trace Edges:
-    Trace edges are edges that connects a node to itself. These edges represent
+    Trace edges are edges that connect a node to itself. These edges represent
     a trace along the given axis. Once again, the axes must be the same
     dimension.
   """
@@ -1208,7 +1214,6 @@ def flatten_edges(edges: List[Edge],
 
   Args:
     edges: A list of edges to flatten.
-    backend: A backend object
     new_edge_name: Optional name to give to the newly created edge.
 
   Returns:
@@ -1301,7 +1306,7 @@ def flatten_edges_between(
 
   Returns:
     The flattened `Edge` object. If there was only one edge between the two
-      nodes, then the original edge is returned. If there where no edges
+      nodes, then the original edge is returned. If there were no edges
       between the nodes, a None is returned.
   """
   shared_edges = get_shared_edges(node1, node2)
@@ -1325,6 +1330,147 @@ def flatten_all_edges(nodes: Iterable[BaseNode]) -> List[Edge]:
   return flattened_edges
 
 
+def _split_trace_edge(edge: Edge,
+                      shape: Tuple[int, ...],
+                      new_edge_names: Optional[List[Text]] = None,
+                      ) -> List[Edge]:
+  """Split trace edges into single edge.
+
+  Args:
+    edge: Trace edge to split.
+    shape: Tuple of integers used to split trace edge into multiple edges.
+    new_edge_names: Optional names of the new edges created.
+
+  Returns:
+    A list of new edges where the product of the dimensions of the new
+    edges corresponds to the dimension of the edge before splitting.
+  """
+  node = edge.node1  # We are in the trace case, so this is the only node.
+  backend = node.backend
+  # Permute until edge axes to be split are at the back and reshape.
+  perm_back = [min(edge.axis1, edge.axis2)]
+  perm_back += [max(edge.axis1, edge.axis2)]
+  perm_front = set(range(len(node.edges))) - set(perm_back)
+  perm_front = sorted(perm_front)
+  node.reorder_axes(perm_front + perm_back)
+  unaffected_shape = backend.shape(node.tensor)[:len(perm_front)]
+  new_shape = backend.concat([unaffected_shape, shape, shape], axis=-1)
+  node.tensor = backend.reshape(node.tensor, new_shape)
+  # Trim edges and add placeholder edges for new axes.
+  node.edges = node.edges[:len(perm_front)] + 2 * len(shape) * [None] 
+  # Create new dangling edges and connect them to each other.
+  new_edges = []
+  for idx in range(len(shape)):
+    edge1 = Edge(node1=node, axis1=len(perm_front) + idx)
+    edge2 = Edge(node1=node, axis1=len(perm_front) + len(shape) + idx)
+    node.edges[len(perm_front) + idx] = edge1
+    node.edges[len(perm_front) + len(shape) + idx] = edge2
+    new_edges.append(connect(edge1, edge2,
+                             new_edge_names[idx] if new_edge_names is not None
+                             else None))
+  # pylint: disable=expression-not-assigned
+  edge.disable() # disable old edge!
+  return new_edges
+
+
+def split_edge(edge: Edge,
+               shape: Tuple[int, ...],
+               new_edge_names: Optional[List[Text]] = None) ->  List[Edge]:
+  """Split an `Edge` into multiple edges according to `shape`. Reshapes
+  the underlying tensors connected to the edge accordingly. 
+  
+  This method acts as the inverse operation of flattening edges and
+  distinguishes between the following edge cases when adding new edges:
+    1) standard edge connecting two different nodes: reshape node dimensions
+    2) dangling edge (node2 is None): reshape node1 dimension
+    3) trace edge (node1 is node2): reshape node1 dimension
+
+  Args:
+    edge: Edge to split.
+    shape: Tuple of integers used to split edge into multiple edges.
+
+  Returns:
+    A list of new edges where the product of the dimensions of the new
+    edges corresponds to the dimension of the edge before splitting.
+
+  Raises:
+    ValueError: If the edge dimension mismatches with the split shape.
+    ValueError: If the edge is connecting nodes with different backends.
+  """
+
+  # Check if reshape operation is possible.
+  if not np.prod(shape) == edge.dimension:
+    raise ValueError(
+        "Edge {} with dimension {} cannot be split according to "
+        "shape {}.".format(edge, edge.dimension, shape))
+  # Check if possible reshape operation is trivial.
+  if len(shape) == 1:
+    return [edge]
+
+  # Handle trace edge case separately.
+  if edge.is_trace():
+    return _split_trace_edge(edge, shape, new_edge_names)
+
+  backends = [node.backend for node in edge.get_nodes() if node is not None]
+  if not all([b.name == backends[0].name for b in backends]):
+    raise ValueError("Not all backends are the same.")
+  backend = backends[0]
+  
+  # Split standard or dangling edge.
+  new_dangling_edges = []
+  expected_nodes = set(edge.get_nodes())
+  for node in expected_nodes:
+    # Required for dangling case.
+    if node is None:
+      continue
+    axis_names = node.axis_names
+    # Permute until edge axes to be split are at the back and reshape.
+    perm_back = [node.edges.index(edge)]
+    perm_front = set(range(len(node.edges))) - set(perm_back)
+    perm_front = sorted(perm_front)
+    node.reorder_axes(perm_front + perm_back)   
+    unaffected_shape = backend.shape(node.tensor)[:len(perm_front)]
+    new_shape = backend.concat([unaffected_shape, shape], axis=-1)
+    node.tensor = backend.reshape(node.tensor, new_shape) # in-place update    
+    # Trim edges.
+    node.edges = node.edges[:len(perm_front)]
+    # Create new dangling edges.    
+    for idx in range(len(shape)):
+      new_dangling_edge = Edge(node1=node, axis1=len(perm_front) + idx,
+                               name=new_edge_names[idx] if new_edge_names
+                               is not None else None)
+      node.edges += [new_dangling_edge]      
+      new_dangling_edges.append(new_dangling_edge)
+    # TODO: Allow renaming of new axes (possibly distinct from new_edge_names).
+    if axis_names:
+      new_axis_names = [axis_names[n] for n in range(len(unaffected_shape))]
+      if new_edge_names:
+        new_axis_names.extend(new_edge_names)
+      else:
+        new_axis_names.extend([str(n) for n in range(len(unaffected_shape),
+                                                     len(node.edges))])
+      node.axis_names = new_axis_names      
+    else:
+      node.axis_names = [str(n) for n in range(len(node.edges))]
+
+  node1, node2 = tuple(expected_nodes)
+  # pylint: disable=expression-not-assigned
+  edge.disable() # disable old edge
+
+  # Return new dangling edges for dangling case.
+  if node1 is None or node2 is None:
+    return new_dangling_edges
+
+  # Create connected edges between nodes for standard case.
+  new_edges = []
+  for idx in range(len(shape)):    
+    new_edges.append(connect(new_dangling_edges[idx],
+                             new_dangling_edges[len(shape) + idx],
+                             new_edge_names[idx] if new_edge_names
+                             is not None else None))  
+  return new_edges
+  
+
 def _remove_trace_edge(edge: Edge, new_node: BaseNode) -> None:
   """Collapse a trace edge. `edge` is disabled before returning.
 
@@ -1596,7 +1742,7 @@ def connect(edge1: Edge, edge2: Edge, name: Optional[Text] = None) -> Edge:
                      "Dimension of edge '{}': {}.".format(
                          edge1, edge1.dimension, edge2, edge2.dimension))
 
-  #edge1 and edg2 are always dangling in this case
+  #edge1 and edge2 are always dangling in this case
   node1 = edge1.node1
   node2 = edge2.node1
   axis1_num = node1.get_axis_number(edge1.axis1)


Test patch:
diff --git a/tensornetwork/tests/network_test.py b/tensornetwork/tests/network_test.py
--- a/tensornetwork/tests/network_test.py
+++ b/tensornetwork/tests/network_test.py
@@ -320,6 +320,86 @@ def test_flatten_edges_different_nodes_value_error(backend):
     tn.flatten_edges([e1, e2])
 
 
+def test_split_trace_edge(backend):
+  a = tn.Node(np.zeros((2, 6, 4, 6, 5, 5)), backend=backend)
+  c = tn.Node(np.zeros((2, 4)), backend=backend)
+  e1 = tn.connect(a[1], a[3])
+  e2 = tn.connect(a[4], a[5])
+  external_1 = tn.connect(a[0], c[0])
+  external_2 = tn.connect(c[1], a[2])
+  shape = (2, 1, 3)
+  new_edge_names = ["New Edge 2", "New Edge 1", "New Edge 3"]
+  new_edges = tn.split_edge(e1, shape, new_edge_names)  
+  assert a.shape == (2, 4, 5, 5) + shape + shape
+  assert a.edges == [external_1, external_2, e2, e2, *new_edges, *new_edges]
+  for new_edge, dim in zip(new_edges, shape):
+    assert new_edge.dimension == dim
+  for new_edge, new_name in zip(new_edges, new_edge_names):
+    assert new_edge.name == new_name
+  tn.check_correct({a, c})
+
+
+def test_split_edges_standard(backend):
+  a = tn.Node(np.zeros((6, 3, 5)), name="A", backend=backend)
+  b = tn.Node(np.zeros((2, 4, 6, 3)), name="B", backend=backend)
+  e1 = tn.connect(a[0], b[2], "Edge_1_1") # to be split
+  e2 = tn.connect(a[1], b[3], "Edge_1_2") # background standard edge
+  edge_a_2 = a[2] # dangling
+  edge_b_0 = b[0] # dangling
+  edge_b_1 = b[1] # dangling
+  shape = (2, 1, 3)
+  new_edge_names = ["New Edge 2", "New Edge 1", "New Edge 3"]
+  new_edges = tn.split_edge(e1, shape, new_edge_names)
+  assert a.shape == (3, 5) + shape
+  assert b.shape == (2, 4, 3) + shape
+  assert a.edges == [e2, edge_a_2, *new_edges]
+  assert b.edges == [edge_b_0, edge_b_1, e2, *new_edges]
+  for new_edge, dim in zip(new_edges, shape):
+    assert new_edge.dimension == dim
+  for new_edge, new_name in zip(new_edges, new_edge_names):
+    assert new_edge.name == new_name
+  tn.check_correct({a, b})
+
+
+def test_split_edges_standard_contract_between(backend):
+  a = tn.Node(np.random.randn(6, 3, 5), name="A", backend=backend)
+  b = tn.Node(np.random.randn(2, 4, 6, 3), name="B", backend=backend)
+  e1 = tn.connect(a[0], b[2], "Edge_1_1") # to be split
+  tn.connect(a[1], b[3], "Edge_1_2") # background standard edge
+  node_dict, _ = tn.copy({a, b})
+  c_prior = node_dict[a] @ node_dict[b]
+  shape = (2, 1, 3)
+  tn.split_edge(e1, shape)
+  tn.check_correct({a, b})
+  c_post = tn.contract_between(a, b)
+  np.testing.assert_allclose(c_prior.tensor, c_post.tensor)
+
+
+def test_split_edges_dangling(backend):
+  a = tn.Node(np.zeros((2, 10, 4, 5)), name="A", backend=backend)
+  e1 = a[0]
+  e2 = a[1]
+  e3 = a[2]
+  e4 = a[3]
+  shape = (2, 5)
+  new_edge_names = ["New Edge 2", "New Edge 5"]
+  new_edges = tn.split_edge(e2, shape, new_edge_names)
+  assert a.shape == (2, 4, 5, 2, 5)
+  assert a.edges == [e1, e3, e4, *new_edges]
+  for new_edge, dim in zip(new_edges, shape):
+    assert new_edge.dimension == dim
+  for new_edge, new_name in zip(new_edges, new_edge_names):
+    assert new_edge.name == new_name
+  tn.check_correct({a})
+
+
+def test_split_edges_dimension_mismatch_value_error(backend):
+  a = tn.Node(np.eye(5), backend=backend)
+  e1 = tn.connect(a[0], a[1])
+  with pytest.raises(ValueError):
+    tn.split_edge(e1, (2, 2))
+
+
 def test_get_shared_edges(backend):
   a = tn.Node(np.ones((2, 2, 2)), backend=backend)
   b = tn.Node(np.ones((2, 2, 2)), backend=backend)