Merge pull request #27 from NatLabRockies/TFv220

Tensorflow v2.20

Author: jlaw9

etc/environment.yml

@@ -8,7 +8,7 @@ dependencies:
   - rdkit
   - pytest
   - tqdm
-  - numpy=1.19
+  - numpy
   - networkx
   - pymatgen
   - pip

nfp/__init__.py

@@ -8,8 +8,11 @@ class MissingDependencyException(RuntimeError):
 
 
 custom_objects = {
+    "RBFExpansion": RBFExpansion,
     "Gather": Gather,
     "Reduce": Reduce,
+    "ConcatDense": ConcatDense,
+    "Tile": Tile,
     "masked_mean_squared_error": masked_mean_squared_error,
     "masked_mean_absolute_error": masked_mean_absolute_error,
     "masked_log_cosh": masked_log_cosh,

nfp/layers/graph_layers.py

@@ -3,8 +3,10 @@
 
 assert tf, "Tensorflow 2.x required for GraphLayers"
 tf_layers = tf.keras.layers
+register_keras_serializable = tf.keras.utils.register_keras_serializable
 
 
+@register_keras_serializable(package="nfp")
 class GraphLayer(tf_layers.Layer):
     """Base class for all GNN layers"""
 
@@ -28,9 +30,16 @@ def build(self, input_shape):
             self.dropout_layer = tf_layers.Dropout(self.dropout)
 
     def get_config(self):
-        return {"dropout": self.dropout}
+        config = super().get_config()
+        config.update({"dropout": self.dropout})
+        return config
+
+    @classmethod
+    def from_config(cls, config):
+        return cls(**config)
 
 
+@register_keras_serializable(package="nfp")
 class EdgeUpdate(GraphLayer):
     def build(self, input_shape):
         """inputs = [atom_state, bond_state, connectivity]
@@ -40,6 +49,10 @@ def build(self, input_shape):
 
         self.gather = nfp.Gather()
         self.concat = nfp.ConcatDense()
+        concat_shapes = [input_shape[1], input_shape[0], input_shape[0]]
+        if self.use_global:
+            concat_shapes.append((input_shape[1][0], input_shape[1][1], input_shape[3][-1]))
+        self.concat.build(concat_shapes)
 
     def call(self, inputs, mask=None, **kwargs):
         """Inputs: [atom_state, bond_state, connectivity]
@@ -77,7 +90,12 @@ def compute_mask(self, inputs, mask=None):
         else:
             return None
 
+    @classmethod
+    def from_config(cls, config):
+        return cls(**config)
+
 
+@register_keras_serializable(package="nfp")
 class NodeUpdate(GraphLayer):
     def build(self, input_shape):
         super().build(input_shape)
@@ -91,6 +109,12 @@ def build(self, input_shape):
 
         self.dense1 = tf_layers.Dense(2 * num_features, activation="relu")
         self.dense2 = tf_layers.Dense(num_features)
+        concat_shapes = [input_shape[0], input_shape[1]]
+        if self.use_global:
+            concat_shapes.append((input_shape[1][0], input_shape[1][1], input_shape[3][-1]))
+        self.concat.build(concat_shapes)
+        self.dense1.build((None, num_features))
+        self.dense2.build((None, 2 * num_features))
 
     def call(self, inputs, mask=None, **kwargs):
         """Inputs: [atom_state, bond_state, connectivity]
@@ -110,7 +134,7 @@ def call(self, inputs, mask=None, **kwargs):
         else:
             messages = self.concat([source_atom, bond_state, global_state])
 
-        if mask is not None:
+        if mask is not None and mask[1] is not None:
             # Only works for sum, max
             messages = tf.where(
                 tf.expand_dims(mask[1], axis=-1), messages, tf.zeros_like(messages)
@@ -136,19 +160,27 @@ def compute_mask(self, inputs, mask=None):
         else:
             return None
 
+    @classmethod
+    def from_config(cls, config):
+        return cls(**config)
 
+
+@register_keras_serializable(package="nfp")
 class GlobalUpdate(GraphLayer):
     def __init__(self, units, num_heads, **kwargs):
         super().__init__(**kwargs)
         self.units = units  # H
         self.num_heads = num_heads  # N
-        self.supports_masking = False
+        self.supports_masking = True
 
     def build(self, input_shape):
         super().build(input_shape)
         dense_units = self.units * self.num_heads  # N*H
+        num_features = input_shape[0][-1]
         self.query_layer = tf_layers.Dense(self.num_heads, name="query")
         self.value_layer = tf_layers.Dense(dense_units, name="value")
+        self.query_layer.build((None, num_features))
+        self.value_layer.build((None, num_features))
 
     def transpose_scores(self, input_tensor):
         input_shape = tf.shape(input_tensor)
@@ -167,7 +199,7 @@ def call(self, inputs, mask=None, **kwargs):
         graph_elements = tf.concat([atom_state, bond_state], axis=1)
         query = self.query_layer(graph_elements)  # [B,N,S,H]
 
-        if mask is not None:
+        if mask is not None and mask[0] is not None and mask[1] is not None:
             graph_element_mask = tf.concat([mask[0], mask[1]], axis=1)
             query = tf.where(
                 tf.expand_dims(graph_element_mask, axis=-1),
@@ -187,7 +219,18 @@ def call(self, inputs, mask=None, **kwargs):
 
         return context
 
+    def compute_output_shape(self, input_shape):
+        # input_shape[0] is atom_state: (batch, N_atoms, features)
+        return (input_shape[0][0], self.num_heads * self.units)
+
+    def compute_mask(self, inputs, mask=None):
+        return None  # 2D output has no sequence mask
+
     def get_config(self):
         config = super(GlobalUpdate, self).get_config()
         config.update({"units": self.units, "num_heads": self.num_heads})
         return config
+
+    @classmethod
+    def from_config(cls, config):
+        return cls(**config)

nfp/layers/layers.py

@@ -2,8 +2,10 @@
 
 assert tf, "Tensorflow 2.x required for GraphLayers"
 tf_layers = tf.keras.layers
+register_keras_serializable = tf.keras.utils.register_keras_serializable
 
 
+@register_keras_serializable(package="nfp")
 class RBFExpansion(tf_layers.Layer):
     def __init__(
         self, dimension=128, init_gap=10, init_max_distance=7, trainable=False
@@ -55,12 +57,18 @@ def compute_mask(self, inputs, mask=None):
         return tf.logical_not(tf.math.is_nan(inputs))
 
     def get_config(self):
-        return {
+        config = super().get_config()
+        config.update({
             "init_gap": self.init_gap,
             "init_max_distance": self.init_max_distance,
             "dimension": self.dimension,
             "trainable": self.trainable,
-        }
+        })
+        return config
+
+    @classmethod
+    def from_config(cls, config):
+        return cls(**config)
 
 
 def batched_segment_op(
@@ -117,16 +125,33 @@ def batched_segment_op(
 #         return cls(**config)
 
 
+@register_keras_serializable(package="nfp")
 class Gather(tf_layers.Layer):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.supports_masking = True
+
     def call(self, inputs, mask=None, **kwargs):
         reference, indices = inputs
         return tf.gather(reference, indices, batch_dims=1)
 
+    def compute_mask(self, inputs, mask=None):
+        return None
 
+    def get_config(self):
+        return super().get_config()
+
+    @classmethod
+    def from_config(cls, config):
+        return cls(**config)
+
+
+@register_keras_serializable(package="nfp")
 class Reduce(tf_layers.Layer):
     def __init__(self, reduction="sum", *args, **kwargs):
         super(Reduce, self).__init__(*args, **kwargs)
         self.reduction = reduction
+        self.supports_masking = True
 
     def compute_output_shape(self, input_shape):
         data_shape, _, target_shape = input_shape
@@ -143,8 +168,17 @@ def call(self, inputs, mask=None, **kwargs):
             reduction=self.reduction,
         )
 
+    def compute_mask(self, inputs, mask=None):
+        return None
+
     def get_config(self):
-        return {"reduction": self.reduction}
+        config = super().get_config()
+        config.update({"reduction": self.reduction})
+        return config
+
+    @classmethod
+    def from_config(cls, config):
+        return cls(**config)
 
     @staticmethod
     def _parse_inputs_and_mask(inputs, mask=None):
@@ -159,6 +193,7 @@ def _parse_inputs_and_mask(inputs, mask=None):
         return data, segment_ids, target, data_mask
 
 
+@register_keras_serializable(package="nfp")
 class ConcatDense(tf_layers.Layer):
     """Layer to combine the concatenation and two dense layers. Just useful as a common
     operation in the graph layers"""
@@ -169,23 +204,41 @@ def __init__(self, **kwargs):
 
     def build(self, input_shape):
         num_features = input_shape[0][-1]
+        concat_features = sum(shape[-1] for shape in input_shape)
         self.concat = tf_layers.Concatenate()
         self.dense1 = tf_layers.Dense(2 * num_features, activation="relu")
         self.dense2 = tf_layers.Dense(num_features)
+        self.dense1.build((None, concat_features))
+        self.dense2.build((None, 2 * num_features))
+        super().build(input_shape)
 
     def call(self, inputs, mask=None, **kwargs):
-        output = self.concat(inputs)
+        # Use tf.concat instead of Keras Concatenate to avoid Keras 3's
+        # compute_mask creating an all-True mask when some inputs lack masks
+        output = tf.concat(inputs, axis=-1)
         output = self.dense1(output)
         output = self.dense2(output)
         return output
 
     def compute_mask(self, inputs, mask=None):
         if mask is None:
             return None
-        else:
-            return tf.math.reduce_all(tf.stack(mask), axis=0)
+        valid_masks = [m for m in mask if m is not None]
+        if not valid_masks:
+            return None
+        if len(valid_masks) == 1:
+            return valid_masks[0]
+        return tf.math.reduce_all(tf.stack(valid_masks), axis=0)
+
+    def get_config(self):
+        return super().get_config()
+
+    @classmethod
+    def from_config(cls, config):
+        return cls(**config)
 
 
+@register_keras_serializable(package="nfp")
 class Tile(tf_layers.Layer):
     def __init__(self, **kwargs):
         super(Tile, self).__init__(**kwargs)
@@ -202,3 +255,10 @@ def compute_mask(self, inputs, mask=None):
             return None
         else:
             return mask[1]
+
+    def get_config(self):
+        return super().get_config()
+
+    @classmethod
+    def from_config(cls, config):
+        return cls(**config)

setup.cfg

@@ -13,10 +13,10 @@ install_requires =
 
 [options.extras_require]
 pymatgen = pymatgen
-tf = tensorflow>=2.4
+tf = tensorflow>=2.20
 rdkit = rdkit-pypi
 test = pytest
 
 [flake8]
 max-line-length = 88
-extend-ignore = E203
+extend-ignore = E203

tests/layers/test_layers.py

@@ -23,7 +23,7 @@
 
 
 def test_gather():
-    in1 = layers.Input(shape=[None], dtype="float", name="data")
+    in1 = layers.Input(shape=[None], dtype="float32", name="data")
     in2 = layers.Input(shape=[None], dtype=tf.int64, name="indices")
 
     gather = nfp.Gather()([in1, in2])
@@ -73,7 +73,7 @@ def test_reduce(smiles_inputs, method):
 
 
 def test_tile():
-    state = layers.Input(shape=[None], dtype="float", name="data")
+    state = layers.Input(shape=[None], dtype="float32", name="data")
     target = layers.Input(shape=[None, 3], dtype=tf.int64, name="indices")
 
     tile = nfp.Tile()([state, target])