Adding MPO layer (#576)

* Adding MPO layer

* Improved network creation and contraction

* Adding get_config to layers and updating tests

* Adding SavedModel test and updating docstrings

* Updating get_config to idiomatic Keras

* Adding decorators for serializable layers

* Fixing errors

* Small fix

* Decorator error

* Investigating global custom objects error

Author: Ben Penchas

tensornetwork/tn_keras/__init__.py

@@ -1 +1,2 @@
-from tensornetwork.tn_keras.dense import DenseDecomp
+# from tensornetwork.tn_keras.dense import DenseDecomp
+# from tensornetwork.tn_keras.mpo import DenseMPO

tensornetwork/tn_keras/dense.py

@@ -7,6 +7,8 @@
 import numpy as np
 
 
+@tf.keras.utils.register_keras_serializable() # type: ignore
+    # package='tensornetwork', name='dense_decomp')
 class DenseDecomp(Layer):
   """TN layer comparable to Dense that carries out matrix multiplication
   with 2 significantly smaller weight matrices instead of 1 large one.
@@ -19,10 +21,10 @@ class DenseDecomp(Layer):
   # as first layer in a sequential model:
   model = Sequential()
   model.add(
-    DenseDecomp(512, decomp_size=128, activation='relu', input_dim=1024))
+    DenseDecomp(512, decomp_size=128, activation='relu', input_shape=(1024,)))
   # now the model will take as input arrays of shape (*, 1024)
-  # and output arrays of shape (*, 512)
-  # after the first layer, you don't need to specify
+  # and output arrays of shape (*, 512).
+  # After the first layer, you don't need to specify
   # the size of the input anymore:
   model.add(DenseDecomp(512, decomp_size=128, activation='relu'))
   ```
@@ -39,7 +41,7 @@ class DenseDecomp(Layer):
     bias_initializer: Initializer for the bias vector.
 
   Input shape:
-    2D tensor with shape: `(batch_size, input_dim)`.
+    2D tensor with shape: `(batch_size, input_shape[-1])`.
 
   Output shape:
     2D tensor with shape: `(batch_size, output_dim)`.
@@ -53,16 +55,21 @@ def __init__(self,
                kernel_initializer: Optional[Text] = 'glorot_uniform',
                bias_initializer: Optional[Text] = 'zeros',
                **kwargs) -> None:
+
+    # Allow specification of input_dim instead of input_shape,
+    # for compatability with Keras layers that support this
     if 'input_shape' not in kwargs and 'input_dim' in kwargs:
       kwargs['input_shape'] = (kwargs.pop('input_dim'),)
+
+    super(DenseDecomp, self).__init__(**kwargs)
+
     self.output_dim = output_dim
     self.decomp_size = decomp_size
 
     self.use_bias = use_bias
     self.activation = activations.get(activation)
     self.kernel_initializer = initializers.get(kernel_initializer)
     self.bias_initializer = initializers.get(bias_initializer)
-    super(DenseDecomp, self).__init__(**kwargs)
 
   def build(self, input_shape: List[int]) -> None:
     # Disable the attribute-defined-outside-init violations in this function
@@ -71,24 +78,24 @@ def build(self, input_shape: List[int]) -> None:
       raise ValueError('The last dimension of the inputs to `Dense` '
                        'should be defined. Found `None`.')
 
-    self.a_var = self.add_weight(
-        name='a',
-        shape=(input_shape[-1], self.decomp_size),
-        trainable=True,
-        initializer=self.kernel_initializer)
-    self.b_var = self.add_weight(
-        name='b',
-        shape=(self.decomp_size, self.output_dim),
-        trainable=True,
-        initializer=self.kernel_initializer)
+    super(DenseDecomp, self).build(input_shape)
+
+    self.a_var = self.add_weight(name='a',
+                                 shape=(input_shape[-1], self.decomp_size),
+                                 trainable=True,
+                                 initializer=self.kernel_initializer)
+    self.b_var = self.add_weight(name='b',
+                                 shape=(self.decomp_size, self.output_dim),
+                                 trainable=True,
+                                 initializer=self.kernel_initializer)
     self.bias_var = self.add_weight(
         name='bias',
         shape=(self.output_dim,),
         trainable=True,
         initializer=self.bias_initializer) if self.use_bias else None
-    super(DenseDecomp, self).build(input_shape)
 
-  def call(self, inputs: tf.Tensor, **kwargs) -> tf.Tensor:  # pylint: disable=unused-argument
+  def call(self, inputs: tf.Tensor, **kwargs) -> tf.Tensor: # pylint: disable=unused-argument
+
 
     def f(x: tf.Tensor, a_var: tf.Tensor, b_var: tf.Tensor, use_bias: bool,
           bias_var: tf.Tensor) -> tf.Tensor:
@@ -124,3 +131,32 @@ def f(x: tf.Tensor, a_var: tf.Tensor, b_var: tf.Tensor, use_bias: bool,
 
   def compute_output_shape(self, input_shape: List[int]) -> Tuple[int, int]:
     return (input_shape[0], self.output_dim)
+
+  def get_config(self) -> dict:
+    """Returns the config of the layer.
+
+    The same layer can be reinstantiated later
+    (without its trained weights) from this configuration.
+
+    Returns:
+      Python dictionary containing the configuration of the layer.
+    """
+    config = {}
+
+    # Include the DenseDecomp-specific arguments
+    decomp_args = ['output_dim', 'decomp_size', 'use_bias']
+    for arg in decomp_args:
+      config[arg] = getattr(self, arg)
+
+    # Serialize the activation
+    config['activation'] = activations.serialize(getattr(self, 'activation'))
+
+    # Serialize the initializers
+    decomp_initializers = ['kernel_initializer', 'bias_initializer']
+    for initializer_arg in decomp_initializers:
+      config[initializer_arg] = initializers.serialize(
+          getattr(self, initializer_arg))
+
+    # Get base config
+    base_config = super(DenseDecomp, self).get_config()
+    return dict(list(base_config.items()) + list(config.items()))

tensornetwork/tn_keras/mpo.py

@@ -0,0 +1,207 @@
+import tensorflow as tf
+from tensorflow.keras.layers import Layer  # type: ignore
+from tensorflow.keras import activations
+from tensorflow.keras import initializers
+from typing import List, Optional, Text, Tuple
+import tensornetwork as tn
+from tensornetwork.network_components import Node
+import numpy as np
+import math
+
+
+@tf.keras.utils.register_keras_serializable() # type: ignore
+    # package='tensornetwork', name='dense_mpo')
+class DenseMPO(Layer):
+  """Matrix Product Operator (MPO) TN layer.
+
+  Example:
+
+  ```python
+  # as first layer in a sequential model:
+  model = Sequential()
+  model.add(
+    DenseMPO(1024, num_nodes=4, bond_dim=8, activation='relu',
+    input_shape=(1024,)))
+  # now the model will take as input arrays of shape (*, 1024)
+  # and output arrays of shape (*, 1024).
+  # After the first layer, you don't need to specify
+  # the size of the input anymore:
+  model.add(DenseMPO(1024, num_nodes=4, bond_dim=8, activation='relu'))
+  ```
+
+  Args:
+    output_dim: Positive integer, dimensionality of the output space.
+    num_nodes: Positive integer, number of nodes in the MPO.
+      Note input_shape[-1]**(1. / num_nodes) and output_dim**(1. / num_nodes)
+      must both be round.
+    bond_dim: Positive integer, size of the intermediate dimension.
+    activation: Activation function to use.
+      If you don't specify anything, no activation is applied
+      (ie. "linear" activation: `a(x) = x`).
+    use_bias: Boolean, whether the layer uses a bias vector.
+    kernel_initializer: Initializer for the node weight matrices.
+    bias_initializer: Initializer for the bias vector.
+
+  Input shape:
+    2D tensor with shape: `(batch_size, input_shape[-1])`.
+
+  Output shape:
+    2D tensor with shape: `(batch_size, output_dim)`.
+  """
+
+  def __init__(self,
+               output_dim: int,
+               num_nodes: int,
+               bond_dim: int,
+               use_bias: Optional[bool] = True,
+               activation: Optional[Text] = None,
+               kernel_initializer: Optional[Text] = 'glorot_uniform',
+               bias_initializer: Optional[Text] = 'zeros',
+               **kwargs) -> None:
+
+    # Allow specification of input_dim instead of input_shape,
+    # for compatability with Keras layers that support this
+    if 'input_shape' not in kwargs and 'input_dim' in kwargs:
+      kwargs['input_shape'] = (kwargs.pop('input_dim'),)
+
+    assert num_nodes > 2, 'Need at least 2 nodes to create MPO'
+
+    super(DenseMPO, self).__init__(**kwargs)
+
+    self.output_dim = output_dim
+    self.num_nodes = num_nodes
+    self.bond_dim = bond_dim
+    self.nodes = []
+    self.use_bias = use_bias
+    self.activation = activations.get(activation)
+    self.kernel_initializer = initializers.get(kernel_initializer)
+    self.bias_initializer = initializers.get(bias_initializer)
+
+  def build(self, input_shape: List[int]) -> None:
+    # Disable the attribute-defined-outside-init violations in this function
+    # pylint: disable=attribute-defined-outside-init
+    if input_shape[-1] is None:
+      raise ValueError('The last dimension of the inputs to `Dense` '
+                       'should be defined. Found `None`.')
+
+    def is_perfect_root(n, n_nodes):
+      root = n**(1. / n_nodes)
+      return round(root)**n_nodes == n
+
+    # Ensure the MPO dimensions will work
+    assert is_perfect_root(input_shape[-1], self.num_nodes), \
+      f'Input dim incorrect.\
+      {input_shape[-1]}**(1. / {self.num_nodes}) must be round.'
+
+    assert is_perfect_root(self.output_dim, self.num_nodes), \
+      f'Output dim incorrect. \
+      {self.output_dim}**(1. / {self.num_nodes}) must be round.'
+
+    super(DenseMPO, self).build(input_shape)
+
+    self.in_leg_dim = math.ceil(input_shape[-1]**(1. / self.num_nodes))
+    self.out_leg_dim = math.ceil(self.output_dim**(1. / self.num_nodes))
+
+    self.nodes.append(
+        self.add_weight(name='end_node_first',
+                        shape=(self.in_leg_dim, self.bond_dim,
+                               self.out_leg_dim),
+                        trainable=True,
+                        initializer=self.kernel_initializer))
+    for i in range(self.num_nodes - 2):
+      self.nodes.append(
+          self.add_weight(name=f'middle_node_{i}',
+                          shape=(self.in_leg_dim, self.bond_dim, self.bond_dim,
+                                 self.out_leg_dim),
+                          trainable=True,
+                          initializer=self.kernel_initializer))
+    self.nodes.append(
+        self.add_weight(name='end_node_last',
+                        shape=(self.in_leg_dim, self.bond_dim,
+                               self.out_leg_dim),
+                        trainable=True,
+                        initializer=self.kernel_initializer))
+
+    self.bias_var = self.add_weight(
+        name='bias',
+        shape=(self.output_dim,),
+        trainable=True,
+        initializer=self.bias_initializer) if self.use_bias else None
+
+  def call(self, inputs: tf.Tensor, **kwargs) -> tf.Tensor: # pylint: disable=unused-argument
+
+
+    def f(x: tf.Tensor, nodes: List[Node], num_nodes: int, in_leg_dim: int,
+          use_bias: bool, bias_var: tf.Tensor) -> tf.Tensor:
+      orig_shape = x.shape
+      l = [in_leg_dim] * num_nodes
+      input_reshaped = tf.reshape(x, tuple(l))
+      x_node = tn.Node(input_reshaped, name='xnode', backend="tensorflow")
+
+      tn_nodes = []
+      for i, v in enumerate(nodes):
+        tn_nodes.append(tn.Node(v, name=f'node_{i}', backend="tensorflow"))
+        # Connect every node to input node
+        x_node[i] ^ tn_nodes[i][0]
+
+      # Connect all core nodes
+      tn_nodes[0][1] ^ tn_nodes[1][1]
+      for i, _ in enumerate(tn_nodes):
+        if len(tn_nodes[i].shape) == 4:
+          tn_nodes[i][2] ^ tn_nodes[i + 1][1]
+
+      # The TN should now look like this
+      #   |     |    |
+      #   1 --- 2 --- ...
+      #    \   /    /
+      #      x
+
+      # Contract TN using zipper algorithm
+      temp = x_node @ tn_nodes[0]
+      for i in range(1, len(tn_nodes)):
+        temp = temp @ tn_nodes[i]
+
+      result = tf.reshape(temp.tensor, orig_shape)
+      if use_bias:
+        result += bias_var
+
+      return result
+
+    result = tf.vectorized_map(
+        lambda vec: f(vec, self.nodes, self.num_nodes, self.in_leg_dim, self.
+                      use_bias, self.bias_var), inputs)
+    if self.activation is not None:
+      result = self.activation(result)
+    return tf.reshape(result, (-1, self.output_dim))
+
+  def compute_output_shape(self, input_shape: List[int]) -> Tuple[int, int]:
+    return (input_shape[0], self.output_dim)
+
+  def get_config(self) -> dict:
+    """Returns the config of the layer.
+
+    The same layer can be reinstantiated later
+    (without its trained weights) from this configuration.
+
+    Returns:
+      Python dictionary containing the configuration of the layer.
+    """
+    config = {}
+
+    # Include the MPO-specific arguments
+    mpo_args = ['output_dim', 'num_nodes', 'bond_dim', 'use_bias']
+    for arg in mpo_args:
+      config[arg] = getattr(self, arg)
+
+    # Serialize the activation
+    config['activation'] = activations.serialize(getattr(self, 'activation'))
+
+    # Serialize the initializers
+    mpo_initializers = ['kernel_initializer', 'bias_initializer']
+    for initializer_arg in mpo_initializers:
+      config[initializer_arg] = initializers.serialize(
+          getattr(self, initializer_arg))
+
+    # Get base config
+    base_config = super(DenseMPO, self).get_config()
+    return dict(list(base_config.items()) + list(config.items()))