Debugging PiNN layers and networks

# Install konda
!pip -q install konda

import konda
konda.install()

# Accept Anaconda Terms of Service
!conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/main
!conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/r

# Create the environment
!konda create -q -n pinn_env python=3.10 -y

!konda activate pinn_env
!konda run "pip -q install 'tensorflow[and-cuda]==2.15.1' 'ase&gt;=3.25' 'PyYAML~=6.0.1' 'numpy&lt;2'"
!konda run "pip -q install git+https://github.com/Teoroo-CMC/PiNN"

!wget -nv -nc https://raw.githubusercontent.com/Teoroo-CMC/PiNN_lab/master/resources/qm9_train.{yml,tfr}

Loading data

For the purpose of testing we download a subset of the QM9 dataset used in PiNN_lab.

%%writefile layer_debug.py

import warnings
from pinn.io import load_tfrecord, sparse_batch
index_warning = 'Converting sparse IndexedSlices'
warnings.filterwarnings('ignore', index_warning)

dataset = load_tfrecord("qm9_train.yml").apply(sparse_batch(10))
for datum in dataset:
    print({k: v.shape for k, v in datum.items()})
    break

Overwriting layer_debug.py

!konda run "python layer_debug.py"

{'A': TensorShape([10]), 'B': TensorShape([10]), 'C': TensorShape([10]), 'Cv': TensorShape([10]), 'G': TensorShape([10]), 'H': TensorShape([10]), 'U': TensorShape([10]), 'U0': TensorShape([10]), 'alpha': TensorShape([10]), 'coord': TensorShape([176, 3]), 'elems': TensorShape([176]), 'gap': TensorShape([10]), 'homo': TensorShape([10]), 'lumo': TensorShape([10]), 'mu': TensorShape([10]), 'r2': TensorShape([10]), 'zpve': TensorShape([10]), 'ind_1': TensorShape([176, 1])}

Using PiNN Layers

PiNN networks and layers are Keras Layers and Models.

To use them, you create an instance of layer, after that, the layer object can be used as a function. Each layer is initialized with different parameters and requires different input tensors, see their individual documentation for the details.

%%writefile -a layer_debug.py

from pinn.layers import CellListNL

nl = CellListNL(rc=5)

for datum in dataset:
    nl(datum)
    break

Appending to layer_debug.py

!konda run "python layer_debug.py"

{'A': TensorShape([10]), 'B': TensorShape([10]), 'C': TensorShape([10]), 'Cv': TensorShape([10]), 'G': TensorShape([10]), 'H': TensorShape([10]), 'U': TensorShape([10]), 'U0': TensorShape([10]), 'alpha': TensorShape([10]), 'coord': TensorShape([176, 3]), 'elems': TensorShape([176]), 'gap': TensorShape([10]), 'homo': TensorShape([10]), 'lumo': TensorShape([10]), 'mu': TensorShape([10]), 'r2': TensorShape([10]), 'zpve': TensorShape([10]), 'ind_1': TensorShape([176, 1])}

The definition of a layer needs three parts:

__init__ defines the layer object;
build creates the necessary variables or sub-layers;
call defines how the input tensors are processed.

The build() method is only called once when the layer is used for the first tiem (e.g. in a loop). See below for an example definition for the PILayer

%%writefile check_PILayer.py

from pinn.networks.pinet import PILayer
import inspect

print(inspect.getsource(PILayer))

Overwriting check_PILayer.py

!konda run "python check_PILayer.py"

2026-09-12 19:20:32.087148: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
2026-09-12 19:20:32.138910: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
2026-09-12 19:20:32.138997: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
2026-09-12 19:20:32.140461: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2026-09-12 19:20:32.148916: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
2026-09-12 19:20:32.149214: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2026-09-12 19:20:33.232936: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
class PILayer(tf.keras.layers.Layer):
    R"""`PILayer` takes the properties ($\mathbb{P}_{i\alpha},
    \mathbb{P}_{j\alpha}$) of a pair of atoms as input and outputs a set of
    interactions for each pair. The inputs will be broadcasted and concatenated
    as the input of a feed-forward neural network (`FFLayer`), and the
    interactions are generated by taking the output of the `FFLayer` as weights
    of radial basis functions, i.e.:

    $$
    \begin{aligned}
    w_{ij(b\beta)} &= \mathrm{FFLayer}\left((\mathbf{1}_{j}\mathbb{P}_{i\alpha})\Vert(\mathbf{1}_{i}\mathbb{P}_{j\alpha})\right) \\
    \mathbb{I}'_{ij\beta} &= \sum_b W_{ij(b\beta)} \, e_{ijb}
    \end{aligned}
    $$

    , where $w_{ij(b\beta)}$ is an intemediate weight tensor for the
    radial basis functions, output by the `FFLayer`; the output channel is
    reshaped into two dimensions, where $b$ is the index for the basis function
    and $d$ is the index for output interaction.


    `n_nodes` specifies the number of nodes in the `FFLayer`. Note that the last
    element of n_nodes specifies the number of output channels after applying
    the basis function ($d$ instead of $bd$), i.e. the output dimension of
    FFLayer is `[n_pairs,n_nodes[-1]*n_basis]`, the output is then summed with
    the basis to form the output interaction.

    """

    def __init__(self, n_nodes=[64], **kwargs):
        """
        Args:
            n_nodes (list of int): number of nodes to use
            **kwargs (dict): keyword arguments will be parsed to the feed forward layers
        """
        super(PILayer, self).__init__()
        self.n_nodes = n_nodes
        self.kwargs = kwargs

    def build(self, shapes):
        """"""
        self.n_basis = shapes[2][-1]
        n_nodes_iter = self.n_nodes.copy()
        n_nodes_iter[-1] *= self.n_basis
        self.ff_layer = FFLayer(n_nodes_iter, **self.kwargs)

    def call(self, tensors):
        """
        PILayer take a list of three tensors as input:

        - ind_2: [sparse indices](layers.md#sparse-indices) of pairs with shape `(n_pairs, 2)`
        - prop: property tensor with shape `(n_atoms, n_prop)`
        - basis: interaction tensor with shape `(n_pairs, n_basis)`

        Args:
            tensors (list of tensors): list of `[ind_2, prop, basis]` tensors

        Returns:
            inter (tensor): interaction tensor with shape `(n_pairs, n_nodes[-1])`
        """
        ind_2, prop, basis = tensors
        ind_i = ind_2[:, 0]
        ind_j = ind_2[:, 1]
        prop_i = tf.gather(prop, ind_i)
        prop_j = tf.gather(prop, ind_j)

        inter = tf.concat([prop_i, prop_j], axis=-1)
        inter = self.ff_layer(inter)
        inter = tf.reshape(inter, [-1, self.n_nodes[-1], self.n_basis])
        inter = tf.einsum("pcb,pb->pc", inter, basis)
        return inter

Using PiNN Networks

network (Keras Models) are defined similarly, but they can be directly used to perform regression task.

By default, network produces per-atom predictions, this can be changed by the out_pool parameter to get some simple per-structure predictions. In that case, the network object can be used to perform trainig directly.

%%writefile -a layer_debug.py

from pinn.networks.pinet import PiNet

def label_data(data):
    # defines the label to train on
    x = data
    y = data['lumo']
    return x, y

train = dataset.map(label_data)
pinet = PiNet(out_pool='min')
pinet.compile(optimizer='Adam', loss='MAE')
pinet.fit(train, epochs=3)

Appending to layer_debug.py

!konda run "python layer_debug.py"

{'A': TensorShape([10]), 'B': TensorShape([10]), 'C': TensorShape([10]), 'Cv': TensorShape([10]), 'G': TensorShape([10]), 'H': TensorShape([10]), 'U': TensorShape([10]), 'U0': TensorShape([10]), 'alpha': TensorShape([10]), 'coord': TensorShape([176, 3]), 'elems': TensorShape([176]), 'gap': TensorShape([10]), 'homo': TensorShape([10]), 'lumo': TensorShape([10]), 'mu': TensorShape([10]), 'r2': TensorShape([10]), 'zpve': TensorShape([10]), 'ind_1': TensorShape([176, 1])}
Epoch 1/3
2000/2000 [==============================] - 62s 27ms/step - loss: 0.0890
Epoch 2/3 
2000/2000 [==============================] - 59s 29ms/step - loss: 0.0296 
Epoch 3/3 
2000/2000 [==============================] - 54s 27ms/step - loss: 0.0248

Further benchmarks

For more advanced usage you are recommended to use the Model API to define the trainig loss, derived predicates. For traininig potential energy surfaces, you are recommended to use pinn.models.potential_model in combination with the command line interface (CLI).

Alternatively, see the Trainig Tips notebook to see how to run the tranining interactively in a notebook.