[docs]classAxis:# MARK: Axis""" Axes are the objects that stick edges to nodes. Each instance of the :class:`AbstractNode` class has a list of :math:`N` axes, each corresponding
@@ -572,7 +575,7 @@
[docs]classAbstractNode(ABC):# MARK: AbstractNode""" Abstract class for all types of nodes. Defines what a node is and most of its properties and methods. Since it is an abstract class, cannot be instantiated.
@@ -670,18 +673,22 @@
Source code for tensorkrowch.components
To learn more about this, see :meth:`~TensorNetwork.set_data_nodes` and :meth:`~TensorNetwork.add_data`.
- * **"virtual_stack"**: Name of the ``virtual`` :class:`ParamStackNode` that
+ * **"virtual_result"**: Name of ``virtual`` nodes that are not explicitly
+ part of the network, but are required for some situations during
+ contraction. For instance, the :class:`ParamStackNode` that results from stacking :class:`ParamNodes <ParamNode>` as the first operation in the network contraction, if ``auto_stack`` mode is set to
- ``True``. There might be as much ``"virtual_stack"`` nodes as stacks are
- created from ``ParamNodes``. To learn more about this, see
- :class:`ParamStackNode`.
+ ``True``. To learn more about this, see :class:`ParamStackNode`. * **"virtual_uniform"**: Name of the ``virtual`` :class:`Node` or :class:`ParamNode` that is used in uniform (translationally invariant) tensor networks to store the tensor that will be shared by all ``leaf`` nodes. There might be as much ``"virtual_uniform"`` nodes as shared memories are used for the ``leaf`` nodes in the network (usually just one).
+
+ For ``"virtual_result"`` and ``"virtual_uniform"``, these special
+ behaviours are not restricted to nodes having those names, but also nodes
+ whose names contain those strings. Although these names can in principle be used for other nodes, this can lead to undesired behaviour.
@@ -751,10 +758,9 @@
Source code for tensorkrowch.components
ifnotisinstance(shape,(tuple,list,Size)):raiseTypeError('`shape` should be tuple[int], list[int] or torch.Size type')
- ifisinstance(shape,(tuple,list)):
- foriinshape:
- ifnotisinstance(i,int):
- raiseTypeError('`shape` elements should be int type')
+ ifisinstance(shape,Sequence):
+ ifany([notisinstance(i,int)foriinshape]):
+ raiseTypeError('`shape` elements should be int type')aux_shape=Size(shape)else:aux_shape=tensor.shape
@@ -769,7 +775,7 @@
Source code for tensorkrowch.components
axes=[Axis(num=i,name=f'axis_{i}',node=self)fori,_inenumerate(aux_shape)]else:
- ifnotisinstance(axes_names,(tuple,list)):
+ ifnotisinstance(axes_names,Sequence):raiseTypeError('`axes_names` should be tuple[str] or list[str] type')iflen(axes_names)!=len(aux_shape):
@@ -999,6 +1005,13 @@
[docs]defis_connected_to(self,other:'AbstractNode')->List[Tuple[Axis]]:
+"""Returns list of tuples of axes where the node is connected to ``other``"""
+ connected_axes=[]
+ fori1,edge1inenumerate(self._edges):
+ fori2,edge2inenumerate(other._edges):
+ if(edge1==edge2)andnotedge1.is_dangling():
+ ifself.is_node1(i1)!=other.is_node1(i2):
+ connected_axes.append((self._axes[i1],
+ other._axes[i2]))
+ returnconnected_axes
foraxinself._axes:ifaxis==ax._num:returnax._num
- raiseIndexError(f'Node {self!s} has no axis with index {axis}')
+ raiseIndexError(f'Node "{self!s}" has no axis with index {axis}')elifisinstance(axis,str):foraxinself._axes:ifaxis==ax._name:returnax._num
- raiseIndexError(f'Node {self!s} has no axis with name {axis}')
+ raiseIndexError(f'Node "{self!s}" has no axis with name "{axis}"')elifisinstance(axis,Axis):foraxinself._axes:ifaxis==ax:returnax._num
- raiseIndexError(f'Node {self!s} has no axis {axis!r}')
+ raiseIndexError(f'Node "{self!s}" has no axis "{axis!r}"')else:raiseTypeError('`axis` should be int, str or Axis type')
[docs]defin_which_axis(self,edge:'Edge')->Union[Axis,Tuple[Axis]]:""" Returns :class:`Axis` given the :class:`Edge` that is attached to the node through it.
@@ -1296,9 +1320,11 @@
Source code for tensorkrowch.components
returnlst[0]else:# Case of trace edges (attached to the node in two axes)
- returnlst
[docs]defreattach_edges(self,
+ axes:Optional[Sequence[Ax]]=None,
+ override:bool=False)->None:""" Substitutes current edges by copies of them that are attached to the node. It can happen that an edge is not attached to the node if it is the result
@@ -1312,7 +1338,10 @@
Source code for tensorkrowch.components
Parameters ----------
- override: bool
+ axis : list[int, str or Axis] or tuple[int, str or Axis], optional
+ The edge attached to these axes will be reattached. If ``None``,
+ all edges will be reattached.
+ override : bool Boolean indicating if the new, reattached edges should also replace the corresponding edges in the node's neighbours (``True``). Otherwise, the neighbours' edges will be pointing to the original nodes from which
@@ -1348,7 +1377,20 @@
Source code for tensorkrowch.components
If ``override`` is ``True``, ``nodeB['right']`` would be replaced by the new ``result['right']``. """
- fori,(edge,node1)inenumerate(zip(self._edges,self.is_node1())):
+ ifaxesisNone:
+ edges=list(enumerate(self._edges))
+ else:
+ edges=[]
+ foraxisinaxes:
+ axis_num=self.get_axis_num(axis)
+ edges.append((axis_num,self._edges[axis_num]))
+
+ skip_edges=[]
+ fori,edgeinedges:
+ ifiinskip_edges:
+ continue
+
+ node1=self._axes[i]._node1node=edge._nodes[1-node1]ifnode!=self:# New edges are always a copy, so that the original
@@ -1361,17 +1403,20 @@
Source code for tensorkrowch.components
# Case of trace edges (attached to the node in two axes)neighbour=new_edge._nodes[node1]
- ifneighbour==node:
- forj,other_edgeinenumerate(self._edges):
- if(other_edge==edge)and(i!=j):
- self._edges[j]=new_edge
- new_edge._nodes[node1]=self
- new_edge._axes[node1]=self._axes[j]
+ ifnotnew_edge.is_dangling():
+ ifneighbour!=self:
+ forj,other_edgeinedges[(i+1):]:
+ ifother_edge==edge:
+ new_edge._nodes[node1]=self
+ new_edge._axes[node1]=self._axes[j]
+ self._edges[j]=new_edge
+ skip_edges.append(j)ifoverride:
- ifnotnew_edge.is_dangling()and(neighbour!=node):
- neighbour._add_edge(
- new_edge,new_edge._axes[node1],notnode1)
"""# If node stores its own tensorifnotself.is_resultant()and(self._tensor_info['address']isnotNone):
+ if(tensorisnotNone)andnotself._compatible_shape(tensor):
+ warnings.warn(f'`tensor` is being cropped to fit the shape of '
+ f'node "{self!s}" at non-batch edges')self._unrestricted_set_tensor(tensor=tensor,init_method=init_method,device=device,
@@ -2135,6 +2186,24 @@
[docs]defnumel(self)->Tensor:
+"""
+ Returns the total number of elements in the node's tensor.
+
+ See also `torch.numel() <https://pytorch.org/docs/stable/generated/torch.numel.html>`_.
+
+ Returns
+ -------
+ int
+
+ Examples
+ --------
+ >>> node = tk.randn(shape=(2, 3), axes_names=('left', 'right'))
+ >>> node.numel()
+ 6
+ """
+ returnself.tensor.numel()
[docs]classNode(AbstractNode):# MARK: Node""" Base class for non-trainable nodes. Should be subclassed by any class of nodes that are not intended to be trained (e.g. :class:`StackNode`).
@@ -2413,10 +2482,58 @@
[docs]defchange_type(self,
+ leaf:bool=False,
+ data:bool=False,
+ virtual:bool=False,)->None:
+"""
+ Changes node type, only if node is not a resultant node.
+
+ Parameters
+ ----------
+ leaf : bool
+ Boolean indicating if the new node type is ``leaf``.
+ data : bool
+ Boolean indicating if the new node type is ``data``.
+ virtual : bool
+ Boolean indicating if the new node type is ``virtual``.
+ """
+ ifself.is_resultant():
+ raiseValueError('Only non-resultant nodes\' types can be changed')
+
+ if(leaf+data+virtual)!=1:
+ raiseValueError('One, and only one, of `leaf`, `data` and `virtual`'
+ ' can be set to True')
+
+ # Unset current type
+ ifself._leafandnotleaf:
+ node_dict=self._network._leaf_nodes
+ self._leaf=False
+ delnode_dict[self._name]
+ elifself._dataandnotdata:
+ node_dict=self._network._data_nodes
+ self._data=False
+ delnode_dict[self._name]
+ elifself._virtualandnotvirtual:
+ node_dict=self._network._virtual_nodes
+ self._virtual=False
+ delnode_dict[self._name]
+
+ # Set new type
+ ifleaf:
+ self._leaf=True
+ self._network._leaf_nodes[self._name]=self
+ elifdata:
+ self._data=True
+ self._network._data_nodes[self._name]=self
+ elifvirtual:
+ self._virtual=True
+ self._network._virtual_nodes[self._name]=self
+
+
+
[docs]classParamNode(AbstractNode):# MARK: ParamNode""" Class for trainable nodes. Should be subclassed by any class of nodes that are intended to be trained (e.g. :class:`ParamStackNode`).
@@ -2633,7 +2750,7 @@
[docs]defchange_type(self,
+ leaf:bool=False,
+ virtual:bool=False,)->None:
+"""
+ Changes node type, only if node is not a resultant node.
+
+ Parameters
+ ----------
+ leaf : bool
+ Boolean indicating if the new node type is ``leaf``.
+ virtual : bool
+ Boolean indicating if the new node type is ``virtual``.
+ """
+ if(leaf+virtual)!=1:
+ raiseValueError('One, and only one, of `leaf`, and `virtual`'
+ ' can be set to True')
+
+ # Unset current type
+ ifself._leafandnotleaf:
+ node_dict=self._network._leaf_nodes
+ self._leaf=False
+ delnode_dict[self._name]
+ elifself._virtualandnotvirtual:
+ node_dict=self._network._virtual_nodes
+ self._virtual=False
+ delnode_dict[self._name]
+
+ # Set new type
+ ifleaf:
+ self._leaf=True
+ self._network._leaf_nodes[self._name]=self
+ elifvirtual:
+ self._virtual=True
+ self._network._virtual_nodes[self._name]=self
[docs]classStackNode(Node):# MARK: StackNode""" Class for stacked nodes. ``StackNodes`` are nodes that store the information of a list of nodes that are stacked via :func:`stack`, although they can also
@@ -2893,12 +3045,10 @@
Source code for tensorkrowch.components
node1_list:Optional[List[bool]]=None)->None:ifnodesisnotNone:
- ifnotisinstance(nodes,(list,tuple)):
+ ifnotisinstance(nodes,Sequence):raiseTypeError('`nodes` should be a list or tuple of nodes')
- fornodeinnodes:
- ifisinstance(node,(StackNode,ParamStackNode)):
- raiseTypeError(
- 'Cannot create a stack using (Param)StackNode\'s')
+ ifany([isinstance(node,(StackNode,ParamStackNode))fornodeinnodes]):
+ raiseTypeError('Cannot create a stack using (Param)StackNode\'s')iftensorisnotNone:raiseValueError('If `nodes` are provided, `tensor` must not be given')
@@ -3012,10 +3162,23 @@
[docs]defreconnect(self,other:Union['StackNode','ParamStackNode'])->None:
+"""
+ Re-connects the ``StackNode`` to another ``(Param)StackNode``, in the
+ axes where the original stacked nodes were already connected.
+ """
+ foraxis1inself._edges_dict:
+ foraxis2inother._edges_dict:
+ ifself._edges_dict[axis1][0]==other._edges_dict[axis2][0]:
+ connect_stack(self.get_edge(axis1),other.get_edge(axis2))
[docs]classParamStackNode(ParamNode):# MARK: ParamStackNode""" Class for parametric stacked nodes. They are essentially the same as :class:`StackNodes <StackNode>` but they are :class:`ParamNodes <ParamNode>`.
@@ -3028,8 +3191,9 @@
Source code for tensorkrowch.components
attribute of the :class:`TensorNetwork` is set to ``True``). Hence, that first :func:`stack` is never actually computed.
- The ``ParamStackNode`` that results from this process uses the reserved name
- ``"virtual_stack"``, as explained :class:`here <AbstractNode>`. This node
+ The ``ParamStackNode`` that results from this process has the name
+ ``"virtual_result_stack"``, which contains the reserved name
+ ``"virtual_result"``, as explained :class:`here <AbstractNode>`. This node stores the tensor from which all the stacked :class:`ParamNodes <ParamNode>` just take one `slice`.
@@ -3109,13 +3273,10 @@
Source code for tensorkrowch.components
virtual:bool=False,override_node:bool=False)->None:
- ifnotisinstance(nodes,(list,tuple)):
+ ifnotisinstance(nodes,Sequence):raiseTypeError('`nodes` should be a list or tuple of nodes')
-
- fornodeinnodes:
- ifisinstance(node,(StackNode,ParamStackNode)):
- raiseTypeError(
- 'Cannot create a stack using (Param)StackNode\'s')
+ ifany([isinstance(node,(StackNode,ParamStackNode))fornodeinnodes]):
+ raiseTypeError('Cannot create a stack using (Param)StackNode\'s')foriinrange(len(nodes[:-1])):ifnotisinstance(nodes[i],type(nodes[i+1])):
@@ -3191,13 +3352,26 @@
[docs]defreconnect(self,other:Union['StackNode','ParamStackNode'])->None:
+"""
+ Re-connects the ``StackNode`` to another ``(Param)StackNode``, in the
+ axes where the original stacked nodes were already connected.
+ """
+ foraxis1inself._edges_dict:
+ foraxis2inother._edges_dict:
+ ifself._edges_dict[axis1][0]==other._edges_dict[axis2][0]:
+ connect_stack(self.get_edge(axis1),other.get_edge(axis2))
[docs]classEdge:# MARK: Edge""" Base class for edges. Should be subclassed by any new class of edges.
@@ -3212,9 +3386,10 @@
Source code for tensorkrowch.components
|
- Furthermore, edges have specific operations like :meth:`contract_` or
- :meth:`svd_` (and its variations) that allow in-place modification of the
- :class:`TensorNetwork`.
+ Furthermore, edges have specific operations like :meth:`contract` or
+ :meth:`svd` (and its variations), as well as in-place versions of them
+ (:meth:`contract_`, :meth:`svd_`, etc.) that allow in-place modification
+ of the :class:`TensorNetwork`. |
@@ -3360,7 +3535,7 @@
Source code for tensorkrowch.components
[docs]defis_batch(self)->bool:"""Returns boolean indicating whether the edge is a batch edge."""
- returnself.axis1._batch
+ returnself._axes[0]._batch
[docs]defis_attached_to(self,node:AbstractNode)->bool:"""Returns boolean indicating whether the edge is attached to ``node``."""
@@ -3456,7 +3631,7 @@
Source code for tensorkrowch.components
Note that this connectes edges from ``leaf`` (or ``data``, ``virtual``) nodes, but never from ``resultant`` nodes. If one tries to connect one of the inherited edges of a ``resultant`` node, the new connected
- edge will be attached to the original ``leaf` nodes from which the
+ edge will be attached to the original ``leaf`` nodes from which the ``resultant`` node inherited its edges. Hence, the ``resultant`` node will not "see" the connection until the :class:`TensorNetwork` is :meth:`~TensorNetwork.reset`.
@@ -3547,7 +3722,7 @@
[docs]classStackEdge(Edge):# MARK: StackEdge""" Class for stacked edges. They are just like :class:`Edges <Edge>` but used when stacking a collection of nodes into a :class:`StackNode`. When doing
@@ -3660,7 +3835,7 @@
Source code for tensorkrowch.components
Note that this connectes edges from ``leaf`` (or ``data``, ``virtual``) nodes, but never from ``resultant`` nodes. If one tries to connect one of the inherited edges of a ``resultant`` node, the new connected edge will be
- attached to the original ``leaf` nodes from which the ``resultant`` node
+ attached to the original ``leaf`` nodes from which the ``resultant`` node inherited its edges. Hence, the ``resultant`` node will not "see" the connection until the :class:`TensorNetwork` is :meth:`~TensorNetwork.reset`.
@@ -3707,17 +3882,17 @@
Source code for tensorkrowch.components
foredgein[edge1,edge2]:ifnotedge.is_dangling():
- raiseValueError(f'Edge {edge!s} is not a dangling edge. '
- f'This edge points to nodes: {edge.node1!s} and '
- f'{edge.node2!s}')
+ raiseValueError(f'Edge "{edge!s}" is not a dangling edge. '
+ f'This edge points to nodes: "{edge.node1!s}" and '
+ f'"{edge.node2!s}"')ifedge.is_batch():
- raiseValueError(f'Edge {edge!s} is a batch edge. Batch edges '
+ raiseValueError(f'Edge "{edge!s}" is a batch edge. Batch edges ''cannot be connected')ifedge1.size()!=edge2.size():raiseValueError(f'Cannot connect edges of unequal size. '
- f'Size of edge {edge1!s}: {edge1.size()}. '
- f'Size of edge {edge2!s}: {edge2.size()}')
+ f'Size of edge "{edge1!s}": {edge1.size()}. '
+ f'Size of edge "{edge2!s}": {edge2.size()}')node1,axis1=edge1.node1,edge1.axis1node2,axis2=edge2.node1,edge2.axis1
@@ -3770,9 +3945,9 @@
Source code for tensorkrowch.components
foredgein[edge1,edge2]:ifnotedge.is_dangling():
- raiseValueError(f'Edge {edge!s} is not a dangling edge. '
- f'This edge points to nodes: {edge.node1!s} and '
- f'{edge.node2!s}')
+ raiseValueError(f'Edge "{edge!s}" is not a dangling edge. '
+ f'This edge points to nodes: "{edge.node1!s}" and '
+ f'"{edge.node2!s}"')node1,axis1=edge1.node1,edge1.axis1node2,axis2=edge2.node1,edge2.axis1
@@ -3847,7 +4022,7 @@
Source code for tensorkrowch.components
################################################################################ SUCCESSOR ################################################################################
-
[docs]classSuccessor:# MARK: Successor""" Class for successors. This is a sort of cache memory for :class:`Operations <Operation>` that have been already computed.
@@ -3926,7 +4101,7 @@
[docs]classTensorNetwork(nn.Module):# MARK: TensorNetwork""" Class for arbitrary Tensor Networks. Subclass of **PyTorch** ``torch.nn.Module``.
@@ -4098,7 +4273,7 @@
Source code for tensorkrowch.components
# TN modes# Auto-management of memory mode
- self._auto_stack=False# train -> True / eval -> True
+ self._auto_stack=True# train -> True / eval -> Trueself._auto_unbind=False# train -> False / eval -> Trueself._tracing=False# Tracing mode (True while calling .trace())self._traced=False# True if .trace() is called, False if reset()
@@ -4167,7 +4342,7 @@
Source code for tensorkrowch.components
defauto_stack(self)->bool:""" Returns boolean indicating whether ``auto_stack`` mode is active. By
- default, it is ``False``.
+ default, it is ``True``. This mode indicates whether the operation :func:`stack` can take control of the memory management of the network to skip some steps in future
@@ -4379,7 +4554,11 @@
"""Updates a single node's name, without taking care of the other names."""# Node is ParamNode and tensor is not Noneifisinstance(node.tensor,Parameter):
- delattr(self,'_'.join(['param',node._name]))
+ ifhasattr(self,'param_'+node._name):
+ delattr(self,'param_'+node._name)foredgeinnode._edges:ifedge.is_attached_to(node):self._remove_edge(edge)
@@ -4429,15 +4609,15 @@
Source code for tensorkrowch.components
node._name=new_nameifisinstance(node.tensor,Parameter):
- ifnothasattr(self,'_'.join(['param',node._name])):
- self.register_parameter(
- '_'.join(['param',node._name]),
- self._memory_nodes[node._name])
- else:
- # Nodes names are never repeated, so it is likely that
- # this case will never occur
- raiseValueError(
- f'Network already has attribute named {node._name}')
+ ifnode._tensor_info['address']isnotNone:
+ ifnothasattr(self,'param_'+node._name):
+ self.register_parameter('param_'+node._name,
+ self._memory_nodes[node._name])
+ else:
+ # Nodes names are never repeated, so it is likely that
+ # this case will never occur
+ raiseValueError(
+ f'Network already has attribute named {node._name}')foredgeinnode._edges:self._add_edge(edge)
@@ -4531,10 +4711,9 @@
Source code for tensorkrowch.components
# Node is ParamNode and tensor is not Noneifisinstance(node.tensor,Parameter):
- ifnothasattr(self,'_'.join(['param',node._name])):
- self.register_parameter(
- '_'.join(['param',node._name]),
- self._memory_nodes[node._name])
+ ifnothasattr(self,'param_'+node._name,):
+ self.register_parameter('param_'+node._name,
+ self._memory_nodes[node._name])else:# Nodes names are never repeated, so it is likely that# this case will never occur
@@ -4574,7 +4753,8 @@
Source code for tensorkrowch.components
"""# Node is ParamNode and tensor is not Noneifisinstance(node.tensor,Parameter):
- delattr(self,'_'.join(['param',node._name]))
+ ifhasattr(self,'param_'+node._name):
+ delattr(self,'param_'+node._name)foredgeinnode._edges:ifedge.is_attached_to(node):self._remove_edge(edge)
@@ -4637,8 +4817,8 @@
Source code for tensorkrowch.components
in-place (``True``) or copied and then parameterized (``False``). """ifself._resultant_nodes:
- warnings.warn('Resultant nodes will be removed before parameterizing'
- ' the TN')
+ warnings.warn(
+ 'Resultant nodes will be removed before parameterizing the TN')self.reset()ifoverride:
@@ -4786,7 +4966,7 @@
Source code for tensorkrowch.components
elifisinstance(edge,Edge):ifedgenotinself._edges:raiseValueError(
- f'Edge {edge!r} should be a dangling edge of the '
+ f'Edge "{edge!r}" should be a dangling edge of the ''Tensor Network')else:raiseTypeError(
@@ -4879,7 +5059,10 @@
* ``virtual``: Only virtual nodes created in :class:`operations <Operation>` are :meth:`deleted <delete_node>`. This only includes
- nodes using the reserved name ``"virtual_stack"``.
+ nodes using the reserved name ``"virtual_result"``. * ``resultant``: These nodes are :meth:`deleted <delete_node>` from the network.
@@ -4931,9 +5114,10 @@
Source code for tensorkrowch.components
aux_dict.update(self._resultant_nodes)aux_dict.update(self._virtual_nodes)fornodeinaux_dict.values():
- ifnode._virtualand('virtual_stack'notinnode._name):
- # Virtual nodes named "virtual_stack" are ParamStackNodes
- # that result from stacking a collection of ParamNodes
+ ifnode._virtualand('virtual_result'notinnode._name):
+ # Virtual nodes named "virtual_result" are nodes that are
+ # required in some situations during contraction, like
+ # ParamStackNodes# This condition is satisfied by the rest of virtual nodescontinue
@@ -4972,13 +5156,13 @@
Source code for tensorkrowch.components
aux_dict.update(self._resultant_nodes)aux_dict.update(self._virtual_nodes)fornodeinlist(aux_dict.values()):
- ifnode._virtualand('virtual_stack'notinnode._name):
+ ifnode._virtualand('virtual_result'notinnode._name):# This condition is satisfied by the rest of virtual nodes
- # (e.g. "virtual_feature", "virtual_n_features")continueself.delete_node(node,False)
[docs]defvec_to_mps(vec:torch.Tensor,
+ n_batches:int=0,
+ rank:Optional[int]=None,
+ cum_percentage:Optional[float]=None,
+ cutoff:Optional[float]=None)->List[torch.Tensor]:
+r"""
+ Splits a vector into a sequence of MPS tensors via consecutive SVD
+ decompositions. The resultant tensors can be used to instantiate a
+ :class:`~tensorkrowch.models.MPS` with ``boundary = "obc"``.
+
+ The number of resultant tensors and their respective physical dimensions
+ depend on the shape of the input vector. That is, if one expects to recover
+ a MPS with physical dimensions
+
+ .. math::
+
+ d_1 \times \cdots \times d_n
+
+ the input vector will have to be provided with that shape. This can be done
+ with `reshape <https://pytorch.org/docs/stable/generated/torch.reshape.html>`_.
+
+ If the input vector has batch dimensions, having as shape
+
+ .. math::
+
+ b_1 \times \cdots \times b_m \times d_1 \times \cdots \times d_n
+
+ the number of batch dimensions :math:`m` can be specified in ``n_batches``.
+ In this case, the resultant tensors will all have the extra batch dimensions.
+ These tensors can be used to instantiate a :class:`~tensorkrowch.models.MPSData`
+ with ``boundary = "obc"``.
+
+ To specify the bond dimension of each cut done via SVD, one can use the
+ arguments ``rank``, ``cum_percentage`` and ``cutoff``. If more than
+ one is specified, the resulting rank will be the one that satisfies all
+ conditions.
+
+ Parameters
+ ----------
+ vec : torch.Tensor
+ Input vector to decompose.
+ n_batches : int
+ Number of batch dimensions of the input vector. Each resultant tensor
+ will have also the corresponding batch dimensions. It should be between
+ 0 and the rank of ``vec``.
+ rank : int, optional
+ Number of singular values to keep.
+ cum_percentage : float, optional
+ Proportion that should be satisfied between the sum of all singular
+ values kept and the total sum of all singular values.
+
+ .. math::
+
+ \frac{\sum_{i \in \{kept\}}{s_i}}{\sum_{i \in \{all\}}{s_i}} \ge
+ cum\_percentage
+ cutoff : float, optional
+ Quantity that lower bounds singular values in order to be kept.
+
+ Returns
+ -------
+ List[torch.Tensor]
+ """
+ ifnotisinstance(vec,torch.Tensor):
+ raiseTypeError('`vec` should be torch.Tensor type')
+
+ ifn_batches>len(vec.shape):
+ raiseValueError(
+ '`n_batches` should be between 0 and the rank of `vec`')
+
+ batches_shape=vec.shape[:n_batches]
+ phys_dims=torch.tensor(vec.shape[n_batches:])
+
+ prev_bond=1
+ tensors=[]
+ foriinrange(len(phys_dims)-1):
+ vec=vec.view(*batches_shape,
+ prev_bond*phys_dims[i],
+ phys_dims[(i+1):].prod())
+
+ u,s,vh=torch.linalg.svd(vec,full_matrices=False)
+
+ lst_ranks=[]
+
+ ifrankisNone:
+ rank=s.shape[-1]
+ lst_ranks.append(rank)
+ else:
+ lst_ranks.append(min(max(1,int(rank)),s.shape[-1]))
+
+ ifcum_percentageisnotNone:
+ s_percentages=s.cumsum(-1)/ \
+ (s.sum(-1,keepdim=True).expand(s.shape)+1e-10)# To avoid having all 0's
+ cum_percentage_tensor=cum_percentage*torch.ones_like(s)
+ cp_rank=torch.lt(
+ s_percentages,
+ cum_percentage_tensor
+ ).view(-1,s.shape[-1]).all(dim=0).sum()
+ lst_ranks.append(max(1,cp_rank.item()+1))
+
+ ifcutoffisnotNone:
+ cutoff_tensor=cutoff*torch.ones_like(s)
+ co_rank=torch.ge(
+ s,
+ cutoff_tensor
+ ).view(-1,s.shape[-1]).all(dim=0).sum()
+ lst_ranks.append(max(1,co_rank.item()))
+
+ # Select rank from specified restrictions
+ rank=min(lst_ranks)
+
+ u=u[...,:rank]
+ ifi>0:
+ u=u.reshape(*batches_shape,prev_bond,phys_dims[i],rank)
+
+ s=s[...,:rank]
+ vh=vh[...,:rank,:]
+ vh=torch.diag_embed(s)@vh
+
+ tensors.append(u)
+ prev_bond=rank
+ vec=torch.diag_embed(s)@vh
+
+ tensors.append(vec)
+ returntensors
+
+
+
[docs]defmat_to_mpo(mat:torch.Tensor,
+ rank:Optional[int]=None,
+ cum_percentage:Optional[float]=None,
+ cutoff:Optional[float]=None)->List[torch.Tensor]:
+r"""
+ Splits a matrix into a sequence of MPO tensors via consecutive SVD
+ decompositions. The resultant tensors can be used to instantiate a
+ :class:`~tensorkrowch.models.MPO` with ``boundary = "obc"``.
+
+ The number of resultant tensors and their respective input/output dimensions
+ depend on the shape of the input matrix. That is, if one expects to recover
+ a MPO with input/output dimensions
+
+ .. math::
+
+ in_1 \times out_1 \times \cdots \times in_n \times out_n
+
+ the input matrix will have to be provided with that shape. Thus it must
+ have an even number of dimensions. To accomplish this, it may happen that
+ some input/output dimensions are 1. This can be done with
+ `reshape <https://pytorch.org/docs/stable/generated/torch.reshape.html>`_.
+
+ To specify the bond dimension of each cut done via SVD, one can use the
+ arguments ``rank``, ``cum_percentage`` and ``cutoff``. If more than
+ one is specified, the resulting rank will be the one that satisfies all
+ conditions.
+
+ Parameters
+ ----------
+ mat : torch.Tensor
+ Input matrix to decompose. It must have an even number of dimensions.
+ rank : int, optional
+ Number of singular values to keep.
+ cum_percentage : float, optional
+ Proportion that should be satisfied between the sum of all singular
+ values kept and the total sum of all singular values.
+
+ .. math::
+
+ \frac{\sum_{i \in \{kept\}}{s_i}}{\sum_{i \in \{all\}}{s_i}} \ge
+ cum\_percentage
+ cutoff : float, optional
+ Quantity that lower bounds singular values in order to be kept.
+
+ Returns
+ -------
+ List[torch.Tensor]
+ """
+ ifnotisinstance(mat,torch.Tensor):
+ raiseTypeError('`mat` should be torch.Tensor type')
+ ifnotlen(mat.shape)%2==0:
+ raiseValueError('`mat` have an even number of dimensions')
+
+ in_out_dims=torch.tensor(mat.shape)
+ iflen(in_out_dims)==2:
+ return[mat]
+
+ prev_bond=1
+ tensors=[]
+ foriinrange(0,len(in_out_dims)-2,2):
+ mat=mat.view(prev_bond*in_out_dims[i]*in_out_dims[i+1],
+ in_out_dims[(i+2):].prod())
+
+ u,s,vh=torch.linalg.svd(mat,full_matrices=False)
+
+ lst_ranks=[]
+
+ ifrankisNone:
+ rank=s.shape[-1]
+ lst_ranks.append(rank)
+ else:
+ lst_ranks.append(min(max(1,int(rank)),s.shape[-1]))
+
+ ifcum_percentageisnotNone:
+ s_percentages=s.cumsum(-1)/ \
+ (s.sum(-1,keepdim=True).expand(s.shape)+1e-10)# To avoid having all 0's
+ cum_percentage_tensor=cum_percentage*torch.ones_like(s)
+ cp_rank=torch.lt(
+ s_percentages,
+ cum_percentage_tensor
+ ).view(-1,s.shape[-1]).all(dim=0).sum()
+ lst_ranks.append(max(1,cp_rank.item()+1))
+
+ ifcutoffisnotNone:
+ cutoff_tensor=cutoff*torch.ones_like(s)
+ co_rank=torch.ge(
+ s,
+ cutoff_tensor
+ ).view(-1,s.shape[-1]).all(dim=0).sum()
+ lst_ranks.append(max(1,co_rank.item()))
+
+ # Select rank from specified restrictions
+ rank=min(lst_ranks)
+
+ u=u[...,:rank]
+ ifi==0:
+ u=u.reshape(in_out_dims[i],in_out_dims[i+1],rank)
+ u=u.permute(0,2,1)# left x input x right
+ else:
+ u=u.reshape(prev_bond,in_out_dims[i],in_out_dims[i+1],rank)
+ u=u.permute(0,1,3,2)# left x input x right x output
+
+ s=s[...,:rank]
+ vh=vh[...,:rank,:]
+ vh=torch.diag_embed(s)@vh
+
+ tensors.append(u)
+ prev_bond=rank
+ mat=torch.diag_embed(s)@vh
+
+ mat=mat.reshape(rank,in_out_dims[-2],in_out_dims[-1])
+ tensors.append(mat)
+ returntensors
[docs]defunit(data:torch.Tensor,dim:int=2,axis:int=-1)->torch.Tensor:r""" Embedds the data tensor using the local feature map defined in the original `paper <https://arxiv.org/abs/1605.05775>`_ by E. Miles Stoudenmire and David
@@ -325,21 +328,25 @@
Source code for tensorkrowch.embeddings
.. math::
- batch\_size \times n_{features}
+ (batch_0 \times \cdots \times batch_n \times) n_{features} That is, ``data`` is a (batch) vector with :math:`n_{features}`
- components. The :math:`batch\_size` is optional.
+ components. The :math:`batch` sizes are optional. dim : int New feature dimension.
+ axis : int
+ Axis where the ``data`` tensor is 'expanded'. Should be between 0 and
+ the rank of ``data``. By default, it is -1, which returns a tensor with
+ shape
+
+ .. math::
+
+ (batch_0 \times \cdots \times batch_n \times) n_{features}
+ \times dim Returns ------- torch.Tensor
- New data tensor with shape
-
- .. math::
-
- batch\_size \times n_{features} \times dim Examples --------
@@ -356,17 +363,20 @@
Parameters ---------- data : torch.Tensor
- Data tensor with shape
+ Data tensor with shape .. math::
- batch\_size \times n_{features}
-
+ (batch_0 \times \cdots \times batch_n \times) n_{features}
+ That is, ``data`` is a (batch) vector with :math:`n_{features}`
- components. The :math:`batch\_size` is optional.
+ components. The :math:`batch` sizes are optional. axis : int
- Axis where the ``data`` tensor is 'expanded' with the 1's. Should be
- between 0 and the rank of ``data``. By default, it is -1, which returns
- a tensor with shape
+ Axis where the ``data`` tensor is 'expanded'. Should be between 0 and
+ the rank of ``data``. By default, it is -1, which returns a tensor with
+ shape .. math::
- batch\_size \times n_{features} \times 2
+ (batch_0 \times \cdots \times batch_n \times) n_{features}
+ \times 2 Returns -------
@@ -435,6 +446,8 @@
Source code for tensorkrowch.embeddings
>>> emb_b.shape torch.Size([100, 5, 2]) """
+ ifnotisinstance(data,torch.Tensor):
+ raiseTypeError('`data` should be torch.Tensor type')returntorch.stack([torch.ones_like(data),data],dim=axis)
@@ -471,20 +484,22 @@
Source code for tensorkrowch.embeddings
.. math::
- batch\_size \times n_{features}
-
+ (batch_0 \times \cdots \times batch_n \times) n_{features}
+ That is, ``data`` is a (batch) vector with :math:`n_{features}`
- components. The :math:`batch\_size` is optional.
+ components. The :math:`batch` sizes are optional. degree : int
- Maximum degree of the monomials.
+ Maximum degree of the monomials. The feature dimension will be
+ ``degree + 1``. axis : int
- Axis where the ``data`` tensor is 'expanded' with monomials. Should be
- between 0 and the rank of ``data``. By default, it is -1, which returns
- a tensor with shape
+ Axis where the ``data`` tensor is 'expanded'. Should be between 0 and
+ the rank of ``data``. By default, it is -1, which returns a tensor with
+ shape .. math::
- batch\_size \times n_{features} \times (degree + 1)
+ (batch_0 \times \cdots \times batch_n \times) n_{features}
+ \times (degree + 1) Returns -------
@@ -505,14 +520,201 @@
[docs]defdiscretize(data:torch.Tensor,
+ level:int,
+ base:int=2,
+ axis:int=-1)->torch.Tensor:
+r"""
+ Embedds the data tensor discretizing each variable in a certain ``basis``
+ and with a certain ``level`` of precision, assuming the values to discretize
+ are all between 0 and 1. That is, given a vector
+
+ .. math::
+
+ x = \begin{bmatrix}
+ x_1\\
+ \vdots\\
+ x_N
+ \end{bmatrix}
+
+ returns a matrix
+
+ .. math::
+
+ \hat{x} = \begin{bmatrix}
+ \lfloor x_1 b^1 \rfloor \mod b & \cdots &
+ \lfloor x_1 b^{l} \rfloor \mod b\\
+ \vdots & \ddots & \vdots\\
+ \lfloor x_N b^1 \rfloor \mod b & \cdots &
+ \lfloor x_N b^{l} \rfloor \mod b
+ \end{bmatrix}
+
+ where :math:`b` stands for ``base``, and :math:`l` for ``level``.
+
+ Parameters
+ ----------
+ data : torch.Tensor
+ Data tensor with shape
+
+ .. math::
+
+ (batch_0 \times \cdots \times batch_n \times) n_{features}
+
+ That is, ``data`` is a (batch) vector with :math:`n_{features}`
+ components. The :math:`batch` sizes are optional. The ``data`` tensor
+ is assumed to have elements between 0 and 1.
+ level : int
+ Level of precision of the discretization. This will be the new feature
+ dimension.
+ base : int
+ The base of the discretization.
+ axis : int
+ Axis where the ``data`` tensor is 'expanded'. Should be between 0 and
+ the rank of ``data``. By default, it is -1, which returns a tensor with
+ shape
+
+ .. math::
+
+ (batch_0 \times \cdots \times batch_n \times) n_{features}
+ \times level
+
+ Returns
+ -------
+ torch.Tensor
+
+ Examples
+ --------
+ >>> a = torch.tensor([0, 0.5, 0.75, 1])
+ >>> a
+ tensor([0.0000, 0.5000, 0.7500, 1.0000])
+
+ >>> emb_a = tk.embeddings.discretize(a, level=3)
+ >>> emb_a
+ tensor([[0., 0., 0.],
+ [1., 0., 0.],
+ [1., 1., 0.],
+ [1., 1., 1.]])
+
+ >>> b = torch.rand(100, 5)
+ >>> emb_b = tk.embeddings.discretize(b, level=3)
+ >>> emb_b.shape
+ torch.Size([100, 5, 3])
+ """
+ ifnotisinstance(data,torch.Tensor):
+ raiseTypeError('`data` should be torch.Tensor type')
+ ifnottorch.ge(data,torch.zeros_like(data)).all():
+ raiseValueError('Elements of `data` should be between 0 and 1')
+ ifnottorch.le(data,torch.ones_like(data)).all():
+ raiseValueError('Elements of `data` should be between 0 and 1')
+
+ max_discr_value=(base-1)*sum([base**-iforiinrange(1,level+1)])
+ data=torch.where(data>max_discr_value,max_discr_value,data)
+
+ base=torch.tensor(base,device=data.device)
+ ids=[torch.remainder((data*base.pow(i)).floor(),base)
+ foriinrange(1,level+1)]
+ ids=torch.stack(ids,dim=axis)
+ returnids
+
+
+
[docs]defbasis(data:torch.Tensor,dim:int=2,axis:int=-1)->torch.Tensor:
+r"""
+ Embedds the data tensor transforming each value, assumed to be an integer
+ between 0 and ``dim - 1``, into the corresponding vector of the
+ computational basis. That is, given a vector
+
+ .. math::
+
+ x = \begin{bmatrix}
+ x_1\\
+ \vdots\\
+ x_N
+ \end{bmatrix}
+
+ returns a matrix
+
+ .. math::
+
+ \hat{x} = \begin{bmatrix}
+ \lvert x_1 \rangle\\
+ \vdots\\
+ \lvert x_N \rangle
+ \end{bmatrix}
+
+ Parameters
+ ----------
+ data : torch.Tensor
+ Data tensor with shape
+
+ .. math::
+
+ (batch_0 \times \cdots \times batch_n \times) n_{features}
+
+ That is, ``data`` is a (batch) vector with :math:`n_{features}`
+ components. The :math:`batch` sizes are optional. The ``data`` tensor
+ is assumed to have integer elements between 0 and ``dim - 1``.
+ dim : int
+ The dimension of the computational basis. This will be the new feature
+ dimension.
+ axis : int
+ Axis where the ``data`` tensor is 'expanded'. Should be between 0 and
+ the rank of ``data``. By default, it is -1, which returns a tensor with
+ shape
+
+ .. math::
+
+ (batch_0 \times \cdots \times batch_n \times) n_{features}
+ \times dim
+
+ Returns
+ -------
+ torch.Tensor
+
+ Examples
+ --------
+ >>> a = torch.arange(5)
+ >>> a
+ tensor([0, 1, 2, 3, 4])
+
+ >>> emb_a = tk.embeddings.basis(a, dim=5)
+ >>> emb_a
+ tensor([[1, 0, 0, 0, 0],
+ [0, 1, 0, 0, 0],
+ [0, 0, 1, 0, 0],
+ [0, 0, 0, 1, 0],
+ [0, 0, 0, 0, 1]])
+
+ >>> b = torch.randint(low=0, high=10, size=(100, 5))
+ >>> emb_b = tk.embeddings.basis(b, dim=10)
+ >>> emb_b.shape
+ torch.Size([100, 5, 10])
+ """
+ ifnotisinstance(data,torch.Tensor):
+ raiseTypeError('`data` should be torch.Tensor type')
+ iftorch.is_floating_point(data):
+ raiseValueError('`data` should be a tensor of integers')
+ ifnottorch.ge(data,torch.zeros_like(data)).all():
+ raiseValueError('Elements of `data` should be between 0 and (dim - 1)')
+ ifnottorch.le(data,torch.ones_like(data)*(dim-1)).all():
+ raiseValueError('Elements of `data` should be between 0 and (dim - 1)')
+
+ ids=torch.arange(dim,device=data.device).repeat(*data.shape,1)
+ ids=torch.where(ids==data.unsqueeze(-1),1,0)
+ ids=ids.movedim(-1,axis)
+ returnids
[docs]classMPO(TensorNetwork):# MARK: MPO
+"""
+ Class for Matrix Product Operators. This is the base class from which
+ :class:`UMPO` inherits.
+
+ Matrix Product Operators are formed by:
+
+ * ``mats_env``: Environment of `matrix` nodes with axes
+ ``("left", "input", "right", "output")``.
+
+ * ``left_node``, ``right_node``: `Vector` nodes with axes ``("right",)``
+ and ``("left",)``, respectively. These are used to close the boundary
+ in the case ``boudary`` is ``"obc"``. Otherwise, both are ``None``.
+
+ In contrast with :class:`MPS`, in ``MPO`` all nodes act both as input and
+ output, with corresponding edges dedicated to that. Thus, data nodes will
+ be connected to the ``"input"`` edge of all nodes. Upon contraction of the
+ whole network, a resultant tensor will be formed, with as many dimensions
+ as nodes were in the MPO.
+
+ If all nodes have the same input dimensions, the input data tensor can be
+ passed as a single tensor. Otherwise, it would have to be passed as a list
+ of tensors with different sizes.
+
+ Parameters
+ ----------
+ n_features : int, optional
+ Number of nodes that will be in ``mats_env``. That is, number of nodes
+ without taking into account ``left_node`` and ``right_node``.
+ in_dim : int, list[int] or tuple[int], optional
+ Input dimension(s). If given as a sequence, its length should be equal
+ to ``n_features``.
+ out_dim : int, list[int] or tuple[int], optional
+ Output dimension(s). If given as a sequence, its length should be equal
+ to ``n_features``.
+ bond_dim : int, list[int] or tuple[int], optional
+ Bond dimension(s). If given as a sequence, its length should be equal
+ to ``n_features`` (if ``boundary = "pbc"``) or ``n_features - 1`` (if
+ ``boundary = "obc"``). The i-th bond dimension is always the dimension
+ of the right edge of the i-th node.
+ boundary : {"obc", "pbc"}
+ String indicating whether periodic or open boundary conditions should
+ be used.
+ tensors: list[torch.Tensor] or tuple[torch.Tensor], optional
+ Instead of providing ``n_features``, ``in_dim``, ``in_dim``, ``bond_dim``
+ and ``boundary``, a list of MPO tensors can be provided. In such case,
+ all mentioned attributes will be inferred from the given tensors. All
+ tensors should be rank-4 tensors, with shape ``(bond_dim, in_dim,
+ bond_dim, out_dim)``. If the first and last elements are rank-3 tensors,
+ with shapes ``(in_dim, bond_dim, out_dim)``, ``(bond_dim, in_dim, out_dim)``,
+ respectively, the inferred boundary conditions will be "obc". Also, if
+ ``tensors`` contains a single element, it can be rank-2 ("obc") or
+ rank-4 ("pbc").
+ n_batches : int
+ Number of batch edges of input ``data`` nodes. Usually ``n_batches = 1``
+ (where the batch edge is used for the data batched) but it could also
+ be ``n_batches = 2`` (one edge for data batched, other edge for image
+ patches in convolutional layers).
+ init_method : {"zeros", "ones", "copy", "rand", "randn"}, optional
+ Initialization method. Check :meth:`initialize` for a more detailed
+ explanation of the different initialization methods.
+ device : torch.device, optional
+ Device where to initialize the tensors if ``init_method`` is provided.
+ kwargs : float
+ Keyword arguments for the different initialization methods. See
+ :meth:`~tensorkrowch.AbstractNode.make_tensor`.
+
+ Examples
+ --------
+ ``MPO`` with same input/output dimensions:
+
+ >>> mpo = tk.models.MPO(n_features=5,
+ ... in_dim=2,
+ ... out_dim=2,
+ ... bond_dim=5)
+ >>> data = torch.ones(20, 5, 2) # batch_size x n_features x feature_size
+ >>> result = mpo(data)
+ >>> result.shape
+ torch.Size([20, 2, 2, 2, 2, 2])
+
+ ``MPO`` with different input/physical dimensions:
+
+ >>> mpo = tk.models.MPO(n_features=5,
+ ... in_dim=list(range(2, 7)),
+ ... out_dim=list(range(7, 2, -1)),
+ ... bond_dim=5)
+ >>> data = [torch.ones(20, i)
+ ... for i in range(2, 7)] # n_features * [batch_size x feature_size]
+ >>> result = mpo(data)
+ >>> result.shape
+ torch.Size([20, 7, 6, 5, 4, 3])
+ """
+
+ def__init__(self,
+ n_features:Optional[int]=None,
+ in_dim:Optional[Union[int,Sequence[int]]]=None,
+ out_dim:Optional[Union[int,Sequence[int]]]=None,
+ bond_dim:Optional[Union[int,Sequence[int]]]=None,
+ boundary:Text='obc',
+ tensors:Optional[Sequence[torch.Tensor]]=None,
+ n_batches:int=1,
+ init_method:Text='randn',
+ device:Optional[torch.device]=None,
+ **kwargs)->None:
+
+ super().__init__(name='mpo')
+
+ iftensorsisNone:
+ # boundary
+ ifboundarynotin['obc','pbc']:
+ raiseValueError('`boundary` should be one of "obc" or "pbc"')
+ self._boundary=boundary
+
+ # n_features
+ ifnotisinstance(n_features,int):
+ raiseTypeError('`n_features` should be int type')
+ elifn_features<1:
+ raiseValueError('`n_features` should be at least 1')
+ self._n_features=n_features
+
+ # in_dim
+ ifisinstance(in_dim,(list,tuple)):
+ iflen(in_dim)!=n_features:
+ raiseValueError('If `in_dim` is given as a sequence of int, '
+ 'its length should be equal to `n_features`')
+ self._in_dim=list(in_dim)
+ elifisinstance(in_dim,int):
+ self._in_dim=[in_dim]*n_features
+ else:
+ raiseTypeError('`in_dim` should be int, tuple[int] or list[int] '
+ 'type')
+
+ # out_dim
+ ifisinstance(out_dim,(list,tuple)):
+ iflen(out_dim)!=n_features:
+ raiseValueError('If `out_dim` is given as a sequence of int, '
+ 'its length should be equal to `n_features`')
+ self._out_dim=list(out_dim)
+ elifisinstance(out_dim,int):
+ self._out_dim=[out_dim]*n_features
+ else:
+ raiseTypeError('`out_dim` should be int, tuple[int] or list[int] '
+ 'type')
+
+ # bond_dim
+ ifisinstance(bond_dim,(list,tuple)):
+ ifboundary=='obc':
+ iflen(bond_dim)!=n_features-1:
+ raiseValueError(
+ 'If `bond_dim` is given as a sequence of int, and '
+ '`boundary` is "obc", its length should be equal '
+ 'to `n_features` - 1')
+ elifboundary=='pbc':
+ iflen(bond_dim)!=n_features:
+ raiseValueError(
+ 'If `bond_dim` is given as a sequence of int, and '
+ '`boundary` is "pbc", its length should be equal '
+ 'to `n_features`')
+ self._bond_dim=list(bond_dim)
+ elifisinstance(bond_dim,int):
+ ifboundary=='obc':
+ self._bond_dim=[bond_dim]*(n_features-1)
+ elifboundary=='pbc':
+ self._bond_dim=[bond_dim]*n_features
+ else:
+ raiseTypeError('`bond_dim` should be int, tuple[int] or list[int]'
+ ' type')
+
+ else:
+ ifnotisinstance(tensors,(list,tuple)):
+ raiseTypeError('`tensors` should be a tuple[torch.Tensor] or '
+ 'list[torch.Tensor] type')
+ else:
+ self._n_features=len(tensors)
+ self._in_dim=[]
+ self._out_dim=[]
+ self._bond_dim=[]
+ fori,tinenumerate(tensors):
+ ifnotisinstance(t,torch.Tensor):
+ raiseTypeError('`tensors` should be a tuple[torch.Tensor]'
+ ' or list[torch.Tensor] type')
+
+ ifi==0:
+ iflen(t.shape)notin[2,3,4]:
+ raiseValueError(
+ 'The first and last elements in `tensors` '
+ 'should be both rank-3 or rank-4 tensors. If'
+ ' the first element is also the last one,'
+ ' it should be a rank-2 tensor')
+ iflen(t.shape)==2:
+ self._boundary='obc'
+ self._in_dim.append(t.shape[0])
+ self._out_dim.append(t.shape[1])
+ eliflen(t.shape)==3:
+ self._boundary='obc'
+ self._in_dim.append(t.shape[0])
+ self._bond_dim.append(t.shape[1])
+ self._out_dim.append(t.shape[2])
+ else:
+ self._boundary='pbc'
+ self._in_dim.append(t.shape[1])
+ self._bond_dim.append(t.shape[2])
+ self._out_dim.append(t.shape[3])
+ elifi==(self._n_features-1):
+ iflen(t.shape)!=len(tensors[0].shape):
+ raiseValueError(
+ 'The first and last elements in `tensors` '
+ 'should have the same rank. Both should be '
+ 'rank-3 or rank-4 tensors. If the first '
+ 'element is also the last one, it should '
+ 'be a rank-2 tensor')
+ iflen(t.shape)==3:
+ self._in_dim.append(t.shape[1])
+ self._out_dim.append(t.shape[2])
+ else:
+ ift.shape[2]!=tensors[0].shape[0]:
+ raiseValueError(
+ 'If the first and last elements in `tensors`'
+ ' are rank-4 tensors, the first dimension '
+ 'of the first element should coincide with'
+ ' the third dimension of the last element')
+ self._in_dim.append(t.shape[1])
+ self._bond_dim.append(t.shape[2])
+ self._out_dim.append(t.shape[3])
+ else:
+ iflen(t.shape)!=4:
+ raiseValueError(
+ 'The elements of `tensors` should be rank-4 '
+ 'tensors, except the first and lest elements'
+ ' if boundary is "obc"')
+ self._in_dim.append(t.shape[1])
+ self._bond_dim.append(t.shape[2])
+ self._out_dim.append(t.shape[3])
+
+ # n_batches
+ ifnotisinstance(n_batches,int):
+ raiseTypeError('`n_batches` should be int type')
+ self._n_batches=n_batches
+
+ # Properties
+ self._left_node=None
+ self._right_node=None
+ self._mats_env=[]
+
+ # Create Tensor Network
+ self._make_nodes()
+ self.initialize(tensors=tensors,
+ init_method=init_method,
+ device=device,
+ **kwargs)
+
+ # ----------
+ # Properties
+ # ----------
+ @property
+ defn_features(self)->int:
+"""Returns number of nodes."""
+ returnself._n_features
+
+ @property
+ defin_dim(self)->List[int]:
+"""Returns input dimensions."""
+ returnself._in_dim
+
+ @property
+ defout_dim(self)->List[int]:
+"""Returns output dimensions."""
+ returnself._out_dim
+
+ @property
+ defbond_dim(self)->List[int]:
+"""Returns bond dimensions."""
+ returnself._bond_dim
+
+ @property
+ defboundary(self)->Text:
+"""Returns boundary condition ("obc" or "pbc")."""
+ returnself._boundary
+
+ @property
+ defn_batches(self)->int:
+"""
+ Returns number of batch edges of the ``data`` nodes. To change this
+ attribute, first call :meth:`~tensorkrowch.TensorNetwork.unset_data_nodes`
+ if there are already data nodes in the network.
+ """
+ returnself._n_batches
+
+ @n_batches.setter
+ defn_batches(self,n_batches:int)->None:
+ ifn_batches!=self._n_batches:
+ ifself._data_nodes:
+ raiseValueError(
+ '`n_batches` cannot be changed if the MPS has data nodes. '
+ 'Use unset_data_nodes first')
+ elifnotisinstance(n_batches,int):
+ raiseTypeError('`n_batches` should be int type')
+ self._n_batches=n_batches
+
+ @property
+ defleft_node(self)->Optional[AbstractNode]:
+"""Returns the ``left_node``."""
+ returnself._left_node
+
+ @property
+ defright_node(self)->Optional[AbstractNode]:
+"""Returns the ``right_node``."""
+ returnself._right_node
+
+ @property
+ defmats_env(self)->List[AbstractNode]:
+"""Returns the list of nodes in ``mats_env``."""
+ returnself._mats_env
+
+ # -------
+ # Methods
+ # -------
+ def_make_nodes(self)->None:
+"""Creates all the nodes of the MPO."""
+ ifself._leaf_nodes:
+ raiseValueError('Cannot create MPO nodes if the MPO already has '
+ 'nodes')
+
+ aux_bond_dim=self._bond_dim
+
+ ifself._boundary=='obc':
+ ifnotaux_bond_dim:
+ aux_bond_dim=[1]
+
+ self._left_node=ParamNode(shape=(aux_bond_dim[0],),
+ axes_names=('right',),
+ name='left_node',
+ network=self)
+ self._right_node=ParamNode(shape=(aux_bond_dim[-1],),
+ axes_names=('left',),
+ name='right_node',
+ network=self)
+
+ aux_bond_dim=aux_bond_dim+[aux_bond_dim[-1]]+[aux_bond_dim[0]]
+
+ foriinrange(self._n_features):
+ node=ParamNode(shape=(aux_bond_dim[i-1],
+ self._in_dim[i],
+ aux_bond_dim[i],
+ self._out_dim[i]),
+ axes_names=('left','input','right','output'),
+ name=f'mats_env_node_({i})',
+ network=self)
+ self._mats_env.append(node)
+
+ ifi!=0:
+ self._mats_env[-2]['right']^self._mats_env[-1]['left']
+
+ ifself._boundary=='pbc':
+ ifi==0:
+ periodic_edge=self._mats_env[-1]['left']
+ ifi==self._n_features-1:
+ self._mats_env[-1]['right']^periodic_edge
+ else:
+ ifi==0:
+ self._left_node['right']^self._mats_env[-1]['left']
+ ifi==self._n_features-1:
+ self._mats_env[-1]['right']^self._right_node['left']
+
+
[docs]definitialize(self,
+ tensors:Optional[Sequence[torch.Tensor]]=None,
+ init_method:Optional[Text]='randn',
+ device:Optional[torch.device]=None,
+ **kwargs:float)->None:
+"""
+ Initializes all the nodes of the :class:`MPO`. It can be called when
+ instantiating the model, or to override the existing nodes' tensors.
+
+ There are different methods to initialize the nodes:
+
+ * ``{"zeros", "ones", "copy", "rand", "randn"}``: Each node is
+ initialized calling :meth:`~tensorkrowch.AbstractNode.set_tensor` with
+ the given method, ``device`` and ``kwargs``.
+
+ Parameters
+ ----------
+ tensors : list[torch.Tensor] or tuple[torch.Tensor], optional
+ Sequence of tensors to set in each of the MPO nodes. If ``boundary``
+ is ``"obc"``, all tensors should be rank-4, except the first and
+ last ones, which can be rank-3, or rank-2 (if the first and last are
+ the same). If ``boundary`` is ``"pbc"``, all tensors should be
+ rank-4.
+ init_method : {"zeros", "ones", "copy", "rand", "randn"}, optional
+ Initialization method.
+ device : torch.device, optional
+ Device where to initialize the tensors if ``init_method`` is provided.
+ kwargs : float
+ Keyword arguments for the different initialization methods. See
+ :meth:`~tensorkrowch.AbstractNode.make_tensor`.
+ """
+ ifself._boundary=='obc':
+ self._left_node.set_tensor(init_method='copy',device=device)
+ self._right_node.set_tensor(init_method='copy',device=device)
+
+ iftensorsisnotNone:
+ iflen(tensors)!=self._n_features:
+ raiseValueError('`tensors` should be a sequence of `n_features`'
+ ' elements')
+
+ ifself._boundary=='obc':
+ tensors=tensors[:]
+ iflen(tensors)==1:
+ tensors[0]=tensors[0].reshape(1,
+ tensors[0].shape[0],
+ 1,
+ tensors[0].shape[1])
+
+ else:
+ # Left node
+ aux_tensor=torch.zeros(*self._mats_env[0].shape,
+ device=tensors[0].device)
+ aux_tensor[0]=tensors[0]
+ tensors[0]=aux_tensor
+
+ # Right node
+ aux_tensor=torch.zeros(*self._mats_env[-1].shape,
+ device=tensors[-1].device)
+ aux_tensor[...,0,:]=tensors[-1]
+ tensors[-1]=aux_tensor
+
+ fortensor,nodeinzip(tensors,self._mats_env):
+ node.tensor=tensor
+
+ elifinit_methodisnotNone:
+
+ fori,nodeinenumerate(self._mats_env):
+ node.set_tensor(init_method=init_method,
+ device=device,
+ **kwargs)
+
+ ifself._boundary=='obc':
+ aux_tensor=torch.zeros(*node.shape,device=device)
+ ifi==0:
+ # Left node
+ aux_tensor[0]=node.tensor[0]
+ elifi==(self._n_features-1):
+ # Right node
+ aux_tensor[...,0,:]=node.tensor[...,0,:]
+ node.tensor=aux_tensor
+
+
[docs]defset_data_nodes(self)->None:
+"""
+ Creates ``data`` nodes and connects each of them to the ``"input"``
+ edge of each node.
+ """
+ input_edges=[node['input']fornodeinself._mats_env]
+ super().set_data_nodes(input_edges=input_edges,
+ num_batch_edges=self._n_batches)
+
+
[docs]defcopy(self,share_tensors:bool=False)->'MPO':
+"""
+ Creates a copy of the :class:`MPO`.
+
+ Parameters
+ ----------
+ share_tensor : bool, optional
+ Boolean indicating whether tensors in the copied MPO should be
+ set as the tensors in the current MPO (``True``), or cloned
+ (``False``). In the former case, tensors in both MPO's will be
+ the same, which might be useful if one needs more than one copy
+ of an MPO, but wants to compute all the gradients with respect
+ to the same, unique, tensors.
+
+ Returns
+ -------
+ MPO
+ """
+ new_mpo=MPO(n_features=self._n_features,
+ in_dim=self._in_dim,
+ out_dim=self._out_dim,
+ bond_dim=self._bond_dim,
+ boundary=self._boundary,
+ tensors=None,
+ n_batches=self._n_batches,
+ init_method=None,
+ device=None)
+ new_mpo.name=self.name+'_copy'
+ ifshare_tensors:
+ fornew_node,nodeinzip(new_mpo._mats_env,self._mats_env):
+ new_node.tensor=node.tensor
+ else:
+ fornew_node,nodeinzip(new_mpo._mats_env,self._mats_env):
+ new_node.tensor=node.tensor.clone()
+
+ returnnew_mpo
+
+
[docs]defparameterize(self,
+ set_param:bool=True,
+ override:bool=False)->'TensorNetwork':
+"""
+ Parameterizes all nodes of the MPO. If there are ``resultant`` nodes
+ in the MPO, it will be first :meth:`~tensorkrowch.TensorNetwork.reset`.
+
+ Parameters
+ ----------
+ set_param : bool
+ Boolean indicating whether the tensor network has to be parameterized
+ (``True``) or de-parameterized (``False``).
+ override : bool
+ Boolean indicating whether the tensor network should be parameterized
+ in-place (``True``) or copied and then parameterized (``False``).
+ """
+ ifself._resultant_nodes:
+ warnings.warn(
+ 'Resultant nodes will be removed before parameterizing the TN')
+ self.reset()
+
+ ifoverride:
+ net=self
+ else:
+ net=self.copy(share_tensors=False)
+
+ foriinrange(self._n_features):
+ net._mats_env[i]=net._mats_env[i].parameterize(set_param)
+
+ ifnet._boundary=='obc':
+ net._left_node=net._left_node.parameterize(set_param)
+ net._right_node=net._right_node.parameterize(set_param)
+
+ returnnet
[docs]defcontract(self,
+ inline_input:bool=False,
+ inline_mats:bool=False,
+ mps:Optional[MPSData]=None)->Node:
+"""
+ Contracts the whole MPO with input data nodes. The input can be in the
+ form of an :class:`MPSData`, which may be convenient for tensorizing
+ vector-matrix multiplication in the form of MPS-MPO contraction.
+
+ If the ``MPO`` is contracted with a ``MPSData``, MPS nodes will become
+ part of the MPO network, and they will be connected to the ``"input"``
+ edges of the MPO. Thus, the MPS and the MPO should have the same number
+ of features (``n_features``).
+
+ Even though it is not necessary to connect the ``MPSData`` nodes to the
+ MPO nodes by hand before contraction, it can be done. However, one
+ should first move the MPS nodes to the MPO network.
+
+ Also, when contracting the MPO with and ``MPSData``, if any of the
+ contraction arguments, ``inline_input`` or ``inline_mats``, is set to
+ ``False``, the MPO (already connected to the MPS) should be
+ :meth:`~tensorkrowch.TensorNetwork.reset` before contraction if new
+ data is set into the ``MPSData`` nodes. This is because :class:`MPSData`
+ admits data tensors with different bond dimensions for each iteration,
+ and this may cause undesired behaviour when reusing some information of
+ previous calls to :func:~tensorkrowch.stack` with the previous data
+ tensors.
+
+ To perform the MPS-MPO contraction, first input data tensors have to
+ be put into the :class:`MPSData` via :meth:`MPSData.add_data`. Then,
+ contraction is carried out by calling ``mpo(mps=mps_data)``, without
+ passing the input data again, as it is already stored in the MPSData
+ nodes.
+
+ Parameters
+ ----------
+ inline_input : bool
+ Boolean indicating whether input ``data`` nodes should be contracted
+ with the ``MPO`` nodes inline (one contraction at a time) or in a
+ single stacked contraction.
+ inline_mats : bool
+ Boolean indicating whether the sequence of matrices (resultant
+ after contracting the input ``data`` nodes) should be contracted
+ inline or as a sequence of pairwise stacked contrations.
+ mps : MPSData, optional
+ MPS that is to be contracted with the MPO. New data can be
+ put into the MPS via :meth:`MPSData.add_data`, and the MPS-MPO
+ contraction is performed by calling ``mpo(mps=mps_data)``, without
+ passing the input data again, as it is already stored in the MPS
+ cores.
+
+ Returns
+ -------
+ Node
+ """
+ ifmpsisnotNone:
+ ifnotisinstance(mps,MPSData):
+ raiseTypeError('`mps` should be MPSData type')
+ ifmps._n_features!=self._n_features:
+ raiseValueError(
+ '`mps` should have as many features as the MPO')
+
+ # Move MPSData ndoes to self
+ mps._mats_env[0].move_to_network(self)
+
+ # Connect mps nodes to mpo nodes
+ formps_node,mpo_nodeinzip(mps._mats_env,self._mats_env):
+ mps_node['feature']^mpo_node['input']
+
+ mats_env=self._input_contraction(self._mats_env,inline_input)
+
+ ifinline_mats:
+ result=self._contract_envs_inline(mats_env,mps)
+ else:
+ result=self._pairwise_contraction(mats_env,mps)
+
+ # Contract periodic edge
+ ifresult.is_connected_to(result):
+ result@=result
+
+ # Put batch edges in first positions
+ batch_edges=[]
+ other_edges=[]
+ fori,edgeinenumerate(result.edges):
+ ifedge.is_batch():
+ batch_edges.append(i)
+ else:
+ other_edges.append(i)
+
+ all_edges=batch_edges+other_edges
+ ifall_edges!=list(range(len(all_edges))):
+ result=op.permute(result,tuple(all_edges))
+
+ returnresult
+
+
+
[docs]classUMPO(MPO):# MARK: UMPO
+"""
+ Class for Uniform (translationally invariant) Matrix Product Operators. It is
+ the uniform version of :class:`MPO`, that is, all nodes share the same
+ tensor. Thus this class cannot have different input/output or bond dimensions
+ for each node, and boundary conditions are always periodic (``"pbc"``).
+
+ |
+
+ For a more detailed list of inherited properties and methods,
+ check :class:`MPO`.
+
+ Parameters
+ ----------
+ n_features : int
+ Number of nodes that will be in ``mats_env``.
+ in_dim : int, optional
+ Input dimension.
+ out_dim : int, optional
+ Output dimension.
+ bond_dim : int, optional
+ Bond dimension.
+ tensor: torch.Tensor, optional
+ Instead of providing ``in_dim``, ``out_dim`` and ``bond_dim``, a single
+ tensor can be provided. ``n_features`` is still needed to specify how
+ many times the tensor should be used to form a finite MPO. The tensor
+ should be rank-4, with its first and third dimensions being equal.
+ n_batches : int
+ Number of batch edges of input ``data`` nodes. Usually ``n_batches = 1``
+ (where the batch edge is used for the data batched) but it could also
+ be ``n_batches = 2`` (one edge for data batched, other edge for image
+ patches in convolutional layers).
+ init_method : {"zeros", "ones", "copy", "rand", "randn"}, optional
+ Initialization method. Check :meth:`initialize` for a more detailed
+ explanation of the different initialization methods.
+ device : torch.device, optional
+ Device where to initialize the tensors if ``init_method`` is provided.
+ kwargs : float
+ Keyword arguments for the different initialization methods. See
+ :meth:`~tensorkrowch.AbstractNode.make_tensor`.
+
+ Examples
+ --------
+ >>> mpo = tk.models.UMPO(n_features=4,
+ ... in_dim=2,
+ ... out_dim=2,
+ ... bond_dim=5)
+ >>> for node in mpo.mats_env:
+ ... assert node.tensor_address() == 'virtual_uniform'
+ ...
+ >>> data = torch.ones(20, 4, 2) # batch_size x n_features x feature_size
+ >>> result = mpo(data)
+ >>> result.shape
+ torch.Size([20, 2, 2, 2, 2])
+ """
+
+ def__init__(self,
+ n_features:int=None,
+ in_dim:Optional[int]=None,
+ out_dim:Optional[int]=None,
+ bond_dim:Optional[int]=None,
+ tensor:Optional[torch.Tensor]=None,
+ n_batches:int=1,
+ init_method:Text='randn',
+ device:Optional[torch.device]=None,
+ **kwargs)->None:
+
+ tensors=None
+
+ # n_features
+ ifnotisinstance(n_features,int):
+ raiseTypeError('`n_features` should be int type')
+ elifn_features<1:
+ raiseValueError('`n_features` should be at least 1')
+
+ iftensorisNone:
+ # in_dim
+ ifnotisinstance(in_dim,int):
+ raiseTypeError('`in_dim` should be int type')
+
+ # out_dim
+ ifnotisinstance(out_dim,int):
+ raiseTypeError('`out_dim` should be int type')
+
+ # bond_dim
+ ifnotisinstance(bond_dim,int):
+ raiseTypeError('`bond_dim` should be int type')
+
+ else:
+ ifnotisinstance(tensor,torch.Tensor):
+ raiseTypeError('`tensor` should be torch.Tensor type')
+ iflen(tensor.shape)!=4:
+ raiseValueError('`tensor` should be a rank-4 tensor')
+ iftensor.shape[0]!=tensor.shape[2]:
+ raiseValueError('`tensor` first and last dimensions should'
+ ' be equal so that the MPS can have '
+ 'periodic boundary conditions')
+
+ tensors=[tensor]*n_features
+
+ super().__init__(n_features=n_features,
+ in_dim=in_dim,
+ out_dim=out_dim,
+ bond_dim=bond_dim,
+ boundary='pbc',
+ tensors=tensors,
+ n_batches=n_batches,
+ init_method=init_method,
+ device=device,
+ **kwargs)
+ self.name='umpo'
+
+ def_make_nodes(self)->None:
+"""Creates all the nodes of the MPO."""
+ super()._make_nodes()
+
+ # Virtual node
+ uniform_memory=ParamNode(shape=(self._bond_dim[0],
+ self._in_dim[0],
+ self._bond_dim[0],
+ self._out_dim[0]),
+ axes_names=('left','input','right','output'),
+ name='virtual_uniform',
+ network=self,
+ virtual=True)
+ self.uniform_memory=uniform_memory
+
+ fornodeinself._mats_env:
+ node.set_tensor_from(uniform_memory)
+
+
[docs]definitialize(self,
+ tensors:Optional[Sequence[torch.Tensor]]=None,
+ init_method:Optional[Text]='randn',
+ device:Optional[torch.device]=None,
+ **kwargs:float)->None:
+"""
+ Initializes the common tensor of the :class:`UMPO`. It can be called
+ when instantiating the model, or to override the existing nodes' tensors.
+
+ There are different methods to initialize the nodes:
+
+ * ``{"zeros", "ones", "copy", "rand", "randn"}``: The tensor is
+ initialized calling :meth:`~tensorkrowch.AbstractNode.set_tensor` with
+ the given method, ``device`` and ``kwargs``.
+
+ Parameters
+ ----------
+ tensors : list[torch.Tensor] or tuple[torch.Tensor], optional
+ Sequence of a single tensor to set in each of the MPO nodes. The
+ tensor should be rank-4, with its first and third dimensions being
+ equal.
+ init_method : {"zeros", "ones", "copy", "rand", "randn"}, optional
+ Initialization method.
+ device : torch.device, optional
+ Device where to initialize the tensors if ``init_method`` is provided.
+ kwargs : float
+ Keyword arguments for the different initialization methods. See
+ :meth:`~tensorkrowch.AbstractNode.make_tensor`.
+ """
+ node=self.uniform_memory
+
+ iftensorsisnotNone:
+ self.uniform_memory.tensor=tensors[0]
+
+ elifinit_methodisnotNone:
+ self.uniform_memory.set_tensor(init_method=init_method,
+ device=device,
+ **kwargs)
+
+
[docs]defcopy(self,share_tensors:bool=False)->'UMPO':
+"""
+ Creates a copy of the :class:`UMPO`.
+
+ Parameters
+ ----------
+ share_tensor : bool, optional
+ Boolean indicating whether the common tensor in the copied UMPO
+ should be set as the tensor in the current UMPO (``True``), or
+ cloned (``False``). In the former case, the tensor in both UMPO's
+ will be the same, which might be useful if one needs more than one
+ copy of a UMPO, but wants to compute all the gradients with respect
+ to the same, unique, tensor.
+
+ Returns
+ -------
+ UMPO
+ """
+ new_mpo=UMPO(n_features=self._n_features,
+ in_dim=self._in_dim[0],
+ out_dim=self._out_dim[0],
+ bond_dim=self._bond_dim[0],
+ tensor=None,
+ n_batches=self._n_batches,
+ init_method=None,
+ device=None)
+ new_mpo.name=self.name+'_copy'
+ ifshare_tensors:
+ new_mpo.uniform_memory.tensor=self.uniform_memory.tensor
+ else:
+ new_mpo.uniform_memory.tensor=self.uniform_memory.tensor.clone()
+ returnnew_mpo
+
+
[docs]defparameterize(self,
+ set_param:bool=True,
+ override:bool=False)->'TensorNetwork':
+"""
+ Parameterizes all nodes of the MPO. If there are ``resultant`` nodes
+ in the MPO, it will be first :meth:`~tensorkrowch.TensorNetwork.reset`.
+
+ Parameters
+ ----------
+ set_param : bool
+ Boolean indicating whether the tensor network has to be parameterized
+ (``True``) or de-parameterized (``False``).
+ override : bool
+ Boolean indicating whether the tensor network should be parameterized
+ in-place (``True``) or copied and then parameterized (``False``).
+ """
+ ifself._resultant_nodes:
+ warnings.warn(
+ 'Resultant nodes will be removed before parameterizing the TN')
+ self.reset()
+
+ ifoverride:
+ net=self
+ else:
+ net=self.copy(share_tensors=False)
+
+ foriinrange(self._n_features):
+ net._mats_env[i]=net._mats_env[i].parameterize(set_param)
+
+ # It is important that uniform_memory is parameterized after the rest
+ # of the nodes
+ net.uniform_memory=net.uniform_memory.parameterize(set_param)
+
+ # Tensor addresses have to be reassigned to reference
+ # the uniform memory
+ fornodeinnet._mats_env:
+ node.set_tensor_from(net.uniform_memory)
+
+ returnnet
[docs]classMPS(TensorNetwork):# MARK: MPS"""
- Class for Matrix Product States, where all nodes are input nodes, that is,
- they are all connected to ``data`` nodes that will store the input data
- tensor(s). When contracting the MPS with new input data, the result will
- be a just a number.
+ Class for Matrix Product States. This is the base class from which
+ :class:`UMPS`, :class:`MPSLayer` and :class:`UMPSLayer` inherit.
- If the input dimensions of all the input nodes are equal, the input data
- tensor can be passed as a single tensor. Otherwise, it would have to be
- passed as a list of tensors with different sizes.
+ Matrix Product States are formed by:
- An ``MPS`` is formed by the following nodes:
-
- * ``left_node``, ``right_node``: `Vector` nodes with axes ``("input", "right")``
- and ``("left", "input")``, respectively. These are the nodes at the
- extremes of the ``MPS``. If ``boundary`` is ``"pbc""``, both are ``None``.
-
- * ``mats_env``: Environment of `matrix` nodes that. These nodes have axes
+ * ``mats_env``: Environment of `matrix` nodes with axes ``("left", "input", "right")``.
+
+ * ``left_node``, ``right_node``: `Vector` nodes with axes ``("right",)``
+ and ``("left",)``, respectively. These are used to close the boundary
+ in the case ``boudary`` is ``"obc"``. Otherwise, both are ``None``.
+
+ The base ``MPS`` class enables setting various nodes as either input or
+ output nodes. This feature proves useful when computing marginal or
+ conditional distributions. The assignment of roles can be altered
+ dynamically, allowing input nodes to transition to output nodes, and vice
+ versa.
+
+ Input nodes will be connected to data nodes at their ``"input"`` edges, and
+ contracted against them when calling :meth:`contract`. Output nodes, on the
+ other hand, will remain disconnected. If ``marginalize_output = True`` in
+ :meth:`contract`, the open indices of the output nodes can be marginalized
+ so that the output is a single scalar (or a vector with only batch
+ dimensions). If ``marginalize_output = False`` the result will be a tensor
+ with as many dimensions as output nodes where in the MPS, plus the
+ corresponding batch dimensions.
+
+ If all input nodes have the same physical dimensions, the input data tensor
+ can be passed as a single tensor. Otherwise, it would have to be passed as
+ a list of tensors with different sizes. Parameters ----------
- n_features : int
- Number of input nodes.
- in_dim : int, list[int] or tuple[int]
- Input dimension(s). Equivalent to the physical dimension. If given as a
- sequence, its length should be equal to ``n_features``.
- bond_dim : int, list[int] or tuple[int]
+ n_features : int, optional
+ Number of nodes that will be in ``mats_env``. That is, number of nodes
+ without taking into account ``left_node`` and ``right_node``.
+ phys_dim : int, list[int] or tuple[int], optional
+ Physical dimension(s). If given as a sequence, its length should be
+ equal to ``n_features``.
+ bond_dim : int, list[int] or tuple[int], optional Bond dimension(s). If given as a sequence, its length should be equal to ``n_features`` (if ``boundary = "pbc"``) or ``n_features - 1`` (if ``boundary = "obc"``). The i-th bond dimension is always the dimension
@@ -338,254 +367,726 @@
Source code for tensorkrowch.models.mps
boundary : {"obc", "pbc"} String indicating whether periodic or open boundary conditions should be used.
+ tensors: list[torch.Tensor] or tuple[torch.Tensor], optional
+ Instead of providing ``n_features``, ``phys_dim``, ``bond_dim`` and
+ ``boundary``, a list of MPS tensors can be provided. In such case, all
+ mentioned attributes will be inferred from the given tensors. All
+ tensors should be rank-3 tensors, with shape ``(bond_dim, phys_dim,
+ bond_dim)``. If the first and last elements are rank-2 tensors, with
+ shapes ``(phys_dim, bond_dim)``, ``(bond_dim, phys_dim)``, respectively,
+ the inferred boundary conditions will be "obc". Also, if ``tensors``
+ contains a single element, it can be rank-1 ("obc") or rank-3 ("pbc").
+ in_features: list[int] or tuple[int], optional
+ List of indices indicating the positions of the MPS nodes that will be
+ considered as input nodes. These nodes will have a neighbouring data
+ node connected to its ``"input"`` edge when the :meth:`set_data_nodes`
+ method is called. ``in_features`` is the complementary set of
+ ``out_features``, so it is only required to specify one of them.
+ out_features: list[int] or tuple[int], optional
+ List of indices indicating the positions of the MPS nodes that will be
+ considered as output nodes. These nodes will be left with their ``"input"``
+ edges open when contrating the network. If ``marginalize_output`` is
+ set to ``True`` in :meth:`contract`, the network will be connected to
+ itself at these nodes, and contracted. ``out_features`` is the
+ complementary set of ``in_features``, so it is only required to specify
+ one of them. n_batches : int Number of batch edges of input ``data`` nodes. Usually ``n_batches = 1`` (where the batch edge is used for the data batched) but it could also be ``n_batches = 2`` (one edge for data batched, other edge for image patches in convolutional layers).
+ init_method : {"zeros", "ones", "copy", "rand", "randn", "randn_eye", "unit"}, optional
+ Initialization method. Check :meth:`initialize` for a more detailed
+ explanation of the different initialization methods.
+ device : torch.device, optional
+ Device where to initialize the tensors if ``init_method`` is provided.
+ kwargs : float
+ Keyword arguments for the different initialization methods. See
+ :meth:`~tensorkrowch.AbstractNode.make_tensor`. Examples --------
- ``MPS`` with same input/physical dimensions:
+ ``MPS`` with the same physical dimensions: >>> mps = tk.models.MPS(n_features=5,
- ... in_dim=2,
+ ... phys_dim=2, ... bond_dim=5) >>> data = torch.ones(20, 5, 2) # batch_size x n_features x feature_size >>> result = mps(data) >>> result.shape torch.Size([20])
- ``MPS`` with different input/physical dimensions:
+ ``MPS`` with different physical dimensions: >>> mps = tk.models.MPS(n_features=5,
- ... in_dim=list(range(2, 7)),
+ ... phys_dim=list(range(2, 7)), ... bond_dim=5) >>> data = [torch.ones(20, i) ... for i in range(2, 7)] # n_features * [batch_size x feature_size] >>> result = mps(data) >>> result.shape torch.Size([20])
+
+ ``MPS`` can also be initialized from a list of tensors:
+
+ >>> tensors = [torch.randn(5, 2, 5) for _ in range(10)]
+ >>> mps = tk.models.MPS(tensors=tensors)
+
+ If ``in_features``/``out_features`` are specified, data will only be
+ connected to the input nodes, leaving output nodes open:
+
+ >>> mps = tk.models.MPS(tensors=tensors,
+ ... out_features=[0, 3, 9])
+ >>> data = torch.ones(20, 7, 2) # batch_size x n_features x feature_size
+ >>> result = mps(data)
+ >>> result.shape
+ torch.Size([20, 2, 2, 2])
+
+ >>> mps.reset()
+ >>> result = mps(data, marginalize_output=True)
+ >>> result.shape
+ torch.Size([20, 20]) """def__init__(self,
- n_features:int,
- in_dim:Union[int,Sequence[int]],
- bond_dim:Union[int,Sequence[int]],
+ n_features:Optional[int]=None,
+ phys_dim:Optional[Union[int,Sequence[int]]]=None,
+ bond_dim:Optional[Union[int,Sequence[int]]]=None,boundary:Text='obc',
- n_batches:int=1)->None:
+ tensors:Optional[Sequence[torch.Tensor]]=None,
+ in_features:Optional[Sequence[int]]=None,
+ out_features:Optional[Sequence[int]]=None,
+ n_batches:int=1,
+ init_method:Text='randn',
+ device:Optional[torch.device]=None,
+ **kwargs)->None:super().__init__(name='mps')
-
- # boundary
- ifboundarynotin['obc','pbc']:
- raiseValueError('`boundary` should be one of "obc" or "pbc"')
- self._boundary=boundary
-
- # n_features
- ifboundary=='obc':
- ifn_features<2:
- raiseValueError('If `boundary` is "obc", at least '
- 'there has to be 2 nodes')
- elifboundary=='pbc':
- ifn_features<1:
- raiseValueError('If `boundary` is "pbc", at least '
- 'there has to be one node')
- else:
- raiseValueError('`boundary` should be one of "obc" or "pbc"')
- self._n_features=n_features
-
- # in_dim
- ifisinstance(in_dim,(list,tuple)):
- iflen(in_dim)!=n_features:
- raiseValueError('If `in_dim` is given as a sequence of int, '
- 'its length should be equal to `n_features`')
- self._in_dim=list(in_dim)
- elifisinstance(in_dim,int):
- self._in_dim=[in_dim]*n_features
+
+ iftensorsisNone:
+ # boundary
+ ifboundarynotin['obc','pbc']:
+ raiseValueError('`boundary` should be one of "obc" or "pbc"')
+ self._boundary=boundary
+
+ # n_features
+ ifnotisinstance(n_features,int):
+ raiseTypeError('`n_features` should be int type')
+ elifn_features<1:
+ raiseValueError('`n_features` should be at least 1')
+ self._n_features=n_features
+
+ # phys_dim
+ ifisinstance(phys_dim,Sequence):
+ iflen(phys_dim)!=n_features:
+ raiseValueError('If `phys_dim` is given as a sequence of int, '
+ 'its length should be equal to `n_features`')
+ self._phys_dim=list(phys_dim)
+ elifisinstance(phys_dim,int):
+ self._phys_dim=[phys_dim]*n_features
+ else:
+ raiseTypeError('`phys_dim` should be int, tuple[int] or list[int] '
+ 'type')
+
+ # bond_dim
+ ifisinstance(bond_dim,Sequence):
+ ifboundary=='obc':
+ iflen(bond_dim)!=n_features-1:
+ raiseValueError(
+ 'If `bond_dim` is given as a sequence of int, and '
+ '`boundary` is "obc", its length should be equal '
+ 'to `n_features` - 1')
+ elifboundary=='pbc':
+ iflen(bond_dim)!=n_features:
+ raiseValueError(
+ 'If `bond_dim` is given as a sequence of int, and '
+ '`boundary` is "pbc", its length should be equal '
+ 'to `n_features`')
+ self._bond_dim=list(bond_dim)
+ elifisinstance(bond_dim,int):
+ ifboundary=='obc':
+ self._bond_dim=[bond_dim]*(n_features-1)
+ elifboundary=='pbc':
+ self._bond_dim=[bond_dim]*n_features
+ else:
+ raiseTypeError('`bond_dim` should be int, tuple[int] or list[int]'
+ ' type')
+
else:
- raiseTypeError('`in_dim` should be int, tuple[int] or list[int] '
- 'type')
-
- # bond_dim
- ifisinstance(bond_dim,(list,tuple)):
- ifboundary=='obc':
- iflen(bond_dim)!=n_features-1:
- raiseValueError('If `bond_dim` is given as a sequence of int,'
- ' and `boundary` is "obc", its length '
- 'should be equal to `n_features` - 1')
- elifboundary=='pbc':
- iflen(bond_dim)!=n_features:
- raiseValueError('If `bond_dim` is given as a sequence of int,'
- ' and `boundary` is "pbc", its length '
- 'should be equal to `n_features`')
- self._bond_dim=list(bond_dim)
- elifisinstance(bond_dim,int):
- ifboundary=='obc':
- self._bond_dim=[bond_dim]*(n_features-1)
- elifboundary=='pbc':
- self._bond_dim=[bond_dim]*n_features
+ ifnotisinstance(tensors,Sequence):
+ raiseTypeError('`tensors` should be a tuple[torch.Tensor] or '
+ 'list[torch.Tensor] type')
+ else:
+ self._n_features=len(tensors)
+ self._phys_dim=[]
+ self._bond_dim=[]
+ fori,tinenumerate(tensors):
+ ifnotisinstance(t,torch.Tensor):
+ raiseTypeError('`tensors` should be a tuple[torch.Tensor]'
+ ' or list[torch.Tensor] type')
+
+ ifi==0:
+ iflen(t.shape)notin[1,2,3]:
+ raiseValueError(
+ 'The first and last elements in `tensors` '
+ 'should be both rank-2 or rank-3 tensors. If'
+ ' the first element is also the last one,'
+ ' it should be a rank-1 tensor')
+ iflen(t.shape)==1:
+ self._boundary='obc'
+ self._phys_dim.append(t.shape[0])
+ eliflen(t.shape)==2:
+ self._boundary='obc'
+ self._phys_dim.append(t.shape[0])
+ self._bond_dim.append(t.shape[1])
+ else:
+ self._boundary='pbc'
+ self._phys_dim.append(t.shape[1])
+ self._bond_dim.append(t.shape[2])
+ elifi==(self._n_features-1):
+ iflen(t.shape)!=len(tensors[0].shape):
+ raiseValueError(
+ 'The first and last elements in `tensors` '
+ 'should have the same rank. Both should be '
+ 'rank-2 or rank-3 tensors. If the first '
+ 'element is also the last one, it should '
+ 'be a rank-1 tensor')
+ iflen(t.shape)==2:
+ self._phys_dim.append(t.shape[1])
+ else:
+ ift.shape[-1]!=tensors[0].shape[0]:
+ raiseValueError(
+ 'If the first and last elements in `tensors`'
+ ' are rank-3 tensors, the first dimension '
+ 'of the first element should coincide with'
+ ' the last dimension of the last element')
+ self._phys_dim.append(t.shape[1])
+ self._bond_dim.append(t.shape[2])
+ else:
+ iflen(t.shape)!=3:
+ raiseValueError(
+ 'The elements of `tensors` should be rank-3 '
+ 'tensors, except the first and lest elements'
+ ' if boundary is "obc"')
+ self._phys_dim.append(t.shape[1])
+ self._bond_dim.append(t.shape[2])
+
+ # in_features and out_features
+ ifin_featuresisNone:
+ ifout_featuresisNone:
+ # By default, all nodes are input nodes
+ self._in_features=list(range(self._n_features))
+ self._out_features=[]
+ else:
+ ifisinstance(out_features,(list,tuple)):
+ forout_finout_features:
+ ifnotisinstance(out_f,int):
+ raiseTypeError('`out_features` should be tuple[int]'
+ ' or list[int] type')
+ if(out_f<0)or(out_f>=self._n_features):
+ raiseValueError('Elements of `out_features` should'
+ ' be between 0 and (`n_features` - 1)')
+ out_features=set(out_features)
+ in_features=set(range(self._n_features)).difference(out_features)
+
+ self._in_features=list(in_features)
+ self._out_features=list(out_features)
+
+ self._in_features.sort()
+ self._out_features.sort()
+ else:
+ raiseTypeError('`out_features` should be tuple[int]'
+ ' or list[int] type')else:
- raiseTypeError('`in_dim` should be int, tuple[int] or list[int] '
- 'type')
-
+ ifisinstance(in_features,(list,tuple)):
+ forin_finin_features:
+ ifnotisinstance(in_f,int):
+ raiseTypeError('`in_features` should be tuple[int]'
+ ' or list[int] type')
+ if(in_f<0)or(in_f>=self._n_features):
+ raiseValueError('Elements in `in_features` should'
+ 'be between 0 and (`n_features` - 1)')
+ in_features=set(in_features)
+ else:
+ raiseTypeError('`in_features` should be tuple[int]'
+ ' or list[int] type')
+
+ ifout_featuresisNone:
+ out_features=set(range(self._n_features)).difference(in_features)
+
+ self._in_features=list(in_features)
+ self._out_features=list(out_features)
+
+ self._in_features.sort()
+ self._out_features.sort()
+ else:
+ out_features=set(out_features)
+ union=in_features.union(out_features)
+ inter=in_features.intersection(out_features)
+
+ if(union==set(range(self._n_features)))and(inter==set([])):
+ self._in_features=list(in_features)
+ self._out_features=list(out_features)
+
+ self._in_features.sort()
+ self._out_features.sort()
+ else:
+ raiseValueError(
+ 'If both `in_features` and `out_features` are provided,'
+ ' they should be complementary. That is, the union should'
+ ' be the total range 0, ..., (`n_features` - 1), and '
+ 'the intersection should be empty')
+
# n_batchesifnotisinstance(n_batches,int):
- raiseTypeError('`n_batches should be int type')
+ raiseTypeError('`n_batches` should be int type')self._n_batches=n_batches
+
+ # Properties
+ self._left_node=None
+ self._right_node=None
+ self._mats_env=[]# Create Tensor Networkself._make_nodes()
- self.initialize()
-
+ self.initialize(tensors=tensors,
+ init_method=init_method,
+ device=device,
+ **kwargs)
+
+ # ----------
+ # Properties
+ # ----------@propertydefn_features(self)->int:"""Returns number of nodes."""returnself._n_features@property
- defboundary(self)->Text:
-"""Returns boundary condition ("obc" or "pbc")."""
- returnself._boundary
-
- @property
- defin_dim(self)->List[int]:
-"""Returns input/physical dimension."""
- returnself._in_dim
+ defphys_dim(self)->List[int]:
+"""Returns physical dimensions."""
+ returnself._phys_dim@propertydefbond_dim(self)->List[int]:
-"""Returns bond dimension."""
+"""Returns bond dimensions."""returnself._bond_dim
+
+ @property
+ defboundary(self)->Text:
+"""Returns boundary condition ("obc" or "pbc")."""
+ returnself._boundary@propertydefn_batches(self)->int:
-"""Returns number of batch edges of the ``data`` nodes."""
+"""
+ Returns number of batch edges of the ``data`` nodes. To change this
+ attribute, first call :meth:`~tensorkrowch.TensorNetwork.unset_data_nodes`
+ if there are already data nodes in the network.
+ """returnself._n_batches
-
+
+ @n_batches.setter
+ defn_batches(self,n_batches:int)->None:
+ ifn_batches!=self._n_batches:
+ ifself._data_nodes:
+ raiseValueError(
+ '`n_batches` cannot be changed if the MPS has data nodes. '
+ 'Use unset_data_nodes first')
+ elifnotisinstance(n_batches,int):
+ raiseTypeError('`n_batches` should be int type')
+ self._n_batches=n_batches
+
+ @property
+ defin_features(self)->List[int]:
+"""
+ Returns list of positions of the input nodes. To change this
+ attribute, first call :meth:`~tensorkrowch.TensorNetwork.unset_data_nodes`
+ if there are already data nodes in the network. When changing it,
+ :attr:`out_features` will change accordingly to be the complementary.
+ """
+ returnself._in_features
+
+ @in_features.setter
+ defin_features(self,in_features)->None:
+ ifself._data_nodes:
+ raiseValueError(
+ '`in_features` cannot be changed if the MPS has data nodes. '
+ 'Use unset_data_nodes first')
+
+ ifisinstance(in_features,(list,tuple)):
+ forin_finin_features:
+ ifnotisinstance(in_f,int):
+ raiseTypeError('`in_features` should be tuple[int]'
+ ' or list[int] type')
+ if(in_f<0)or(in_f>=self._n_features):
+ raiseValueError('Elements in `in_features` should'
+ 'be between 0 and (`n_features` - 1)')
+ in_features=set(in_features)
+ out_features=set(range(self._n_features)).difference(in_features)
+
+ self._in_features=list(in_features)
+ self._out_features=list(out_features)
+
+ self._in_features.sort()
+ self._out_features.sort()
+ else:
+ raiseTypeError(
+ '`in_features` should be tuple[int] or list[int] type')
+
+ @property
+ defout_features(self)->List[int]:
+"""
+ Returns list of positions of the output nodes. To change this
+ attribute, first call :meth:`~tensorkrowch.TensorNetwork.unset_data_nodes`
+ if there are already data nodes in the network. When changing it,
+ :attr:`in_features` will change accordingly to be the complementary.
+ """
+ returnself._out_features
+
+ @out_features.setter
+ defout_features(self,out_features)->None:
+ ifself._data_nodes:
+ raiseValueError(
+ '`out_features` cannot be changed if the MPS has data nodes. '
+ 'Use unset_data_nodes first')
+
+ ifisinstance(out_features,(list,tuple)):
+ forout_finout_features:
+ ifnotisinstance(out_f,int):
+ raiseTypeError('`out_features` should be tuple[int]'
+ ' or list[int] type')
+ if(out_f<0)or(out_f>=self._n_features):
+ raiseValueError('Elements in `out_features` should'
+ 'be between 0 and (`n_features` - 1)')
+ out_features=set(out_features)
+ in_features=set(range(self._n_features)).difference(out_features)
+
+ self._in_features=list(in_features)
+ self._out_features=list(out_features)
+
+ self._in_features.sort()
+ self._out_features.sort()
+ else:
+ raiseTypeError(
+ '`out_features` should be tuple[int] or list[int] type')
+
+ @property
+ defin_regions(self)->List[List[int]]:
+""" Returns a list of lists of consecutive input positions."""
+ returnsplit_sequence_into_regions(self._in_features)
+
+ @property
+ defout_regions(self)->List[List[int]]:
+""" Returns a list of lists of consecutive output positions."""
+ returnsplit_sequence_into_regions(self._out_features)
+
+ @property
+ defleft_node(self)->Optional[AbstractNode]:
+"""Returns the ``left_node``."""
+ returnself._left_node
+
+ @property
+ defright_node(self)->Optional[AbstractNode]:
+"""Returns the ``right_node``."""
+ returnself._right_node
+
+ @property
+ defmats_env(self)->List[AbstractNode]:
+"""Returns the list of nodes in ``mats_env``."""
+ returnself._mats_env
+
+ @property
+ defin_env(self)->List[AbstractNode]:
+"""Returns the list of input nodes."""
+ return[self._mats_env[i]foriinself._in_features]
+
+ @property
+ defout_env(self)->List[AbstractNode]:
+"""Returns the list of output nodes."""
+ return[self._mats_env[i]foriinself._out_features]
+
+ # -------
+ # Methods
+ # -------def_make_nodes(self)->None:"""Creates all the nodes of the MPS."""
- ifself.leaf_nodes:
+ ifself._leaf_nodes:raiseValueError('Cannot create MPS nodes if the MPS already has ''nodes')
+
+ aux_bond_dim=self._bond_dim
+
+ ifself._boundary=='obc':
+ ifnotaux_bond_dim:
+ aux_bond_dim=[1]
+
+ self._left_node=ParamNode(shape=(aux_bond_dim[0],),
+ axes_names=('right',),
+ name='left_node',
+ network=self)
+ self._right_node=ParamNode(shape=(aux_bond_dim[-1],),
+ axes_names=('left',),
+ name='right_node',
+ network=self)
+
+ aux_bond_dim=aux_bond_dim+[aux_bond_dim[-1]]+[aux_bond_dim[0]]
+
+ foriinrange(self._n_features):
+ node=ParamNode(shape=(aux_bond_dim[i-1],
+ self._phys_dim[i],
+ aux_bond_dim[i]),
+ axes_names=('left','input','right'),
+ name=f'mats_env_node_({i})',
+ network=self)
+ self._mats_env.append(node)
- self.left_node=None
- self.right_node=None
- self.mats_env=[]
-
- ifself.boundary=='obc':
- self.left_node=ParamNode(shape=(self.in_dim[0],self.bond_dim[0]),
- axes_names=('input','right'),
- name='left_node',
- network=self)
-
- foriinrange(self._n_features-2):
- node=ParamNode(shape=(self.bond_dim[i],
- self.in_dim[i+1],
- self.bond_dim[i+1]),
- axes_names=('left','input','right'),
- name=f'mats_env_node_({i})',
- network=self)
- self.mats_env.append(node)
+ ifi!=0:
+ self._mats_env[-2]['right']^self._mats_env[-1]['left']
+ ifself._boundary=='pbc':ifi==0:
- self.left_node['right']^self.mats_env[-1]['left']
- else:
- self.mats_env[-2]['right']^self.mats_env[-1]['left']
-
- self.right_node=ParamNode(shape=(self.bond_dim[-1],
- self.in_dim[-1]),
- axes_names=('left','input'),
- name='right_node',
- network=self)
-
- ifself._n_features>2:
- self.mats_env[-1]['right']^self.right_node['left']
+ periodic_edge=self._mats_env[-1]['left']
+ ifi==self._n_features-1:
+ self._mats_env[-1]['right']^periodic_edgeelse:
- self.left_node['right']^self.right_node['left']
-
- else:
- foriinrange(self._n_features):
- node=ParamNode(shape=(self.bond_dim[i-1],
- self.in_dim[i],
- self.bond_dim[i]),
- axes_names=('left','input','right'),
- name=f'mats_env_node_({i})',
- network=self)
- self.mats_env.append(node)
-
ifi==0:
- periodic_edge=self.mats_env[-1]['left']
- else:
- self.mats_env[-2]['right']^self.mats_env[-1]['left']
-
+ self._left_node['right']^self._mats_env[-1]['left']ifi==self._n_features-1:
- self.mats_env[-1]['right']^periodic_edge
-
-
[docs]definitialize(self,std:float=1e-9)->None:
+ self._mats_env[-1]['right']^self._right_node['left']
+
+ def_make_unitaries(self,device:Optional[torch.device]=None)->List[torch.Tensor]:
+"""
+ Creates random unitaries to initialize the MPS in canonical form with
+ orthogonality center at the rightmost node."""
+ tensors=[]
+ fori,nodeinenumerate(self._mats_env):
+ ifself._boundary=='obc':
+ ifi==0:
+ node_shape=node.shape[1:]
+ aux_shape=node_shape
+ elifi==(self._n_features-1):
+ node_shape=node.shape[:2]
+ aux_shape=node_shape
+ else:
+ node_shape=node.shape
+ aux_shape=(node.shape[:2].numel(),node.shape[2])
+ else:
+ node_shape=node.shape
+ aux_shape=(node.shape[:2].numel(),node.shape[2])
+ size=max(aux_shape[0],aux_shape[1])
+
+ tensor=random_unitary(size,device=device)
+ tensor=tensor[:min(aux_shape[0],size),:min(aux_shape[1],size)]
+ tensor=tensor.reshape(*node_shape)
+
+ tensors.append(tensor)
+
+ ifself._boundary=='obc':
+ tensors[-1]=tensors[-1]/tensors[-1].norm()
+ returntensors
+
+
[docs]definitialize(self,
+ tensors:Optional[Sequence[torch.Tensor]]=None,
+ init_method:Optional[Text]='randn',
+ device:Optional[torch.device]=None,
+ **kwargs:float)->None:"""
- Initializes all the nodes as explained `here <https://arxiv.org/abs/1605.03795>`_.
- It can be overriden for custom initializations.
+ Initializes all the nodes of the :class:`MPS`. It can be called when
+ instantiating the model, or to override the existing nodes' tensors.
+
+ There are different methods to initialize the nodes:
+
+ * ``{"zeros", "ones", "copy", "rand", "randn"}``: Each node is
+ initialized calling :meth:`~tensorkrowch.AbstractNode.set_tensor` with
+ the given method, ``device`` and ``kwargs``.
+
+ * ``"randn_eye"``: Nodes are initialized as in this
+ `paper <https://arxiv.org/abs/1605.03795>`_, adding identities at the
+ top of random gaussian tensors. In this case, ``std`` should be
+ specified with a low value, e.g., ``std = 1e-9``.
+
+ * ``"unit"``: Nodes are initialized as random unitaries, so that the
+ MPS is in canonical form, with the orthogonality center at the
+ rightmost node.
+
+ Parameters
+ ----------
+ tensors : list[torch.Tensor] or tuple[torch.Tensor], optional
+ Sequence of tensors to set in each of the MPS nodes. If ``boundary``
+ is ``"obc"``, all tensors should be rank-3, except the first and
+ last ones, which can be rank-2, or rank-1 (if the first and last are
+ the same). If ``boundary`` is ``"pbc"``, all tensors should be
+ rank-3.
+ init_method : {"zeros", "ones", "copy", "rand", "randn", "randn_eye", "unit"}, optional
+ Initialization method.
+ device : torch.device, optional
+ Device where to initialize the tensors if ``init_method`` is provided.
+ kwargs : float
+ Keyword arguments for the different initialization methods. See
+ :meth:`~tensorkrowch.AbstractNode.make_tensor`. """
- # Left node
- ifself.left_nodeisnotNone:
- tensor=torch.randn(self.left_node.shape)*std
- aux=torch.zeros(tensor.shape[1])*std
- aux[0]=1.
- tensor[0,:]=aux
- self.left_node.tensor=tensor
-
- # Right node
- ifself.right_nodeisnotNone:
- tensor=torch.randn(self.right_node.shape)*std
- aux=torch.zeros(tensor.shape[0])*std
- aux[0]=1.
- tensor[:,0]=aux
- self.right_node.tensor=tensor
-
- # Mats env
- fornodeinself.mats_env:
- tensor=torch.randn(node.shape)*std
- aux=torch.eye(tensor.shape[0],tensor.shape[2])
- tensor[:,0,:]=aux
- node.tensor=tensor
[docs]defset_data_nodes(self)->None:"""
- Creates ``data`` nodes and connects each of them to the input/physical
+ Creates ``data`` nodes and connects each of them to the ``"input"`` edge of each input node.
+ """
+ input_edges=[node['input']fornodeinself.in_env]
+ super().set_data_nodes(input_edges=input_edges,
+ num_batch_edges=self._n_batches)
+
+
[docs]defcopy(self,share_tensors:bool=False)->'MPS':
+"""
+ Creates a copy of the :class:`MPS`.
+
+ Parameters
+ ----------
+ share_tensor : bool, optional
+ Boolean indicating whether tensors in the copied MPS should be
+ set as the tensors in the current MPS (``True``), or cloned
+ (``False``). In the former case, tensors in both MPS's will be
+ the same, which might be useful if one needs more than one copy
+ of an MPS, but wants to compute all the gradients with respect
+ to the same, unique, tensors.
+
+ Returns
+ -------
+ MPS """
- input_edges=[]
- ifself.left_nodeisnotNone:
- input_edges.append(self.left_node['input'])
- input_edges+=list(map(lambdanode:node['input'],self.mats_env))
- ifself.right_nodeisnotNone:
- input_edges.append(self.right_node['input'])
+ new_mps=MPS(n_features=self._n_features,
+ phys_dim=self._phys_dim,
+ bond_dim=self._bond_dim,
+ boundary=self._boundary,
+ tensors=None,
+ in_features=self._in_features,
+ out_features=self._out_features,
+ n_batches=self._n_batches,
+ init_method=None,
+ device=None)
+ new_mps.name=self.name+'_copy'
+ ifshare_tensors:
+ fornew_node,nodeinzip(new_mps._mats_env,self._mats_env):
+ new_node.tensor=node.tensor
+ else:
+ fornew_node,nodeinzip(new_mps._mats_env,self._mats_env):
+ new_node.tensor=node.tensor.clone()
+ returnnew_mps
+
+
[docs]defparameterize(self,
+ set_param:bool=True,
+ override:bool=False)->'TensorNetwork':
+"""
+ Parameterizes all nodes of the MPS. If there are ``resultant`` nodes
+ in the MPS, it will be first :meth:`~tensorkrowch.TensorNetwork.reset`.
- super().set_data_nodes(input_edges=input_edges,
- num_batch_edges=self.n_batches)
+ Parameters
+ ----------
+ set_param : bool
+ Boolean indicating whether the tensor network has to be parameterized
+ (``True``) or de-parameterized (``False``).
+ override : bool
+ Boolean indicating whether the tensor network should be parameterized
+ in-place (``True``) or copied and then parameterized (``False``).
+ """
+ ifself._resultant_nodes:
+ warnings.warn(
+ 'Resultant nodes will be removed before parameterizing the TN')
+ self.reset()
- ifself.mats_env:
- self.mats_env_data=list(map(lambdanode:node.neighbours('input'),
- self.mats_env))
returnaux_nodes,leftoverreturnnodes,[]
- def_pairwise_contraction(self,mats_nodes:List[Node])->Node:
-"""Contracts the left and right environments pairwise."""
- length=len(mats_nodes)
- aux_nodes=mats_nodes
+ def_pairwise_contraction(self,mats_env:List[AbstractNode])->Node:
+"""Contracts nodes environments pairwise."""
+ length=len(mats_env)
+ aux_nodes=mats_enviflength>1:leftovers=[]whilelength>1:
@@ -657,108 +1163,642 @@
Source code for tensorkrowch.models.mps
[docs]defcontract(self,inline_input:bool=False,
- inline_mats:bool=False)->Node:
+ inline_mats:bool=False,
+ marginalize_output:bool=False,
+ embedding_matrices:Optional[
+ Union[torch.Tensor,
+ Sequence[torch.Tensor]]]=None,
+ mpo:Optional[MPO]=None
+ )->Node:""" Contracts the whole MPS.
+ If the MPS has input nodes, these are contracted against input ``data``
+ nodes.
+
+ If the MPS has output nodes, these can be left with their ``"input"``
+ edges open, or can be marginalized, contracting the remaining output
+ nodes with themselves, if the argument ``"marginalize_output"`` is set
+ to ``True``.
+
+ In the latter case, one can add additional nodes in between the MPS-MPS
+ contraction:
+
+ * ``embedding_matrices``: A list of matrices with appropiate physical
+ dimensions can be passed, one for each output node. These matrices
+ will connect the two ``"input"`` edges of the corresponding nodes.
+
+ * ``mpo``: If an :class:`MPO` is passed, when calling
+ ``mps(marginalize_output=True, mpo=mpo)``, this will perform the
+ MPS-MPO-MPS contraction at the output nodes of the MPS. Therefore,
+ the MPO should have as many nodes as output nodes are in the MPS.
+
+ After contraction, the MPS will still be connected to the MPO nodes
+ until these are manually disconnected.
+
+ The provided MPO can also be already connected to the MPS before
+ contraction. In this case, it is assumed that the output nodes of the
+ MPS are connected to the ``"output"`` edges of the MPO nodes, and
+ that the MPO nodes have been moved to the MPS, so that all nodes
+ belong to the MPS network. In this case, each MPO node will connect
+ the two ``"input"`` edges of the corresponding MPS nodes.
+
+ If the MPO nodes are not trainable, they can be de-parameterized
+ by doing ``mpo = mpo.parameterize(set_param=False, override=True)``.
+ This should be done before the contraction, or before connecting
+ the MPO nodes to the MPS, since the de-parameterized nodes are not
+ the same nodes as the original ``ParamNodes`` of the MPO.
+
+ When ``marginalize_output = True``, the contracted input nodes are
+ duplicated using different batch dimensions. That is, if the MPS
+ is contracted with input data with ``batch_size = 100``, and some
+ other (output) nodes are marginalized, the result will be a tensor
+ with shape ``(100, 100)`` rather than just ``(100,)``.
+ Parameters ---------- inline_input : bool Boolean indicating whether input ``data`` nodes should be contracted
- with the ``MPS`` nodes inline (one contraction at a time) or in a
- single stacked contraction.
+ with the ``MPS`` input nodes inline (one contraction at a time) or
+ in a single stacked contraction. inline_mats : bool Boolean indicating whether the sequence of matrices (resultant after contracting the input ``data`` nodes) should be contracted inline or as a sequence of pairwise stacked contrations.
+ marginalize_output : bool
+ Boolean indicating whether output nodes should be marginalized. If
+ ``True``, after contracting all the input nodes with their
+ neighbouring data nodes, this resultant network is contracted with
+ itself connecting output nodes to itselves at ``"input"`` edges. If
+ ``False``, output nodes are left with their ``"input"`` edges
+ disconnected.
+ embedding_matrices : torch.Tensor, list[torch.Tensor] or tuple[torch.Tensor], optional
+ If ``marginalize_output = True``, a matrix can be introduced
+ between each output node and its copy, connecting the ``"input"``
+ edges. This can be useful when data vectors are not represented
+ as qubits in the computational basis, but are transformed via
+ some :ref:`Embeddings` function.
+ mpo : MPO, optional
+ MPO that is to be contracted with the MPS at the output nodes, if
+ ``marginalize_output = True``. In this case, the ``"output"`` edges
+ of the MPO nodes will be connected to the ``"input"`` edges of the
+ MPS output nodes. If there are no input nodes, the MPS-MPO-MPS
+ is performed by calling ``mps(marginalize_output=True, mpo=mpo)``,
+ without passing extra data tensors. Returns ------- Node """
- mats_env=self._input_contraction(inline_input)
-
- ifinline_mats:
- result=self._contract_envs_inline(mats_env)
+ ifembedding_matricesisnotNone:
+ ifisinstance(embedding_matrices,Sequence):
+ iflen(embedding_matrices)!=len(self._out_features):
+ raiseValueError(
+ '`embedding_matrices` should have the same amount of '
+ 'elements as output nodes are in the MPS')
+ else:
+ embedding_matrices=[embedding_matrices]*len(self._out_features)
+
+ fori,(mat,node)inenumerate(zip(embedding_matrices,
+ self.out_env)):
+ ifnotisinstance(mat,torch.Tensor):
+ raiseTypeError(
+ '`embedding_matrices` should be torch.Tensor type')
+ iflen(mat.shape)!=2:
+ raiseValueError(
+ '`embedding_matrices should ne rank-2 tensors')
+ ifmat.shape[0]!=mat.shape[1]:
+ raiseValueError(
+ '`embedding_matrices` should have equal dimensions')
+ ifnode['input'].size()!=mat.shape[0]:
+ raiseValueError(
+ '`embedding_matrices` dimensions should be equal '
+ 'to the input dimensions of the corresponding MPS '
+ 'output nodes')
+ elifmpoisnotNone:
+ ifnotisinstance(mpo,MPO):
+ raiseTypeError('`mpo` should be MPO type')
+ ifmpo._n_features!=len(self._out_features):
+ raiseValueError(
+ '`mpo` should have as many features as output nodes are '
+ 'in the MPS')
+
+ in_regions=self.in_regions
+ out_regions=self.out_regions
+
+ mats_in_env=self._input_contraction(
+ nodes_env=self.in_env,
+ input_nodes=[node.neighbours('input')fornodeinself.in_env],
+ inline_input=inline_input)
+
+ in_results=[]
+ forregioninin_regions:
+ ifinline_mats:
+ result=self._contract_envs_inline(
+ mats_env=mats_in_env[:len(region)])
+ else:
+ result=self._pairwise_contraction(
+ mats_env=mats_in_env[:len(region)])
+
+ mats_in_env=mats_in_env[len(region):]
+ in_results.append(result)
+
+ ifnotout_regions:
+ # If there is only input region, in_results has only 1 node
+ result=in_results[0]
+
else:
- result=self._pairwise_contraction(mats_env)
-
+ # Contract each in_result with the next output node
+ nodes_out_env=[]
+ out_first=out_regions[0][0]==0
+ out_last=out_regions[-1][-1]==(self._n_features-1)
+
+ foriinrange(len(out_regions)):
+ aux_out_env=[self._mats_env[j]forjinout_regions[i]]
+
+ if(i==0)andout_first:
+ ifself._boundary=='obc':
+ aux_out_env[0]=self._left_node@aux_out_env[0]
+ else:
+ aux_out_env[0]=in_results[i-out_first]@aux_out_env[0]
+
+ nodes_out_env+=aux_out_env
+
+ ifout_last:
+ if(self._boundary=='obc'):
+ nodes_out_env[-1]=nodes_out_env[-1]@self._right_node
+ else:
+ nodes_out_env[-1]=nodes_out_env[-1]@in_results[-1]
+
+ ifnotmarginalize_output:
+ # Contract all output nodes sequentially
+ result=self._inline_contraction(mats_env=nodes_out_env)
+
+ else:
+ # Copy output nodes sharing tensors
+ copied_nodes=[]
+ fornodeinnodes_out_env:
+ copied_node=node.__class__(shape=node._shape,
+ axes_names=node.axes_names,
+ name='virtual_result_copy',
+ network=self,
+ virtual=True)
+ copied_node.set_tensor_from(node)
+ copied_nodes.append(copied_node)
+
+ # Change batch names so that they not coincide with
+ # original batches, which gives dupliicate output batches
+ foraxincopied_node.axes:
+ ifax._batch:
+ ax.name=ax.name+'_copy'
+
+ # Connect copied nodes with neighbours
+ foriinrange(len(copied_nodes)):
+ if(i==0)and(self._boundary=='pbc'):
+ ifnodes_out_env[i-1].is_connected_to(nodes_out_env[i]):
+ copied_nodes[i-1]['right']^copied_nodes[i]['left']
+ elifi>0:
+ copied_nodes[i-1]['right']^copied_nodes[i]['left']
+
+ # Contract with embedding matrices
+ ifembedding_matricesisnotNone:
+ mats_nodes=[]
+ fori,nodeinenumerate(nodes_out_env):
+ # Reattach input edges
+ node.reattach_edges(axes=['input'])
+
+ # Create matrices
+ mat_node=Node(tensor=embedding_matrices[i],
+ axes_names=('input','output'),
+ name='virtual_result_mat',
+ network=self,
+ virtual=True)
+
+ # Connect matrices to output nodes
+ mat_node['output']^node['input']
+ mats_nodes.append(mat_node)
+
+ # Connect matrices to copies
+ format_node,copied_nodeinzip(mats_nodes,copied_nodes):
+ copied_node['input']^mat_node['input']
+
+ # Contract output nodes with matrices
+ nodes_out_env=self._input_contraction(
+ nodes_env=nodes_out_env,
+ input_nodes=mats_nodes,
+ inline_input=True)
+
+ # Contract with mpo
+ elifmpoisnotNone:
+ # Move all the connected component to the MPS network
+ mpo._mats_env[0].move_to_network(self)
+
+ # Move uniform memory
+ ifisinstance(mpo,UMPO):
+ mpo.uniform_memory.move_to_network(self)
+ fornodeinmpo._mats_env:
+ node.set_tensor_from(mpo.uniform_memory)
+
+ # Connect MPO to MPS
+ formps_node,mpo_nodeinzip(nodes_out_env,mpo._mats_env):
+ # Reattach input edges
+ mps_node.reattach_edges(axes=['input'])
+ mpo_node['output']^mps_node['input']
+
+ # Connect MPO to copies
+ forcopied_node,mpo_nodeinzip(copied_nodes,mpo._mats_env):
+ copied_node['input']^mpo_node['input']
+
+ # Contract MPO with MPS
+ nodes_out_env=self._input_contraction(
+ nodes_env=nodes_out_env,
+ input_nodes=mpo._mats_env,
+ inline_input=True)
+
+ # Contract MPO left and right nodes
+ ifmpo._boundary=='obc':
+ nodes_out_env[0]=mpo._left_node@nodes_out_env[0]
+ nodes_out_env[-1]=nodes_out_env[-1]@mpo._right_node
+
+ else:
+ # Reattach input edges of resultant output nodes and connect
+ # with copied nodes
+ fornode,copied_nodeinzip(nodes_out_env,copied_nodes):
+ # Reattach input edges
+ node.reattach_edges(axes=['input'])
+
+ # Connect copies directly to output nodes
+ copied_node['input']^node['input']
+
+ # Contract output nodes with copies
+ mats_out_env=self._input_contraction(
+ nodes_env=nodes_out_env,
+ input_nodes=copied_nodes,
+ inline_input=True)
+
+ # Contract resultant matrices
+ result=self._inline_contraction(mats_env=mats_out_env)
+
+ # Contract periodic edge
+ ifresult.is_connected_to(result):
+ result@=result
+
+ # Put batch edges in first positions
+ batch_edges=[]
+ other_edges=[]
+ fori,edgeinenumerate(result.edges):
+ ifedge.is_batch():
+ batch_edges.append(i)
+ else:
+ other_edges.append(i)
+
+ all_edges=batch_edges+other_edges
+ ifall_edges!=list(range(len(all_edges))):
+ result=op.permute(result,tuple(all_edges))
+
+ returnresult
+
+
[docs]defnorm(self)->torch.Tensor:
+"""
+ Computes the norm of the MPS.
+
+ This method internally sets ``in_features = []``, and calls the
+ :meth:`~tensorkrowch.TensorNetwork.forward` method with
+ ``marginalize_output = True``. Therefore, it may alter the behaviour
+ of the MPS if it is not :meth:`~tensorkrowch.TensorNetwork.reset`
+ afterwards. Also, if the MPS was contracted before with other arguments,
+ it should be ``reset`` before calling ``norm`` to avoid undesired
+ behaviour.
+ """
+ ifself._data_nodes:
+ self.unset_data_nodes()
+ self.in_features=[]
+
+ result=self.forward(marginalize_output=True)
+ result=result.sqrt()returnresult
-
[docs]defcanonicalize(self,
- oc:Optional[int]=None,
- mode:Text='svd',
- rank:Optional[int]=None,
- cum_percentage:Optional[float]=None,
- cutoff:Optional[float]=None)->None:
-r"""
- Turns MPS into canonical form via local SVD/QR decompositions.
+
[docs]defpartial_density(self,trace_sites:Sequence[int]=[])->torch.Tensor:
+"""
+ Returns de partial density matrix, tracing out the sites specified
+ by ``trace_sites``.
+
+ This method internally sets ``out_features = trace_sites``, and calls
+ the :meth:`~tensorkrowch.TensorNetwork.forward` method with
+ ``marginalize_output = True``. Therefore, it may alter the behaviour
+ of the MPS if it is not :meth:`~tensorkrowch.TensorNetwork.reset`
+ afterwards. Also, if the MPS was contracted before with other arguments,
+ it should be ``reset`` before calling ``partial_density`` to avoid
+ undesired behaviour.
+
+ Since the density matrix is computed by contracting the MPS, it means
+ one can take gradients of it with respect to the MPS tensors, if it
+ is needed.
+
+ This method may also alter the attribute :attr:`n_batches` of the
+ :class:`MPS`. Parameters ----------
- oc : int
- Position of the orthogonality center. It should be between 0 and
- ``n_features -1``.
- mode : {"svd", "svdr", "qr"}
- Indicates which decomposition should be used to split a node after
- contracting it. See more at :func:`svd_`, :func:`svdr_`, :func:`qr_`.
- If mode is "qr", operation :func:`qr_` will be performed on nodes at
- the left of the output node, whilst operation :func:`rq_` will be
- used for nodes at the right.
- rank : int, optional
- Number of singular values to keep.
- cum_percentage : float, optional
- Proportion that should be satisfied between the sum of all singular
- values kept and the total sum of all singular values.
-
- .. math::
-
- \frac{\sum_{i \in \{kept\}}{s_i}}{\sum_{i \in \{all\}}{s_i}} \ge
- cum\_percentage
- cutoff : float, optional
- Quantity that lower bounds singular values in order to be kept.
-
+ trace_sites : list[int] or tuple[int]
+ Sequence of nodes' indices in the MPS. These indices specify the
+ nodes that should be traced to compute the density matrix. If
+ it is empty ``[]``, the total density matrix will be returned,
+ though this may be costly if :attr:`n_features` is big.
+ Examples -------- >>> mps = tk.models.MPS(n_features=4,
- ... in_dim=2,
+ ... phys_dim=[2, 3, 4, 5], ... bond_dim=5)
- >>> mps.canonicalize(rank=3)
- >>> mps.bond_dim
- [2, 3, 2]
+ >>> density = mps.partial_density(trace_sites=[0, 2])
+ >>> density.shape
+ torch.Size([3, 5, 3, 5]) """
- self.reset()
-
- prev_auto_stack=self._auto_stack
- self.auto_stack=False
-
- ifocisNone:
- oc=self._n_features-1
- elifoc>=self._n_features:
- raiseValueError(f'Orthogonality center position `oc` should be '
- f'between 0 and {self._n_features-1}')
-
- nodes=self.mats_env
- ifself.boundary=='obc':
- nodes=[self.left_node]+nodes+[self.right_node]
-
- foriinrange(oc):
- ifmode=='svd':
- result1,result2=nodes[i]['right'].svd_(
- side='right',
- rank=rank,
- cum_percentage=cum_percentage,
- cutoff=cutoff)
+ ifnotisinstance(trace_sites,Sequence):
+ raiseTypeError(
+ '`trace_sites` should be list[int] or tuple[int] type')
+
+ forsiteintrace_sites:
+ ifnotisinstance(site,int):
+ raiseTypeError(
+ 'elements of `trace_sites` should be int type')
+ if(site<0)or(site>=self._n_features):
+ raiseValueError(
+ 'Elements of `trace_sites` should be between 0 and '
+ '(`n_features` - 1)')
+
+ ifself._data_nodes:
+ self.unset_data_nodes()
+ self.out_features=trace_sites
+
+ # Create dataset with all possible combinations for the input nodes
+ # so that they are kept sort of "open"
+ dims=torch.tensor([self._phys_dim[i]foriinself._in_features])
+
+ data=[]
+ foriinrange(len(self._in_features)):
+ aux=torch.arange(dims[i]).view(-1,1)
+ aux=aux.repeat(1,dims[(i+1):].prod()).flatten().view(-1,1)
+ aux=aux.repeat(dims[:i].prod(),1)
+
+ data.append(aux.reshape(*dims,1))
+
+ n_dims=len(set(dims))
+ ifn_dims>=1:
+ ifn_dims==1:
+ data=torch.cat(data,dim=-1)
+ data=basis(data,dim=dims[0]).float().to(self.in_env[0].device)
+ elifn_dims>1:
+ data=[
+ basis(dat,dim=dim).squeeze(-2).float().to(self.in_env[0].device)
+ fordat,diminzip(data,dims)
+ ]
+
+ self.n_batches=len(dims)
+ result=self.forward(data,marginalize_output=True)
+
+ else:
+ result=self.forward(marginalize_output=True)
+
+ returnresult
+
+
[docs]@torch.no_grad()
+ defmi(self,
+ middle_site:int,
+ renormalize:bool=False)->Union[float,Tuple[float]]:
+r"""
+ Computes the Mutual Information between subsystems :math:`A` and
+ :math:`B`, :math:`\textrm{MI}(A:B)`, where :math:`A` goes from site
+ 0 to ``middle_site``, and :math:`B` goes from ``middle_site + 1`` to
+ ``n_features - 1``.
+
+ To compute the mutual information, the MPS is put into canonical form
+ with orthogonality center at ``middle_site``. Bond dimensions are not
+ changed if possible. Only when the bond dimension is bigger than the
+ physical dimension multiplied by the other bond dimension of the node,
+ it will be cropped to that size.
+
+ If the MPS is not normalized, it may happen that the computation of the
+ mutual information fails due to errors in the Singular Value
+ Decompositions. To avoid this, it is recommended to set
+ ``renormalize = True``. In this case, the norm of each node after the
+ SVD is extracted in logarithmic form, and accumulated. As a result,
+ the function will return the tuple ``(mi, log_norm)``, which is a sort
+ of `scaled` mutual information. The actual mutual information could be
+ obtained as ``exp(log_norm) * mi``.
+
+ Parameters
+ ----------
+ middle_site : int
+ Position that separates regios :math:`A` and :math:`B`. It should
+ be between 0 and ``n_features - 2``.
+ renormalize : bool
+ Indicates whether nodes should be renormalized after SVD/QR
+ decompositions. If not, it may happen that the norm explodes as it
+ is being accumulated from all nodes. Renormalization aims to avoid
+ this undesired behavior by extracting the norm of each node on a
+ logarithmic scale after SVD/QR decompositions are computed. Finally,
+ the normalization factor is evenly distributed among all nodes of
+ the MPS.
+
+ Returns
+ -------
+ float or tuple[float, float]
+ """
+ self.reset()
+
+ prev_auto_stack=self._auto_stack
+ self.auto_stack=False
+
+ if(middle_site<0)or(middle_site>(self._n_features-2)):
+ raiseValueError(
+ '`middle_site` should be between 0 and `n_features` - 2')
+
+ log_norm=0
+
+ nodes=self._mats_env[:]
+ ifself._boundary=='obc':
+ nodes[0].tensor[1:]=torch.zeros_like(
+ nodes[0].tensor[1:])
+ nodes[-1].tensor[...,1:]=torch.zeros_like(
+ nodes[-1].tensor[...,1:])
+
+ foriinrange(middle_site):
+ result1,result2=nodes[i]['right'].svd_(
+ side='right',
+ rank=nodes[i]['right'].size())
+
+ ifrenormalize:
+ aux_norm=result2.norm()/sqrt(result2.shape[0])
+ ifnotaux_norm.isinf()and(aux_norm>0):
+ result2.tensor=result2.tensor/aux_norm
+ log_norm+=aux_norm.log()
+
+ result1=result1.parameterize()
+ nodes[i]=result1
+ nodes[i+1]=result2
+
+ foriinrange(len(nodes)-1,middle_site,-1):
+ result1,result2=nodes[i]['left'].svd_(
+ side='left',
+ rank=nodes[i]['left'].size())
+
+ ifrenormalize:
+ aux_norm=result1.norm()/sqrt(result1.shape[0])
+ ifnotaux_norm.isinf()and(aux_norm>0):
+ result1.tensor=result1.tensor/aux_norm
+ log_norm+=aux_norm.log()
+
+ result2=result2.parameterize()
+ nodes[i]=result2
+ nodes[i-1]=result1
+
+ nodes[middle_site]=nodes[middle_site].parameterize()
+
+ # Compute mutual information
+ middle_tensor=nodes[middle_site].tensor.clone()
+ _,s,_=torch.linalg.svd(
+ middle_tensor.reshape(middle_tensor.shape[:-1].numel(),# left x input
+ middle_tensor.shape[-1]),# right
+ full_matrices=False)
+
+ s=s[s>0]
+ mutual_info=-(s*(s.log()+log_norm)).sum()
+
+ # Rescale
+ iflog_norm!=0:
+ rescale=(log_norm/len(nodes)).exp()
+
+ ifrenormalizeand(log_norm!=0):
+ fornodeinnodes:
+ node.tensor=node.tensor*rescale
+
+ # Update variables
+ ifself._boundary=='obc':
+ self._bond_dim=[node['right'].size()fornodeinnodes[:-1]]
+ else:
+ self._bond_dim=[node['right'].size()fornodeinnodes]
+ self._mats_env=nodes
+
+ self.auto_stack=prev_auto_stack
+
+ ifrenormalize:
+ returnmutual_info,log_norm
+ else:
+ returnmutual_info
+
+
[docs]@torch.no_grad()
+ defcanonicalize(self,
+ oc:Optional[int]=None,
+ mode:Text='svd',
+ rank:Optional[int]=None,
+ cum_percentage:Optional[float]=None,
+ cutoff:Optional[float]=None,
+ renormalize:bool=False)->None:
+r"""
+ Turns MPS into canonical form via local SVD/QR decompositions.
+
+ To specify the new bond dimensions, the arguments ``rank``,
+ ``cum_percentage`` or ``cutoff`` can be specified. These will be used
+ equally for all SVD computations.
+
+ If none of them are specified, the bond dimensions won't be modified
+ if possible. Only when the bond dimension is bigger than the physical
+ dimension multiplied by the other bond dimension of the node, it will
+ be cropped to that size.
+
+ Parameters
+ ----------
+ oc : int
+ Position of the orthogonality center. It should be between 0 and
+ ``n_features - 1``.
+ mode : {"svd", "svdr", "qr"}
+ Indicates which decomposition should be used to split a node after
+ contracting it. See more at :func:`~tensorkrowch.svd_`,
+ :func:`~tensorkrowch.svdr_`, :func:`~tensorkrowch.qr_`.
+ If mode is "qr", operation :func:`~tensorkrowch.qr_` will be
+ performed on nodes at the left of the output node, whilst operation
+ :func:`~tensorkrowch.rq_` will be used for nodes at the right.
+ rank : int, optional
+ Number of singular values to keep.
+ cum_percentage : float, optional
+ Proportion that should be satisfied between the sum of all singular
+ values kept and the total sum of all singular values.
+
+ .. math::
+
+ \frac{\sum_{i \in \{kept\}}{s_i}}{\sum_{i \in \{all\}}{s_i}} \ge
+ cum\_percentage
+ cutoff : float, optional
+ Quantity that lower bounds singular values in order to be kept.
+ renormalize : bool
+ Indicates whether nodes should be renormalized after SVD/QR
+ decompositions. If not, it may happen that the norm explodes as it
+ is being accumulated from all nodes. Renormalization aims to avoid
+ this undesired behavior by extracting the norm of each node on a
+ logarithmic scale after SVD/QR decompositions are computed. Finally,
+ the normalization factor is evenly distributed among all nodes of
+ the MPS.
+
+ Examples
+ --------
+ >>> mps = tk.models.MPS(n_features=4,
+ ... phys_dim=2,
+ ... bond_dim=5)
+ >>> mps.canonicalize(rank=3)
+ >>> mps.bond_dim
+ [3, 3, 3]
+ """
+ self.reset()
+
+ prev_auto_stack=self._auto_stack
+ self.auto_stack=False
+
+ ifocisNone:
+ oc=self._n_features-1
+ elif(oc<0)or(oc>=self._n_features):
+ raiseValueError('Orthogonality center position `oc` should be '
+ 'between 0 and `n_features` - 1')
+
+ log_norm=0
+
+ nodes=self._mats_env[:]
+ ifself._boundary=='obc':
+ nodes[0].tensor[1:]=torch.zeros_like(
+ nodes[0].tensor[1:])
+ nodes[-1].tensor[...,1:]=torch.zeros_like(
+ nodes[-1].tensor[...,1:])
+
+ # If mode is svd or svr and none of the args is provided, the ranks are
+ # kept as they were originally
+ keep_rank=False
+ if(rankisNone)and(cum_percentageisNone)and(cutoffisNone):
+ keep_rank=True
+
+ foriinrange(oc):
+ ifmode=='svd':
+ result1,result2=nodes[i]['right'].svd_(
+ side='right',
+ rank=nodes[i]['right'].size()ifkeep_rankelserank,
+ cum_percentage=cum_percentage,
+ cutoff=cutoff)elifmode=='svdr':result1,result2=nodes[i]['right'].svdr_(side='right',
- rank=rank,
+ rank=nodes[i]['right'].size()ifkeep_rankelserank,cum_percentage=cum_percentage,cutoff=cutoff)elifmode=='qr':result1,result2=nodes[i]['right'].qr_()else:raiseValueError('`mode` can only be "svd", "svdr" or "qr"')
+
+ ifrenormalize:
+ aux_norm=result2.norm()/sqrt(result2.shape[0])
+ ifnotaux_norm.isinf()and(aux_norm>0):
+ result2.tensor=result2.tensor/aux_norm
+ log_norm+=aux_norm.log()result1=result1.parameterize()nodes[i]=result1
@@ -768,38 +1808,46 @@
nodes:List[AbstractNode],idx:int,left_nodeL:AbstractNode):
-"""Returns canonicalize version of the tensor at site ``idx``."""
+"""Returns canonicalized version of the tensor at site ``idx``."""L=nodes[idx]# left x input x rightleft_nodeC=None
@@ -897,17 +1945,17 @@
[docs]@torch.no_grad()
+ defcanonicalize_univocal(self):""" Turns MPS into the univocal canonical form defined `here <https://arxiv.org/abs/2202.12319>`_.
@@ -949,10 +1998,17 @@
[docs]classUMPS(MPS):# MARK: UMPS"""
- Class for Uniform (translationally invariant) Matrix Product States where
- all nodes are input nodes. It is the uniform version of :class:`MPS`, that
- is, all nodes share the same tensor. Thus this class cannot have different
- input or bond dimensions for each node, and boundary conditions are
- always periodic (``"pbc"``).
+ Class for Uniform (translationally invariant) Matrix Product States. It is
+ the uniform version of :class:`MPS`, that is, all nodes share the same
+ tensor. Thus this class cannot have different physical or bond dimensions
+ for each node, and boundary conditions are always periodic (``"pbc"``).
- A ``UMPS`` is formed by the following nodes:
-
- * ``mats_env``: Environment of `matrix` nodes that. These nodes have axes
- ``("left", "input", "right")``.
+ |
+
+ For a more detailed list of inherited properties and methods,
+ check :class:`MPS`. Parameters ---------- n_features : int
- Number of input nodes.
- in_dim : int
- Input dimension. Equivalent to the physical dimension.
- bond_dim : int
+ Number of nodes that will be in ``mats_env``.
+ phys_dim : int, optional
+ Physical dimension.
+ bond_dim : int, optional Bond dimension.
+ tensor: torch.Tensor, optional
+ Instead of providing ``phys_dim`` and ``bond_dim``, a single tensor
+ can be provided. ``n_features`` is still needed to specify how many
+ times the tensor should be used to form a finite MPS. The tensor
+ should be rank-3, with its first and last dimensions being equal.
+ in_features: list[int] or tuple[int], optional
+ List of indices indicating the positions of the MPS nodes that will be
+ considered as input nodes. These nodes will have a neighbouring data
+ node connected to its ``"input"`` edge when the :meth:`set_data_nodes`
+ method is called. ``in_features`` is the complementary set of
+ ``out_features``, so it is only required to specify one of them.
+ out_features: list[int] or tuple[int], optional
+ List of indices indicating the positions of the MPS nodes that will be
+ considered as output nodes. These nodes will be left with their ``"input"``
+ edges open when contrating the network. If ``marginalize_output`` is
+ set to ``True`` in :meth:`contract`, the network will be connected to
+ itself at these nodes, and contracted. ``out_features`` is the
+ complementary set of ``in_features``, so it is only required to specify
+ one of them. n_batches : int Number of batch edges of input ``data`` nodes. Usually ``n_batches = 1`` (where the batch edge is used for the data batched) but it could also be ``n_batches = 2`` (one edge for data batched, other edge for image patches in convolutional layers).
+ init_method : {"zeros", "ones", "copy", "rand", "randn", "randn_eye", "unit"}, optional
+ Initialization method. Check :meth:`initialize` for a more detailed
+ explanation of the different initialization methods.
+ device : torch.device, optional
+ Device where to initialize the tensors if ``init_method`` is provided.
+ kwargs : float
+ Keyword arguments for the different initialization methods. See
+ :meth:`~tensorkrowch.AbstractNode.make_tensor`. Examples -------- >>> mps = tk.models.UMPS(n_features=4,
- ... in_dim=2,
+ ... phys_dim=2, ... bond_dim=5) >>> for node in mps.mats_env: ... assert node.tensor_address() == 'virtual_uniform'
@@ -1022,245 +2105,1438 @@
Source code for tensorkrowch.models.mps
def__init__(self,n_features:int,
- in_dim:int,
- bond_dim:int,
- n_batches:int=1)->None:
-
- super().__init__(name='mps')
-
+ phys_dim:Optional[int]=None,
+ bond_dim:Optional[int]=None,
+ tensor:Optional[torch.Tensor]=None,
+ in_features:Optional[Sequence[int]]=None,
+ out_features:Optional[Sequence[int]]=None,
+ n_batches:int=1,
+ init_method:Text='randn',
+ device:Optional[torch.device]=None,
+ **kwargs)->None:
+
+ tensors=None
+
# n_features
- ifn_features<1:
- raiseValueError('`n_features` cannot be lower than 1')
- self._n_features=n_features
-
- # in_dim
- ifisinstance(in_dim,int):
- self._in_dim=in_dim
- else:
- raiseTypeError('`in_dim` should be int type')
-
- # bond_dim
- ifisinstance(bond_dim,int):
- self._bond_dim=bond_dim
+ ifnotisinstance(n_features,int):
+ raiseTypeError('`n_features` should be int type')
+ elifn_features<1:
+ raiseValueError('`n_features` should be at least 1')
+
+ iftensorisNone:
+ # phys_dim
+ ifnotisinstance(phys_dim,int):
+ raiseTypeError('`phys_dim` should be int type')
+
+ # bond_dim
+ ifnotisinstance(bond_dim,int):
+ raiseTypeError('`bond_dim` should be int type')
+
else:
- raiseTypeError('`bond_dim` should be int type')
-
- # n_batches
- ifnotisinstance(n_batches,int):
- raiseTypeError('`n_batches should be int type')
- self._n_batches=n_batches
-
- # Create Tensor Network
- self._make_nodes()
- self.initialize()
-
- @property
- defn_features(self)->int:
-"""Returns number of nodes."""
- returnself._n_features
-
- @property
- defin_dim(self)->int:
-"""Returns input/physical dimension."""
- returnself._in_dim
-
- @property
- defbond_dim(self)->int:
-"""Returns bond dimension."""
- returnself._bond_dim
-
- @property
- defn_batches(self)->int:
-"""Returns number of batch edges of the ``data`` nodes."""
- returnself._n_batches
+ ifnotisinstance(tensor,torch.Tensor):
+ raiseTypeError('`tensor` should be torch.Tensor type')
+ iflen(tensor.shape)!=3:
+ raiseValueError('`tensor` should be a rank-3 tensor')
+ iftensor.shape[0]!=tensor.shape[2]:
+ raiseValueError('`tensor` first and last dimensions should'
+ ' be equal so that the MPS can have '
+ 'periodic boundary conditions')
+
+ tensors=[tensor]*n_features
+
+ super().__init__(n_features=n_features,
+ phys_dim=phys_dim,
+ bond_dim=bond_dim,
+ boundary='pbc',
+ tensors=tensors,
+ in_features=in_features,
+ out_features=out_features,
+ n_batches=n_batches,
+ init_method=init_method,
+ device=device,
+ **kwargs)
+ self.name='umps'def_make_nodes(self)->None:"""Creates all the nodes of the MPS."""
- ifself._leaf_nodes:
- raiseValueError('Cannot create MPS nodes if the MPS already has '
- 'nodes')
-
- self.left_node=None
- self.right_node=None
- self.mats_env=[]
-
- foriinrange(self._n_features):
- node=ParamNode(shape=(self.bond_dim,self.in_dim,self.bond_dim),
- axes_names=('left','input','right'),
- name=f'mats_env_node_({i})',
- network=self)
- self.mats_env.append(node)
-
- ifi==0:
- periodic_edge=self.mats_env[-1]['left']
- else:
- self.mats_env[-2]['right']^self.mats_env[-1]['left']
-
- ifi==self._n_features-1:
- self.mats_env[-1]['right']^periodic_edge
-
+ super()._make_nodes()
+
# Virtual node
- uniform_memory=ParamNode(shape=(self.bond_dim,self.in_dim,self.bond_dim),
+ uniform_memory=ParamNode(shape=(self._bond_dim[0],
+ self._phys_dim[0],
+ self._bond_dim[0]),axes_names=('left','input','right'),name='virtual_uniform',network=self,virtual=True)self.uniform_memory=uniform_memory
-
-
[docs]definitialize(self,
+ tensors:Optional[Sequence[torch.Tensor]]=None,
+ init_method:Optional[Text]='randn',
+ device:Optional[torch.device]=None,
+ **kwargs:float)->None:"""
- Initializes output and uniform nodes as explained `here
- <https://arxiv.org/abs/1605.03795>`_.
- It can be overriden for custom initializations.
+ Initializes the common tensor of the :class:`UMPS`. It can be called
+ when instantiating the model, or to override the existing nodes' tensors.
+
+ There are different methods to initialize the nodes:
+
+ * ``{"zeros", "ones", "copy", "rand", "randn"}``: The tensor is
+ initialized calling :meth:`~tensorkrowch.AbstractNode.set_tensor` with
+ the given method, ``device`` and ``kwargs``.
+
+ * ``"randn_eye"``: Tensor is initialized as in this
+ `paper <https://arxiv.org/abs/1605.03795>`_, adding identities at the
+ top of a random gaussian tensor. In this case, ``std`` should be
+ specified with a low value, e.g., ``std = 1e-9``.
+
+ * ``"unit"``: Tensor is initialized as a random unitary, so that the
+ MPS is in canonical form.
+
+ Parameters
+ ----------
+ tensors : list[torch.Tensor] or tuple[torch.Tensor], optional
+ Sequence of a single tensor to set in each of the MPS nodes. The
+ tensor should be rank-3, with its first and last dimensions being
+ equal.
+ init_method : {"zeros", "ones", "copy", "rand", "randn", "randn_eye", "unit"}, optional
+ Initialization method.
+ device : torch.device, optional
+ Device where to initialize the tensors if ``init_method`` is provided.
+ kwargs : float
+ Keyword arguments for the different initialization methods. See
+ :meth:`~tensorkrowch.AbstractNode.make_tensor`. """
- # Virtual node
- tensor=torch.randn(self.uniform_memory.shape)*std
- random_eye=torch.randn(tensor.shape[0],tensor.shape[2])*std
- random_eye=random_eye+torch.eye(tensor.shape[0],tensor.shape[2])
- tensor[:,0,:]=random_eye
-
- self.uniform_memory.tensor=tensor
-
- fornodeinself.mats_env:
- node.set_tensor_from(self.uniform_memory)
[docs]defcopy(self,share_tensors:bool=False)->'UMPS':"""
- Creates ``data`` nodes and connects each of them to the input/physical
- edge of each input node.
- """
- input_edges=[]
- ifself.left_nodeisnotNone:
- input_edges.append(self.left_node['input'])
- input_edges+=list(map(lambdanode:node['input'],self.mats_env))
- ifself.right_nodeisnotNone:
- input_edges.append(self.right_node['input'])
-
- super().set_data_nodes(input_edges=input_edges,
- num_batch_edges=self._n_batches)
-
- ifself.mats_env:
- self.mats_env_data=list(map(lambdanode:node.neighbours('input'),
- self.mats_env))
+ Creates a copy of the :class:`UMPS`.
- def_input_contraction(self,
- inline_input:bool=False)->Tuple[
- Optional[List[Node]],
- Optional[List[Node]]]:
-"""Contracts input data nodes with MPS nodes."""
- ifinline_input:
- mats_result=[]
- fornodeinself.mats_env:
- mats_result.append(node@node.neighbours('input'))
- returnmats_result
+ Parameters
+ ----------
+ share_tensor : bool, optional
+ Boolean indicating whether the common tensor in the copied UMPS
+ should be set as the tensor in the current UMPS (``True``), or
+ cloned (``False``). In the former case, the tensor in both UMPS's
+ will be the same, which might be useful if one needs more than one
+ copy of a UMPS, but wants to compute all the gradients with respect
+ to the same, unique, tensor.
+ Returns
+ -------
+ UMPS
+ """
+ new_mps=UMPS(n_features=self._n_features,
+ phys_dim=self._phys_dim[0],
+ bond_dim=self._bond_dim[0],
+ tensor=None,
+ in_features=self._in_features,
+ out_features=self._out_features,
+ n_batches=self._n_batches,
+ init_method=None,
+ device=None)
+ new_mps.name=self.name+'_copy'
+ ifshare_tensors:
+ new_mps.uniform_memory.tensor=self.uniform_memory.tensorelse:
- ifself.mats_env:
- stack=op.stack(self.mats_env)
- stack_data=op.stack(self.mats_env_data)
-
- stack['input']^stack_data['feature']
+ new_mps.uniform_memory.tensor=self.uniform_memory.tensor.clone()
+ returnnew_mps
+
+
[docs]defparameterize(self,
+ set_param:bool=True,
+ override:bool=False)->'TensorNetwork':
+"""
+ Parameterizes all nodes of the MPS. If there are ``resultant`` nodes
+ in the MPS, it will be first :meth:`~tensorkrowch.TensorNetwork.reset`.
- result=stack_data@stack
- mats_result=op.unbind(result)
- returnmats_result
- else:
- return[]
+ Parameters
+ ----------
+ set_param : bool
+ Boolean indicating whether the tensor network has to be parameterized
+ (``True``) or de-parameterized (``False``).
+ override : bool
+ Boolean indicating whether the tensor network should be parameterized
+ in-place (``True``) or copied and then parameterized (``False``).
+ """
+ ifself._resultant_nodes:
+ warnings.warn(
+ 'Resultant nodes will be removed before parameterizing the TN')
+ self.reset()
- @staticmethod
- def_inline_contraction(nodes:List[Node])->Node:
-"""Contracts sequence of MPS nodes (matrices) inline."""
- result_node=nodes[0]
- fornodeinnodes[1:]:
- result_node@=node
- returnresult_node
-
- def_contract_envs_inline(self,mats_env:List[Node])->Node:
-"""Contracts the left and right environments inline."""
- iflen(mats_env)>1:
- contract_lst=mats_env
+ ifoverride:
+ net=selfelse:
- returnmats_env[0]@mats_env[0]
-
- returnself._inline_contraction(contract_lst)
-
- def_aux_pairwise(self,nodes:List[Node])->Tuple[List[Node],
- List[Node]]:
-"""Contracts a sequence of MPS nodes (matrices) pairwise."""
- length=len(nodes)
- aux_nodes=nodes
- iflength>1:
- half_length=length//2
- nice_length=2*half_length
-
- even_nodes=aux_nodes[0:nice_length:2]
- odd_nodes=aux_nodes[1:nice_length:2]
- leftover=aux_nodes[nice_length:]
-
- stack1=op.stack(even_nodes)
- stack2=op.stack(odd_nodes)
+ net=self.copy(share_tensors=False)
+
+ foriinrange(self._n_features):
+ net._mats_env[i]=net._mats_env[i].parameterize(set_param)
+
+ # It is important that uniform_memory is parameterized after the rest
+ # of the nodes
+ net.uniform_memory=net.uniform_memory.parameterize(set_param)
+
+ # Tensor addresses have to be reassigned to reference
+ # the uniform memory
+ fornodeinnet._mats_env:
+ node.set_tensor_from(net.uniform_memory)
+
+ returnnet
+
+ defcanonicalize(self,
+ oc:Optional[int]=None,
+ mode:Text='svd',
+ rank:Optional[int]=None,
+ cum_percentage:Optional[float]=None,
+ cutoff:Optional[float]=None,
+ renormalize:bool=False)->None:
+""":meta private:"""
+ raiseNotImplementedError(
+ '`canonicalize` not implemented for UMPS')
+
+ defcanonicalize_univocal(self):
+""":meta private:"""
+ raiseNotImplementedError(
+ '`canonicalize_univocal` not implemented for UMPS')
[docs]classMPSLayer(MPS):# MARK: MPSLayer
+"""
+ Class for Matrix Product States with a single output node. That is, this
+ MPS has :math:`n` nodes, being :math:`n-1` input nodes connected to ``data``
+ nodes (nodes that will contain the data tensors), and one output node,
+ whose physical dimension (``out_dim``) is used as the label (for
+ classification tasks).
+
+ Besides, since this class has an output edge, when contracting the whole
+ tensor network (with input data), the result will be a vector that can be
+ plugged into the next layer (being this other tensor network or a neural
+ network layer).
+
+ If the physical dimensions of all the input nodes (``in_dim``) are equal,
+ the input data tensor can be passed as a single tensor. Otherwise, it would
+ have to be passed as a list of tensors with different sizes.
+
+ |
+
+ That is, ``MPSLayer`` is equivalent to :class:`MPS` with
+ ``out_features = [out_position]``. However, ``in_features`` and
+ ``out_features`` are still free to be changed if necessary, even though
+ this may change the expected behaviour of the ``MPSLayer``. The expected
+ behaviour can be recovered by setting ``out_features = [out_position]``
+ again.
+
+ |
+
+ For a more detailed list of inherited properties and methods,
+ check :class:`MPS`.
- returnaux_nodes,leftover
- returnnodes,[]
+ Parameters
+ ----------
+ n_features : int, optional
+ Number of nodes that will be in ``mats_env``. That is, number of nodes
+ without taking into account ``left_node`` and ``right_node``. This also
+ includes the output node, so if one wants to instantiate an ``MPSLayer``
+ for a dataset with ``n`` features, it should be ``n_features = n + 1``,
+ to account for the output node.
+ in_dim : int, list[int] or tuple[int], optional
+ Input dimension(s). Equivalent to the physical dimension(s) but only
+ for input nodes. If given as a sequence, its length should be equal to
+ ``n_features - 1``, since these are the input dimensions of the input
+ nodes.
+ out_dim : int, optional
+ Output dimension (labels) for the output node. Equivalent to the
+ physical dimension of the output node.
+ bond_dim : int, list[int] or tuple[int], optional
+ Bond dimension(s). If given as a sequence, its length should be equal
+ to ``n_features`` (if ``boundary = "pbc"``) or ``n_features - 1`` (if
+ ``boundary = "obc"``). The i-th bond dimension is always the dimension
+ of the right edge of the i-th node (including output node).
+ out_position : int, optional
+ Position of the output node (label). Should be between 0 and
+ ``n_features - 1``. If ``None``, the output node will be located at the
+ middle of the MPS.
+ boundary : {"obc", "pbc"}
+ String indicating whether periodic or open boundary conditions should
+ be used.
+ tensors: list[torch.Tensor] or tuple[torch.Tensor], optional
+ Instead of providing ``n_features``, ``in_dim``, ``out_dim``,
+ ``bond_dim`` and ``boundary``, a list of MPS tensors can be provided.
+ In such case, all mentioned attributes will be inferred from the given
+ tensors. All tensors should be rank-3 tensors, with shape ``(bond_dim,
+ phys_dim, bond_dim)``. If the first and last elements are rank-2 tensors,
+ with shapes ``(phys_dim, bond_dim)``, ``(bond_dim, phys_dim)``,
+ respectively, the inferred boundary conditions will be "obc". Also, if
+ ``tensors`` contains a single element, it can be rank-1 ("obc") or
+ rank-3 ("pbc").
+ n_batches : int
+ Number of batch edges of input ``data`` nodes. Usually ``n_batches = 1``
+ (where the batch edge is used for the data batched) but it could also
+ be ``n_batches = 2`` (e.g. one edge for data batched, other edge for
+ image patches in convolutional layers).
+ init_method : {"zeros", "ones", "copy", "rand", "randn", "randn_eye", "unit"}, optional
+ Initialization method. Check :meth:`initialize` for a more detailed
+ explanation of the different initialization methods.
+ device : torch.device, optional
+ Device where to initialize the tensors if ``init_method`` is provided.
+ kwargs : float
+ Keyword arguments for the different initialization methods. See
+ :meth:`~tensorkrowch.AbstractNode.make_tensor`.
+
+ Examples
+ --------
+ ``MPSLayer`` with same input dimensions:
+
+ >>> mps_layer = tk.models.MPSLayer(n_features=4,
+ ... in_dim=2,
+ ... out_dim=10,
+ ... bond_dim=5)
+ >>> data = torch.ones(20, 3, 2) # batch_size x (n_features - 1) x feature_size
+ >>> result = mps_layer(data)
+ >>> result.shape
+ torch.Size([20, 10])
+
+ ``MPSLayer`` with different input dimensions:
+
+ >>> mps_layer = tk.models.MPSLayer(n_features=4,
+ ... in_dim=list(range(2, 5)),
+ ... out_dim=10,
+ ... bond_dim=5)
+ >>> data = [torch.ones(20, i)
+ ... for i in range(2, 5)] # (n_features - 1) * [batch_size x feature_size]
+ >>> result = mps_layer(data)
+ >>> result.shape
+ torch.Size([20, 10])
+ """
- def_pairwise_contraction(self,mats_nodes:List[Node])->Node:
-"""Contracts the left and right environments pairwise."""
- length=len(mats_nodes)
- aux_nodes=mats_nodes
- iflength>1:
- leftovers=[]
- whilelength>1:
- aux1,aux2=self._aux_pairwise(aux_nodes)
- aux_nodes=aux1
- leftovers=aux2+leftovers
- length=len(aux1)
+ def__init__(self,
+ n_features:Optional[int]=None,
+ in_dim:Optional[Union[int,Sequence[int]]]=None,
+ out_dim:Optional[int]=None,
+ bond_dim:Optional[Union[int,Sequence[int]]]=None,
+ out_position:Optional[int]=None,
+ boundary:Text='obc',
+ tensors:Optional[Sequence[torch.Tensor]]=None,
+ n_batches:int=1,
+ init_method:Text='randn',
+ device:Optional[torch.device]=None,
+ **kwargs)->None:
+
+ phys_dim=None
+
+ iftensorsisnotNone:
+ ifnotisinstance(tensors,(list,tuple)):
+ raiseTypeError('`tensors` should be a tuple[torch.Tensor] or '
+ 'list[torch.Tensor] type')
+ n_features=len(tensors)
+ else:
+ ifnotisinstance(n_features,int):
+ raiseTypeError('`n_features` should be int type')
+
+ # out_position
+ ifout_positionisNone:
+ out_position=n_features//2
+ if(out_position<0)or(out_position>n_features):
+ raiseValueError(
+ f'`out_position` should be between 0 and {n_features}')
+ self._out_position=out_position
+
+ iftensorsisNone:
+ # in_dim
+ ifisinstance(in_dim,(list,tuple)):
+ iflen(in_dim)!=(n_features-1):
+ raiseValueError(
+ 'If `in_dim` is given as a sequence of int, its '
+ 'length should be equal to `n_features` - 1')
+ else:
+ fordiminin_dim:
+ ifnotisinstance(dim,int):
+ raiseTypeError(
+ '`in_dim` should be int, tuple[int] or '
+ 'list[int] type')
+ in_dim=list(in_dim)
+ elifisinstance(in_dim,int):
+ in_dim=[in_dim]*(n_features-1)
+ else:
+ ifn_features==1:
+ in_dim=[]
+ else:
+ raiseTypeError(
+ '`in_dim` should be int, tuple[int] or list[int] type')
+
+ # out_dim
+ ifnotisinstance(out_dim,int):
+ raiseTypeError('`out_dim` should be int type')
+
+ # phys_dim
+ phys_dim=in_dim[:out_position]+[out_dim]+in_dim[out_position:]
+
+ super().__init__(n_features=n_features,
+ phys_dim=phys_dim,
+ bond_dim=bond_dim,
+ boundary=boundary,
+ tensors=tensors,
+ in_features=None,
+ out_features=[out_position],
+ n_batches=n_batches,
+ init_method=init_method,
+ device=device,
+ **kwargs)
+ self.name='mpslayer'
+ self._in_dim=self._phys_dim[:out_position]+ \
+ self._phys_dim[(out_position+1):]
+ self._out_dim=self._phys_dim[out_position]
+
+ @property
+ defin_dim(self)->List[int]:
+"""Returns input dimensions."""
+ returnself._in_dim
+
+ @property
+ defout_dim(self)->int:
+"""
+ Returns the output dimension, that is, the number of labels in the
+ output node. Same as ``in_dim`` for input nodes.
+ """
+ returnself._out_dim
+
+ @property
+ defout_position(self)->int:
+"""Returns position of the output node (label)."""
+ returnself._out_position
+
+ @property
+ defout_node(self)->ParamNode:
+"""Returns the output node."""
+ returnself._mats_env[self._out_position]
+
+
[docs]definitialize(self,
+ tensors:Optional[Sequence[torch.Tensor]]=None,
+ init_method:Optional[Text]='randn',
+ device:Optional[torch.device]=None,
+ **kwargs:float)->None:
+"""
+ Initializes all the nodes of the :class:`MPSLayer`. It can be called when
+ instantiating the model, or to override the existing nodes' tensors.
+
+ There are different methods to initialize the nodes:
+
+ * ``{"zeros", "ones", "copy", "rand", "randn"}``: Each node is
+ initialized calling :meth:`~tensorkrowch.AbstractNode.set_tensor` with
+ the given method, ``device`` and ``kwargs``.
+
+ * ``"randn_eye"``: Nodes are initialized as in this
+ `paper <https://arxiv.org/abs/1605.03795>`_, adding identities at the
+ top of random gaussian tensors. In this case, ``std`` should be
+ specified with a low value, e.g., ``std = 1e-9``.
+
+ * ``"unit"``: Nodes are initialized as random unitaries, so that the
+ MPS is in canonical form, with the orthogonality center at the
+ rightmost node.
+
+ Parameters
+ ----------
+ tensors : list[torch.Tensor] or tuple[torch.Tensor], optional
+ Sequence of tensors to set in each of the MPS nodes. If ``boundary``
+ is ``"obc"``, all tensors should be rank-3, except the first and
+ last ones, which can be rank-2, or rank-1 (if the first and last are
+ the same). If ``boundary`` is ``"pbc"``, all tensors should be
+ rank-3.
+ init_method : {"zeros", "ones", "copy", "rand", "randn", "randn_eye", "unit"}, optional
+ Initialization method.
+ device : torch.device, optional
+ Device where to initialize the tensors if ``init_method`` is provided.
+ kwargs : float
+ Keyword arguments for the different initialization methods. See
+ :meth:`~tensorkrowch.AbstractNode.make_tensor`.
+ """
+ ifself._boundary=='obc':
+ self._left_node.set_tensor(init_method='copy',device=device)
+ self._right_node.set_tensor(init_method='copy',device=device)
+
+ ifinit_method=='unit':
+ tensors=self._make_unitaries(device=device)
+
+ iftensorsisnotNone:
+ iflen(tensors)!=self._n_features:
+ raiseValueError('`tensors` should be a sequence of `n_features`'
+ ' elements')
+
+ ifself._boundary=='obc':
+ tensors=tensors[:]
+ iflen(tensors)==1:
+ tensors[0]=tensors[0].reshape(1,-1,1)
+ else:
+ # Left node
+ aux_tensor=torch.zeros(*self._mats_env[0].shape,
+ device=tensors[0].device)
+ aux_tensor[0]=tensors[0]
+ tensors[0]=aux_tensor
+
+ # Right node
+ aux_tensor=torch.zeros(*self._mats_env[-1].shape,
+ device=tensors[-1].device)
+ aux_tensor[...,0]=tensors[-1]
+ tensors[-1]=aux_tensor
+
+ fortensor,nodeinzip(tensors,self._mats_env):
+ node.tensor=tensor
+
+ elifinit_methodisnotNone:
+ add_eye=False
+ ifinit_method=='randn_eye':
+ init_method='randn'
+ add_eye=True
+
+ fori,nodeinenumerate(self._mats_env):
+ node.set_tensor(init_method=init_method,
+ device=device,
+ **kwargs)
+ ifadd_eye:
+ aux_tensor=node.tensor.detach()
+ eye_tensor=torch.eye(node.shape[0],
+ node.shape[2],
+ device=device)
+ ifi==self._out_position:
+ eye_tensor=eye_tensor.unsqueeze(1)
+ eye_tensor=eye_tensor.expand(node.shape)
+ aux_tensor+=eye_tensor
+ else:
+ aux_tensor[:,0,:]+=eye_tensor
+ node.tensor=aux_tensor
+
+ ifself._boundary=='obc':
+ aux_tensor=torch.zeros(*node.shape,device=device)
+ ifi==0:
+ # Left node
+ aux_tensor[0]=node.tensor[0]
+ node.tensor=aux_tensor
+ elifi==(self._n_features-1):
+ # Right node
+ aux_tensor[...,0]=node.tensor[...,0]
+ node.tensor=aux_tensor
+
+
[docs]defcopy(self,share_tensors:bool=False)->'MPSLayer':
+"""
+ Creates a copy of the :class:`MPSLayer`.
+
+ Parameters
+ ----------
+ share_tensor : bool, optional
+ Boolean indicating whether tensors in the copied MPSLayer should be
+ set as the tensors in the current MPSLayer (``True``), or cloned
+ (``False``). In the former case, tensors in both MPSLayer's will be
+ the same, which might be useful if one needs more than one copy
+ of an MPSLayer, but wants to compute all the gradients with respect
+ to the same, unique, tensors.
+
+ Returns
+ -------
+ MPSLayer
+ """
+ new_mps=MPSLayer(n_features=self._n_features,
+ in_dim=self._in_dim,
+ out_dim=self._out_dim,
+ bond_dim=self._bond_dim,
+ out_position=self._out_position,
+ boundary=self._boundary,
+ tensors=None,
+ n_batches=self._n_batches,
+ init_method=None,
+ device=None)
+ new_mps.name=self.name+'_copy'
+ ifshare_tensors:
+ fornew_node,nodeinzip(new_mps._mats_env,self._mats_env):
+ new_node.tensor=node.tensor
+ else:
+ fornew_node,nodeinzip(new_mps._mats_env,self._mats_env):
+ new_node.tensor=node.tensor.clone()
+
+ returnnew_mps
+
+
+
[docs]classUMPSLayer(MPS):# MARK: UMPSLayer
+"""
+ Class for Uniform (translationally invariant) Matrix Product States with an
+ output node. It is the uniform version of :class:`MPSLayer`, with all input
+ nodes sharing the same tensor, but with a different node for the output node.
+ Thus this class cannot have different input or bond dimensions for each node,
+ and boundary conditions are always periodic (``"pbc"``).
+
+ |
+
+ For a more detailed list of inherited properties and methods,
+ check :class:`MPS`.
+
+ Parameters
+ ----------
+ n_features : int
+ Number of nodes that will be in ``mats_env``. This also includes the
+ output node, so if one wants to instantiate a ``UMPSLayer`` for a
+ dataset with ``n`` features, it should be ``n_features = n + 1``, to
+ account for the output node.
+ in_dim : int, optional
+ Input dimension. Equivalent to the physical dimension but only for
+ input nodes.
+ out_dim : int, optional
+ Output dimension (labels) for the output node.
+ bond_dim : int, optional
+ Bond dimension.
+ out_position : int, optional
+ Position of the output node (label). Should be between 0 and
+ ``n_features - 1``. If ``None``, the output node will be located at the
+ middle of the MPS.
+ tensors: list[torch.Tensor] or tuple[torch.Tensor], optional
+ Instead of providing ``in_dim``, ``out_dim`` and ``bond_dim``, a
+ sequence of 2 tensors can be provided, the first one will be the uniform
+ tensor, and the second one will be the output node's tensor.
+ ``n_features`` is still needed to specify how many times the uniform
+ tensor should be used to form a finite MPS. In this case, since the
+ output node will have a different tensor, the uniform tensor will be
+ used in the remaining ``n_features - 1`` input nodes. Both tensors
+ should be rank-3, with all their first and last dimensions being equal.
+ n_batches : int
+ Number of batch edges of input ``data`` nodes. Usually ``n_batches = 1``
+ (where the batch edge is used for the data batched) but it could also
+ be ``n_batches = 2`` (one edge for data batched, other edge for image
+ patches in convolutional layers).
+ init_method : {"zeros", "ones", "copy", "rand", "randn", "randn_eye", "unit"}, optional
+ Initialization method. Check :meth:`initialize` for a more detailed
+ explanation of the different initialization methods.
+ device : torch.device, optional
+ Device where to initialize the tensors if ``init_method`` is provided.
+ kwargs : float
+ Keyword arguments for the different initialization methods. See
+ :meth:`~tensorkrowch.AbstractNode.make_tensor`.
+
+ Examples
+ --------
+ >>> mps_layer = tk.models.UMPSLayer(n_features=4,
+ ... in_dim=2,
+ ... out_dim=10,
+ ... bond_dim=5)
+ >>> for i, node in enumerate(mps_layer.mats_env):
+ ... if i != mps_layer.out_position:
+ ... assert node.tensor_address() == 'virtual_uniform'
+ ...
+ >>> data = torch.ones(20, 3, 2) # batch_size x (n_features - 1) x feature_size
+ >>> result = mps_layer(data)
+ >>> result.shape
+ torch.Size([20, 10])
+ """
+
+ def__init__(self,
+ n_features:int,
+ in_dim:Optional[int]=None,
+ out_dim:Optional[int]=None,
+ bond_dim:Optional[int]=None,
+ out_position:Optional[int]=None,
+ tensors:Optional[Sequence[torch.Tensor]]=None,
+ n_batches:int=1,
+ init_method:Text='randn',
+ device:Optional[torch.device]=None,
+ **kwargs)->None:
+
+ phys_dim=None
+
+ # n_features
+ ifnotisinstance(n_features,int):
+ raiseTypeError('`n_features` should be int type')
+ elifn_features<1:
+ raiseValueError('`n_features` should be at least 1')
+
+ # out_position
+ ifout_positionisNone:
+ out_position=n_features//2
+ if(out_position<0)or(out_position>n_features):
+ raiseValueError(
+ f'`out_position` should be between 0 and {n_features}')
+ self._out_position=out_position
+
+ iftensorsisNone:
+ # in_dim
+ ifisinstance(in_dim,int):
+ in_dim=[in_dim]*(n_features-1)
+ else:
+ ifn_features==1:
+ in_dim=[]
+ else:
+ raiseTypeError(
+ '`in_dim` should be int, tuple[int] or list[int] type')
+
+ # out_dim
+ ifnotisinstance(out_dim,int):
+ raiseTypeError('`out_dim` should be int type')
+
+ # phys_dim
+ phys_dim=in_dim[:out_position]+[out_dim]+in_dim[out_position:]
+
+ else:
+ ifnotisinstance(tensors,Sequence):
+ raiseTypeError('`tensors` should be a tuple[torch.Tensor] or '
+ 'list[torch.Tensor] type')
+ iflen(tensors)!=2:
+ raiseValueError('`tensors` should have 2 elements, the first'
+ ' corresponding to the common input tensor, '
+ 'and another one for the output node')
+ fortintensors:
+ ifnotisinstance(t,torch.Tensor):
+ raiseTypeError(
+ 'Elements of `tensors` should be torch.Tensor type')
+ iflen(t.shape)!=3:
+ raiseValueError(
+ 'Elements of `tensors` should be a rank-3 tensor')
+ ift.shape[0]!=t.shape[2]:
+ raiseValueError(
+ 'Elements of `tensors` should have equal first and last'
+ ' dimensions so that the MPS can have periodic boundary'
+ ' conditions')
+
+ ifn_features==1:
+ # Only output node is used, uniform memory will
+ # take that tensor too
+ tensors=[tensors[1]]
+ else:
+ tensors=[tensors[0]]*out_position+[tensors[1]]+ \
+ [tensors[0]]*(n_features-1-out_position)
+
+ super().__init__(n_features=n_features,
+ phys_dim=phys_dim,
+ bond_dim=bond_dim,
+ boundary='pbc',
+ tensors=tensors,
+ in_features=None,
+ out_features=[out_position],
+ n_batches=n_batches,
+ init_method=init_method,
+ device=device,
+ **kwargs)
+ self.name='umpslayer'
+ self._in_dim=self._phys_dim[:out_position]+ \
+ self._phys_dim[(out_position+1):]
+ self._out_dim=self._phys_dim[out_position]
+
+ @property
+ defin_dim(self)->List[int]:
+"""Returns input dimensions."""
+ returnself._in_dim
+
+ @property
+ defout_dim(self)->int:
+"""
+ Returns the output dimension, that is, the number of labels in the
+ output node. Same as ``in_dim`` for input nodes.
+ """
+ returnself._out_dim
+
+ @property
+ defout_position(self)->int:
+"""Returns position of the output node (label)."""
+ returnself._out_position
+
+ @property
+ defout_node(self)->ParamNode:
+"""Returns the output node."""
+ returnself._mats_env[self._out_position]
+
+ def_make_nodes(self)->None:
+"""Creates all the nodes of the MPS."""
+ super()._make_nodes()
+
+ # Virtual node
+ uniform_memory=ParamNode(shape=(self._bond_dim[0],
+ self._phys_dim[0],
+ self._bond_dim[0]),
+ axes_names=('left','input','right'),
+ name='virtual_uniform',
+ network=self,
+ virtual=True)
+ self.uniform_memory=uniform_memory
+
+ in_nodes=self._mats_env[:self._out_position]+ \
+ self._mats_env[(self._out_position+1):]
+ fornodeinin_nodes:
+ node.set_tensor_from(uniform_memory)
+
+ def_make_unitaries(self,device:Optional[torch.device]=None)->List[torch.Tensor]:
+"""Initializes MPS in canonical form."""
+ tensors=[]
+ fornodein[self.uniform_memory,self.out_node]:
+ node_shape=node.shape
+ aux_shape=(node.shape[:2].numel(),node.shape[2])
+
+ size=max(aux_shape[0],aux_shape[1])
+
+ tensor=random_unitary(size,device=device)
+ tensor=tensor[:min(aux_shape[0],size),:min(aux_shape[1],size)]
+ tensor=tensor.reshape(*node_shape)
+
+ tensors.append(tensor)
+ returntensors
+
+
[docs]definitialize(self,
+ tensors:Optional[Sequence[torch.Tensor]]=None,
+ init_method:Optional[Text]='randn',
+ device:Optional[torch.device]=None,
+ **kwargs:float)->None:
+"""
+ Initializes the common tensor of the :class:`UMPSLayer`. It can be called
+ when instantiating the model, or to override the existing nodes' tensors.
+
+ There are different methods to initialize the nodes:
+
+ * ``{"zeros", "ones", "copy", "rand", "randn"}``: The tensor is
+ initialized calling :meth:`~tensorkrowch.AbstractNode.set_tensor` with
+ the given method, ``device`` and ``kwargs``.
+
+ * ``"randn_eye"``: Tensor is initialized as in this
+ `paper <https://arxiv.org/abs/1605.03795>`_, adding identities at the
+ top of a random gaussian tensor. In this case, ``std`` should be
+ specified with a low value, e.g., ``std = 1e-9``.
+
+ * ``"unit"``: Tensor is initialized as a random unitary, so that the
+ MPS is in canonical form.
+
+ Parameters
+ ----------
+ tensors : list[torch.Tensor] or tuple[torch.Tensor], optional
+ Sequence of a 2 tensors, the first one will be the uniform tensor
+ that will be set in all input nodes, and the second one will be the
+ output node's tensor. Both tensors should be rank-3, with all their
+ first and last dimensions being equal.
+ init_method : {"zeros", "ones", "copy", "rand", "randn", "randn_eye", "unit"}, optional
+ Initialization method.
+ device : torch.device, optional
+ Device where to initialize the tensors if ``init_method`` is provided.
+ kwargs : float
+ Keyword arguments for the different initialization methods. See
+ :meth:`~tensorkrowch.AbstractNode.make_tensor`.
+ """
+ ifinit_method=='unit':
+ tensors=self._make_unitaries(device=device)
+
+ iftensorsisnotNone:
+ self.uniform_memory.tensor=tensors[0]
+ self.out_node.tensor=tensors[-1]
+
+ elifinit_methodisnotNone:
+ fori,nodeinenumerate([self.uniform_memory,self.out_node]):
+ add_eye=False
+ ifinit_method=='randn_eye':
+ init_method='randn'
+ add_eye=True
+
+ node.set_tensor(init_method=init_method,
+ device=device,
+ **kwargs)
+ ifadd_eye:
+ aux_tensor=node.tensor.detach()
+ eye_tensor=torch.eye(node.shape[0],
+ node.shape[2],
+ device=device)
+ ifi==0:
+ aux_tensor[:,0,:]+=eye_tensor
+ else:
+ eye_tensor=eye_tensor.unsqueeze(1)
+ eye_tensor=eye_tensor.expand(node.shape)
+ aux_tensor+=eye_tensor
+ node.tensor=aux_tensor
+
+
[docs]defcopy(self,share_tensors:bool=False)->'UMPSLayer':
+"""
+ Creates a copy of the :class:`UMPSLayer`.
+
+ Parameters
+ ----------
+ share_tensor : bool, optional
+ Boolean indicating whether tensors in the copied UMPSLayer should be
+ set as the tensors in the current UMPSLayer (``True``), or cloned
+ (``False``). In the former case, tensors in both UMPSLayer's will be
+ the same, which might be useful if one needs more than one copy
+ of an UMPSLayer, but wants to compute all the gradients with respect
+ to the same, unique, tensors.
+
+ Returns
+ -------
+ UMPSLayer
+ """
+ new_mps=UMPSLayer(n_features=self._n_features,
+ in_dim=self._in_dim[0]ifself._in_dimelseNone,
+ out_dim=self._out_dim,
+ bond_dim=self._bond_dim,
+ out_position=self._out_position,
+ tensor=None,
+ n_batches=self._n_batches,
+ init_method=None,
+ device=None)
+ new_mps.name=self.name+'_copy'
+ ifshare_tensors:
+ new_mps.uniform_memory.tensor=self.uniform_memory.tensor
+ new_mps.out_node.tensor=self.out_node.tensor
+ else:
+ new_mps.uniform_memory.tensor=self.uniform_memory.tensor.clone()
+ new_mps.out_node.tensor=self.out_node.tensor.clone()
+ returnnew_mps
+
+
[docs]defparameterize(self,
+ set_param:bool=True,
+ override:bool=False)->'TensorNetwork':
+"""
+ Parameterizes all nodes of the MPS. If there are ``resultant`` nodes
+ in the MPS, it will be first :meth:`~tensorkrowch.TensorNetwork.reset`.
+
+ Parameters
+ ----------
+ set_param : bool
+ Boolean indicating whether the tensor network has to be parameterized
+ (``True``) or de-parameterized (``False``).
+ override : bool
+ Boolean indicating whether the tensor network should be parameterized
+ in-place (``True``) or copied and then parameterized (``False``).
+ """
+ ifself._resultant_nodes:
+ warnings.warn(
+ 'Resultant nodes will be removed before parameterizing the TN')
+ self.reset()
+
+ ifoverride:
+ net=self
+ else:
+ net=self.copy(share_tensors=False)
+
+ foriinrange(self._n_features):
+ net._mats_env[i]=net._mats_env[i].parameterize(set_param)
+
+ # It is important that uniform_memory is parameterized after the rest
+ # of the nodes
+ net.uniform_memory=net.uniform_memory.parameterize(set_param)
+
+ # Tensor addresses have to be reassigned to reference
+ # the uniform memory
+ fornodeinnet._mats_env:
+ node.set_tensor_from(net.uniform_memory)
+
+ returnnet
+
+ defcanonicalize(self,
+ oc:Optional[int]=None,
+ mode:Text='svd',
+ rank:Optional[int]=None,
+ cum_percentage:Optional[float]=None,
+ cutoff:Optional[float]=None,
+ renormalize:bool=False)->None:
+""":meta private:"""
+ raiseNotImplementedError(
+ '`canonicalize` not implemented for UMPSLayer')
+
+ defcanonicalize_univocal(self):
+""":meta private:"""
+ raiseNotImplementedError(
+ '`canonicalize_univocal` not implemented for UMPSLayer')
+
+
+###############################################################################
+# CONV MODELS #
+###############################################################################
+classAbstractConvClass(ABC):# MARK: AbstractConvClass
+
+ @abstractmethod
+ def__init__(self):
+ pass
+
+ def_set_attributes(self,
+ in_channels:int,
+ kernel_size:Union[int,Sequence[int]],
+ stride:int,
+ padding:int,
+ dilation:int)->nn.Module:
+
+ ifisinstance(kernel_size,int):
+ kernel_size=(kernel_size,kernel_size)
+ elifnotisinstance(kernel_size,Sequence):
+ raiseTypeError('`kernel_size` must be int, list[int] or tuple[int]')
+
+ ifisinstance(stride,int):
+ stride=(stride,stride)
+ elifnotisinstance(stride,Sequence):
+ raiseTypeError('`stride` must be int, list[int] or tuple[int]')
+
+ ifisinstance(padding,int):
+ padding=(padding,padding)
+ elifnotisinstance(padding,Sequence):
+ raiseTypeError('`padding` must be int, list[int] or tuple[int]')
+
+ ifisinstance(dilation,int):
+ dilation=(dilation,dilation)
+ elifnotisinstance(dilation,Sequence):
+ raiseTypeError('`dilation` must be int , list[int] or tuple[int]')
+
+ self._in_channels=in_channels
+ self._kernel_size=kernel_size
+ self._stride=stride
+ self._padding=padding
+ self._dilation=dilation
+
+ unfold=nn.Unfold(kernel_size=kernel_size,
+ stride=stride,
+ padding=padding,
+ dilation=dilation)
+ returnunfold
+
+ defforward(self,image,mode='flat',*args,**kwargs):
+r"""
+ Overrides :meth:`~tensorkrowch.TensorNetwork.forward` to compute a
+ convolution on the input image.
+
+ Parameters
+ ----------
+ image : torch.Tensor
+ Input batch of images with shape
+
+ .. math::
+
+ batch\_size \times in\_channels \times height \times width
+ mode : {"flat", "snake"}
+ Indicates the order in which MPS should take the pixels in the image.
+ When ``"flat"``, the image is flattened putting one row of the image
+ after the other. When ``"snake"``, its row is put in the opposite
+ orientation as the previous row (like a snake running through the
+ image).
+ args :
+ Arguments that might be used in :meth:`~MPS.contract`.
+ kwargs :
+ Keyword arguments that might be used in :meth:`~MPS.contract`,
+ like ``inline_input`` or ``inline_mats``.
+ """
+ # Input image shape: batch_size x in_channels x height x width
+
+ patches=self.unfold(image).transpose(1,2)
+ # batch_size x nb_windows x (in_channels * nb_pixels)
+
+ patches=patches.view(*patches.shape[:-1],self.in_channels,-1)
+ # batch_size x nb_windows x in_channels x nb_pixels
+
+ patches=patches.transpose(2,3)
+ # batch_size x nb_windows x nb_pixels x in_channels
+
+ ifmode=='snake':
+ new_patches=patches[...,:self._kernel_size[1],:]
+ foriinrange(1,self._kernel_size[0]):
+ ifi%2==0:
+ aux=patches[...,(i*self._kernel_size[1]):
+ ((i+1)*self._kernel_size[1]),:]
+ else:
+ aux=patches[...,
+ (i*self._kernel_size[1]):
+ ((i+1)*self._kernel_size[1]),:].flip(dims=[0])
+ new_patches=torch.cat([new_patches,aux],dim=2)
+
+ patches=new_patches
+
+ elifmode!='flat':
+ raiseValueError('`mode` can only be "flat" or "snake"')
+
+ h_in=image.shape[2]
+ w_in=image.shape[3]
+
+ h_out=int((h_in+2*self.padding[0]-self.dilation[0]*
+ (self.kernel_size[0]-1)-1)/self.stride[0]+1)
+ w_out=int((w_in+2*self.padding[1]-self.dilation[1]*
+ (self.kernel_size[1]-1)-1)/self.stride[1]+1)
+
+ result=super().forward(patches,*args,**kwargs)
+ # batch_size x nb_windows (x out_channels ...)
+
+ iflen(result.shape)==3:
+ result=result.movedim(1,-1)
+ # batch_size (x out_channels ...) x nb_windows
+
+ result=result.view(*result.shape[:-1],h_out,w_out)
+ # batch_size (x out_channels ...) x height_out x width_out
+
+ returnresult
+
+
+
[docs]classConvMPS(AbstractConvClass,MPS):# MARK: ConvMPS
+"""
+ Convolutional version of :class:`MPS`, where the input data is assumed to
+ be a batch of images.
+
+ Input data as well as initialization parameters are described in `torch.nn.Conv2d
+ <https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html>`_.
+
+ Parameters
+ ----------
+ in_channels : int
+ Input channels. Same as ``phys_dim`` in :class:`MPS`.
+ bond_dim : int, list[int] or tuple[int]
+ Bond dimension(s). If given as a sequence, its length should be equal
+ to :math:`kernel\_size_0 \cdot kernel\_size_1` (if ``boundary = "pbc"``)
+ or :math:`kernel\_size_0 \cdot kernel\_size_1 - 1` (if
+ ``boundary = "obc"``). The i-th bond dimension is always the dimension
+ of the right edge of the i-th node.
+ kernel_size : int, list[int] or tuple[int]
+ Kernel size used in `torch.nn.Unfold
+ <https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html#torch.nn.Unfold>`_.
+ If given as an ``int``, the actual kernel size will be
+ ``(kernel_size, kernel_size)``.
+ stride : int
+ Stride used in `torch.nn.Unfold
+ <https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html#torch.nn.Unfold>`_.
+ padding : int
+ Padding used in `torch.nn.Unfold
+ <https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html#torch.nn.Unfold>`_.
+ If given as an ``int``, the actual kernel size will be
+ ``(kernel_size, kernel_size)``.
+ dilation : int
+ Dilation used in `torch.nn.Unfold
+ <https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html#torch.nn.Unfold>`_.
+ If given as an ``int``, the actual kernel size will be
+ ``(kernel_size, kernel_size)``.
+ boundary : {"obc", "pbc"}
+ String indicating whether periodic or open boundary conditions should
+ be used.
+ tensors: list[torch.Tensor] or tuple[torch.Tensor], optional
+ To initialize MPS nodes, a list of MPS tensors can be provided. All
+ tensors should be rank-3 tensors, with shape ``(bond_dim, in_channels,
+ bond_dim)``. If the first and last elements are rank-2 tensors, with
+ shapes ``(in_channels, bond_dim)``, ``(bond_dim, in_channels)``,
+ respectively, the inferred boundary conditions will be "obc". Also, if
+ ``tensors`` contains a single element, it can be rank-1 ("obc") or
+ rank-3 ("pbc").
+ init_method : {"zeros", "ones", "copy", "rand", "randn", "randn_eye", "unit"}, optional
+ Initialization method. Check :meth:`~MPS.initialize` for a more detailed
+ explanation of the different initialization methods.
+ device : torch.device, optional
+ Device where to initialize the tensors if ``init_method`` is provided.
+ kwargs : float
+ Keyword arguments for the different initialization methods. See
+ :meth:`~tensorkrowch.AbstractNode.make_tensor`.
+
+ Examples
+ --------
+ >>> conv_mps = tk.models.ConvMPS(in_channels=2,
+ ... bond_dim=5,
+ ... kernel_size=2)
+ >>> data = torch.ones(20, 2, 2, 2) # batch_size x in_channels x height x width
+ >>> result = conv_mps(data)
+ >>> result.shape
+ torch.Size([20, 1, 1])
+ """
+
+ def__init__(self,
+ in_channels:int,
+ bond_dim:Union[int,Sequence[int]],
+ kernel_size:Union[int,Sequence],
+ stride:int=1,
+ padding:int=0,
+ dilation:int=1,
+ boundary:Text='obc',
+ tensors:Optional[Sequence[torch.Tensor]]=None,
+ init_method:Text='randn',
+ device:Optional[torch.device]=None,
+ **kwargs):
+
+ unfold=self._set_attributes(in_channels=in_channels,
+ kernel_size=kernel_size,
+ stride=stride,
+ padding=padding,
+ dilation=dilation)
+
+ MPS.__init__(self,
+ n_features=self._kernel_size[0]*self._kernel_size[1],
+ phys_dim=in_channels,
+ bond_dim=bond_dim,
+ boundary=boundary,
+ tensors=tensors,
+ n_batches=2,
+ init_method=init_method,
+ device=device,
+ **kwargs)
+
+ self.unfold=unfold
+
+ @property
+ defin_channels(self)->int:
+"""Returns ``in_channels``. Same as ``phys_dim`` in :class:`MPS`."""
+ returnself._in_channels
+
+ @property
+ defkernel_size(self)->Tuple[int,int]:
+"""
+ Returns ``kernel_size``. Number of nodes is given by
+ :math:`kernel\_size_0 \cdot kernel\_size_1`.
+ """
+ returnself._kernel_size
+
+ @property
+ defstride(self)->Tuple[int,int]:
+"""
+ Returns stride used in `torch.nn.Unfold
+ <https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html#torch.nn.Unfold>`_.
+ """
+ returnself._stride
+
+ @property
+ defpadding(self)->Tuple[int,int]:
+"""
+ Returns padding used in `torch.nn.Unfold
+ <https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html#torch.nn.Unfold>`_.
+ """
+ returnself._padding
+
+ @property
+ defdilation(self)->Tuple[int,int]:
+"""
+ Returns dilation used in `torch.nn.Unfold
+ <https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html#torch.nn.Unfold>`_.
+ """
+ returnself._dilation
+
+
[docs]defcopy(self,share_tensors:bool=False)->'ConvMPS':
+"""
+ Creates a copy of the :class:`ConvMPS`.
+
+ Parameters
+ ----------
+ share_tensor : bool, optional
+ Boolean indicating whether tensors in the copied ConvMPS should be
+ set as the tensors in the current ConvMPS (``True``), or cloned
+ (``False``). In the former case, tensors in both ConvMPS's will be
+ the same, which might be useful if one needs more than one copy
+ of a ConvMPS, but wants to compute all the gradients with respect
+ to the same, unique, tensors.
+
+ Returns
+ -------
+ ConvMPS
+ """
+ new_mps=ConvMPS(in_channels=self._in_channels,
+ bond_dim=self._bond_dim,
+ kernel_size=self._kernel_size,
+ stride=self._stride,
+ padding=self._padding,
+ dilation=self.dilation,
+ boundary=self._boundary,
+ tensors=None,
+ init_method=None,
+ device=None)
+ new_mps.name=self.name+'_copy'
+ ifshare_tensors:
+ fornew_node,nodeinzip(new_mps._mats_env,self._mats_env):
+ new_node.tensor=node.tensor
+ else:
+ fornew_node,nodeinzip(new_mps._mats_env,self._mats_env):
+ new_node.tensor=node.tensor.clone()
+
+ returnnew_mps
+
+
+
[docs]classConvUMPS(AbstractConvClass,UMPS):# MARK: ConvUMPS
+"""
+ Convolutional version of :class:`UMPS`, where the input data is assumed to
+ be a batch of images.
+
+ Input data as well as initialization parameters are described in `torch.nn.Conv2d
+ <https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html>`_.
+
+ Parameters
+ ----------
+ in_channels : int
+ Input channels. Same as ``phys_dim`` in :class:`UMPS`.
+ bond_dim : int
+ Bond dimension.
+ kernel_size : int, list[int] or tuple[int]
+ Kernel size used in `torch.nn.Unfold
+ <https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html#torch.nn.Unfold>`_.
+ If given as an ``int``, the actual kernel size will be
+ ``(kernel_size, kernel_size)``.
+ stride : int
+ Stride used in `torch.nn.Unfold
+ <https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html#torch.nn.Unfold>`_.
+ padding : int
+ Padding used in `torch.nn.Unfold
+ <https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html#torch.nn.Unfold>`_.
+ If given as an ``int``, the actual kernel size will be
+ ``(kernel_size, kernel_size)``.
+ dilation : int
+ Dilation used in `torch.nn.Unfold
+ <https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html#torch.nn.Unfold>`_.
+ If given as an ``int``, the actual kernel size will be
+ ``(kernel_size, kernel_size)``.
+ tensor: torch.Tensor, optional
+ To initialize MPS nodes, a MPS tensor can be provided. The tensor
+ should be rank-3, with its first and last dimensions being equal.
+ init_method : {"zeros", "ones", "copy", "rand", "randn", "randn_eye", "unit"}, optional
+ Initialization method. Check :meth:`~UMPS.initialize` for a more detailed
+ explanation of the different initialization methods.
+ device : torch.device, optional
+ Device where to initialize the tensors if ``init_method`` is provided.
+ kwargs : float
+ Keyword arguments for the different initialization methods. See
+ :meth:`~tensorkrowch.AbstractNode.make_tensor`.
+
+
+ Examples
+ --------
+ >>> conv_mps = tk.models.ConvUMPS(in_channels=2,
+ ... bond_dim=5,
+ ... kernel_size=2)
+ >>> for node in conv_mps.mats_env:
+ ... assert node.tensor_address() == 'virtual_uniform'
+ ...
+ >>> data = torch.ones(20, 2, 2, 2) # batch_size x in_channels x height x width
+ >>> result = conv_mps(data)
+ >>> result.shape
+ torch.Size([20, 1, 1])
+ """
+
+ def__init__(self,
+ in_channels:int,
+ bond_dim:int,
+ kernel_size:Union[int,Sequence],
+ stride:int=1,
+ padding:int=0,
+ dilation:int=1,
+ tensor:Optional[torch.Tensor]=None,
+ init_method:Text='randn',
+ device:Optional[torch.device]=None,
+ **kwargs):
+
+ unfold=self._set_attributes(in_channels=in_channels,
+ kernel_size=kernel_size,
+ stride=stride,
+ padding=padding,
+ dilation=dilation)
+
+ UMPS.__init__(self,
+ n_features=self._kernel_size[0]*self._kernel_size[1],
+ phys_dim=in_channels,
+ bond_dim=bond_dim,
+ tensor=tensor,
+ n_batches=2,
+ init_method=init_method,
+ device=device,
+ **kwargs)
+
+ self.unfold=unfold
+
+ @property
+ defin_channels(self)->int:
+"""Returns ``in_channels``. Same as ``phys_dim`` in :class:`MPS`."""
+ returnself._in_channels
+
+ @property
+ defkernel_size(self)->Tuple[int,int]:
+"""
+ Returns ``kernel_size``. Number of nodes is given by
+ :math:`kernel\_size_0 \cdot kernel\_size_1`.
+ """
+ returnself._kernel_size
- aux_nodes=aux_nodes+leftovers
- returnself._pairwise_contraction(aux_nodes)
+ @property
+ defstride(self)->Tuple[int,int]:
+"""
+ Returns stride used in `torch.nn.Unfold
+ <https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html#torch.nn.Unfold>`_.
+ """
+ returnself._stride
- returnself._contract_envs_inline(aux_nodes)
+ @property
+ defpadding(self)->Tuple[int,int]:
+"""
+ Returns padding used in `torch.nn.Unfold
+ <https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html#torch.nn.Unfold>`_.
+ """
+ returnself._padding
-
[docs]defcontract(self,
- inline_input:bool=False,
- inline_mats:bool=False)->Node:
+ @property
+ defdilation(self)->Tuple[int,int]:"""
- Contracts the whole MPS.
-
+ Returns dilation used in `torch.nn.Unfold
+ <https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html#torch.nn.Unfold>`_.
+ """
+ returnself._dilation
+
+
[docs]defcopy(self,share_tensors:bool=False)->'ConvUMPS':
+"""
+ Creates a copy of the :class:`ConvUMPS`.
+
Parameters ----------
- inline_input : bool
- Boolean indicating whether input ``data`` nodes should be contracted
- with the ``MPS`` nodes inline (one contraction at a time) or in a
- single stacked contraction.
- inline_mats : bool
- Boolean indicating whether the sequence of matrices (resultant
- after contracting the input ``data`` nodes) should be contracted
- inline or as a sequence of pairwise stacked contrations.
+ share_tensor : bool, optional
+ Boolean indicating whether the common tensor in the copied ConvUMPS
+ should be set as the tensor in the current ConvUMPS (``True``), or
+ cloned (``False``). In the former case, the tensor in both ConvUMPS's
+ will be the same, which might be useful if one needs more than one
+ copy of a ConvUMPS, but wants to compute all the gradients with respect
+ to the same, unique, tensor. Returns -------
- Node
+ ConvUMPS """
- mats_env=self._input_contraction(inline_input)
-
- ifinline_mats:
- result=self._contract_envs_inline(mats_env)
+ new_mps=ConvUMPS(in_channels=self._in_channels,
+ bond_dim=self._bond_dim[0],
+ kernel_size=self._kernel_size,
+ stride=self._stride,
+ padding=self._padding,
+ dilation=self.dilation,
+ tensor=None,
+ init_method=None,
+ device=None)
+ new_mps.name=self.name+'_copy'
+ ifshare_tensors:
+ new_mps.uniform_memory.tensor=self.uniform_memory.tensorelse:
- result=self._pairwise_contraction(mats_env)
-
- returnresult
[docs]classConvMPSLayer(AbstractConvClass,MPSLayer):# MARK: ConvMPSLayer"""
- Class for Matrix Product States, where all nodes are input nodes, and where
- the input data is a batch of images. It is the convolutional version of
- :class:`MPS`.
+ Convolutional version of :class:`MPSLayer`, where the input data is assumed to
+ be a batch of images. Input data as well as initialization parameters are described in `torch.nn.Conv2d <https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html>`_.
@@ -1268,13 +3544,15 @@
Source code for tensorkrowch.models.mps
Parameters ---------- in_channels : int
- Input channels. Same as ``in_dim`` in :class:`MPS`.
+ Input channels. Same as ``in_dim`` in :class:`MPSLayer`.
+ out_channels : int
+ Output channels. Same as ``out_dim`` in :class:`MPSLayer`. bond_dim : int, list[int] or tuple[int] Bond dimension(s). If given as a sequence, its length should be equal
- to :math:`kernel\_size_0 \cdot kernel\_size_1` (if ``boundary = "pbc"``)
- or :math:`kernel\_size_0 \cdot kernel\_size_1 - 1` (if
- ``boundary = "obc"``). The i-th bond dimension is always the dimension
- of the right edge of the i-th node.
+ to :math:`kernel\_size_0 \cdot kernel\_size_1 + 1`
+ (if ``boundary = "pbc"``) or :math:`kernel\_size_0 \cdot kernel\_size_1`
+ (if ``boundary = "obc"``). The i-th bond dimension is always the dimension
+ of the right edge of the i-th node (including output node). kernel_size : int, list[int] or tuple[int] Kernel size used in `torch.nn.Unfold <https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html#torch.nn.Unfold>`_.
@@ -1293,77 +3571,95 @@
Source code for tensorkrowch.models.mps
<https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html#torch.nn.Unfold>`_. If given as an ``int``, the actual kernel size will be ``(kernel_size, kernel_size)``.
+ out_position : int, optional
+ Position of the output node (label). Should be between 0 and
+ :math:`kernel\_size_0 \cdot kernel\_size_1`. If ``None``, the output node
+ will be located at the middle of the MPS. boundary : {"obc", "pbc"} String indicating whether periodic or open boundary conditions should be used.
+ tensors: list[torch.Tensor] or tuple[torch.Tensor], optional
+ To initialize MPS nodes, a list of MPS tensors can be provided. All
+ tensors should be rank-3 tensors, with shape ``(bond_dim, in_channels,
+ bond_dim)``. If the first and last elements are rank-2 tensors, with
+ shapes ``(in_channels, bond_dim)``, ``(bond_dim, in_channels)``,
+ respectively, the inferred boundary conditions will be "obc". Also, if
+ ``tensors`` contains a single element, it can be rank-1 ("obc") or
+ rank-3 ("pbc").
+ init_method : {"zeros", "ones", "copy", "rand", "randn", "randn_eye", "unit"}, optional
+ Initialization method. Check :meth:`~MPSLayer.initialize` for a more detailed
+ explanation of the different initialization methods.
+ device : torch.device, optional
+ Device where to initialize the tensors if ``init_method`` is provided.
+ kwargs : float
+ Keyword arguments for the different initialization methods. See
+ :meth:`~tensorkrowch.AbstractNode.make_tensor`. Examples --------
- >>> conv_mps = tk.models.ConvMPS(in_channels=2,
- ... bond_dim=5,
- ... kernel_size=2)
+ >>> conv_mps_layer = tk.models.ConvMPSLayer(in_channels=2,
+ ... out_channels=10,
+ ... bond_dim=5,
+ ... kernel_size=2) >>> data = torch.ones(20, 2, 2, 2) # batch_size x in_channels x height x width
- >>> result = conv_mps(data)
- >>> print(result.shape)
- torch.Size([20, 1, 1])
+ >>> result = conv_mps_layer(data)
+ >>> result.shape
+ torch.Size([20, 10, 1, 1]) """def__init__(self,in_channels:int,
+ out_channels:int,bond_dim:Union[int,Sequence[int]],kernel_size:Union[int,Sequence],stride:int=1,padding:int=0,dilation:int=1,
- boundary:Text='obc'):
-
- ifisinstance(kernel_size,int):
- kernel_size=(kernel_size,kernel_size)
- elifnotisinstance(kernel_size,Sequence):
- raiseTypeError('`kernel_size` must be int or Sequence')
-
- ifisinstance(stride,int):
- stride=(stride,stride)
- elifnotisinstance(stride,Sequence):
- raiseTypeError('`stride` must be int or Sequence')
-
- ifisinstance(padding,int):
- padding=(padding,padding)
- elifnotisinstance(padding,Sequence):
- raiseTypeError('`padding` must be int or Sequence')
-
- ifisinstance(dilation,int):
- dilation=(dilation,dilation)
- elifnotisinstance(dilation,Sequence):
- raiseTypeError('`dilation` must be int or Sequence')
-
- self._in_channels=in_channels
- self._kernel_size=kernel_size
- self._stride=stride
- self._padding=padding
- self._dilation=dilation
-
- super().__init__(n_features=kernel_size[0]*kernel_size[1],
- in_dim=in_channels,
- bond_dim=bond_dim,
- boundary=boundary,
- n_batches=2)
-
- self.unfold=nn.Unfold(kernel_size=kernel_size,
- stride=stride,
- padding=padding,
- dilation=dilation)
-
+ out_position:Optional[int]=None,
+ boundary:Text='obc',
+ tensors:Optional[Sequence[torch.Tensor]]=None,
+ init_method:Text='randn',
+ device:Optional[torch.device]=None,
+ **kwargs):
+
+ unfold=self._set_attributes(in_channels=in_channels,
+ kernel_size=kernel_size,
+ stride=stride,
+ padding=padding,
+ dilation=dilation)
+
+ MPSLayer.__init__(self,
+ n_features=self._kernel_size[0]* \
+ self._kernel_size[1]+1,
+ in_dim=in_channels,
+ out_dim=out_channels,
+ bond_dim=bond_dim,
+ out_position=out_position,
+ boundary=boundary,
+ tensors=tensors,
+ n_batches=2,
+ init_method=init_method,
+ device=device,
+ **kwargs)
+
+ self._out_channels=out_channels
+ self.unfold=unfold
+
@propertydefin_channels(self)->int:
-"""Returns ``in_channels``. Same as ``in_dim`` in :class:`MPS`."""
+"""Returns ``in_channels``. Same as ``in_dim`` in :class:`MPSLayer`."""returnself._in_channels
+ @property
+ defout_channels(self)->int:
+"""Returns ``out_channels``. Same as ``out_dim`` in :class:`MPSLayer`."""
+ returnself._out_channels
+
@propertydefkernel_size(self)->Tuple[int,int]:""" Returns ``kernel_size``. Number of nodes is given by
- :math:`kernel\_size_0 \cdot kernel\_size_1`.
+ :math:`kernel\_size_0 \cdot kernel\_size_1 + 1`. """returnself._kernel_size
@@ -1390,83 +3686,51 @@
[docs]defcopy(self,share_tensors:bool=False)->'ConvMPSLayer':
+"""
+ Creates a copy of the :class:`ConvMPSLayer`.
-
[docs]defforward(self,image,mode='flat',*args,**kwargs):
-r"""
- Overrides ``torch.nn.Module``'s forward to compute a convolution on the
- input image.
- Parameters ----------
- image : torch.Tensor
- Input batch of images with shape
-
- .. math::
-
- batch\_size \times in\_channels \times height \times width
- mode : {"flat", "snake"}
- Indicates the order in which MPS should take the pixels in the image.
- When ``"flat"``, the image is flattened putting one row of the image
- after the other. When ``"snake"``, its row is put in the opposite
- orientation as the previous row (like a snake running through the
- image).
- args :
- Arguments that might be used in :meth:`~MPS.contract`.
- kwargs :
- Keyword arguments that might be used in :meth:`~MPS.contract`,
- like ``inline_input`` or ``inline_mats``.
+ share_tensor : bool, optional
+ Boolean indicating whether tensors in the copied ConvMPSLayer should
+ be set as the tensors in the current ConvMPSLayer (``True``), or
+ cloned (``False``). In the former case, tensors in both ConvMPSLayer's
+ will be the same, which might be useful if one needs more than one
+ copy of an ConvMPSLayer, but wants to compute all the gradients with
+ respect to the same, unique, tensors.
+
+ Returns
+ -------
+ ConvMPSLayer """
- # Input image shape: batch_size x in_channels x height x width
-
- patches=self.unfold(image).transpose(1,2)
- # batch_size x nb_windows x (in_channels * nb_pixels)
-
- patches=patches.view(*patches.shape[:-1],self.in_channels,-1)
- # batch_size x nb_windows x in_channels x nb_pixels
-
- patches=patches.transpose(2,3)
- # batch_size x nb_windows x nb_pixels x in_channels
-
- ifmode=='snake':
- new_patches=patches[...,:self._kernel_size[1],:]
- foriinrange(1,self._kernel_size[0]):
- ifi%2==0:
- aux=patches[...,(i*self._kernel_size[1]):
- ((i+1)*self._kernel_size[1]),:]
- else:
- aux=patches[...,
- (i*self._kernel_size[1]):
- ((i+1)*self._kernel_size[1]),:].flip(dims=[0])
- new_patches=torch.cat([new_patches,aux],dim=2)
-
- patches=new_patches
-
- elifmode!='flat':
- raiseValueError('`mode` can only be "flat" or "snake"')
-
- result=super().forward(patches,*args,**kwargs)
- # batch_size x nb_windows
-
- h_in=image.shape[2]
- w_in=image.shape[3]
-
- h_out=int((h_in+2*self.padding[0]-self.dilation[0]*
- (self.kernel_size[0]-1)-1)/self.stride[0]+1)
- w_out=int((w_in+2*self.padding[1]-self.dilation[1]*
- (self.kernel_size[1]-1)-1)/self.stride[1]+1)
-
- result=result.view(*result.shape[:-1],h_out,w_out)
- # batch_size x height_out x width_out
-
- returnresult
[docs]classConvUMPSLayer(AbstractConvClass,UMPSLayer):# MARK: ConvUMPSLayer"""
- Class for Uniform Matrix Product States, where all nodes are input nodes,
- and where the input data is a batch of images. It is the convolutional
- version of :class:`UMPS`. This class cannot have different bond dimensions
- for each site and boundary conditions are always periodic.
+ Convolutional version of :class:`UMPSLayer`, where the input data is assumed to
+ be a batch of images. Input data as well as initialization parameters are described in `torch.nn.Conv2d <https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html>`_.
@@ -1474,7 +3738,9 @@
Source code for tensorkrowch.models.mps
Parameters ---------- in_channels : int
- Input channels. Same as ``in_dim`` in :class:`UMPS`.
+ Input channels. Same as ``in_dim`` in :class:`UMPSLayer`.
+ out_channels : int
+ Output channels. Same as ``out_dim`` in :class:`UMPSLayer`. bond_dim : int Bond dimension. kernel_size : int, list[int] or tuple[int]
@@ -1495,75 +3761,92 @@
Source code for tensorkrowch.models.mps
<https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html#torch.nn.Unfold>`_. If given as an ``int``, the actual kernel size will be ``(kernel_size, kernel_size)``.
-
+ out_position : int, optional
+ Position of the output node (label). Should be between 0 and
+ :math:`kernel\_size_0 \cdot kernel\_size_1`. If ``None``, the output node
+ will be located at the middle of the MPS.
+
+ tensors: list[torch.Tensor] or tuple[torch.Tensor], optional
+ To initialize MPS nodes, a sequence of 2 tensors can be provided, the
+ first one will be the uniform tensor, and the second one will be the
+ output node's tensor. Both tensors should be rank-3, with all their
+ first and last dimensions being equal.
+ init_method : {"zeros", "ones", "copy", "rand", "randn", "randn_eye", "unit"}, optional
+ Initialization method. Check :meth:`~UMPSLayer.initialize` for a more
+ detailed explanation of the different initialization methods.
+ device : torch.device, optional
+ Device where to initialize the tensors if ``init_method`` is provided.
+ kwargs : float
+ Keyword arguments for the different initialization methods. See
+ :meth:`~tensorkrowch.AbstractNode.make_tensor`.
+ Examples --------
- >>> conv_mps = tk.models.ConvMPS(in_channels=2,
- ... bond_dim=5,
- ... kernel_size=2)
- >>> for node in mps.mats_env:
- ... assert node.tensor_address() == 'virtual_uniform'
+ >>> conv_mps_layer = tk.models.ConvUMPSLayer(in_channels=2,
+ ... out_channels=10,
+ ... bond_dim=5,
+ ... kernel_size=2)
+ >>> for i, node in enumerate(conv_mps_layer.mats_env):
+ ... if i != conv_mps_layer.out_position:
+ ... assert node.tensor_address() == 'virtual_uniform' ... >>> data = torch.ones(20, 2, 2, 2) # batch_size x in_channels x height x width
- >>> result = conv_mps(data)
- >>> print(result.shape)
- torch.Size([20, 1, 1])
+ >>> result = conv_mps_layer(data)
+ >>> result.shape
+ torch.Size([20, 10, 1, 1]) """def__init__(self,in_channels:int,
- bond_dim:int,
+ out_channels:int,
+ bond_dim:Union[int,Sequence[int]],kernel_size:Union[int,Sequence],stride:int=1,padding:int=0,
- dilation:int=1):
-
- ifisinstance(kernel_size,int):
- kernel_size=(kernel_size,kernel_size)
- elifnotisinstance(kernel_size,Sequence):
- raiseTypeError('`kernel_size` must be int or Sequence')
-
- ifisinstance(stride,int):
- stride=(stride,stride)
- elifnotisinstance(stride,Sequence):
- raiseTypeError('`stride` must be int or Sequence')
-
- ifisinstance(padding,int):
- padding=(padding,padding)
- elifnotisinstance(padding,Sequence):
- raiseTypeError('`padding` must be int or Sequence')
-
- ifisinstance(dilation,int):
- dilation=(dilation,dilation)
- elifnotisinstance(dilation,Sequence):
- raiseTypeError('`dilation` must be int or Sequence')
-
- self._in_channels=in_channels
- self._kernel_size=kernel_size
- self._stride=stride
- self._padding=padding
- self._dilation=dilation
-
- super().__init__(n_features=kernel_size[0]*kernel_size[1],
- in_dim=in_channels,
- bond_dim=bond_dim,
- n_batches=2)
-
- self.unfold=nn.Unfold(kernel_size=kernel_size,
- stride=stride,
- padding=padding,
- dilation=dilation)
-
+ dilation:int=1,
+ out_position:Optional[int]=None,
+ tensors:Optional[Sequence[torch.Tensor]]=None,
+ init_method:Text='randn',
+ device:Optional[torch.device]=None,
+ **kwargs):
+
+ unfold=self._set_attributes(in_channels=in_channels,
+ kernel_size=kernel_size,
+ stride=stride,
+ padding=padding,
+ dilation=dilation)
+
+ UMPSLayer.__init__(self,
+ n_features=self._kernel_size[0]* \
+ self._kernel_size[1]+1,
+ in_dim=in_channels,
+ out_dim=out_channels,
+ bond_dim=bond_dim,
+ out_position=out_position,
+ tensors=tensors,
+ n_batches=2,
+ init_method=init_method,
+ device=device,
+ **kwargs)
+
+ self._out_channels=out_channels
+ self.unfold=unfold
+
@propertydefin_channels(self)->int:
-"""Returns ``in_channels``. Same as ``in_dim`` in :class:`UMPS`."""
+"""Returns ``in_channels``. Same as ``in_dim`` in :class:`UMPSLayer`."""returnself._in_channels
+ @property
+ defout_channels(self)->int:
+"""Returns ``out_channels``. Same as ``out_dim`` in :class:`UMPSLayer`."""
+ returnself._out_channels
+
@propertydefkernel_size(self)->Tuple[int,int]:""" Returns ``kernel_size``. Number of nodes is given by
- :math:`kernel\_size_0 \cdot kernel\_size_1`.
+ :math:`kernel\_size_0 \cdot kernel\_size_1 + 1`. """returnself._kernel_size
@@ -1590,75 +3873,49 @@
Source code for tensorkrowch.models.mps
<https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html#torch.nn.Unfold>`_. """returnself._dilation
+
+ @property
+ defout_channels(self)->int:
+"""Returns ``out_channels``. Same as ``phys_dim`` in :class:`MPS`."""
+ returnself._out_channels
+
+
[docs]defcopy(self,share_tensors:bool=False)->'ConvUMPSLayer':
+"""
+ Creates a copy of the :class:`ConvUMPSLayer`.
-
[docs]defforward(self,image,mode='flat',*args,**kwargs):
-r"""
- Overrides ``torch.nn.Module``'s forward to compute a convolution on the
- input image.
- Parameters ----------
- image : torch.Tensor
- Input batch of images with shape
-
- .. math::
-
- batch\_size \times in\_channels \times height \times width
- mode : {"flat", "snake"}
- Indicates the order in which MPS should take the pixels in the image.
- When ``"flat"``, the image is flattened putting one row of the image
- after the other. When ``"snake"``, its row is put in the opposite
- orientation as the previous row (like a snake running through the
- image).
- args :
- Arguments that might be used in :meth:`~UMPS.contract`.
- kwargs :
- Keyword arguments that might be used in :meth:`~UMPS.contract`,
- like ``inline_input`` or ``inline_mats``.
+ share_tensor : bool, optional
+ Boolean indicating whether tensors in the copied ConvUMPSLayer should
+ be set as the tensors in the current ConvUMPSLayer (``True``), or
+ cloned (``False``). In the former case, tensors in both ConvUMPSLayer's
+ will be the same, which might be useful if one needs more than one
+ copy of an ConvUMPSLayer, but wants to compute all the gradients with
+ respect to the same, unique, tensors.
+
+ Returns
+ -------
+ ConvUMPSLayer """
- # Input image shape: batch_size x in_channels x height x width
-
- patches=self.unfold(image).transpose(1,2)
- # batch_size x nb_windows x (in_channels * nb_pixels)
-
- patches=patches.view(*patches.shape[:-1],self.in_channels,-1)
- # batch_size x nb_windows x in_channels x nb_pixels
-
- patches=patches.transpose(2,3)
- # batch_size x nb_windows x nb_pixels x in_channels
-
- ifmode=='snake':
- new_patches=patches[...,:self._kernel_size[1],:]
- foriinrange(1,self._kernel_size[0]):
- ifi%2==0:
- aux=patches[...,(i*self._kernel_size[1]):
- ((i+1)*self._kernel_size[1]),:]
- else:
- aux=patches[...,
- (i*self._kernel_size[1]):
- ((i+1)*self._kernel_size[1]),:].flip(dims=[0])
- new_patches=torch.cat([new_patches,aux],dim=2)
-
- patches=new_patches
-
- elifmode!='flat':
- raiseValueError('`mode` can only be "flat" or "snake"')
-
- result=super().forward(patches,*args,**kwargs)
- # batch_size x nb_windows
-
- h_in=image.shape[2]
- w_in=image.shape[3]
-
- h_out=int((h_in+2*self.padding[0]-self.dilation[0]*
- (self.kernel_size[0]-1)-1)/self.stride[0]+1)
- w_out=int((w_in+2*self.padding[1]-self.dilation[1]*
- (self.kernel_size[1]-1)-1)/self.stride[1]+1)
-
- result=result.view(*result.shape[:-1],h_out,w_out)
- # batch_size x height_out x width_out
-
- returnresult
[docs]classMPSData(TensorNetwork):# MARK: MPSData
+"""
+ Class for data vectors in the form of Matrix Product States. That is, this
+ is a class similar to :class:`MPS`, but where all nodes can have additional
+ batch edges.
+
+ Besides, since this class is intended to store data vectors, all
+ :class:`nodes <tensorkrowch.Node>` are non-parametric, and are ``data``
+ nodes. Also, this class does not have an inherited ``contract`` method,
+ since it is not intended to be contracted with input data, but rather
+ act itself as input data of another tensor network model.
+
+ Similar to :class:`MPS`, ``MPSData`` is formed by:
+
+ * ``mats_env``: Environment of `matrix` nodes with axes
+ ``("batch_0", ..., "batch_n", "left", "feature", "right")``.
+
+ * ``left_node``, ``right_node``: `Vector` nodes with axes ``("right",)``
+ and ``("left",)``, respectively. These are used to close the boundary
+ in the case ``boudary`` is ``"obc"``. Otherwise, both are ``None``.
+
+ Since ``MPSData`` is designed to store input data vectors, this can be
+ accomplished by calling the custom :meth:`add_data` method with a given
+ list of tensors. This is in contrast to the usual way of setting nodes'
+ tensors in :class:`MPS` and its derived classes via :meth:`MPS.initialize`.
+
+ Parameters
+ ----------
+ n_features : int, optional
+ Number of nodes that will be in ``mats_env``. That is, number of nodes
+ without taking into account ``left_node`` and ``right_node``.
+ phys_dim : int, list[int] or tuple[int], optional
+ Physical dimension(s). If given as a sequence, its length should be
+ equal to ``n_features``.
+ bond_dim : int, list[int] or tuple[int], optional
+ Bond dimension(s). If given as a sequence, its length should be equal
+ to ``n_features`` (if ``boundary = "pbc"``) or ``n_features - 1`` (if
+ ``boundary = "obc"``). The i-th bond dimension is always the dimension
+ of the right edge of the i-th node.
+ boundary : {"obc", "pbc"}
+ String indicating whether periodic or open boundary conditions should
+ be used.
+ n_batches : int
+ Number of batch edges of the MPS nodes. Usually ``n_batches = 1``
+ (where the batch edge is used for the data batched) but it could also
+ be ``n_batches = 2`` (one edge for data batched, other edge for image
+ patches in convolutional layers).
+ tensors: list[torch.Tensor] or tuple[torch.Tensor], optional
+ Instead of providing ``n_features``, ``phys_dim``, ``bond_dim`` and
+ ``boundary``, a list of MPS tensors can be provided. In such case, all
+ mentioned attributes will be inferred from the given tensors. All
+ tensors should be rank-(n+3) tensors, with shape
+ ``(batch_1, ..., batch_n, bond_dim, phys_dim, bond_dim)``. If the first
+ and last elements are rank-(n+2) tensors, with shapes
+ ``(batch_1, ..., batch_n, phys_dim, bond_dim)``,
+ ``(batch_1, ..., batch_n, bond_dim, phys_dim)``, respectively,
+ the inferred boundary conditions will be "obc". Also, if ``tensors``
+ contains a single element, it can be rank-(n+1) ("obc") or rank-(n+3)
+ ("pbc").
+ init_method : {"zeros", "ones", "copy", "rand", "randn"}, optional
+ Initialization method. Check :meth:`initialize` for a more detailed
+ explanation of the different initialization methods. By default it is
+ ``None``, since ``MPSData`` is intended to store input data vectors in
+ MPS form, rather than initializing its own random tensors. Check
+ :meth:`add_data` to see how to initialize MPS nodes with data tensors.
+ device : torch.device, optional
+ Device where to initialize the tensors if ``init_method`` is provided.
+ kwargs : float
+ Keyword arguments for the different initialization methods. See
+ :meth:`~tensorkrowch.AbstractNode.make_tensor`.
+
+ Examples
+ --------
+ >>> mps = tk.models.MPSData(n_features=5,
+ ... phys_dim=2,
+ ... bond_dim=5,
+ ... boundary="pbc")
+ >>> # n_features * (batch_size x bond_dim x feature_size x bond_dim)
+ >>> data = [torch.ones(20, 5, 2, 5) for _ in range(5)]
+ >>> mps.add_data(data)
+ >>> for node in mps.mats_env:
+ ... assert node.shape == (20, 5, 2, 5)
+ """
+
+ def__init__(self,
+ n_features:Optional[int]=None,
+ phys_dim:Optional[Union[int,Sequence[int]]]=None,
+ bond_dim:Optional[Union[int,Sequence[int]]]=None,
+ boundary:Text='obc',
+ n_batches:int=1,
+ tensors:Optional[Sequence[torch.Tensor]]=None,
+ init_method:Optional[Text]=None,
+ device:Optional[torch.device]=None,
+ **kwargs)->None:
+
+ super().__init__(name='mps_data')
+
+ # n_batches
+ ifnotisinstance(n_batches,int):
+ raiseTypeError('`n_batches` should be int type')
+ self._n_batches=n_batches
+
+ iftensorsisNone:
+ # boundary
+ ifboundarynotin['obc','pbc']:
+ raiseValueError('`boundary` should be one of "obc" or "pbc"')
+ self._boundary=boundary
+
+ # n_features
+ ifnotisinstance(n_features,int):
+ raiseTypeError('`n_features` should be int type')
+ elifn_features<1:
+ raiseValueError('`n_features` should be at least 1')
+ self._n_features=n_features
+
+ # phys_dim
+ ifisinstance(phys_dim,(list,tuple)):
+ iflen(phys_dim)!=n_features:
+ raiseValueError('If `phys_dim` is given as a sequence of int, '
+ 'its length should be equal to `n_features`')
+ self._phys_dim=list(phys_dim)
+ elifisinstance(phys_dim,int):
+ self._phys_dim=[phys_dim]*n_features
+ else:
+ raiseTypeError('`phys_dim` should be int, tuple[int] or list[int] '
+ 'type')
+
+ # bond_dim
+ ifisinstance(bond_dim,(list,tuple)):
+ ifboundary=='obc':
+ iflen(bond_dim)!=n_features-1:
+ raiseValueError(
+ 'If `bond_dim` is given as a sequence of int, and '
+ '`boundary` is "obc", its length should be equal '
+ 'to `n_features` - 1')
+ elifboundary=='pbc':
+ iflen(bond_dim)!=n_features:
+ raiseValueError(
+ 'If `bond_dim` is given as a sequence of int, and '
+ '`boundary` is "pbc", its length should be equal '
+ 'to `n_features`')
+ self._bond_dim=list(bond_dim)
+ elifisinstance(bond_dim,int):
+ ifboundary=='obc':
+ self._bond_dim=[bond_dim]*(n_features-1)
+ elifboundary=='pbc':
+ self._bond_dim=[bond_dim]*n_features
+ else:
+ raiseTypeError('`bond_dim` should be int, tuple[int] or list[int]'
+ ' type')
+
+ else:
+ ifnotisinstance(tensors,Sequence):
+ raiseTypeError('`tensors` should be a tuple[torch.Tensor] or '
+ 'list[torch.Tensor] type')
+ else:
+ self._n_features=len(tensors)
+ self._phys_dim=[]
+ self._bond_dim=[]
+ fori,tinenumerate(tensors):
+ ifnotisinstance(t,torch.Tensor):
+ raiseTypeError('`tensors` should be a tuple[torch.Tensor]'
+ ' or list[torch.Tensor] type')
+
+ ifi==0:
+ iflen(t.shape)notin[n_batches+1,
+ n_batches+2,
+ n_batches+3]:
+ raiseValueError(
+ 'The first and last elements in `tensors` '
+ 'should be both rank-(n+2) or rank-(n+3) tensors.'
+ ' If the first element is also the last one,'
+ ' it should be a rank-(n+1) tensor')
+ iflen(t.shape)==n_batches+1:
+ self._boundary='obc'
+ self._phys_dim.append(t.shape[-1])
+ eliflen(t.shape)==n_batches+2:
+ self._boundary='obc'
+ self._phys_dim.append(t.shape[-2])
+ self._bond_dim.append(t.shape[-1])
+ else:
+ self._boundary='pbc'
+ self._phys_dim.append(t.shape[-2])
+ self._bond_dim.append(t.shape[-1])
+ elifi==(self._n_features-1):
+ iflen(t.shape)!=len(tensors[0].shape):
+ raiseValueError(
+ 'The first and last elements in `tensors` '
+ 'should have the same rank. Both should be '
+ 'rank-(n+2) or rank-(n+3) tensors. If the first'
+ ' element is also the last one, it should '
+ 'be a rank-(n+1) tensor')
+ iflen(t.shape)==n_batches+2:
+ self._phys_dim.append(t.shape[-1])
+ else:
+ ift.shape[-1]!=tensors[0].shape[-3]:
+ raiseValueError(
+ 'If the first and last elements in `tensors`'
+ ' are rank-(n+3) tensors, the first dimension'
+ ' of the first element should coincide with'
+ ' the last dimension of the last element '
+ '(ignoring batch dimensions)')
+ self._phys_dim.append(t.shape[-2])
+ self._bond_dim.append(t.shape[-1])
+ else:
+ iflen(t.shape)!=n_batches+3:
+ raiseValueError(
+ 'The elements of `tensors` should be rank-(n+3) '
+ 'tensors, except the first and lest elements'
+ ' if boundary is "obc"')
+ self._phys_dim.append(t.shape[-2])
+ self._bond_dim.append(t.shape[-1])
+
+ # Properties
+ self._left_node=None
+ self._right_node=None
+ self._mats_env=[]
+
+ # Create Tensor Network
+ self._make_nodes()
+ self.initialize(init_method=init_method,
+ device=device,
+ **kwargs)
+
+ iftensorsisnotNone:
+ self.add_data(data=tensors)
+
+ # ----------
+ # Properties
+ # ----------
+ @property
+ defn_features(self)->int:
+"""Returns number of nodes."""
+ returnself._n_features
+
+ @property
+ defphys_dim(self)->List[int]:
+"""Returns physical dimensions."""
+ returnself._phys_dim
+
+ @property
+ defbond_dim(self)->List[int]:
+"""Returns bond dimensions."""
+ returnself._bond_dim
+
+ @property
+ defboundary(self)->Text:
+"""Returns boundary condition ("obc" or "pbc")."""
+ returnself._boundary
+
+ @property
+ defn_batches(self)->int:
+"""Returns number of batch edges of the MPS nodes."""
+ returnself._n_batches
+
+ @property
+ defleft_node(self)->Optional[AbstractNode]:
+"""Returns the ``left_node``."""
+ returnself._left_node
+
+ @property
+ defright_node(self)->Optional[AbstractNode]:
+"""Returns the ``right_node``."""
+ returnself._right_node
+
+ @property
+ defmats_env(self)->List[AbstractNode]:
+"""Returns the list of nodes in ``mats_env``."""
+ returnself._mats_env
+
+ # -------
+ # Methods
+ # -------
+ def_make_nodes(self)->None:
+"""Creates all the nodes of the MPS."""
+ ifself._leaf_nodes:
+ raiseValueError('Cannot create MPS nodes if the MPS already has '
+ 'nodes')
+
+ aux_bond_dim=self._bond_dim
+
+ ifself._boundary=='obc':
+ ifnotaux_bond_dim:
+ aux_bond_dim=[1]
+
+ self._left_node=ParamNode(shape=(aux_bond_dim[0],),
+ axes_names=('right',),
+ name='left_node',
+ network=self)
+ self._right_node=ParamNode(shape=(aux_bond_dim[-1],),
+ axes_names=('left',),
+ name='right_node',
+ network=self)
+
+ aux_bond_dim=aux_bond_dim+[aux_bond_dim[-1]]+[aux_bond_dim[0]]
+
+ foriinrange(self._n_features):
+ node=Node(shape=(*([1]*self._n_batches),
+ aux_bond_dim[i-1],
+ self._phys_dim[i],
+ aux_bond_dim[i]),
+ axes_names=(*(['batch']*self._n_batches),
+ 'left',
+ 'feature',
+ 'right'),
+ name=f'mats_env_data_node_({i})',
+ data=True,
+ network=self)
+ self._mats_env.append(node)
+
+ ifi!=0:
+ self._mats_env[-2]['right']^self._mats_env[-1]['left']
+
+ ifself._boundary=='pbc':
+ ifi==0:
+ periodic_edge=self._mats_env[-1]['left']
+ ifi==self._n_features-1:
+ self._mats_env[-1]['right']^periodic_edge
+ else:
+ ifi==0:
+ self._left_node['right']^self._mats_env[-1]['left']
+ ifi==self._n_features-1:
+ self._mats_env[-1]['right']^self._right_node['left']
+
+
[docs]definitialize(self,
+ init_method:Optional[Text]='randn',
+ device:Optional[torch.device]=None,
+ **kwargs:float)->None:
+"""
+ Initializes all the nodes of the :class:`MPSData`. It can be called when
+ instantiating the model, or to override the existing nodes' tensors.
+
+ There are different methods to initialize the nodes:
+
+ * ``{"zeros", "ones", "copy", "rand", "randn"}``: Each node is
+ initialized calling :meth:`~tensorkrowch.AbstractNode.set_tensor` with
+ the given method, ``device`` and ``kwargs``.
+
+ Parameters
+ ----------
+ init_method : {"zeros", "ones", "copy", "rand", "randn"}, optional
+ Initialization method. Check :meth:`add_data` to see how to
+ initialize MPS nodes with data tensors.
+ device : torch.device, optional
+ Device where to initialize the tensors if ``init_method`` is provided.
+ kwargs : float
+ Keyword arguments for the different initialization methods. See
+ :meth:`~tensorkrowch.AbstractNode.make_tensor`.
+ """
+ ifself._boundary=='obc':
+ self._left_node.set_tensor(init_method='copy',device=device)
+ self._right_node.set_tensor(init_method='copy',device=device)
+
+ ifinit_methodisnotNone:
+
+ fori,nodeinenumerate(self._mats_env):
+ node.set_tensor(init_method=init_method,
+ device=device,
+ **kwargs)
+
+ ifself._boundary=='obc':
+ aux_tensor=torch.zeros(*node.shape,device=device)
+ ifi==0:
+ # Left node
+ aux_tensor[...,0,:,:]=node.tensor[...,0,:,:]
+ node.tensor=aux_tensor
+ elifi==(self._n_features-1):
+ # Right node
+ aux_tensor[...,0]=node.tensor[...,0]
+ node.tensor=aux_tensor
+
+
[docs]defadd_data(self,data:Sequence[torch.Tensor])->None:
+"""
+ Adds data to MPS data nodes. Input is a list of mps tensors.
+
+ The physical dimensions of the given data tensors should coincide with
+ the physical dimensions of the MPS. The bond dimensions can be different.
+
+ Parameters
+ ----------
+ data : list[torch.Tensor] or tuple[torch.Tensor]
+ A sequence of tensors, one for each of the MPS nodes. If ``boundary``
+ is ``"pbc"``, all tensors should have the same rank, with shapes
+ ``(batch_0, ..., batch_n, bond_dim, phys_dim, bond_dim)``. If
+ ``boundary`` is ``"obc"``, the first and last tensors should have
+ shapes ``(batch_0, ..., batch_n, phys_dim, bond_dim)`` and
+ ``(batch_0, ..., batch_n, bond_dim, phys_dim)``, respectively.
+ """
+ ifnotisinstance(data,Sequence):
+ raiseTypeError(
+ '`data` should be list[torch.Tensor] or tuple[torch.Tensor] type')
+ iflen(data)!=self._n_features:
+ raiseValueError('`data` should be a sequence of tensors of length'
+ ' equal to `n_features`')
+ ifany([notisinstance(x,torch.Tensor)forxindata]):
+ raiseTypeError(
+ '`data` should be list[torch.Tensor] or tuple[torch.Tensor] type')
+
+ # Check physical dimensions coincide
+ fori,(data_tensor,node)inenumerate(zip(data,self._mats_env)):
+ if(self._boundary=='obc')and(i==(self._n_features-1)):
+ ifdata_tensor.shape[-1]!=node.shape[-2]:
+ raiseValueError(
+ f'Physical dimension {data_tensor.shape[-1]} of '
+ f'data tensor at position {i} does not coincide '
+ f'with the corresponding physical dimension '
+ f'{node.shape[-2]} of the MPS')
+ else:
+ ifdata_tensor.shape[-2]!=node.shape[-2]:
+ raiseValueError(
+ f'Physical dimension {data_tensor.shape[-2]} of '
+ f'data tensor at position {i} does not coincide '
+ f'with the corresponding physical dimension '
+ f'{node.shape[-2]} of the MPS')
+
+ data=data[:]
+ fori,nodeinenumerate(self._mats_env):
+ ifself._boundary=='obc':
+ aux_tensor=torch.zeros(*node.shape,
+ device=data[i].device)
+ if(i==0)and(i==(self._n_features-1)):
+ aux_tensor=torch.zeros(*data[i].shape[:-1],
+ *node.shape[-3:],
+ device=data[i].device)
+ aux_tensor[...,0,:,0]=data[i]
+ data[i]=aux_tensor
+ elifi==0:
+ aux_tensor=torch.zeros(*data[i].shape[:-2],
+ *node.shape[-3:-1],
+ data[i].shape[-1],
+ device=data[i].device)
+ aux_tensor[...,0,:,:]=data[i]
+ data[i]=aux_tensor
+ elifi==(self._n_features-1):
+ aux_tensor=torch.zeros(*data[i].shape[:-1],
+ *node.shape[-2:],
+ device=data[i].device)
+ aux_tensor[...,0]=data[i]
+ data[i]=aux_tensor
+
+ node._direct_set_tensor(data[i])
[docs]classMPSLayer(TensorNetwork):
-"""
- Class for Matrix Product States with an extra node that is dedicated to the
- output. That is, this MPS has :math:`n` nodes, being :math:`n-1` input nodes
- connected to ``data`` nodes (nodes that will contain the data tensors), and
- one output node, whose physical dimension (``out_dim``) is used as the label
- (for classification tasks).
-
- Besides, since this class has an output edge, when contracting the whole
- tensor network (with input data), the result will be a vector that can be
- plugged into the next layer (being this other tensor network or a neural
- network layer).
-
- If the physical dimensions of all the input nodes (``in_dim``) are equal,
- the input data tensor can be passed as a single tensor. Otherwise, it would
- have to be passed as a list of tensors with different sizes.
-
- An ``MPSLayer`` is formed by the following nodes:
-
- * ``left_node``, ``right_node``: `Vector` nodes with axes ``("input", "right")``
- and ``("left", "input")``, respectively. These are the nodes at the
- extremes of the ``MPSLayer``. If ``boundary`` is ``"pbc""``, both are
- ``None``.
-
- * ``left_env``, ``right_env``: Environments of `matrix` nodes that are at
- the left or right side of the ``output_node``. These nodes have axes
- ``("left", "input", "right")``.
-
- * ``output_node``: Node dedicated to the output. It has axes
- ``("left", "output", "right")``.
-
- Parameters
- ----------
- n_features : int
- Number of input nodes. The total number of nodes (including the output
- node) will be ``n_features + 1``.
- in_dim : int, list[int] or tuple[int]
- Input dimension. Equivalent to the physical dimension. If given as a
- sequence, its length should be equal to ``n_features``, since these are
- the input dimensions of the input nodes.
- out_dim : int
- Output dimension (labels) for the output node. Plays the same role as
- ``in_dim`` for input nodes.
- bond_dim : int, list[int] or tuple[int]
- Bond dimension(s). If given as a sequence, its length should be equal
- to ``n_features + 1`` (if ``boundary = "pbc"``) or ``n_features`` (if
- ``boundary = "obc"``). The i-th bond dimension is always the dimension
- of the right edge of the i-th node (including output node).
- out_position : int, optional
- Position of the output node (label). Should be between 0 and
- ``n_features``. If ``None``, the output node will be located at the
- middle of the MPS.
- boundary : {"obc", "pbc"}
- String indicating whether periodic or open boundary conditions should
- be used.
- n_batches : int
- Number of batch edges of input ``data`` nodes. Usually ``n_batches = 1``
- (where the batch edge is used for the data batched) but it could also
- be ``n_batches = 2`` (e.g. one edge for data batched, other edge for
- image patches in convolutional layers).
-
- Examples
- --------
- ``MPSLayer`` with same input dimensions:
-
- >>> mps_layer = tk.models.MPSLayer(n_features=4,
- ... in_dim=2,
- ... out_dim=10,
- ... bond_dim=5)
- >>> data = torch.ones(20, 4, 2) # batch_size x n_features x feature_size
- >>> result = mps_layer(data)
- >>> result.shape
- torch.Size([20, 10])
-
- ``MPSLayer`` with different input dimensions:
-
- >>> mps_layer = tk.models.MPSLayer(n_features=4,
- ... in_dim=list(range(2, 6)),
- ... out_dim=10,
- ... bond_dim=5)
- >>> data = [torch.ones(20, i)
- ... for i in range(2, 6)] # n_features * [batch_size x feature_size]
- >>> result = mps_layer(data)
- >>> result.shape
- torch.Size([20, 10])
- """
-
- def__init__(self,
- n_features:int,
- in_dim:Union[int,Sequence[int]],
- out_dim:int,
- bond_dim:Union[int,Sequence[int]],
- out_position:Optional[int]=None,
- boundary:Text='obc',
- n_batches:int=1)->None:
-
- super().__init__(name='mps')
-
- # boundary
- ifboundarynotin['obc','pbc']:
- raiseValueError('`boundary` should be one of "obc" or "pbc"')
- self._boundary=boundary
-
- # out_position
- ifout_positionisNone:
- out_position=(n_features+1)//2
- elif(out_position<0)or(out_position>n_features):
- raiseValueError('`out_position` should be between 0 and '
- f'{n_features}')
- self._out_position=out_position
-
- # n_features
- ifn_features<0:
- raiseValueError('`n_features` cannot be lower than 0')
- elif(boundary=='obc')and(n_features<1):
- raiseValueError('If `boundary` is "obc", at least '
- 'there has to be 1 input node')
- self._n_features=n_features
-
- # in_dim
- ifisinstance(in_dim,(list,tuple)):
- iflen(in_dim)!=n_features:
- raiseValueError('If `in_dim` is given as a sequence of int, '
- 'its length should be equal to `n_features`')
- else:
- fordiminin_dim:
- ifnotisinstance(dim,int):
- raiseTypeError('`in_dim` should be int, tuple[int] or '
- 'list[int] type')
- self._in_dim=list(in_dim)
- elifisinstance(in_dim,int):
- self._in_dim=[in_dim]*n_features
- else:
- raiseTypeError('`in_dim` should be int, tuple[int] or list[int] '
- 'type')
-
- # out_dim
- ifnotisinstance(out_dim,int):
- raiseTypeError('`out_dim` should be int type')
- self._out_dim=out_dim
-
- # phys_dim
- ifisinstance(in_dim,(list,tuple)):
- self._phys_dim=list(in_dim[:out_position])+[out_dim]+ \
- list(in_dim[out_position:])
- elifisinstance(in_dim,int):
- self._phys_dim=[in_dim]*out_position+[out_dim]+ \
- [in_dim]*(n_features-out_position)
-
- # bond_dim
- ifisinstance(bond_dim,(list,tuple)):
- ifboundary=='obc':
- iflen(bond_dim)!=n_features:
- raiseValueError('If `bond_dim` is given as a sequence of '
- 'int, and `boundary` is "obc", its length '
- 'should be equal to `n_features`')
- elifboundary=='pbc':
- iflen(bond_dim)!=(n_features+1):
- raiseValueError('If `bond_dim` is given as a sequence of '
- 'int, and `boundary` is "pbc", its length '
- 'should be equal to `n_features + 1`')
- self._bond_dim=list(bond_dim)
- elifisinstance(bond_dim,int):
- self._bond_dim=[bond_dim]*(n_features+(boundary=='pbc'))
- else:
- raiseTypeError('`bond_dim` should be int, tuple[int] or list[int]'
- ' type')
-
- # n_batches
- ifnotisinstance(n_batches,int):
- raiseTypeError('`n_batches should be int type')
- self._n_batches=n_batches
-
- # Create Tensor Network
- self._make_nodes()
- self.initialize()
-
- @property
- defn_features(self)->int:
-"""
- Returns number of input nodes. The total number of nodes (including the
- output node) will be ``n_features + 1``.
- """
- returnself._n_features
-
- @property
- defin_dim(self)->List[int]:
-"""Returns input dimensions."""
- returnself._in_dim
-
- @property
- defout_dim(self)->int:
-"""
- Returns the output dimension, that is, the number of labels in the
- output node. Same as ``in_dim`` for input nodes.
- """
- returnself._out_dim
-
- @property
- defphys_dim(self)->List[int]:
-"""Returns ``in_dim`` list with ``out_dim`` in the ``out_position``."""
- returnself._phys_dim
-
- @property
- defbond_dim(self)->List[int]:
-"""Returns bond dimensions."""
- returnself._bond_dim
-
- @property
- defout_position(self)->int:
-"""Returns position of the output node (label)."""
- returnself._out_position
-
- @property
- defboundary(self)->Text:
-"""Returns boundary condition ("obc" or "pbc")."""
- returnself._boundary
-
- @property
- defn_batches(self)->int:
-"""Returns number of batch edges of the ``data`` nodes."""
- returnself._n_batches
-
- def_make_nodes(self)->None:
-"""Creates all the nodes of the MPS."""
- ifself.leaf_nodes:
- raiseValueError('Cannot create MPS nodes if the MPS already has '
- 'nodes')
-
- self.left_node=None
- self.right_node=None
- self.left_env=[]
- self.right_env=[]
-
- # Open Boundary Conditions
- ifself.boundary=='obc':
- # Left node
- ifself.out_position>0:
- self.left_node=ParamNode(shape=(self.in_dim[0],
- self.bond_dim[0]),
- axes_names=('input','right'),
- name='left_node',
- network=self)
-
- # Left environment
- ifself.out_position>1:
- foriinrange(1,self.out_position):
- node=ParamNode(shape=(self.bond_dim[i-1],
- self.in_dim[i],
- self.bond_dim[i]),
- axes_names=('left','input','right'),
- name=f'left_env_node_({i-1})',
- network=self)
- self.left_env.append(node)
- ifi==1:
- self.left_node['right']^self.left_env[-1]['left']
- else:
- self.left_env[-2]['right']^self.left_env[-1]['left']
-
- # Output node
- ifself.out_position==0:
- self.output_node=ParamNode(shape=(self.out_dim,
- self.bond_dim[0]),
- axes_names=('output','right'),
- name='output_node',
- network=self)
-
- if(self.out_position>0)and(self.out_position<self.n_features):
- self.output_node=ParamNode(
- shape=(self.bond_dim[self.out_position-1],
- self.out_dim,
- self.bond_dim[self.out_position]),
- axes_names=('left','output','right'),
- name='output_node',
- network=self)
- ifself.left_env:
- self.left_env[-1]['right']^self.output_node['left']
- else:
- self.left_node['right']^self.output_node['left']
-
- ifself.out_position==self.n_features:
- self.output_node=ParamNode(shape=(self.bond_dim[-1],
- self.out_dim),
- axes_names=('left',
- 'output'),
- name='output_node',
- network=self)
- ifself.left_env:
- self.left_env[-1]['right']^self.output_node['left']
- elifself.left_node:
- self.left_node['right']^self.output_node['left']
-
- # Right environment
- ifself.out_position<self.n_features-1:
- foriinrange(self.out_position+1,self.n_features):
- node=ParamNode(shape=(self.bond_dim[i-1],
- self.in_dim[i-1],
- self.bond_dim[i]),
- axes_names=('left','input','right'),
- name=f'right_env_node_({i-self.out_position-1})',
- network=self)
- self.right_env.append(node)
- ifi==self.out_position+1:
- self.output_node['right']^self.right_env[-1]['left']
- else:
- self.right_env[-2]['right']^self.right_env[-1]['left']
-
- # Right node
- ifself.out_position<self.n_features:
- self.right_node=ParamNode(shape=(self.bond_dim[-1],
- self.in_dim[-1]),
- axes_names=('left','input'),
- name='right_node',
- network=self)
- ifself.right_env:
- self.right_env[-1]['right']^self.right_node['left']
- else:
- self.output_node['right']^self.right_node['left']
-
- # Periodic Boundary Conditions
- else:
- # Left environment
- ifself.out_position>0:
- foriinrange(self.out_position):
- node=ParamNode(shape=(self.bond_dim[i-1],
- self.in_dim[i],
- self.bond_dim[i]),
- axes_names=('left','input','right'),
- name=f'left_env_node_({i})',
- network=self)
- self.left_env.append(node)
- ifi==0:
- periodic_edge=self.left_env[-1]['left']
- else:
- self.left_env[-2]['right']^self.left_env[-1]['left']
-
- # Output node
- self.output_node=ParamNode(
- shape=(self.bond_dim[self.out_position-1],
- self.out_dim,
- self.bond_dim[self.out_position]),
- axes_names=('left','output','right'),
- name='output_node',
- network=self)
- ifself.left_env:
- self.left_env[-1]['right']^self.output_node['left']
- else:
- periodic_edge=self.output_node['left']
- ifself.out_position==self.n_features:
- self.output_node['right']^periodic_edge
-
- # Right environment
- ifself.out_position<self.n_features:
- foriinrange(self.out_position+1,self.n_features+1):
- node=ParamNode(shape=(self.bond_dim[i-1],
- self.in_dim[i-1],
- self.bond_dim[i]),
- axes_names=('left','input','right'),
- name=f'right_env_node_({i-self.out_position-1})',
- network=self)
- self.right_env.append(node)
- ifi==self.out_position+1:
- self.output_node['right']^self.right_env[-1]['left']
- else:
- self.right_env[-2]['right']^self.right_env[-1]['left']
- ifi==self.n_features:
- self.right_env[-1]['right']^periodic_edge
-
-
[docs]definitialize(self,std:float=1e-9)->None:
-"""
- Initializes all the nodes as explained `here <https://arxiv.org/abs/1605.03795>`_.
- It can be overriden for custom initializations.
- """
- # Left node
- ifself.left_nodeisnotNone:
- tensor=torch.randn(self.left_node.shape)*std
- aux=torch.zeros(tensor.shape[1])*std
- aux[0]=1.
- tensor[0,:]=aux
- self.left_node.tensor=tensor
-
- # Right node
- ifself.right_nodeisnotNone:
- tensor=torch.randn(self.right_node.shape)*std
- aux=torch.zeros(tensor.shape[0])*std
- aux[0]=1.
- tensor[:,0]=aux
- self.right_node.tensor=tensor
-
- # Left env + Right env
- fornodeinself.left_env+self.right_env:
- tensor=torch.randn(node.shape)*std
- aux=torch.eye(tensor.shape[0],tensor.shape[2])
- tensor[:,0,:]=aux
- node.tensor=tensor
-
- # Output node
- if(self.boundary=='obc')and(self.out_position==0):
- eye_tensor=torch.eye(self.output_node.shape[1])[0,:]
- eye_tensor=eye_tensor.view([1,self.output_node.shape[1]])
- eye_tensor=eye_tensor.expand(self.output_node.shape)
- elif(self.boundary=='obc')and(self.out_position==self.n_features):
- eye_tensor=torch.eye(self.output_node.shape[0])[0,:]
- eye_tensor=eye_tensor.view([self.output_node.shape[0],1])
- eye_tensor=eye_tensor.expand(self.output_node.shape)
- else:
- eye_tensor=torch.eye(self.output_node.shape[0],
- self.output_node.shape[2])
- eye_tensor=eye_tensor.view([self.output_node.shape[0],1,
- self.output_node.shape[2]])
- eye_tensor=eye_tensor.expand(self.output_node.shape)
-
- # Add on a bit of random noise
- tensor=eye_tensor+std*torch.randn(self.output_node.shape)
- self.output_node.tensor=tensor
-
-
[docs]defset_data_nodes(self)->None:
-"""
- Creates ``data`` nodes and connects each of them to the input edge of
- each input node.
- """
- input_edges=[]
- ifself.left_nodeisnotNone:
- input_edges.append(self.left_node['input'])
- input_edges+=list(map(lambdanode:node['input'],
- self.left_env+self.right_env))
- ifself.right_nodeisnotNone:
- input_edges.append(self.right_node['input'])
-
- super().set_data_nodes(input_edges=input_edges,
- num_batch_edges=self.n_batches)
-
- ifself.left_env+self.right_env:
- self.lr_env_data=list(map(lambdanode:node.neighbours('input'),
- self.left_env+self.right_env))
[docs]classUMPSLayer(TensorNetwork):
-"""
- Class for Uniform (translationally invariant) Matrix Product States with an
- extra node that is dedicated to the output. It is the uniform version of
- :class:`MPSLayer`, that is, all input nodes share the same tensor. Thus
- this class cannot have different input or bond dimensions for each node,
- and boundary conditions are always periodic (``"pbc"``).
-
- A ``UMPSLayer`` is formed by the following nodes:
-
- * ``left_env``, ``right_env``: Environments of `matrix` nodes that are at
- the left or right side of the ``output_node``. These nodes have axes
- ``("left", "input", "right")``.
-
- * ``output_node``: Node dedicated to the output. It has axes
- ``("left", "output", "right")``.
-
- Parameters
- ----------
- n_features : int
- Number of input nodes. The total number of nodes (including the output
- node) will be ``n_features + 1``
- in_dim : int
- Input dimension. Equivalent to the physical dimension.
- out_dim : int
- Output dimension (labels) for the output node. Plays the same role as
- ``in_dim`` for input nodes.
- bond_dim : int
- Bond dimension.
- out_position : int, optional
- Position of the output node (label). Should be between 0 and
- ``n_features``. If ``None``, the output node will be located at the
- middle of the MPS.
- n_batches : int
- Number of batch edges of input ``data`` nodes. Usually ``n_batches = 1``
- (where the batch edge is used for the data batched) but it could also
- be ``n_batches = 2`` (one edge for data batched, other edge for image
- patches in convolutional layers).
-
- Examples
- --------
- >>> mps_layer = tk.models.UMPSLayer(n_features=4,
- ... in_dim=2,
- ... out_dim=10,
- ... bond_dim=5)
- >>> for node in mps_layer.left_env + mps_layer.right_env:
- ... assert node.tensor_address() == 'virtual_uniform'
- ...
- >>> data = torch.ones(20, 4, 2) # batch_size x n_features x feature_size
- >>> result = mps_layer(data)
- >>> result.shape
- torch.Size([20, 10])
- """
-
- def__init__(self,
- n_features:int,
- in_dim:int,
- out_dim:int,
- bond_dim:int,
- out_position:Optional[int]=None,
- n_batches:int=1)->None:
-
- super().__init__(name='mps')
-
- # n_features
- ifn_features<0:
- raiseValueError('`n_features` cannot be lower than 0')
- self._n_features=n_features
-
- # out_position
- ifout_positionisNone:
- out_position=n_features//2
- self._out_position=out_position
-
- # in_dim
- ifisinstance(in_dim,int):
- self._in_dim=in_dim
- else:
- raiseTypeError('`in_dim` should be int type')
-
- # out_dim
- ifisinstance(out_dim,int):
- self._out_dim=out_dim
- else:
- raiseTypeError('`out_dim` should be int type')
-
- # bond_dim
- ifisinstance(bond_dim,int):
- self._bond_dim=bond_dim
- else:
- raiseTypeError('`bond_dim` should be int type')
-
- # n_batches
- ifnotisinstance(n_batches,int):
- raiseTypeError('`n_batches should be int type')
- self._n_batches=n_batches
-
- # Create Tensor Network
- self._make_nodes()
- self.initialize()
-
- @property
- defn_features(self)->int:
-"""
- Returns number of input nodes. The total number of nodes (including the
- output node) will be ``n_features + 1``.
- """
- returnself._n_features
-
- @property
- defin_dim(self)->int:
-"""Returns input/physical dimension."""
- returnself._in_dim
-
- @property
- defout_dim(self)->int:
-"""
- Returns the output dimension, that is, the number of labels in the
- output node. Same as ``in_dim`` for input nodes.
- """
- returnself._out_dim
-
- @property
- defbond_dim(self)->int:
-"""Returns bond dimensions."""
- returnself._bond_dim
-
- @property
- defout_position(self)->int:
-"""Returns position of the output node (label)."""
- returnself._out_position
-
- @property
- defboundary(self)->Text:
-"""Returns boundary condition ("obc" or "pbc")."""
- returnself._boundary
-
- @property
- defn_batches(self)->int:
-"""Returns number of batch edges of the ``data`` nodes."""
- returnself._n_batches
-
- def_make_nodes(self)->None:
-"""Creates all the nodes of the MPS."""
- ifself.leaf_nodes:
- raiseValueError('Cannot create MPS nodes if the MPS already has '
- 'nodes')
-
- self.left_env=[]
- self.right_env=[]
-
- # Left environment
- ifself.out_position>0:
- foriinrange(self.out_position):
- node=ParamNode(shape=(self.bond_dim,
- self.in_dim,
- self.bond_dim),
- axes_names=('left','input','right'),
- name=f'left_env_node_({i})',
- network=self)
- self.left_env.append(node)
- ifi==0:
- periodic_edge=node['left']
- else:
- self.left_env[-2]['right']^self.left_env[-1]['left']
-
- # Output
- ifself.out_position==0:
- self.output_node=ParamNode(shape=(self.bond_dim,
- self.out_dim,
- self.bond_dim),
- axes_names=('left',
- 'output',
- 'right'),
- name='output_node',
- network=self)
- periodic_edge=self.output_node['left']
-
- ifself.out_position==self.n_features:
- ifself.n_features!=0:
- self.output_node=ParamNode(shape=(self.bond_dim,
- self.out_dim,
- self.bond_dim),
- axes_names=('left',
- 'output',
- 'right'),
- name='output_node',
- network=self)
- self.output_node['right']^periodic_edge
-
- ifself.left_env:
- self.left_env[-1]['right']^self.output_node['left']
-
- if(self.out_position>0)and(self.out_position<self.n_features):
- self.output_node=ParamNode(shape=(self.bond_dim,
- self.out_dim,
- self.bond_dim),
- axes_names=('left',
- 'output',
- 'right'),
- name='output_node',
- network=self)
- ifself.left_env:
- self.left_env[-1]['right']^self.output_node['left']
-
- # Right environment
- ifself.out_position<self.n_features:
- foriinrange(self.out_position+1,self.n_features+1):
- node=ParamNode(shape=(self.bond_dim,
- self.in_dim,
- self.bond_dim),
- axes_names=('left','input','right'),
- name=f'right_env_node_({i-self.out_position-1})',
- network=self)
- self.right_env.append(node)
- ifi==self.out_position+1:
- self.output_node['right']^self.right_env[-1]['left']
- else:
- self.right_env[-2]['right']^self.right_env[-1]['left']
-
- ifi==self.n_features:
- self.right_env[-1]['right']^periodic_edge
-
- # Virtual node
- uniform_memory=ParamNode(shape=(self.bond_dim,
- self.in_dim,
- self.bond_dim),
- axes_names=('left',
- 'input',
- 'right'),
- name='virtual_uniform',
- network=self,
- virtual=True)
- self.uniform_memory=uniform_memory
-
-
[docs]definitialize(self,std:float=1e-9)->None:
-"""
- Initializes output and uniform nodes as explained `here
- <https://arxiv.org/abs/1605.03795>`_.
- It can be overriden for custom initializations.
- """
- # Virtual node
- tensor=torch.randn(self.uniform_memory.shape)*std
- random_eye=torch.randn(tensor.shape[0],tensor.shape[2])*std
- random_eye=random_eye+torch.eye(tensor.shape[0],tensor.shape[2])
- tensor[:,0,:]=random_eye
-
- self.uniform_memory.tensor=tensor
-
- fornodeinself.left_env+self.right_env:
- node.set_tensor_from(self.uniform_memory)
-
- # Output node
- eye_tensor=torch.eye(self.output_node.shape[0],
- self.output_node.shape[2])
- eye_tensor=eye_tensor.view([self.output_node.shape[0],1,
- self.output_node.shape[2]])
- eye_tensor=eye_tensor.expand(self.output_node.shape)
-
- # Add on a bit of random noise
- tensor=eye_tensor+std*torch.randn(self.output_node.shape)
- self.output_node.tensor=tensor
-
-
[docs]defset_data_nodes(self)->None:
-"""
- Creates ``data`` nodes and connects each of them to the physical edge of
- each input node.
- """
- input_edges=list(map(lambdanode:node['input'],
- self.left_env+self.right_env))
-
- super().set_data_nodes(input_edges=input_edges,
- num_batch_edges=self._n_batches)
-
- ifself.left_env+self.right_env:
- self.lr_env_data=list(map(lambdanode:node.neighbours('input'),
- self.left_env+self.right_env))
[docs]defcontract(self,
- inline_input:bool=False,
- inline_mats:bool=False)->Node:
-"""
- Contracts the whole MPS.
-
- Parameters
- ----------
- inline_input : bool
- Boolean indicating whether input ``data`` nodes should be contracted
- with the ``MPS`` nodes inline (one contraction at a time) or in a
- single stacked contraction.
- inline_mats : bool
- Boolean indicating whether the sequence of matrices (resultant
- after contracting the input ``data`` nodes) should be contracted
- inline or as a sequence of pairwise stacked contrations.
-
- Returns
- -------
- Node
- """
- left_env,right_env=self._input_contraction(inline_input)
-
- ifinline_mats:
- left_env_contracted,right_env_contracted= \
- self._contract_envs_inline(left_env,right_env)
- else:
- left_env_contracted,right_env_contracted= \
- self._pairwise_contraction(left_env,right_env)
-
- result=self.output_node
- ifleft_env_contractedandright_env_contracted:
- result=left_env_contracted[0]@result@right_env_contracted[0]
- elifleft_env_contracted:
- result=left_env_contracted[0]@result
- elifright_env_contracted:
- result=right_env_contracted[0]@result
- else:
- result@=result
- returnresult
-
-
-
[docs]classConvMPSLayer(MPSLayer):
-"""
- Class for Matrix Product States with an extra node that is dedicated to the
- output, and where the input data is a batch of images. It is the convolutional
- version of :class:`MPSLayer`.
-
- Input data as well as initialization parameters are described in `torch.nn.Conv2d
- <https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html>`_.
-
- Parameters
- ----------
- in_channels : int
- Input channels. Same as ``in_dim`` in :class:`MPSLayer`.
- out_channels : int
- Output channels. Same as ``out_dim`` in :class:`MPSLayer`.
- bond_dim : int, list[int] or tuple[int]
- Bond dimension(s). If given as a sequence, its length should be equal
- to :math:`kernel\_size_0 \cdot kernel\_size_1 + 1`
- (if ``boundary = "pbc"``) or :math:`kernel\_size_0 \cdot kernel\_size_1`
- (if ``boundary = "obc"``). The i-th bond dimension is always the dimension
- of the right edge of the i-th node (including output node).
- kernel_size : int, list[int] or tuple[int]
- Kernel size used in `torch.nn.Unfold
- <https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html#torch.nn.Unfold>`_.
- If given as an ``int``, the actual kernel size will be
- ``(kernel_size, kernel_size)``.
- stride : int
- Stride used in `torch.nn.Unfold
- <https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html#torch.nn.Unfold>`_.
- padding : int
- Padding used in `torch.nn.Unfold
- <https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html#torch.nn.Unfold>`_.
- If given as an ``int``, the actual kernel size will be
- ``(kernel_size, kernel_size)``.
- dilation : int
- Dilation used in `torch.nn.Unfold
- <https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html#torch.nn.Unfold>`_.
- If given as an ``int``, the actual kernel size will be
- ``(kernel_size, kernel_size)``.
- out_position : int, optional
- Position of the output node (label). Should be between 0 and
- :math:`kernel\_size_0 \cdot kernel\_size_1`. If ``None``, the output node
- will be located at the middle of the MPS.
- boundary : {"obc", "pbc"}
- String indicating whether periodic or open boundary conditions should
- be used.
-
- Examples
- --------
- >>> conv_mps_layer = tk.models.ConvMPSLayer(in_channels=2,
- ... out_channels=10,
- ... bond_dim=5,
- ... kernel_size=2)
- >>> data = torch.ones(20, 2, 2, 2) # batch_size x in_channels x height x width
- >>> result = conv_mps_layer(data)
- >>> result.shape
- torch.Size([20, 10, 1, 1])
- """
-
- def__init__(self,
- in_channels:int,
- out_channels:int,
- bond_dim:Union[int,Sequence[int]],
- kernel_size:Union[int,Sequence],
- stride:int=1,
- padding:int=0,
- dilation:int=1,
- out_position:Optional[int]=None,
- boundary:Text='obc'):
-
- ifisinstance(kernel_size,int):
- kernel_size=(kernel_size,kernel_size)
- elifnotisinstance(kernel_size,Sequence):
- raiseTypeError('`kernel_size` must be int or Sequence')
-
- ifisinstance(stride,int):
- stride=(stride,stride)
- elifnotisinstance(stride,Sequence):
- raiseTypeError('`stride` must be int or Sequence')
-
- ifisinstance(padding,int):
- padding=(padding,padding)
- elifnotisinstance(padding,Sequence):
- raiseTypeError('`padding` must be int or Sequence')
-
- ifisinstance(dilation,int):
- dilation=(dilation,dilation)
- elifnotisinstance(dilation,Sequence):
- raiseTypeError('`dilation` must be int or Sequence')
-
- self._in_channels=in_channels
- self._kernel_size=kernel_size
- self._stride=stride
- self._padding=padding
- self._dilation=dilation
-
- super().__init__(n_features=kernel_size[0]*kernel_size[1],
- in_dim=in_channels,
- out_dim=out_channels,
- bond_dim=bond_dim,
- out_position=out_position,
- boundary=boundary,
- n_batches=2)
-
- self.unfold=nn.Unfold(kernel_size=kernel_size,
- stride=stride,
- padding=padding,
- dilation=dilation)
-
- @property
- defin_channels(self)->int:
-"""Returns ``in_channels``. Same as ``in_dim`` in :class:`MPSLayer`."""
- returnself._in_channels
-
- @property
- defkernel_size(self)->Tuple[int,int]:
-"""
- Returns ``kernel_size``. Number of nodes is given by
- :math:`kernel\_size_0 \cdot kernel\_size_1 + 1`.
- """
- returnself._kernel_size
-
- @property
- defstride(self)->Tuple[int,int]:
-"""
- Returns stride used in `torch.nn.Unfold
- <https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html#torch.nn.Unfold>`_.
- """
- returnself._stride
-
- @property
- defpadding(self)->Tuple[int,int]:
-"""
- Returns padding used in `torch.nn.Unfold
- <https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html#torch.nn.Unfold>`_.
- """
- returnself._padding
-
- @property
- defdilation(self)->Tuple[int,int]:
-"""
- Returns dilation used in `torch.nn.Unfold
- <https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html#torch.nn.Unfold>`_.
- """
- returnself._dilation
-
-
[docs]defforward(self,image,mode='flat',*args,**kwargs):
-r"""
- Overrides ``torch.nn.Module``'s forward to compute a convolution on the
- input image.
-
- Parameters
- ----------
- image : torch.Tensor
- Input batch of images with shape
-
- .. math::
-
- batch\_size \times in\_channels \times height \times width
- mode : {"flat", "snake"}
- Indicates the order in which MPS should take the pixels in the image.
- When ``"flat"``, the image is flattened putting one row of the image
- after the other. When ``"snake"``, its row is put in the opposite
- orientation as the previous row (like a snake running through the
- image).
- args :
- Arguments that might be used in :meth:`~MPSLayer.contract`.
- kwargs :
- Keyword arguments that might be used in :meth:`~MPSLayer.contract`,
- like ``inline_input`` or ``inline_mats``.
- """
- # Input image shape: batch_size x in_channels x height x width
-
- patches=self.unfold(image).transpose(1,2)
- # batch_size x nb_windows x (in_channels * nb_pixels)
-
- patches=patches.view(*patches.shape[:-1],self.in_channels,-1)
- # batch_size x nb_windows x in_channels x nb_pixels
-
- patches=patches.transpose(2,3)
- # batch_size x nb_windows x nb_pixels x in_channels
-
- ifmode=='snake':
- new_patches=patches[...,:self._kernel_size[1],:]
- foriinrange(1,self._kernel_size[0]):
- ifi%2==0:
- aux=patches[...,(i*self._kernel_size[1]):
- ((i+1)*self._kernel_size[1]),:]
- else:
- aux=patches[...,
- (i*self._kernel_size[1]):
- ((i+1)*self._kernel_size[1]),:].flip(dims=[0])
- new_patches=torch.cat([new_patches,aux],dim=2)
-
- patches=new_patches
-
- elifmode!='flat':
- raiseValueError('`mode` can only be "flat" or "snake"')
-
- result=super().forward(patches,*args,**kwargs)
- # batch_size x nb_windows x out_channels
-
- result=result.transpose(1,2)
- # batch_size x out_channels x nb_windows
-
- h_in=image.shape[2]
- w_in=image.shape[3]
-
- h_out=int((h_in+2*self.padding[0]-self.dilation[0]*
- (self.kernel_size[0]-1)-1)/self.stride[0]+1)
- w_out=int((w_in+2*self.padding[1]-self.dilation[1]*
- (self.kernel_size[1]-1)-1)/self.stride[1]+1)
-
- result=result.view(*result.shape[:-1],h_out,w_out)
- # batch_size x out_channels x height_out x width_out
-
- returnresult
-
-
-
[docs]classConvUMPSLayer(UMPSLayer):
-"""
- Class for Uniform Matrix Product States with an extra node that is dedicated
- to the output, and where the input data is a batch of images. It is the
- convolutional version of :class:`UMPSLayer`. This class cannot have different
- bond dimensions for each site and boundary conditions are always periodic.
-
- Input data as well as initialization parameters are described in `torch.nn.Conv2d
- <https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html>`_.
-
- Parameters
- ----------
- in_channels : int
- Input channels. Same as ``in_dim`` in :class:`UMPSLayer`.
- out_channels : int
- Output channels. Same as ``out_dim`` in :class:`UMPSLayer`.
- bond_dim : int
- Bond dimension.
- kernel_size : int, list[int] or tuple[int]
- Kernel size used in `torch.nn.Unfold
- <https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html#torch.nn.Unfold>`_.
- If given as an ``int``, the actual kernel size will be
- ``(kernel_size, kernel_size)``.
- stride : int
- Stride used in `torch.nn.Unfold
- <https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html#torch.nn.Unfold>`_.
- padding : int
- Padding used in `torch.nn.Unfold
- <https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html#torch.nn.Unfold>`_.
- If given as an ``int``, the actual kernel size will be
- ``(kernel_size, kernel_size)``.
- dilation : int
- Dilation used in `torch.nn.Unfold
- <https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html#torch.nn.Unfold>`_.
- If given as an ``int``, the actual kernel size will be
- ``(kernel_size, kernel_size)``.
- out_position : int, optional
- Position of the output node (label). Should be between 0 and
- :math:`kernel\_size_0 \cdot kernel\_size_1`. If ``None``, the output node
- will be located at the middle of the MPS.
-
- Examples
- --------
- >>> conv_mps_layer = tk.models.ConvUMPSLayer(in_channels=2,
- ... out_channels=10,
- ... bond_dim=5,
- ... kernel_size=2)
- >>> for node in conv_mps_layer.left_env + conv_mps_layer.right_env:
- ... assert node.tensor_address() == 'virtual_uniform'
- ...
- >>> data = torch.ones(20, 2, 2, 2) # batch_size x in_channels x height x width
- >>> result = conv_mps_layer(data)
- >>> result.shape
- torch.Size([20, 10, 1, 1])
- """
-
- def__init__(self,
- in_channels:int,
- out_channels:int,
- bond_dim:int,
- kernel_size:Union[int,Sequence],
- stride:int=1,
- padding:int=0,
- dilation:int=1,
- out_position:Optional[int]=None):
-
- ifisinstance(kernel_size,int):
- kernel_size=(kernel_size,kernel_size)
- elifnotisinstance(kernel_size,Sequence):
- raiseTypeError('`kernel_size` must be int or Sequence')
-
- ifisinstance(stride,int):
- stride=(stride,stride)
- elifnotisinstance(stride,Sequence):
- raiseTypeError('`stride` must be int or Sequence')
-
- ifisinstance(padding,int):
- padding=(padding,padding)
- elifnotisinstance(padding,Sequence):
- raiseTypeError('`padding` must be int or Sequence')
-
- ifisinstance(dilation,int):
- dilation=(dilation,dilation)
- elifnotisinstance(dilation,Sequence):
- raiseTypeError('`dilation` must be int or Sequence')
-
- self._in_channels=in_channels
- self._kernel_size=kernel_size
- self._stride=stride
- self._padding=padding
- self._dilation=dilation
-
- super().__init__(n_features=kernel_size[0]*kernel_size[1],
- in_dim=in_channels,
- out_dim=out_channels,
- bond_dim=bond_dim,
- out_position=out_position,
- n_batches=2)
-
- self.unfold=nn.Unfold(kernel_size=kernel_size,
- stride=stride,
- padding=padding,
- dilation=dilation)
-
- @property
- defin_channels(self)->int:
-"""Returns ``in_channels``. Same as ``in_dim`` in :class:`UMPSLayer`."""
- returnself._in_channels
-
- @property
- defkernel_size(self)->Tuple[int,int]:
-"""
- Returns ``kernel_size``. Number of nodes is given by
- :math:`kernel\_size_0 \cdot kernel\_size_1 + 1`.
- """
- returnself._kernel_size
-
- @property
- defstride(self)->Tuple[int,int]:
-"""
- Returns stride used in `torch.nn.Unfold
- <https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html#torch.nn.Unfold>`_.
- """
- returnself._stride
-
- @property
- defpadding(self)->Tuple[int,int]:
-"""
- Returns padding used in `torch.nn.Unfold
- <https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html#torch.nn.Unfold>`_.
- """
- returnself._padding
-
- @property
- defdilation(self)->Tuple[int,int]:
-"""
- Returns dilation used in `torch.nn.Unfold
- <https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html#torch.nn.Unfold>`_.
- """
- returnself._dilation
-
-
[docs]defforward(self,image,mode='flat',*args,**kwargs):
-r"""
- Overrides ``torch.nn.Module``'s forward to compute a convolution on the
- input image.
-
- Parameters
- ----------
- image : torch.Tensor
- Input batch of images with shape
-
- .. math::
-
- batch\_size \times in\_channels \times height \times width
- mode : {"flat", "snake"}
- Indicates the order in which MPS should take the pixels in the image.
- When ``"flat"``, the image is flattened putting one row of the image
- after the other. When ``"snake"``, its row is put in the opposite
- orientation as the previous row (like a snake running through the
- image).
- args :
- Arguments that might be used in :meth:`~UMPSLayer.contract`.
- kwargs :
- Keyword arguments that might be used in :meth:`~UMPSLayer.contract`,
- like ``inline_input`` or ``inline_mats``.
- """
- # Input image shape: batch_size x in_channels x height x width
-
- patches=self.unfold(image).transpose(1,2)
- # batch_size x nb_windows x (in_channels * nb_pixels)
-
- patches=patches.view(*patches.shape[:-1],self.in_channels,-1)
- # batch_size x nb_windows x in_channels x nb_pixels
-
- patches=patches.transpose(2,3)
- # batch_size x nb_windows x nb_pixels x in_channels
-
- ifmode=='snake':
- new_patches=patches[...,:self._kernel_size[1],:]
- foriinrange(1,self._kernel_size[0]):
- ifi%2==0:
- aux=patches[...,(i*self._kernel_size[1]):
- ((i+1)*self._kernel_size[1]),:]
- else:
- aux=patches[...,
- (i*self._kernel_size[1]):
- ((i+1)*self._kernel_size[1]),:].flip(dims=[0])
- new_patches=torch.cat([new_patches,aux],dim=2)
-
- patches=new_patches
-
- elifmode!='flat':
- raiseValueError('`mode` can only be "flat" or "snake"')
-
- result=super().forward(patches,*args,**kwargs)
- # batch_size x nb_windows x out_channels
-
- result=result.transpose(1,2)
- # batch_size x out_channels x nb_windows
-
- h_in=image.shape[2]
- w_in=image.shape[3]
-
- h_out=int((h_in+2*self.padding[0]-self.dilation[0]*
- (self.kernel_size[0]-1)-1)/self.stride[0]+1)
- w_out=int((w_in+2*self.padding[1]-self.dilation[1]*
- (self.kernel_size[1]-1)-1)/self.stride[1]+1)
-
- result=result.view(*result.shape[:-1],h_out,w_out)
- # batch_size x out_channels x height_out x width_out
-
- returnresult
fromtensorkrowch.utilsimport(inverse_permutation,is_permutation,list2slice,permute_list)
-Ax=Union[int,Text,Axis]
-
defcopy_func(f):"""Returns a function with the same code, defaults, closure and name."""
@@ -344,7 +350,7 @@
[docs]classOperation:# MARK: Operation""" Class for node operations.
@@ -370,7 +376,11 @@
Source code for tensorkrowch.operations
Function that is called the next times the operation is performed. """
- def__init__(self,name:Text,check_first,fn_first,fn_next):
+ def__init__(self,
+ name:Text,
+ check_first:Callable,
+ fn_first:Callable,
+ fn_next:Callable)->None:assertisinstance(check_first,Callable)assertisinstance(fn_first,Callable)assertisinstance(fn_next,Callable)
@@ -395,9 +405,10 @@
Permutes the nodes' tensor, as well as its axes and edges to match the new shape.
- See `permute <https://pytorch.org/docs/stable/generated/torch.permute.html>`_.
+ See `permute <https://pytorch.org/docs/stable/generated/torch.permute.html>`_
+ in the **PyTorch** documentation. Nodes ``resultant`` from this operation are called ``"permute"``. The node that keeps information about the :class:`Successor` is ``node``.
+
+ This operation is the same as :meth:`~AbstractNode.permute`. Parameters ----------
@@ -506,7 +520,8 @@
Source code for tensorkrowch.operations
Permutes the nodes' tensor, as well as its axes and edges to match the new shape.
- See `permute <https://pytorch.org/docs/stable/generated/torch.permute.html>`_.
+ See `permute <https://pytorch.org/docs/stable/generated/torch.permute.html>`_
+ in the **PyTorch** documentation. Nodes ``resultant`` from this operation are called ``"permute"``. The node that keeps information about the :class:`Successor` is ``self``.
@@ -542,6 +557,8 @@
Source code for tensorkrowch.operations
See `permute <https://pytorch.org/docs/stable/generated/torch.permute.html>`_. Nodes ``resultant`` from this operation use the same name as ``node``.
+
+ This operation is the same as :meth:`~AbstractNode.permute_`. Parameters ----------
@@ -613,6 +630,7 @@
rank:Optional[int]=None,cum_percentage:Optional[float]=None,cutoff:Optional[float]=None)->Tuple[Node,Node]:
- ifnotisinstance(node1_axes,(list,tuple)):
+ ifnotisinstance(node1_axes,Sequence):raiseTypeError('`node1_edges` should be list or tuple type')
- ifnotisinstance(node2_axes,(list,tuple)):
+ ifnotisinstance(node2_axes,Sequence):raiseTypeError('`node2_edges` should be list or tuple type')args=(node,
@@ -1254,48 +1276,39 @@
Source code for tensorkrowch.operations
if(mode=='svd')or(mode=='svdr'):u,s,vh=torch.linalg.svd(node_tensor,full_matrices=False)
-
- ifcum_percentageisnotNone:
- if(rankisnotNone)or(cutoffisnotNone):
- raiseValueError('Only one of `rank`, `cum_percentage` and '
- '`cutoff` should be provided')
-
- percentages=s.cumsum(-1)/s.sum(-1) \
- .view(*s.shape[:-1],1).expand(s.shape)
- cum_percentage_tensor=torch.tensor(cum_percentage) \
- .repeat(percentages.shape[:-1])
- rank=0
- foriinrange(percentages.shape[-1]):
- p=percentages[...,i]
- rank+=1
- # Cut when ``cum_percentage`` is exceeded in all batches
- iftorch.ge(p,cum_percentage_tensor).all():
- break
-
- elifcutoffisnotNone:
- ifrankisnotNone:
- raiseValueError('Only one of `rank`, `cum_percentage` and '
- '`cutoff` should be provided')
-
- cutoff_tensor=torch.tensor(cutoff).repeat(s.shape[:-1])
- rank=0
- foriinrange(s.shape[-1]):
- # Cut when ``cutoff`` is exceeded in all batches
- iftorch.le(s[...,i],cutoff_tensor).all():
- break
- rank+=1
- ifrank==0:
- rank=1
-
+
+ lst_ranks=[]
+
ifrankisNone:rank=s.shape[-1]
+ lst_ranks.append(rank)else:
- ifrank<s.shape[-1]:
- u=u[...,:rank]
- s=s[...,:rank]
- vh=vh[...,:rank,:]
- else:
- rank=s.shape[-1]
+ lst_ranks.append(min(max(1,int(rank)),s.shape[-1]))
+
+ ifcum_percentageisnotNone:
+ s_percentages=s.cumsum(-1)/ \
+ (s.sum(-1,keepdim=True).expand(s.shape)+1e-10)# To avoid having all 0's
+ cum_percentage_tensor=cum_percentage*torch.ones_like(s)
+ cp_rank=torch.lt(
+ s_percentages,
+ cum_percentage_tensor
+ ).view(-1,s.shape[-1]).all(dim=0).sum()
+ lst_ranks.append(max(1,cp_rank.item()+1))
+
+ ifcutoffisnotNone:
+ cutoff_tensor=cutoff*torch.ones_like(s)
+ co_rank=torch.ge(
+ s,
+ cutoff_tensor
+ ).view(-1,s.shape[-1]).all(dim=0).sum()
+ lst_ranks.append(max(1,co_rank.item()))
+
+ # Select rank from specified restrictions
+ rank=min(lst_ranks)
+
+ u=u[...,:rank]
+ s=s[...,:rank]
+ vh=vh[...,:rank,:]ifmode=='svdr':phase=torch.sign(torch.randn(s.shape))
@@ -1462,48 +1475,39 @@
Source code for tensorkrowch.operations
if(mode=='svd')or(mode=='svdr'):u,s,vh=torch.linalg.svd(node_tensor,full_matrices=False)
-
- ifcum_percentageisnotNone:
- if(rankisnotNone)or(cutoffisnotNone):
- raiseValueError('Only one of `rank`, `cum_percentage` and '
- '`cutoff` should be provided')
-
- percentages=s.cumsum(-1)/s.sum(-1) \
- .view(*s.shape[:-1],1).expand(s.shape)
- cum_percentage_tensor=torch.tensor(
- cum_percentage).repeat(percentages.shape[:-1])
- rank=0
- foriinrange(percentages.shape[-1]):
- p=percentages[...,i]
- rank+=1
- # Cut when ``cum_percentage`` is exceeded in all batches
- iftorch.ge(p,cum_percentage_tensor).all():
- break
-
- elifcutoffisnotNone:
- ifrankisnotNone:
- raiseValueError('Only one of `rank`, `cum_percentage` and '
- '`cutoff` should be provided')
-
- cutoff_tensor=torch.tensor(cutoff).repeat(s.shape[:-1])
- rank=0
- foriinrange(s.shape[-1]):
- # Cut when ``cutoff`` is exceeded in all batches
- iftorch.le(s[...,i],cutoff_tensor).all():
- break
- rank+=1
- ifrank==0:
- rank=1
-
+
+ lst_ranks=[]
+
ifrankisNone:rank=s.shape[-1]
+ lst_ranks.append(rank)else:
- ifrank<s.shape[-1]:
- u=u[...,:rank]
- s=s[...,:rank]
- vh=vh[...,:rank,:]
- else:
- rank=s.shape[-1]
+ lst_ranks.append(min(max(1,rank),s.shape[-1]))
+
+ ifcum_percentageisnotNone:
+ s_percentages=s.cumsum(-1)/ \
+ (s.sum(-1,keepdim=True).expand(s.shape)+1e-10)# To avoid having all 0's
+ cum_percentage_tensor=cum_percentage*torch.ones_like(s)
+ cp_rank=torch.lt(
+ s_percentages,
+ cum_percentage_tensor
+ ).view(-1,s.shape[-1]).all(dim=0).sum()
+ lst_ranks.append(max(1,cp_rank.item()+1))
+
+ ifcutoffisnotNone:
+ cutoff_tensor=cutoff*torch.ones_like(s)
+ co_rank=torch.ge(
+ s,
+ cutoff_tensor
+ ).view(-1,s.shape[-1]).all(dim=0).sum()
+ lst_ranks.append(max(1,co_rank.item()))
+
+ # Select rank from specified restrictions
+ rank=min(lst_ranks)
+
+ u=u[...,:rank]
+ s=s[...,:rank]
+ vh=vh[...,:rank,:]ifmode=='svdr':phase=torch.sign(torch.randn(s.shape))
@@ -1571,7 +1575,7 @@
Source code for tensorkrowch.operations
r""" Splits one node in two via the decomposition specified in ``mode``. To perform this operation the set of edges has to be split in two sets,
- corresponding to the edges of the first and second ``resultant nodes``.
+ corresponding to the edges of the first and second ``resultant`` nodes. Batch edges that don't appear in any of the lists will be repeated in both nodes, and will appear as the first edges of the ``resultant`` nodes, in the order they appeared in ``node``.
@@ -1617,14 +1621,18 @@
Source code for tensorkrowch.operations
where R is a lower triangular matrix and Q is unitary.
- If ``mode`` is "svd" or "svdr", ``side`` must be provided. Besides, one
- (and only one) of ``rank``, ``cum_percentage`` and ``cutoff`` is required.
+ If ``mode`` is "svd" or "svdr", ``side`` must be provided. Besides, at least
+ one of ``rank``, ``cum_percentage`` and ``cutoff`` is required. If more than
+ one is specified, the resulting rank will be the one that satisfies all
+ conditions. Since the node is `split` in two, a new edge appears connecting both nodes. The axis that corresponds to this edge has the name ``"split"``. Nodes ``resultant`` from this operation are called ``"split"``. The node that keeps information about the :class:`Successor` is ``node``.
+
+ This operation is the same as :meth:`~AbstractNode.split`. Parameters ----------
@@ -1765,6 +1773,8 @@
Source code for tensorkrowch.operations
nodes. The axis that corresponds to this edge has the name ``"split"``. Nodes ``resultant`` from this operation are called ``"split_ip"``.
+
+ This operation is the same as :meth:`~AbstractNode.split_`. Parameters ----------
@@ -1828,8 +1838,8 @@
split_node_=copy_func(split_)split_node_.__doc__= \
r"""
- In-place version of :func:`~AbstractNode.split`.
+ In-place version of :meth:`~AbstractNode.split`. Following the **PyTorch** convention, names of functions ended with an underscore indicate **in-place** operations.
@@ -1910,10 +1920,9 @@
[docs]defsvd(edge:Edge,
+ side:Text='left',
+ rank:Optional[int]=None,
+ cum_percentage:Optional[float]=None,
+ cutoff:Optional[float]=None)->Tuple[Node,Node]:r"""
- Contracts an edge in-place via :func:`contract_` and splits it in-place via
- :func:`split_` using ``mode = "svd"``. See :func:`split` for a more complete
- explanation.
-
- Following the **PyTorch** convention, names of functions ended with an
- underscore indicate **in-place** operations.
+ Contracts an edge via :func:`contract` and splits it via :func:`split`
+ using ``mode = "svd"``. See :func:`split` for a more complete explanation.
- Nodes ``resultant`` from this operation use the same names as the original
- nodes connected by ``edge``.
+ This operation is the same as :meth:`~Edge.svd`. Parameters ----------
@@ -1988,25 +1992,31 @@
-svd_edge_=copy_func(svd_)
-svd_edge_.__doc__= \
+svd_edge=copy_func(svd)
+svd_edge.__doc__= \
r"""
- Contracts an edge in-place via :func:`~Edge.contract_` and splits
- it in-place via :func:`~AbstractNode.split_` using ``mode = "svd"``. See
- :func:`split` for a more complete explanation.
-
- Following the **PyTorch** convention, names of functions ended with an
- underscore indicate **in-place** operations.
-
- Nodes ``resultant`` from this operation use the same names as the original
- nodes connected by ``self``.
+ Contracts an edge via :meth:`~Edge.contract` and splits it via
+ :meth:`~AbstractNode.split` using ``mode = "svd"``. See :func:`split` for
+ a more complete explanation. Parameters ----------
@@ -2089,9 +2090,9 @@
Source code for tensorkrowch.operations
cum_percentage : float, optional Proportion that should be satisfied between the sum of all singular values kept and the total sum of all singular values.
-
+
.. math::
-
+
\frac{\sum_{i \in \{kept\}}{s_i}}{\sum_{i \in \{all\}}{s_i}} \ge cum\_percentage cutoff : float, optional
@@ -2111,32 +2112,39 @@
[docs]defsvd_(edge:Edge,
+ side:Text='left',
+ rank:Optional[int]=None,
+ cum_percentage:Optional[float]=None,
+ cutoff:Optional[float]=None)->Tuple[Node,Node]:r"""
+ In-place version of :func:`svd`.
+ Contracts an edge in-place via :func:`contract_` and splits it in-place via
- :func:`split_` using ``mode = "svdr"``. See :func:`split` for a more complete
+ :func:`split_` using ``mode = "svd"``. See :func:`split` for a more complete explanation. Following the **PyTorch** convention, names of functions ended with an
@@ -2144,6 +2152,8 @@
Source code for tensorkrowch.operations
Nodes ``resultant`` from this operation use the same names as the original nodes connected by ``edge``.
+
+ This operation is the same as :meth:`~Edge.svd_`. Parameters ----------
@@ -2182,7 +2192,7 @@
"""ifedge.is_dangling():raiseValueError('Edge should be connected to perform SVD')
+ ifedge.node1isedge.node2:
+ raiseValueError('Edge should connect different nodes')node1,node2=edge.node1,edge.node2node1_name,node2_name=node1._name,node2._name
@@ -2220,14 +2232,14 @@
-svdr_edge_=copy_func(svdr_)
-svdr_edge_.__doc__= \
+svd_edge_=copy_func(svd_)
+svd_edge_.__doc__= \
r"""
- Contracts an edge in-place via :func:`~Edge.contract_` and splits
- it in-place via :func:`~AbstractNode.split_` using ``mode = "svdr"``. See
+ In-place version of :meth:`~Edge.svd`.
+
+ Contracts an edge in-place via :meth:`~Edge.contract_` and splits
+ it in-place via :meth:`~AbstractNode.split_` using ``mode = "svd"``. See :func:`split` for a more complete explanation. Following the **PyTorch** convention, names of functions ended with an
@@ -2305,7 +2319,7 @@
[docs]defsvdr(edge:Edge,
+ side:Text='left',
+ rank:Optional[int]=None,
+ cum_percentage:Optional[float]=None,
+ cutoff:Optional[float]=None)->Tuple[Node,Node]:r"""
- Contracts an edge in-place via :func:`contract_` and splits it in-place via
- :func:`split_` using ``mode = "qr"``. See :func:`split` for a more complete
- explanation.
-
- Following the **PyTorch** convention, names of functions ended with an
- underscore indicate **in-place** operations.
+ Contracts an edge via :func:`contract` and splits it via :func:`split`
+ using ``mode = "svdr"``. See :func:`split` for a more complete explanation.
- Nodes ``resultant`` from this operation use the same names as the original
- nodes connected by ``edge``.
+ This operation is the same as :meth:`~Edge.svdr`. Parameters ---------- edge : Edge Edge whose nodes are to be contracted and split.
+ side : str, optional
+ Indicates the side to which the diagonal matrix :math:`S` should be
+ contracted. If "left", the first resultant node's tensor will be
+ :math:`US`, and the other node's tensor will be :math:`V^{\dagger}`.
+ If "right", their tensors will be :math:`U` and :math:`SV^{\dagger}`,
+ respectively.
+ rank : int, optional
+ Number of singular values to keep.
+ cum_percentage : float, optional
+ Proportion that should be satisfied between the sum of all singular
+ values kept and the total sum of all singular values.
+
+ .. math::
+
+ \frac{\sum_{i \in \{kept\}}{s_i}}{\sum_{i \in \{all\}}{s_i}} \ge
+ cum\_percentage
+ cutoff : float, optional
+ Quantity that lower bounds singular values in order to be kept. Returns -------
@@ -2354,48 +2385,58 @@
-qr_edge_=copy_func(qr_)
-qr_edge_.__doc__= \
+svdr_edge=copy_func(svdr)
+svdr_edge.__doc__= \
r"""
- Contracts an edge in-place via :func:`~Edge.contract_` and splits
- it in-place via :func:`~AbstractNode.split_` using ``mode = "qr"``. See
- :func:`split` for a more complete explanation.
-
- Following the **PyTorch** convention, names of functions ended with an
- underscore indicate **in-place** operations.
-
- Nodes ``resultant`` from this operation use the same names as the original
- nodes connected by ``self``.
+ Contracts an edge via :meth:`~Edge.contract` and splits it via
+ :meth:`~AbstractNode.split` using ``mode = "svdr"``. See :func:`split` for
+ a more complete explanation.
+
+ Parameters
+ ----------
+ side : str, optional
+ Indicates the side to which the diagonal matrix :math:`S` should be
+ contracted. If "left", the first resultant node's tensor will be
+ :math:`US`, and the other node's tensor will be :math:`V^{\dagger}`.
+ If "right", their tensors will be :math:`U` and :math:`SV^{\dagger}`,
+ respectively.
+ rank : int, optional
+ Number of singular values to keep.
+ cum_percentage : float, optional
+ Proportion that should be satisfied between the sum of all singular
+ values kept and the total sum of all singular values.
+
+ .. math::
+
+ \frac{\sum_{i \in \{kept\}}{s_i}}{\sum_{i \in \{all\}}{s_i}} \ge
+ cum\_percentage
+ cutoff : float, optional
+ Quantity that lower bounds singular values in order to be kept. Returns -------
@@ -2452,28 +2505,39 @@
[docs]defsvdr_(edge:Edge,
+ side:Text='left',
+ rank:Optional[int]=None,
+ cum_percentage:Optional[float]=None,
+ cutoff:Optional[float]=None)->Tuple[Node,Node]:r"""
+ In-place version of :func:`svdr`.
+ Contracts an edge in-place via :func:`contract_` and splits it in-place via
- :func:`split_` using ``mode = "rq"``. See :func:`split` for a more complete
+ :func:`split_` using ``mode = "svdr"``. See :func:`split` for a more complete explanation. Following the **PyTorch** convention, names of functions ended with an
@@ -2481,16 +2545,659 @@
Source code for tensorkrowch.operations
Nodes ``resultant`` from this operation use the same names as the original nodes connected by ``edge``.
+
+ This operation is the same as :meth:`~Edge.svdr_`. Parameters ---------- edge : Edge Edge whose nodes are to be contracted and split.
+ side : str, optional
+ Indicates the side to which the diagonal matrix :math:`S` should be
+ contracted. If "left", the first resultant node's tensor will be
+ :math:`US`, and the other node's tensor will be :math:`V^{\dagger}`.
+ If "right", their tensors will be :math:`U` and :math:`SV^{\dagger}`,
+ respectively.
+ rank : int, optional
+ Number of singular values to keep.
+ cum_percentage : float, optional
+ Proportion that should be satisfied between the sum of all singular
+ values kept and the total sum of all singular values.
- Returns
- -------
- tuple[Node, Node]
-
+ .. math::
+
+ \frac{\sum_{i \in \{kept\}}{s_i}}{\sum_{i \in \{all\}}{s_i}} \ge
+ cum\_percentage
+ cutoff : float, optional
+ Quantity that lower bounds singular values in order to be kept.
+
+ Returns
+ -------
+ tuple[Node, Node]
+
+ Examples
+ --------
+ >>> nodeA = tk.randn(shape=(10, 15, 100),
+ ... axes_names=('left', 'right', 'batch'),
+ ... name='nodeA')
+ >>> nodeB = tk.randn(shape=(15, 20, 100),
+ ... axes_names=('left', 'right', 'batch'),
+ ... name='nodeB')
+ ...
+ >>> new_edge = nodeA['right'] ^ nodeB['left']
+ >>> nodeA, nodeB = tk.svdr_(new_edge, rank=7)
+ ...
+ >>> nodeA.shape
+ torch.Size([10, 7, 100])
+
+ >>> nodeB.shape
+ torch.Size([7, 20, 100])
+
+ >>> print(nodeA.axes_names)
+ ['left', 'right', 'batch']
+
+ >>> print(nodeB.axes_names)
+ ['left', 'right', 'batch']
+ """
+ ifedge.is_dangling():
+ raiseValueError('Edge should be connected to perform SVD')
+ ifedge.node1isedge.node2:
+ raiseValueError('Edge should connect different nodes')
+
+ node1,node2=edge.node1,edge.node2
+ node1_name,node2_name=node1._name,node2._name
+ axis1,axis2=edge.axis1,edge.axis2
+
+ batch_axes=[]
+ foraxisinnode1._axes:
+ ifaxis._batchand(axis._nameinnode2.axes_names):
+ batch_axes.append(axis)
+
+ n_batches=len(batch_axes)
+ n_axes1=len(node1._axes)-n_batches-1
+ n_axes2=len(node2._axes)-n_batches-1
+
+ contracted=edge.contract_()
+ new_node1,new_node2=split_(node=contracted,
+ node1_axes=list(
+ range(n_batches,
+ n_batches+n_axes1)),
+ node2_axes=list(
+ range(n_batches+n_axes1,
+ n_batches+n_axes1+n_axes2)),
+ mode='svdr',
+ side=side,
+ rank=rank,
+ cum_percentage=cum_percentage,
+ cutoff=cutoff)
+
+ # new_node1
+ prev_nums=[ax._numforaxinbatch_axes]
+ foriinrange(new_node1.rank):
+ if(inotinprev_nums)and(i!=axis1._num):
+ prev_nums.append(i)
+ prev_nums+=[axis1._num]
+
+ ifprev_nums!=list(range(new_node1.rank)):
+ permutation=inverse_permutation(prev_nums)
+ new_node1=new_node1.permute_(permutation)
+
+ # new_node2
+ prev_nums=[node2.get_axis_num(node1.get_axis(ax)._name)
+ foraxinbatch_axes]+[axis2._num]
+ foriinrange(new_node2.rank):
+ ifinotinprev_nums:
+ prev_nums.append(i)
+
+ ifprev_nums!=list(range(new_node2.rank)):
+ permutation=inverse_permutation(prev_nums)
+ new_node2=new_node2.permute_(permutation)
+
+ new_node1.name=node1_name
+ new_node1.get_axis(axis1._num).name=axis1._name
+
+ new_node2.name=node2_name
+ new_node2.get_axis(axis2._num).name=axis2._name
+
+ returnnew_node1,new_node2
+
+
+svdr_edge_=copy_func(svdr_)
+svdr_edge_.__doc__= \
+r"""
+ In-place version of :meth:`~Edge.svdr`.
+
+ Contracts an edge in-place via :meth:`~Edge.contract_` and splits
+ it in-place via :meth:`~AbstractNode.split_` using ``mode = "svdr"``. See
+ :func:`split` for a more complete explanation.
+
+ Following the **PyTorch** convention, names of functions ended with an
+ underscore indicate **in-place** operations.
+
+ Nodes ``resultant`` from this operation use the same names as the original
+ nodes connected by ``self``.
+
+ Parameters
+ ----------
+ side : str, optional
+ Indicates the side to which the diagonal matrix :math:`S` should be
+ contracted. If "left", the first resultant node's tensor will be
+ :math:`US`, and the other node's tensor will be :math:`V^{\dagger}`.
+ If "right", their tensors will be :math:`U` and :math:`SV^{\dagger}`,
+ respectively.
+ rank : int, optional
+ Number of singular values to keep.
+ cum_percentage : float, optional
+ Proportion that should be satisfied between the sum of all singular
+ values kept and the total sum of all singular values.
+
+ .. math::
+
+ \frac{\sum_{i \in \{kept\}}{s_i}}{\sum_{i \in \{all\}}{s_i}} \ge
+ cum\_percentage
+ cutoff : float, optional
+ Quantity that lower bounds singular values in order to be kept.
+
+ Returns
+ -------
+ tuple[Node, Node]
+
+ Examples
+ --------
+ >>> nodeA = tk.randn(shape=(10, 15, 100),
+ ... axes_names=('left', 'right', 'batch'),
+ ... name='nodeA')
+ >>> nodeB = tk.randn(shape=(15, 20, 100),
+ ... axes_names=('left', 'right', 'batch'),
+ ... name='nodeB')
+ ...
+ >>> new_edge = nodeA['right'] ^ nodeB['left']
+ >>> nodeA, nodeB = new_edge.svdr_(rank=7)
+ ...
+ >>> nodeA.shape
+ torch.Size([10, 7, 100])
+
+ >>> nodeB.shape
+ torch.Size([7, 20, 100])
+
+ >>> print(nodeA.axes_names)
+ ['left', 'right', 'batch']
+
+ >>> print(nodeB.axes_names)
+ ['left', 'right', 'batch']
+ """
+
+Edge.svdr_=svdr_edge_
+
+
+
[docs]defqr(edge:Edge)->Tuple[Node,Node]:
+r"""
+ Contracts an edge via :func:`contract` and splits it via :func:`split`
+ using ``mode = "qr"``. See :func:`split` for a more complete explanation.
+
+ This operation is the same as :meth:`~Edge.qr`.
+
+ Parameters
+ ----------
+ edge : Edge
+ Edge whose nodes are to be contracted and split.
+
+ Returns
+ -------
+ tuple[Node, Node]
+
+ Examples
+ --------
+ >>> nodeA = tk.randn(shape=(10, 15, 100),
+ ... axes_names=('left', 'right', 'batch'),
+ ... name='nodeA')
+ >>> nodeB = tk.randn(shape=(15, 20, 100),
+ ... axes_names=('left', 'right', 'batch'),
+ ... name='nodeB')
+ ...
+ >>> new_edge = nodeA['right'] ^ nodeB['left']
+ >>> new_nodeA, new_nodeB = tk.qr(new_edge)
+ ...
+ >>> new_nodeA.shape
+ torch.Size([10, 10, 100])
+
+ >>> new_nodeB.shape
+ torch.Size([10, 20, 100])
+
+ >>> print(new_nodeA.axes_names)
+ ['left', 'right', 'batch']
+
+ >>> print(new_nodeB.axes_names)
+ ['left', 'right', 'batch']
+
+ Original nodes still exist in the network
+
+ >>> assert nodeA.network == new_nodeA.network
+ >>> assert nodeB.network == new_nodeB.network
+ """
+ ifedge.is_dangling():
+ raiseValueError('Edge should be connected to perform SVD')
+ ifedge.node1isedge.node2:
+ raiseValueError('Edge should connect different nodes')
+
+ node1,node2=edge.node1,edge.node2
+ axis1,axis2=edge.axis1,edge.axis2
+
+ batch_axes=[]
+ foraxisinnode1._axes:
+ ifaxis._batchand(axis._nameinnode2.axes_names):
+ batch_axes.append(axis)
+
+ n_batches=len(batch_axes)
+ n_axes1=len(node1._axes)-n_batches-1
+ n_axes2=len(node2._axes)-n_batches-1
+
+ contracted=edge.contract()
+ new_node1,new_node2=split(node=contracted,
+ node1_axes=list(
+ range(n_batches,
+ n_batches+n_axes1)),
+ node2_axes=list(
+ range(n_batches+n_axes1,
+ n_batches+n_axes1+n_axes2)),
+ mode='qr')
+
+ # new_node1
+ prev_nums=[ax.numforaxinbatch_axes]
+ foriinrange(new_node1.rank):
+ if(inotinprev_nums)and(i!=axis1._num):
+ prev_nums.append(i)
+ prev_nums+=[axis1._num]
+
+ ifprev_nums!=list(range(new_node1.rank)):
+ permutation=inverse_permutation(prev_nums)
+ new_node1=new_node1.permute(permutation)
+
+ # new_node2
+ prev_nums=[node2.get_axis_num(node1.get_axis(ax)._name)
+ foraxinbatch_axes]+[axis2._num]
+ foriinrange(new_node2.rank):
+ ifinotinprev_nums:
+ prev_nums.append(i)
+
+ ifprev_nums!=list(range(new_node2.rank)):
+ permutation=inverse_permutation(prev_nums)
+ new_node2=new_node2.permute(permutation)
+
+ new_node1.get_axis(axis1._num).name=axis1._name
+ new_node2.get_axis(axis2._num).name=axis2._name
+
+ returnnew_node1,new_node2
[docs]defrq_(edge)->Tuple[Node,Node]:
+r"""
+ In-place version of :func:`rq`.
+
+ Contracts an edge in-place via :func:`contract_` and splits it in-place via
+ :func:`split_` using ``mode = "rq"``. See :func:`split` for a more complete
+ explanation.
+
+ Following the **PyTorch** convention, names of functions ended with an
+ underscore indicate **in-place** operations.
+
+ Nodes ``resultant`` from this operation use the same names as the original
+ nodes connected by ``edge``.
+
+ This operation is the same as :meth:`~Edge.rq_`.
+
+ Parameters
+ ----------
+ edge : Edge
+ Edge whose nodes are to be contracted and split.
+
+ Returns
+ -------
+ tuple[Node, Node]
+ Examples -------- >>> nodeA = tk.randn(shape=(10, 15, 100),
@@ -2517,6 +3224,8 @@
Source code for tensorkrowch.operations
"""ifedge.is_dangling():raiseValueError('Edge should be connected to perform SVD')
+ ifedge.node1isedge.node2:
+ raiseValueError('Edge should connect different nodes')node1,node2=edge.node1,edge.node2node1_name,node2_name=node1._name,node2._name
@@ -2524,7 +3233,7 @@
rq_edge_=copy_func(rq_)rq_edge_.__doc__= \
r"""
- Contracts an edge in-place via :func:`~Edge.contract_` and splits
- it in-place via :func:`~AbstractNode.split_` using ``mode = "qr"``. See
+ In-place version of :meth:`~Edge.rq`.
+
+ Contracts an edge in-place via :meth:`~Edge.contract_` and splits
+ it in-place via :meth:`~AbstractNode.split_` using ``mode = "qr"``. See :func:`split` for a more complete explanation. Following the **PyTorch** convention, names of functions ended with an
@@ -2618,6 +3329,7 @@
[docs]defcontract(edge:Edge)->Node:
+"""
+ Contracts the nodes that are connected through the edge.
+
+ Nodes ``resultant`` from this operation are called ``"contract_edges"``.
+ The node that keeps information about the :class:`Successor` is
+ ``edge.node1``.
+
+ This operation is the same as :meth:`~Edge.contract`.
+
+ Parameters
+ ----------
+ edge : Edge
+ Edge that is to be contracted. Batch contraction is automatically
+ performed when both nodes have batch edges with the same names.
+
+ Returns
+ -------
+ Node
+
+ Examples
+ --------
+ >>> nodeA = tk.randn(shape=(10, 15, 20),
+ ... axes_names=('one', 'two', 'three'),
+ ... name='nodeA')
+ >>> nodeB = tk.randn(shape=(10, 15, 20),
+ ... axes_names=('one', 'two', 'three'),
+ ... name='nodeB')
+ ...
+ >>> _ = nodeA['one'] ^ nodeB['one']
+ >>> _ = nodeA['two'] ^ nodeB['two']
+ >>> _ = nodeA['three'] ^ nodeB['three']
+ >>> result = tk.contract(nodeA['one'])
+ >>> result.shape
+ torch.Size([15, 20, 15, 20])
+ """
+ returncontract_edges_op([edge],edge.node1,edge.node2)
+
+contract_edge=copy_func(contract)
+contract_edge.__doc__= \
+"""
+ Contracts the nodes that are connected through the edge.
+
+ Nodes ``resultant`` from this operation are called ``"contract_edges"``.
+ The node that keeps information about the :class:`Successor` is
+ ``self.node1``.
+
+ Returns
+ -------
+ Node
+
+ Examples
+ --------
+ >>> nodeA = tk.randn(shape=(10, 15, 20),
+ ... axes_names=('one', 'two', 'three'),
+ ... name='nodeA')
+ >>> nodeB = tk.randn(shape=(10, 15, 20),
+ ... axes_names=('one', 'two', 'three'),
+ ... name='nodeB')
+ ...
+ >>> _ = nodeA['one'] ^ nodeB['one']
+ >>> _ = nodeA['two'] ^ nodeB['two']
+ >>> _ = nodeA['three'] ^ nodeB['three']
+ >>> result = nodeA['one'].contract()
+ >>> result.shape
+ torch.Size([15, 20, 15, 20])
+ """
+
+Edge.contract=contract_edge
+
+
[docs]defcontract_(edge:Edge)->Node:"""
- Contracts in-place the nodes that are connected through the edge. See
- :func:`contract` for a more complete explanation.
+ In-place version of :func:`contract`. Following the **PyTorch** convention, names of functions ended with an underscore indicate **in-place** operations. Nodes ``resultant`` from this operation are called ``"contract_edges_ip"``.
+
+ This operation is the same as :meth:`~Edge.contract_`. Parameters ----------
@@ -3017,14 +3801,17 @@
Source code for tensorkrowch.operations
>>> del nodeA >>> del nodeB """
- result=contract_edges([edge],edge.node1,edge.node2)
- result.reattach_edges(True)
+ nodes=[edge.node1,edge.node2]
+ result=contract_edges_op([edge],nodes[0],nodes[1])
+ result.reattach_edges(override=True)result._unrestricted_set_tensor(result.tensor.detach())
+
+ nodes=set(nodes)# Delete nodes (and their edges) from the TNnet=result.network
- net.delete_node(edge.node1)
- net.delete_node(edge.node2)
+ fornodeinnodes:
+ net.delete_node(node)# Add edges of result to the TNforres_edgeinresult._edges:
@@ -3034,9 +3821,9 @@
-Edge.contract_=contract_
+contract_edge_=copy_func(contract_)
+contract_edge_.__doc__= \
+"""
+ In-place version of :meth:`~Edge.contract`.
+
+ Following the **PyTorch** convention, names of functions ended with an
+ underscore indicate **in-place** operations.
+
+ Nodes ``resultant`` from this operation are called ``"contract_edges_ip"``.
+
+ Returns
+ -------
+ Node
+
+ Examples
+ --------
+ >>> nodeA = tk.randn(shape=(10, 15, 20),
+ ... axes_names=('one', 'two', 'three'),
+ ... name='nodeA')
+ >>> nodeB = tk.randn(shape=(10, 15, 20),
+ ... axes_names=('one', 'two', 'three'),
+ ... name='nodeB')
+ ...
+ >>> _ = nodeA['one'] ^ nodeB['one']
+ >>> _ = nodeA['two'] ^ nodeB['two']
+ >>> _ = nodeA['three'] ^ nodeB['three']
+ >>> result = nodeA['one'].contract_()
+ >>> result.shape
+ torch.Size([15, 20, 15, 20])
+
+ ``nodeA`` and ``nodeB`` have been removed from the network.
+
+ >>> nodeA.network is None
+ True
+
+ >>> nodeB.network is None
+ True
+
+ >>> del nodeA
+ >>> del nodeB
+ """
+
+Edge.contract_=contract_edge_defget_shared_edges(node1:AbstractNode,
@@ -3066,10 +3895,12 @@
Source code for tensorkrowch.operations
""" Contracts all edges shared between two nodes. Batch contraction is automatically performed when both nodes have batch edges with the same
- names.
+ names. It can also be performed using the operator ``@``. Nodes ``resultant`` from this operation are called ``"contract_edges"``. The node that keeps information about the :class:`Successor` is ``node1``.
+
+ This operation is the same as :meth:`~AbstractNode.contract_between`. Parameters ----------
@@ -3098,7 +3929,7 @@
underscore indicate **in-place** operations. Nodes ``resultant`` from this operation are called ``"contract_edges_ip"``.
+
+ This operation is the same as :meth:`~AbstractNode.contract_between_`. Parameters ----------
@@ -3189,13 +4022,15 @@
Source code for tensorkrowch.operations
>>> del nodeB """result=contract_between(node1,node2)
- result.reattach_edges(True)
+ result.reattach_edges(override=True)result._unrestricted_set_tensor(result.tensor.detach())
+
+ nodes=set([node1,node2])# Delete nodes (and their edges) from the TNnet=result.network
- net.delete_node(node1)
- net.delete_node(node2)
+ fornodeinnodes:
+ net.delete_node(node)# Add edges of result to the TNforres_edgeinresult._edges:
@@ -3205,9 +4040,9 @@
##################################### STACK ###############################
+# MARK: stackdef_check_first_stack(nodes:Sequence[AbstractNode])->Optional[Successor]:ifnotnodes:raiseValueError('`nodes` should be a non-empty sequence of nodes')
@@ -3288,7 +4124,7 @@
Source code for tensorkrowch.operations
stack_indices=[]# In the case above, stack indices of each node in# the reference node's memory
- ifnot(isinstance(nodes,(list,tuple))andisinstance(nodes[0],AbstractNode)):
+ ifnot(isinstance(nodes,Sequence)andisinstance(nodes[0],AbstractNode)):raiseTypeError('`nodes` should be a list or tuple of AbstractNodes')net=nodes[0]._network
@@ -3316,15 +4152,19 @@
Nodes ``resultant`` from this operation are called ``"stack"``. If this operation returns a ``virtual`` :class:`ParamStackNode`, it will be called
- ``"virtual_stack"``. See :class:AbstractNode` to learn about this **reserved
- name**. The node that keeps information about the :class:`Successor` is
- ``nodes[0]``, the first stacked node.
+ ``"virtual_result_stack"``. See :class:AbstractNode` to learn about this
+ **reserved name**. The node that keeps information about the
+ :class:`Successor` is ``nodes[0]``, the first stacked node. Parameters ----------
- nodes : list[Node or ParamNode] or tuple[Node or ParamNode]
+ nodes : list[AbstractNode] or tuple[AbstractNode] Sequence of nodes that are to be stacked. They must be of the same type, have the same rank and axes names, be in the same tensor network, and have edges with the same types.
@@ -3510,6 +4348,7 @@
# If node is indexing from the original stackforj,(aux_slice,dim)inenumerate(zip(node_index[1:],new_node._shape)):
- ifnew_node._axes[j].is_batch():
+ ifnew_node._axes[j]._batch:# Admit any size in batch edgesindex.append(slice(0,None))else:
@@ -3624,15 +4467,17 @@
Source code for tensorkrowch.operations
else:# If node has the same shape as the original stack
- forj,(max_dim,dim)inenumerate(zip(node.shape[1:],
+ forj,(max_dim,dim)inenumerate(zip(node._shape[1:],new_node._shape)):
- ifnew_node._axes[j].is_batch():
+ ifnew_node._axes[j]._batch:# Admit any size in batch edgesindex.append(slice(0,None))else:index.append(slice(max_dim-dim,max_dim))
- new_node._tensor_info['index']=index
+ iflen(index)==1:
+ index=index[0]
+ new_node._tensor_info['index']=index# Create successorargs=(node,)
@@ -3761,7 +4606,64 @@
Source code for tensorkrowch.operations
returnunbind_op(node)
+unbind_node=copy_func(unbind)
+unbind_node.__doc__= \
+"""
+ Unbinds a :class:`StackNode` or :class:`ParamStackNode`, where the first
+ dimension is assumed to be the stack dimension.
+
+ If :meth:`~TensorNetwork.auto_unbind` is set to ``False``, each resultant
+ node will store its own tensor. Otherwise, they will have only a reference
+ to the corresponding slice of the ``(Param)StackNode``.
+
+ See :class:`TensorNetwork` to learn how the ``auto_unbind`` mode affects
+ the computation of :func:`unbind`.
+
+ Nodes ``resultant`` from this operation are called ``"unbind"``. The node
+ that keeps information about the :class:`Successor` is ``self``.
+
+ Returns
+ -------
+ list[Node]
+
+ Examples
+ --------
+ >>> net = tk.TensorNetwork()
+ >>> nodes = [tk.randn(shape=(2, 4, 2),
+ ... axes_names=('left', 'input', 'right'),
+ ... network=net)
+ ... for _ in range(10)]
+ >>> data = [tk.randn(shape=(4,),
+ ... axes_names=('feature',),
+ ... network=net)
+ ... for _ in range(10)]
+ ...
+ >>> for i in range(10):
+ ... _ = nodes[i]['input'] ^ data[i]['feature']
+ ...
+ >>> stack_nodes = tk.stack(nodes)
+ >>> stack_data = tk.stack(data)
+ ...
+ >>> # It is necessary to re-connect stacks
+ >>> _ = stack_nodes['input'] ^ stack_data['feature']
+ >>> result = stack_nodes @ stack_data
+ >>> result = result.unbind()
+ >>> print(result[0].name)
+ unbind_0
+
+ >>> result[0].axes
+ [Axis( left (0) ), Axis( right (1) )]
+
+ >>> result[0].shape
+ torch.Size([2, 2])
+ """
+
+StackNode.unbind=unbind_node
+ParamStackNode.unbind=unbind_node
+
+
################################## EINSUM #################################
+# MARK: einsumdef_check_first_einsum(string:Text,*nodes:AbstractNode)->Optional[Successor]:ifnotnodes:
diff --git a/docs/_build/html/_sources/api.rst.txt b/docs/_build/html/_sources/api.rst.txt
index 8938481..696b184 100644
--- a/docs/_build/html/_sources/api.rst.txt
+++ b/docs/_build/html/_sources/api.rst.txt
@@ -8,4 +8,5 @@ API Reference
operations
models
initializers
- embeddings
\ No newline at end of file
+ embeddings
+ decompositions
\ No newline at end of file
diff --git a/docs/_build/html/_sources/decompositions.rst.txt b/docs/_build/html/_sources/decompositions.rst.txt
new file mode 100644
index 0000000..bd31fa9
--- /dev/null
+++ b/docs/_build/html/_sources/decompositions.rst.txt
@@ -0,0 +1,12 @@
+Decompositions
+==============
+
+.. currentmodule:: tensorkrowch.decompositions
+
+vec_to_mps
+^^^^^^^^^^
+.. autofunction:: vec_to_mps
+
+mat_to_mpo
+^^^^^^^^^^
+.. autofunction:: mat_to_mpo
\ No newline at end of file
diff --git a/docs/_build/html/_sources/embeddings.rst.txt b/docs/_build/html/_sources/embeddings.rst.txt
index 17dcb74..fb1e511 100644
--- a/docs/_build/html/_sources/embeddings.rst.txt
+++ b/docs/_build/html/_sources/embeddings.rst.txt
@@ -1,3 +1,5 @@
+.. _Embeddings:
+
Embeddings
==========
@@ -13,4 +15,12 @@ add_ones
poly
^^^^
-.. autofunction:: poly
\ No newline at end of file
+.. autofunction:: poly
+
+discretize
+^^^^^^^^^^
+.. autofunction:: discretize
+
+basis
+^^^^^
+.. autofunction:: basis
\ No newline at end of file
diff --git a/docs/_build/html/_sources/models.rst.txt b/docs/_build/html/_sources/models.rst.txt
index 1491129..5c51f80 100644
--- a/docs/_build/html/_sources/models.rst.txt
+++ b/docs/_build/html/_sources/models.rst.txt
@@ -4,8 +4,18 @@ Models
.. currentmodule:: tensorkrowch.models
-MPSLayer
---------
+MPS
+---
+
+MPS
+^^^
+.. autoclass:: MPS
+ :members:
+
+UMPS
+^^^^
+.. autoclass:: UMPS
+ :members:
MPSLayer
^^^^^^^^
@@ -17,38 +27,55 @@ UMPSLayer
.. autoclass:: UMPSLayer
:members:
+ConvMPS
+^^^^^^^
+.. autoclass:: ConvMPS
+ :members:
+
+ .. automethod:: forward
+
+ConvUMPS
+^^^^^^^^
+.. autoclass:: ConvUMPS
+ :members:
+
+ .. automethod:: forward
+
ConvMPSLayer
^^^^^^^^^^^^
.. autoclass:: ConvMPSLayer
:members:
+ .. automethod:: forward
+
ConvUMPSLayer
^^^^^^^^^^^^^
.. autoclass:: ConvUMPSLayer
:members:
+ .. automethod:: forward
-MPS
----
-MPS
-^^^
-.. autoclass:: MPS
- :members:
+MPSData
+-------
-UMPS
-^^^^
-.. autoclass:: UMPS
+MPSData
+^^^^^^^
+.. autoclass:: MPSData
:members:
-ConvMPS
-^^^^^^^
-.. autoclass:: ConvMPS
+
+MPO
+---
+
+MPO
+^^^
+.. autoclass:: MPO
:members:
-ConvUMPS
-^^^^^^^^
-.. autoclass:: ConvUMPS
+UMPO
+^^^^
+.. autoclass:: UMPO
:members:
diff --git a/docs/_build/html/_sources/operations.rst.txt b/docs/_build/html/_sources/operations.rst.txt
index 7c670f4..6643a96 100644
--- a/docs/_build/html/_sources/operations.rst.txt
+++ b/docs/_build/html/_sources/operations.rst.txt
@@ -19,22 +19,42 @@ disconnect
^^^^^^^^^^
.. autofunction:: disconnect
+svd
+^^^
+.. autofunction:: svd
+
svd\_
^^^^^
.. autofunction:: svd_
+svdr
+^^^^
+.. autofunction:: svdr
+
svdr\_
^^^^^^
.. autofunction:: svdr_
+qr
+^^
+.. autofunction:: qr
+
qr\_
^^^^
.. autofunction:: qr_
+rq
+^^
+.. autofunction:: rq
+
rq\_
^^^^
.. autofunction:: rq_
+contract
+^^^^^^^^
+.. autofunction:: contract
+
contract\_
^^^^^^^^^^
.. autofunction:: contract_
diff --git a/docs/_build/html/_sources/tutorials/0_first_steps.rst.txt b/docs/_build/html/_sources/tutorials/0_first_steps.rst.txt
index 120ada5..05997e1 100644
--- a/docs/_build/html/_sources/tutorials/0_first_steps.rst.txt
+++ b/docs/_build/html/_sources/tutorials/0_first_steps.rst.txt
@@ -73,7 +73,8 @@ First of all, we need to import the necessary libraries::
# Model parameters
image_size = (28, 28)
- in_channels = 3
+ in_dim = 3
+ out_dim = 10
bond_dim = 10
@@ -126,7 +127,7 @@ Put **MNIST** into ``DataLoaders``::
We are going to train a Matrix Product State (MPS) model. ``TensorKrowch`` comes
with some built-in models like ``MPSLayer``, which is a MPS with one output node
-with a dangling edge. Hence, when the whole tensor netwok gets contracted, we
+with a dangling edge. Hence, when the whole tensor network gets contracted, we
obtain a vector with the probabilities that an image belongs to one of the 10
possible classes.
@@ -136,10 +137,12 @@ possible classes.
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# Instantiate model
- mps = tk.models.MPSLayer(n_features=image_size[0] * image_size[1],
+ mps = tk.models.MPSLayer(n_features=image_size[0] * image_size[1] + 1,
in_dim=in_dim,
- out_dim=10,
- bond_dim=bond_dim)
+ out_dim=out_dim,
+ bond_dim=bond_dim,
+ init_method='randn_eye',
+ std=1e-9)
# Send model to GPU
mps = mps.to(device)
@@ -224,16 +227,16 @@ We use a common training loop used when training neural networks in ``PyTorch``:
f'Train. Acc.: {running_train_acc / num_batches["train"]:.4f}, '
f'Test Acc.: {running_test_acc / num_batches["test"]:.4f}')
- # * Epoch 1: Train. Loss: 0.9456, Train. Acc.: 0.6752, Test Acc.: 0.8924
- # * Epoch 2: Train. Loss: 0.2921, Train. Acc.: 0.9122, Test Acc.: 0.9360
- # * Epoch 3: Train. Loss: 0.2066, Train. Acc.: 0.9378, Test Acc.: 0.9443
- # * Epoch 4: Train. Loss: 0.1642, Train. Acc.: 0.9502, Test Acc.: 0.9595
- # * Epoch 5: Train. Loss: 0.1317, Train. Acc.: 0.9601, Test Acc.: 0.9632
- # * Epoch 6: Train. Loss: 0.1135, Train. Acc.: 0.9654, Test Acc.: 0.9655
- # * Epoch 7: Train. Loss: 0.1046, Train. Acc.: 0.9687, Test Acc.: 0.9669
- # * Epoch 8: Train. Loss: 0.0904, Train. Acc.: 0.9720, Test Acc.: 0.9723
- # * Epoch 9: Train. Loss: 0.0836, Train. Acc.: 0.9740, Test Acc.: 0.9725
- # * Epoch 10: Train. Loss: 0.0751, Train. Acc.: 0.9764, Test Acc.: 0.9748
+ # * Epoch 1: Train. Loss: 1.1955, Train. Acc.: 0.5676, Test Acc.: 0.8820
+ # * Epoch 2: Train. Loss: 0.3083, Train. Acc.: 0.9053, Test Acc.: 0.9371
+ # * Epoch 3: Train. Loss: 0.1990, Train. Acc.: 0.9396, Test Acc.: 0.9509
+ # * Epoch 4: Train. Loss: 0.1573, Train. Acc.: 0.9526, Test Acc.: 0.9585
+ # * Epoch 5: Train. Loss: 0.1308, Train. Acc.: 0.9600, Test Acc.: 0.9621
+ # * Epoch 6: Train. Loss: 0.1123, Train. Acc.: 0.9668, Test Acc.: 0.9625
+ # * Epoch 7: Train. Loss: 0.0998, Train. Acc.: 0.9696, Test Acc.: 0.9677
+ # * Epoch 8: Train. Loss: 0.0913, Train. Acc.: 0.9721, Test Acc.: 0.9729
+ # * Epoch 9: Train. Loss: 0.0820, Train. Acc.: 0.9743, Test Acc.: 0.9736
+ # * Epoch 10: Train. Loss: 0.0728, Train. Acc.: 0.9775, Test Acc.: 0.9734
7. Prune the Model
@@ -250,8 +253,8 @@ Let's count how many parameters our model has before pruning::
print(f'Nº params: {n_params}')
print(f'Memory module: {memory / 1024**2:.4f} MB') # MegaBytes
- # Nº params: 235660
- # Memory module: 0.8990 MB
+ # Nº params: 236220
+ # Memory module: 0.9011 MB
To prune the model, we take the **canonical form** of the ``MPSLayer``. In this
process, many Singular Value Decompositions are performed in the network. By
@@ -262,10 +265,19 @@ that we are losing a lot of uninformative (useless) parameters.
# Canonicalize SVD
# ----------------
- mps.canonicalize(cum_percentage=0.98)
- mps.trace(torch.zeros(1, image_size[0] * image_size[1], in_dim).to(device))
+ mps.canonicalize(cum_percentage=0.99)
+
+ # Number of parametrs
+ n_params = 0
+ memory = 0
+ for p in mps.parameters():
+ n_params += p.nelement()
+ memory += p.nelement() * p.element_size() # Bytes
+ print(f'Nº params: {n_params}')
+ print(f'Memory module: {memory / 1024**2:.4f} MB\n') # MegaBytes
# New test accuracy
+ mps.trace(torch.zeros(1, image_size[0] * image_size[1], in_dim).to(device))
with torch.no_grad():
running_acc = 0.0
@@ -282,19 +294,10 @@ that we are losing a lot of uninformative (useless) parameters.
print(f'Test Acc.: {running_acc / num_batches["test"]:.4f}\n')
- # Number of parametrs
- n_params = 0
- memory = 0
- for p in mps.parameters():
- n_params += p.nelement()
- memory += p.nelement() * p.element_size() # Bytes
- print(f'Nº params: {n_params}')
- print(f'Memory module: {memory / 1024**2:.4f} MB\n') # MegaBytes
-
- # Test Acc.: 0.9194
+ # Nº params: 161204
+ # Memory module: 0.6149 MB
- # Nº params: 150710
- # Memory module: 0.5749 MB
+ # Test Acc.: 0.9400
We could continue training to improve performance after pruning, and pruning
again, until we reach an `equilibrium` point::
@@ -349,9 +352,9 @@ again, until we reach an `equilibrium` point::
f'Train. Acc.: {running_train_acc / num_batches["train"]:.4f}, '
f'Test Acc.: {running_test_acc / num_batches["test"]:.4f}')
- # * Epoch 1: Train. Loss: 0.1018, Train. Acc.: 0.9684, Test Acc.: 0.9693
- # * Epoch 2: Train. Loss: 0.0815, Train. Acc.: 0.9745, Test Acc.: 0.9699
- # * Epoch 3: Train. Loss: 0.0716, Train. Acc.: 0.9777, Test Acc.: 0.9721
+ # * Epoch 1: Train. Loss: 0.0983, Train. Acc.: 0.9700, Test Acc.: 0.9738
+ # * Epoch 2: Train. Loss: 0.0750, Train. Acc.: 0.9768, Test Acc.: 0.9743
+ # * Epoch 3: Train. Loss: 0.0639, Train. Acc.: 0.9793, Test Acc.: 0.9731
-After all the pruning an re-training, we have reduced around 36% of the total
-amount of parameters losing less than 0.3% in accuracy.
+After all the pruning an re-training, we have reduced around 32% of the total
+amount of parameters without losing accuracy.
diff --git a/docs/_build/html/_sources/tutorials/2_contracting_tensor_network.rst.txt b/docs/_build/html/_sources/tutorials/2_contracting_tensor_network.rst.txt
index 85939b5..f42b503 100644
--- a/docs/_build/html/_sources/tutorials/2_contracting_tensor_network.rst.txt
+++ b/docs/_build/html/_sources/tutorials/2_contracting_tensor_network.rst.txt
@@ -192,7 +192,8 @@ Regarding the **node-like** operations, these are:
stack_node = tk.stack(nodes)
stack_data_node = tk.stack(data_nodes)
- stack_node['input'] ^ stack_data_node['feature']
+ # reconnect stacks
+ stack_node ^ stack_data_node
4) :func:`unbind`: Unbinds a :class:`StackNode` and returns a list of nodes that
are already connected to the corresponding neighbours::
@@ -281,7 +282,7 @@ contractions, which will save us some time::
stack_node = tk.stack(nodes)
stack_data_node = tk.stack(data_nodes)
- stack_node['input'] ^ stack_data_node['feature']
+ stack_node ^ stack_data_node
stack_result = stack_node @ stack_data_node
unbind_result = tk.unbind(stack_result)
diff --git a/docs/_build/html/_sources/tutorials/3_memory_management.rst.txt b/docs/_build/html/_sources/tutorials/3_memory_management.rst.txt
index 3abcb7c..c36aa04 100644
--- a/docs/_build/html/_sources/tutorials/3_memory_management.rst.txt
+++ b/docs/_build/html/_sources/tutorials/3_memory_management.rst.txt
@@ -129,8 +129,8 @@ Product State::
assert node.tensor_address() == 'uniform'
-2. How TensorKrowch skipps Operations to run faster
----------------------------------------------------
+2. How TensorKrowch skips Operations to run faster
+--------------------------------------------------
The main purpose of ``TensorKrowch`` is enabling you to experiment
creating and training different tensor networks, only having to worry about
diff --git a/docs/_build/html/_sources/tutorials/4_types_of_nodes.rst.txt b/docs/_build/html/_sources/tutorials/4_types_of_nodes.rst.txt
index a25b488..a3c5ff6 100644
--- a/docs/_build/html/_sources/tutorials/4_types_of_nodes.rst.txt
+++ b/docs/_build/html/_sources/tutorials/4_types_of_nodes.rst.txt
@@ -123,7 +123,7 @@ different roles in the ``TensorNetwork``:
for node in nodes:
assert node.tensor_address() == 'virtual_uniform'
- giving the ``uniform_node`` the role of ``virtual`` makes more sense,
+ Giving the ``uniform_node`` the role of ``virtual`` makes more sense,
since it is a node that one wouldn't desire to see as a ``leaf`` node
of the network. Instead it is `hidden`.
@@ -139,8 +139,7 @@ different roles in the ``TensorNetwork``:
They are intermediate nodes that (almost always) inherit edges from ``leaf``
and ``data`` nodes, the ones that really form the network. These nodes can
store their own tensors or use other node's tensor. The names of the
- ``resultant`` nodes are the name of the ``Operation`` that originated
- it::
+ ``resultant`` nodes are the name of the ``Operation`` that originated it::
node1 = tk.randn(shape=(2, 3))
node2 = tk.randn(shape=(3, 4))
@@ -191,11 +190,12 @@ Other thing one should take into account are **reserved nodes' names**:
# Batch edge has size 1 when created
assert net['stack_data_memory'].shape == (100, 1, 5)
-* **"virtual_stack"**: Name of the ``virtual`` :class:`ParamStackNode` that
- results from stacking ``ParamNodes`` as the first operation in the network
- contraction, if ``auto_stack`` mode is set to ``True``. There might be as
- much ``"virtual_stack"`` nodes as stacks are created from ``ParamNodes``. To
- learn more about this, see :class:`ParamStackNode`.
+* **"virtual_result"**: Name of ``virtual`` nodes that are not explicitly
+ part of the network, but are required for some situations during contraction.
+ For instance, the :class:`ParamStackNode` that results from stacking
+ :class:`ParamNodes ` as the first operation in the network
+ contraction, if ``auto_stack`` mode is set to ``True``. To learn more about
+ this, see :class:`ParamStackNode`.
::
@@ -213,7 +213,7 @@ Other thing one should take into account are **reserved nodes' names**:
# All ParamNodes use a slice of the tensor in stack_node
for node in nodes:
- assert node.tensor_address() == 'virtual_stack'
+ assert node.tensor_address() == 'virtual_result_stack'
* **"virtual_uniform"**: Name of the ``virtual`` ``Node`` or ``ParamNode`` that
is used in uniform (translationally invariant) tensor networks to store the
@@ -221,6 +221,10 @@ Other thing one should take into account are **reserved nodes' names**:
``"virtual_uniform"`` nodes as shared memories are used for the ``leaf``
nodes in the network (usually just one). An example of this can be seen in
the previous section, when ``virtual`` nodes were defined.
+
+For ``"virtual_result"`` and ``"virtual_uniform"``, these special behaviours
+are not restricted to nodes having those names, but also nodes whose names
+contain those strings.
Although these names can in principle be used for other nodes, this can lead
to undesired behaviour.
diff --git a/docs/_build/html/_sources/tutorials/5_subclass_tensor_network.rst.txt b/docs/_build/html/_sources/tutorials/5_subclass_tensor_network.rst.txt
index 1e0d8e5..24b4148 100644
--- a/docs/_build/html/_sources/tutorials/5_subclass_tensor_network.rst.txt
+++ b/docs/_build/html/_sources/tutorials/5_subclass_tensor_network.rst.txt
@@ -207,7 +207,7 @@ images::
stack_input = tk.stack(self.input_nodes)
stack_data = tk.stack(list(self.data_nodes.values()))
- stack_input['input'] ^ stack_data['feature']
+ stack_input ^ stack_data
stack_result = stack_input @ stack_data
stack_result = tk.unbind(stack_result)
@@ -231,7 +231,7 @@ Now we can instantiate our model::
Since our model is a subclass of ``torch.nn.Module``, we can take advantage of
its methods. For instance, we can easily send the model to the GPU::
- device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+ device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
mps = mps.to(device)
Note that only the ``ParamNodes`` are parameters of the model. Thus if your
@@ -313,6 +313,7 @@ already comes with a handful of widely-known models that you can use:
* :class:`MPS`
* :class:`MPSLayer`
+* :class:`MPO`
* :class:`PEPS`
* :class:`Tree`
diff --git a/docs/_build/html/_sources/tutorials/6_mix_with_pytorch.rst.txt b/docs/_build/html/_sources/tutorials/6_mix_with_pytorch.rst.txt
index ba2f645..0788167 100644
--- a/docs/_build/html/_sources/tutorials/6_mix_with_pytorch.rst.txt
+++ b/docs/_build/html/_sources/tutorials/6_mix_with_pytorch.rst.txt
@@ -57,11 +57,11 @@ Now we can define the model::
in_channels=7,
bond_dim=bond_dim,
out_channels=10,
- kernel_size=image_size[0] // 2)
+ kernel_size=image_size[0] // 2,
+ init_method='randn_eye',
+ std=1e-9)
self.layers.append(mps)
- self.softmax = nn.Softmax(dim=1)
-
@staticmethod
def embedding(x):
ones = torch.ones_like(x[:, 0]).unsqueeze(1)
@@ -99,7 +99,7 @@ Now we set the parameters for the training algorithm and our model::
Initialize our model and send it to the appropiate device::
- device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+ device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
cnn_snake = CNN_SnakeSBS(in_channels, bond_dim, image_size)
cnn_snake = cnn_snake.to(device)
@@ -210,15 +210,15 @@ Let the training begin!
f'Train. Loss: {running_train_loss / num_batches["train"]:.4f}, '
f'Train. Acc.: {running_train_acc / num_batches["train"]:.4f}, '
f'Test Acc.: {running_test_acc / num_batches["test"]:.4f}')
-
- # * Epoch 10: Train. Loss: 0.3824, Train. Acc.: 0.8599, Test Acc.: 0.8570
- # * Epoch 20: Train. Loss: 0.3245, Train. Acc.: 0.8814, Test Acc.: 0.8758
- # * Epoch 30: Train. Loss: 0.2924, Train. Acc.: 0.8915, Test Acc.: 0.8829
- # * Epoch 40: Train. Loss: 0.2694, Train. Acc.: 0.8993, Test Acc.: 0.8884
- # * Epoch 50: Train. Loss: 0.2463, Train. Acc.: 0.9078, Test Acc.: 0.8860
- # * Epoch 60: Train. Loss: 0.2257, Train. Acc.: 0.9163, Test Acc.: 0.8958
- # * Epoch 70: Train. Loss: 0.2083, Train. Acc.: 0.9219, Test Acc.: 0.8969
- # * Epoch 80: Train. Loss: 0.2013, Train. Acc.: 0.9226, Test Acc.: 0.8979
+
+ # * Epoch 10: Train. Loss: 0.3714, Train. Acc.: 0.8627, Test Acc.: 0.8502
+ # * Epoch 20: Train. Loss: 0.3149, Train. Acc.: 0.8851, Test Acc.: 0.8795
+ # * Epoch 30: Train. Loss: 0.2840, Train. Acc.: 0.8948, Test Acc.: 0.8848
+ # * Epoch 40: Train. Loss: 0.2618, Train. Acc.: 0.9026, Test Acc.: 0.8915
+ # * Epoch 50: Train. Loss: 0.2357, Train. Acc.: 0.9125, Test Acc.: 0.8901
+ # * Epoch 60: Train. Loss: 0.2203, Train. Acc.: 0.9174, Test Acc.: 0.9009
+ # * Epoch 70: Train. Loss: 0.2052, Train. Acc.: 0.9231, Test Acc.: 0.8984
+ # * Epoch 80: Train. Loss: 0.1878, Train. Acc.: 0.9284, Test Acc.: 0.9011
Wow! That's almost 90% accuracy with just the first model we try!
@@ -227,13 +227,13 @@ Let's check how many parameters our model has::
# Original number of parametrs
n_params = 0
memory = 0
- for p in mps.parameters():
+ for p in cnn_snake.parameters():
n_params += p.nelement()
memory += p.nelement() * p.element_size() # Bytes
print(f'Nº params: {n_params}')
print(f'Memory module: {memory / 1024**2:.4f} MB') # MegaBytes
- # Nº params: 136940
+ # Nº params: 553186
# Memory module: 0.5224 MB
Since we are using tensor networks we can **prune** our model in 4 lines of
@@ -250,7 +250,21 @@ singular values, reducing the sizes of the edges in our network::
Let's see how much our model has changed after pruning with **canonical forms**::
+ # Number of parametrs
+ n_params = 0
+ memory = 0
+ for p in mps.parameters():
+ n_params += p.nelement()
+ memory += p.nelement() * p.element_size() # Bytes
+ print(f'Nº params: {n_params}')
+ print(f'Memory module: {memory / 1024**2:.4f} MB\n') # MegaBytes
+
# New test accuracy
+ for mps in cnn_snake.layers:
+ # Since the nodes are different now, we have to re-trace
+ mps.trace(torch.zeros(
+ 1, 7, image_size[0]//2, image_size[1]//2).to(device))
+
with torch.no_grad():
running_test_acc = 0.0
@@ -267,19 +281,7 @@ Let's see how much our model has changed after pruning with **canonical forms**:
print(f'Test Acc.: {running_test_acc / num_batches["test"]:.4f}\n')
- # Test Acc.: 0.8908
-
- # Number of parametrs
- n_params = 0
- memory = 0
- for p in mps.parameters():
- n_params += p.nelement()
- memory += p.nelement() * p.element_size() # Bytes
- print(f'Nº params: {n_params}')
- print(f'Memory module: {memory / 1024**2:.4f} MB\n') # MegaBytes
-
- # Nº params: 110753
- # Memory module: 0.4225 MB
+ # Nº params: 499320
+ # Memory module: 1.9048 MB
-We have reduced around 20% of the total amount of parameters losing less than
-1% in accuracy.
+ # Test Acc.: 0.8968
diff --git a/docs/_build/html/_static/_sphinx_javascript_frameworks_compat.js b/docs/_build/html/_static/_sphinx_javascript_frameworks_compat.js
deleted file mode 100644
index 8549469..0000000
--- a/docs/_build/html/_static/_sphinx_javascript_frameworks_compat.js
+++ /dev/null
@@ -1,134 +0,0 @@
-/*
- * _sphinx_javascript_frameworks_compat.js
- * ~~~~~~~~~~
- *
- * Compatability shim for jQuery and underscores.js.
- *
- * WILL BE REMOVED IN Sphinx 6.0
- * xref RemovedInSphinx60Warning
- *
- */
-
-/**
- * select a different prefix for underscore
- */
-$u = _.noConflict();
-
-
-/**
- * small helper function to urldecode strings
- *
- * See https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/decodeURIComponent#Decoding_query_parameters_from_a_URL
- */
-jQuery.urldecode = function(x) {
- if (!x) {
- return x
- }
- return decodeURIComponent(x.replace(/\+/g, ' '));
-};
-
-/**
- * small helper function to urlencode strings
- */
-jQuery.urlencode = encodeURIComponent;
-
-/**
- * This function returns the parsed url parameters of the
- * current request. Multiple values per key are supported,
- * it will always return arrays of strings for the value parts.
- */
-jQuery.getQueryParameters = function(s) {
- if (typeof s === 'undefined')
- s = document.location.search;
- var parts = s.substr(s.indexOf('?') + 1).split('&');
- var result = {};
- for (var i = 0; i < parts.length; i++) {
- var tmp = parts[i].split('=', 2);
- var key = jQuery.urldecode(tmp[0]);
- var value = jQuery.urldecode(tmp[1]);
- if (key in result)
- result[key].push(value);
- else
- result[key] = [value];
- }
- return result;
-};
-
-/**
- * highlight a given string on a jquery object by wrapping it in
- * span elements with the given class name.
- */
-jQuery.fn.highlightText = function(text, className) {
- function highlight(node, addItems) {
- if (node.nodeType === 3) {
- var val = node.nodeValue;
- var pos = val.toLowerCase().indexOf(text);
- if (pos >= 0 &&
- !jQuery(node.parentNode).hasClass(className) &&
- !jQuery(node.parentNode).hasClass("nohighlight")) {
- var span;
- var isInSVG = jQuery(node).closest("body, svg, foreignObject").is("svg");
- if (isInSVG) {
- span = document.createElementNS("http://www.w3.org/2000/svg", "tspan");
- } else {
- span = document.createElement("span");
- span.className = className;
- }
- span.appendChild(document.createTextNode(val.substr(pos, text.length)));
- node.parentNode.insertBefore(span, node.parentNode.insertBefore(
- document.createTextNode(val.substr(pos + text.length)),
- node.nextSibling));
- node.nodeValue = val.substr(0, pos);
- if (isInSVG) {
- var rect = document.createElementNS("http://www.w3.org/2000/svg", "rect");
- var bbox = node.parentElement.getBBox();
- rect.x.baseVal.value = bbox.x;
- rect.y.baseVal.value = bbox.y;
- rect.width.baseVal.value = bbox.width;
- rect.height.baseVal.value = bbox.height;
- rect.setAttribute('class', className);
- addItems.push({
- "parent": node.parentNode,
- "target": rect});
- }
- }
- }
- else if (!jQuery(node).is("button, select, textarea")) {
- jQuery.each(node.childNodes, function() {
- highlight(this, addItems);
- });
- }
- }
- var addItems = [];
- var result = this.each(function() {
- highlight(this, addItems);
- });
- for (var i = 0; i < addItems.length; ++i) {
- jQuery(addItems[i].parent).before(addItems[i].target);
- }
- return result;
-};
-
-/*
- * backward compatibility for jQuery.browser
- * This will be supported until firefox bug is fixed.
- */
-if (!jQuery.browser) {
- jQuery.uaMatch = function(ua) {
- ua = ua.toLowerCase();
-
- var match = /(chrome)[ \/]([\w.]+)/.exec(ua) ||
- /(webkit)[ \/]([\w.]+)/.exec(ua) ||
- /(opera)(?:.*version|)[ \/]([\w.]+)/.exec(ua) ||
- /(msie) ([\w.]+)/.exec(ua) ||
- ua.indexOf("compatible") < 0 && /(mozilla)(?:.*? rv:([\w.]+)|)/.exec(ua) ||
- [];
-
- return {
- browser: match[ 1 ] || "",
- version: match[ 2 ] || "0"
- };
- };
- jQuery.browser = {};
- jQuery.browser[jQuery.uaMatch(navigator.userAgent).browser] = true;
-}
diff --git a/docs/_build/html/_static/jquery-3.6.0.js b/docs/_build/html/_static/jquery-3.6.0.js
deleted file mode 100644
index fc6c299..0000000
--- a/docs/_build/html/_static/jquery-3.6.0.js
+++ /dev/null
@@ -1,10881 +0,0 @@
-/*!
- * jQuery JavaScript Library v3.6.0
- * https://jquery.com/
- *
- * Includes Sizzle.js
- * https://sizzlejs.com/
- *
- * Copyright OpenJS Foundation and other contributors
- * Released under the MIT license
- * https://jquery.org/license
- *
- * Date: 2021-03-02T17:08Z
- */
-( function( global, factory ) {
-
- "use strict";
-
- if ( typeof module === "object" && typeof module.exports === "object" ) {
-
- // For CommonJS and CommonJS-like environments where a proper `window`
- // is present, execute the factory and get jQuery.
- // For environments that do not have a `window` with a `document`
- // (such as Node.js), expose a factory as module.exports.
- // This accentuates the need for the creation of a real `window`.
- // e.g. var jQuery = require("jquery")(window);
- // See ticket #14549 for more info.
- module.exports = global.document ?
- factory( global, true ) :
- function( w ) {
- if ( !w.document ) {
- throw new Error( "jQuery requires a window with a document" );
- }
- return factory( w );
- };
- } else {
- factory( global );
- }
-
-// Pass this if window is not defined yet
-} )( typeof window !== "undefined" ? window : this, function( window, noGlobal ) {
-
-// Edge <= 12 - 13+, Firefox <=18 - 45+, IE 10 - 11, Safari 5.1 - 9+, iOS 6 - 9.1
-// throw exceptions when non-strict code (e.g., ASP.NET 4.5) accesses strict mode
-// arguments.callee.caller (trac-13335). But as of jQuery 3.0 (2016), strict mode should be common
-// enough that all such attempts are guarded in a try block.
-"use strict";
-
-var arr = [];
-
-var getProto = Object.getPrototypeOf;
-
-var slice = arr.slice;
-
-var flat = arr.flat ? function( array ) {
- return arr.flat.call( array );
-} : function( array ) {
- return arr.concat.apply( [], array );
-};
-
-
-var push = arr.push;
-
-var indexOf = arr.indexOf;
-
-var class2type = {};
-
-var toString = class2type.toString;
-
-var hasOwn = class2type.hasOwnProperty;
-
-var fnToString = hasOwn.toString;
-
-var ObjectFunctionString = fnToString.call( Object );
-
-var support = {};
-
-var isFunction = function isFunction( obj ) {
-
- // Support: Chrome <=57, Firefox <=52
- // In some browsers, typeof returns "function" for HTML
+
+
+