mpnn.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MPNN\n",
    "\n",
    "### Input preparation\n",
    "- Search bonds : [see this](search_bonds.ipynb)\n",
    "- Search angles : [see this](search_angles.ipynb)\n",
    "- Generate input array : [see this](gen_input_graph.ipynb)\n",
    "    \n",
    "### Prepared files:\n",
    "- Train\n",
    "    - Nodes features\n",
    "    - Edges features input\n",
    "    - Edges output\n",
    "- Test\n",
    "    - Nodes features\n",
    "    - Edges features input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/javascript": "IPython.notebook.set_autosave_interval(60000)"
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Autosaving every 60 seconds\n"
     ]
    }
   ],
   "source": [
    "%autosave 60"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "use_multiGPU = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Import done!\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "# os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\" # turn on/off GPU run\n",
    "# os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' # disable debugging logs\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import tensorflow as tf\n",
    "import networkx\n",
    "import ase\n",
    "import ase.visualize\n",
    "\n",
    "from tqdm import tqdm\n",
    "from sklearn.utils import shuffle\n",
    "\n",
    "print(\"Import done!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /tmp/ipykernel_19985/501681495.py:1: is_gpu_available (from tensorflow.python.framework.test_util) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `tf.config.list_physical_devices('GPU')` instead.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-07-16 10:55:17.076979: I tensorflow/compiler/jit/xla_gpu_device.cc:99] Not creating XLA devices, tf_xla_enable_xla_devices not set\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-07-16 10:55:17.092241: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcuda.so.1\n",
      "2021-07-16 10:55:17.433709: E tensorflow/stream_executor/cuda/cuda_driver.cc:328] failed call to cuInit: CUDA_ERROR_NO_DEVICE: no CUDA-capable device is detected\n",
      "2021-07-16 10:55:17.433797: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:156] kernel driver does not appear to be running on this host (jwlogin23.juwels): /proc/driver/nvidia/version does not exist\n"
     ]
    }
   ],
   "source": [
    "tf.test.is_gpu_available(\n",
    "    cuda_only=False, min_cuda_compute_capability=None\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-07-16 10:55:17.489464: I tensorflow/compiler/jit/xla_cpu_device.cc:41] Not creating XLA devices, tf_xla_enable_xla_devices not set\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# see a list of GPU devices\n",
    "devices = tf.config.experimental.list_physical_devices(\"GPU\")\n",
    "devices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:Some requested devices in `tf.distribute.Strategy` are not visible to TensorFlow: /job:localhost/replica:0/task:0/device:GPU:1,/job:localhost/replica:0/task:0/device:GPU:0,/job:localhost/replica:0/task:0/device:GPU:3,/job:localhost/replica:0/task:0/device:GPU:2\n",
      "INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1', '/job:localhost/replica:0/task:0/device:GPU:2', '/job:localhost/replica:0/task:0/device:GPU:3')\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-07-16 10:55:17.502964: I tensorflow/compiler/jit/xla_gpu_device.cc:99] Not creating XLA devices, tf_xla_enable_xla_devices not set\n"
     ]
    }
   ],
   "source": [
    "strategy = tf.distribute.MirroredStrategy(devices=['GPU:0', 'GPU:1', 'GPU:2', 'GPU:3'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# dir_dataset = \"test/j-coupling-dataset/\"\n",
    "dir_dataset = \"../from-kaggle/\"\n",
    "nodes_train = np.load(dir_dataset + \"nodes_train.npz\")['arr_0'][:1000]\n",
    "in_edges_train = np.load(dir_dataset + \"in_edges_train.npz\")['arr_0'][:1000]\n",
    "# out_edges_train = np.load(dir_dataset + \"out_edges_train.npz\")['arr_0'][:1000]\n",
    "\n",
    "nodes_test = np.load(dir_dataset + \"nodes_test.npz\")['arr_0'][:500]\n",
    "in_edges_test = np.load(dir_dataset + \"in_edges_test.npz\")['arr_0'][:500]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1000, 29, 5)\n",
      "(1000, 29, 29, 16)\n",
      "(500, 29, 5)\n",
      "(500, 29, 29, 16)\n"
     ]
    }
   ],
   "source": [
    "print(nodes_train.shape)\n",
    "print(in_edges_train.shape)\n",
    "# print(out_edges_train.shape)\n",
    "print(nodes_test.shape)\n",
    "print(in_edges_test.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "out_labels = np.random.randn(1000, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1000, 1)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out_labels.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# out_labels = out_edges_train.reshape(-1, out_edges_train.shape[1]*out_edges_train.shape[2], 1)\n",
    "in_edges_train = in_edges_train.reshape(-1, in_edges_train.shape[1]*in_edges_train.shape[2], in_edges_train.shape[3])\n",
    "in_edges_test = in_edges_test.reshape(-1, in_edges_test.shape[1]*in_edges_test.shape[2], in_edges_test.shape[3])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "nodes_train, in_edges_train, out_labels = shuffle(nodes_train, in_edges_train, out_labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Message Passing Neural Network\n",
    "\n",
    "Implement according to Gilmer et al. https://arxiv.org/abs/1704.01212"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build message parser\n",
    "\n",
    "Define message functions $M_{t}$ in $m^{t+1}_{v} =\\sum_{w \\in N(v)} M_{t}(h^{t}_{v}, h^{t}_{w}, e_{v,w})$\n",
    "\n",
    "$\n",
    "\\begin{align}\n",
    "M_{t}(h^{t}_{v}, h^{t}_{w}, e_{v,w}) &= NN(e_{v w}) h^{t}_{w} \\\\\n",
    "&= A_{e_{v w}} h^{t}_{w}\n",
    "\\end{align}\n",
    "$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Message_Passer(tf.keras.layers.Layer):\n",
    "    def __init__(self, node_dim):\n",
    "        super(Message_Passer, self).__init__()\n",
    "        self.node_dim = node_dim\n",
    "        self.nn = tf.keras.layers.Dense(units=self.node_dim*self.node_dim, activation=tf.nn.relu)\n",
    "      \n",
    "    def call(self, node_j, edge_ij):\n",
    "        # Embed the edge as a matrix\n",
    "        A = self.nn(edge_ij)\n",
    "        \n",
    "        # Reshape so matrix mult can be done\n",
    "        A = tf.reshape(A, [-1, self.node_dim, self.node_dim])\n",
    "        node_j = tf.reshape(node_j, [-1, self.node_dim, 1])\n",
    "        \n",
    "        # Multiply edge matrix by node and shape into message list\n",
    "        messages = tf.linalg.matmul(A, node_j)\n",
    "        messages = tf.reshape(messages, [-1, tf.shape(edge_ij)[1], self.node_dim])\n",
    "\n",
    "        return messages"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build aggregator\n",
    "\n",
    "Aggregate all $M_{t}$ to update message being sent to node $v$ at state $t+1$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Message_Agg(tf.keras.layers.Layer):\n",
    "    def __init__(self):\n",
    "        super(Message_Agg, self).__init__()\n",
    "    \n",
    "    def call(self, messages):\n",
    "        return tf.math.reduce_sum(messages, 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build update function - GRU\n",
    "\n",
    "$\n",
    "h^{t+1}_{v} = U_{t}(h^{t}_{v}, m^{t+1}_{v}) \\\\\n",
    "U_{t}(h^{t}_{v}, m^{t+1}_{v}) = GRU(h^{t}_{v}, m^{t+1}_{v})\n",
    "$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Update_Func_GRU(tf.keras.layers.Layer):\n",
    "    def __init__(self, state_dim):\n",
    "        super(Update_Func_GRU, self).__init__()\n",
    "        self.concat_layer = tf.keras.layers.Concatenate(axis=1)\n",
    "        self.GRU = tf.keras.layers.GRU(state_dim)\n",
    "        \n",
    "    def call(self, old_state, agg_messages):\n",
    "        # Remember node dim\n",
    "        n_nodes  = tf.shape(old_state)[1]\n",
    "        node_dim = tf.shape(old_state)[2]\n",
    "        \n",
    "        # Reshape so GRU can be applied, concat so old_state and messages are in sequence\n",
    "        old_state = tf.reshape(old_state, [-1, 1, tf.shape(old_state)[-1]])\n",
    "        agg_messages = tf.reshape(agg_messages, [-1, 1, tf.shape(agg_messages)[-1]])\n",
    "        concat = self.concat_layer([old_state, agg_messages])\n",
    "        \n",
    "        # Apply GRU and then reshape so it can be returned\n",
    "        activation = self.GRU(concat)\n",
    "        new_state = tf.reshape(activation, [-1, n_nodes, node_dim])\n",
    "        \n",
    "        return new_state"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Output layer\n",
    "Define readout function ($R$) as a neural network:\n",
    "\n",
    "$\\hat{y} = R(\\{h^{T}_{v} | v \\in G \\})$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Node_Regressor(tf.keras.layers.Layer):\n",
    "    def __init__(self, intermediate_dim):\n",
    "        super(Node_Regressor, self).__init__()\n",
    "        self.concat_layer = tf.keras.layers.Concatenate()\n",
    "        self.hidden_layer_1 = tf.keras.layers.Dense(units=intermediate_dim, activation=tf.nn.relu)\n",
    "        self.hidden_layer_2 = tf.keras.layers.Dense(units=intermediate_dim, activation=tf.nn.relu)\n",
    "        self.output_layer = tf.keras.layers.Dense(units=1, activation=None)\n",
    "        \n",
    "    def call(self, nodes):\n",
    "        n_nodes = tf.shape(nodes)[1]\n",
    "        node_dim = tf.shape(nodes)[2]\n",
    "        \n",
    "        # Tile and reshape to match edges\n",
    "        state_i = tf.tile(nodes, [1, 1, n_nodes])\n",
    "        state_i = tf.reshape(state_i, [-1, n_nodes*n_nodes, node_dim])\n",
    "        state_j = tf.tile(nodes, [1, n_nodes, 1])\n",
    "        \n",
    "        concat = self.concat_layer([state_i, state_j])\n",
    "        act_1 = self.hidden_layer_1(concat)\n",
    "        act_2 = self.hidden_layer_2(act_1)\n",
    "        \n",
    "        return self.output_layer(act_2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Edge_Regressor(tf.keras.layers.Layer):\n",
    "    def __init__(self, intermediate_dim):\n",
    "        super(Edge_Regressor, self).__init__()\n",
    "        self.concat_layer = tf.keras.layers.Concatenate()\n",
    "        self.hidden_layer_1 = tf.keras.layers.Dense(units=intermediate_dim, activation=tf.nn.relu)\n",
    "        self.hidden_layer_2 = tf.keras.layers.Dense(units=intermediate_dim, activation=tf.nn.relu)\n",
    "        self.output_layer = tf.keras.layers.Dense(units=1, activation=None)\n",
    "\n",
    "    def call(self, nodes, edges):\n",
    "        print(nodes.shape)\n",
    "        n_nodes = tf.shape(nodes)[1]\n",
    "        node_dim = tf.shape(nodes)[2]\n",
    "        \n",
    "        # Tile and reshape to match edges\n",
    "        state_i = tf.tile(nodes, [1, 1, n_nodes])\n",
    "        state_i = tf.reshape(state_i, [-1, n_nodes*n_nodes, node_dim])\n",
    "        state_j = tf.tile(nodes, [1, n_nodes, 1])\n",
    "        \n",
    "        # concat nodes at state i and j, and edges, and then apply MLP\n",
    "        print(state_i.shape)\n",
    "        print(edges.shape)\n",
    "        print(state_j.shape)\n",
    "        concat = self.concat_layer([state_i, edges, state_j])\n",
    "        print(concat.shape)\n",
    "        act_1 = self.hidden_layer_1(concat)\n",
    "        print(act_1.shape)\n",
    "        act_2 = self.hidden_layer_2(act_1)\n",
    "        print(act_2.shape)\n",
    "\n",
    "        return self.output_layer(act_2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 114,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(10, 29, 5)"
      ]
     },
     "execution_count": 114,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nodes_train[:10].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 115,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(10, 29, 5)\n",
      "(10, 29, 145)\n",
      "(10, 841, 5)\n",
      "(10, 841, 16)\n",
      "(10, 841, 5)\n",
      "(10, 841, 26)\n",
      "(10, 841, 512)\n",
      "(10, 841, 512)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(10, 841, 1), dtype=float32, numpy=\n",
       "array([[[-0.04952043],\n",
       "        [-0.02957275],\n",
       "        [-0.02791196],\n",
       "        ...,\n",
       "        [ 0.        ],\n",
       "        [ 0.        ],\n",
       "        [ 0.        ]],\n",
       "\n",
       "       [[ 0.02576611],\n",
       "        [ 0.01177338],\n",
       "        [ 0.05215303],\n",
       "        ...,\n",
       "        [ 0.        ],\n",
       "        [ 0.        ],\n",
       "        [ 0.        ]],\n",
       "\n",
       "       [[ 0.02598085],\n",
       "        [-0.04252946],\n",
       "        [ 0.01289575],\n",
       "        ...,\n",
       "        [ 0.        ],\n",
       "        [ 0.        ],\n",
       "        [ 0.        ]],\n",
       "\n",
       "       ...,\n",
       "\n",
       "       [[ 0.02598085],\n",
       "        [-0.0425198 ],\n",
       "        [ 0.01286001],\n",
       "        ...,\n",
       "        [ 0.        ],\n",
       "        [ 0.        ],\n",
       "        [ 0.        ]],\n",
       "\n",
       "       [[ 0.02598085],\n",
       "        [-0.04249534],\n",
       "        [-0.00568492],\n",
       "        ...,\n",
       "        [ 0.        ],\n",
       "        [ 0.        ],\n",
       "        [ 0.        ]],\n",
       "\n",
       "       [[ 0.02598085],\n",
       "        [-0.04252305],\n",
       "        [ 0.01288661],\n",
       "        ...,\n",
       "        [ 0.        ],\n",
       "        [ 0.        ],\n",
       "        [ 0.        ]]], dtype=float32)>"
      ]
     },
     "execution_count": 115,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ok = Edge_Regressor(512)\n",
    "ok.call(nodes_train[:10], in_edges_train[:10])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build Single Message Passing Layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MP_Layer(tf.keras.layers.Layer):\n",
    "    def __init__(self, state_dim):\n",
    "        super(MP_Layer, self).__init__(self)\n",
    "        self.message_passer = Message_Passer(node_dim=state_dim) \n",
    "        self.message_aggs = Message_Agg()\n",
    "        self.update_functions = Update_Func_GRU(state_dim=state_dim)\n",
    "        self.state_dim = state_dim\n",
    "\n",
    "    def call(self, nodes, edges, mask):\n",
    "        n_nodes = tf.shape(nodes)[1]\n",
    "        node_dim = tf.shape(nodes)[2]\n",
    "\n",
    "        state_j = tf.tile(nodes, [1, n_nodes, 1])\n",
    "\n",
    "        messages = self.message_passer(state_j, edges)\n",
    "\n",
    "        # Do this to ignore messages from non-existant nodes\n",
    "        masked = tf.math.multiply(messages, mask)\n",
    "        masked = tf.reshape(masked, [tf.shape(messages)[0], n_nodes, n_nodes, node_dim])\n",
    "        agg_m = self.message_aggs(masked)\n",
    "        \n",
    "        updated_nodes = self.update_functions(nodes, agg_m)\n",
    "        \n",
    "        nodes_out = updated_nodes\n",
    "        # Batch norm seems not to work. \n",
    "        # nodes_out = self.batch_norm(updated_nodes)\n",
    "        \n",
    "        return nodes_out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Formulate MPNN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "metadata": {},
   "outputs": [],
   "source": [
    "nod_input = tf.keras.Input(shape=(None,), name='nod_input')\n",
    "adj_input = tf.keras.Input(shape=(None,), name='adj_input')\n",
    "\n",
    "class MPNN(tf.keras.Model):\n",
    "    def __init__(self, out_int_dim, state_dim, T):\n",
    "        super(MPNN, self).__init__(self)   \n",
    "        self.T = T\n",
    "        self.embed = tf.keras.layers.Dense(units=state_dim, activation=tf.nn.relu)\n",
    "        self.MP = MP_Layer(state_dim)\n",
    "        self.edge_regressor  = Edge_Regressor(out_int_dim)\n",
    "        #self.batch_norm = tf.keras.layers.BatchNormalization()\n",
    "\n",
    "    def call(self, inputs=[nod_input, adj_input]):\n",
    "        nodes = inputs[0]\n",
    "        edges = inputs[1]\n",
    "        # Get distances, and create mask wherever 0 (i.e. non-existant nodes)\n",
    "        # This also masks node self-interactions...\n",
    "        # This assumes distance is last\n",
    "        len_edges = tf.shape(edges)[-1]\n",
    "\n",
    "        _, x = tf.split(edges, [len_edges - 1, 1], 2)\n",
    "        mask = tf.where(tf.equal(x, 0), x, tf.ones_like(x))\n",
    "\n",
    "        # Embed node to be of the chosen node dimension\n",
    "        nodes = self.embed(nodes)\n",
    "        \n",
    "        # nodes = self.batch_norm(nodes)\n",
    "        # Run the T message passing steps\n",
    "        for mp in tqdm(range(self.T)):\n",
    "            nodes = self.MP(nodes, edges, mask)\n",
    "            \n",
    "        # Regress the output values\n",
    "        con_edges = self.edge_regressor(nodes, edges)\n",
    "        \n",
    "        return con_edges"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define metrics (loss)\n",
    "\n",
    "Supported now:\n",
    "- MSE\n",
    "- Log MSE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mse(orig , preds):\n",
    "    # Mask values for which no scalar coupling exists\n",
    "    mask  = tf.where(tf.equal(orig, 0), orig, tf.ones_like(orig))\n",
    "\n",
    "    nums  = tf.boolean_mask(orig,  mask)\n",
    "    preds = tf.boolean_mask(preds,  mask)\n",
    "\n",
    "    reconstruction_error = tf.reduce_mean(tf.square(tf.subtract(nums, preds)))\n",
    "\n",
    "    return reconstruction_error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "metadata": {},
   "outputs": [],
   "source": [
    "def log_mse(orig , preds):\n",
    "    # Mask values for which no scalar coupling exists\n",
    "    mask  = tf.where(tf.equal(orig, 0), orig, tf.ones_like(orig))\n",
    "\n",
    "    nums  = tf.boolean_mask(orig,  mask)\n",
    "    preds = tf.boolean_mask(preds,  mask)\n",
    "\n",
    "    reconstruction_error = tf.math.log(tf.reduce_mean(tf.square(tf.subtract(nums, preds))))\n",
    "\n",
    "    return reconstruction_error"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define callback"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "metadata": {},
   "outputs": [],
   "source": [
    "learning_rate = 0.001\n",
    "def step_decay(epoch):\n",
    "    initial_lrate = learning_rate\n",
    "    drop = 0.1\n",
    "    epochs_drop = 20.0\n",
    "    lrate = initial_lrate * np.power(drop, np.floor((epoch)/epochs_drop))\n",
    "    tf.print(\"Learning rate: \", lrate)\n",
    "    return lrate\n",
    "\n",
    "lrate = tf.keras.callbacks.LearningRateScheduler(step_decay)\n",
    "stop_early = tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience = 15, restore_best_weights=True)\n",
    "\n",
    "# lrate  =  tf.keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=5, min_lr=0.00001, verbose = 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define optimizer\n",
    "For training model on multi-gpu, optimizer needs to be defined under the scope as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not use_multiGPU:\n",
    "    opt = tf.optimizers.Adam(learning_rate=learning_rate)\n",
    "else:\n",
    "    with strategy.scope():\n",
    "        opt = tf.optimizers.Adam(learning_rate=learning_rate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Construct a model and compile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not use_multiGPU:\n",
    "    mpnn = MPNN(out_int_dim=512, state_dim=128, T=4)\n",
    "    mpnn.compile(opt, mse, metrics = [mse, log_mse])\n",
    "else:\n",
    "    with strategy.scope():\n",
    "        mpnn = MPNN(out_int_dim=512, state_dim=128, T=4)\n",
    "        mpnn.compile(opt, mse, metrics=[mse, log_mse])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 4/4 [00:00<00:00, 16.47it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(10, 29, 128)\n",
      "(10, 841, 128)\n",
      "(10, 841, 16)\n",
      "(10, 841, 128)\n",
      "(10, 841, 272)\n",
      "(10, 841, 512)\n",
      "(10, 841, 512)\n",
      "(10, 841, 1)\n",
      "ok\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(10, 841, 1), dtype=float32, numpy=\n",
       "array([[[ 1.8868537e-04],\n",
       "        [ 2.4864132e-02],\n",
       "        [ 5.6891907e-03],\n",
       "        ...,\n",
       "        [ 0.0000000e+00],\n",
       "        [ 0.0000000e+00],\n",
       "        [ 0.0000000e+00]],\n",
       "\n",
       "       [[ 4.6886897e-05],\n",
       "        [ 2.4823228e-02],\n",
       "        [ 6.0076960e-03],\n",
       "        ...,\n",
       "        [ 0.0000000e+00],\n",
       "        [ 0.0000000e+00],\n",
       "        [ 0.0000000e+00]],\n",
       "\n",
       "       [[ 6.7952863e-04],\n",
       "        [ 3.6223266e-02],\n",
       "        [ 5.0405585e-03],\n",
       "        ...,\n",
       "        [ 0.0000000e+00],\n",
       "        [ 0.0000000e+00],\n",
       "        [ 0.0000000e+00]],\n",
       "\n",
       "       ...,\n",
       "\n",
       "       [[-9.5222786e-06],\n",
       "        [ 3.5439461e-02],\n",
       "        [ 4.3027271e-03],\n",
       "        ...,\n",
       "        [ 0.0000000e+00],\n",
       "        [ 0.0000000e+00],\n",
       "        [ 0.0000000e+00]],\n",
       "\n",
       "       [[ 5.0220609e-04],\n",
       "        [ 3.5291258e-02],\n",
       "        [ 5.3370581e-03],\n",
       "        ...,\n",
       "        [ 0.0000000e+00],\n",
       "        [ 0.0000000e+00],\n",
       "        [ 0.0000000e+00]],\n",
       "\n",
       "       [[ 7.2502252e-04],\n",
       "        [ 3.5441000e-02],\n",
       "        [ 6.3309888e-03],\n",
       "        ...,\n",
       "        [ 0.0000000e+00],\n",
       "        [ 0.0000000e+00],\n",
       "        [ 0.0000000e+00]]], dtype=float32)>"
      ]
     },
     "execution_count": 105,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# test call\n",
    "mpnn.call([nodes_train[:10], in_edges_train[:10]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Start training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_size = int(len(out_labels)*0.8)\n",
    "batch_size = 16\n",
    "epochs = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data_x = tf.data.Dataset.from_tensor_slices((nodes_train[:train_size], in_edges_train[:train_size]))\n",
    "train_data_y = tf.data.Dataset.from_tensor_slices(out_labels[:train_size])\n",
    "\n",
    "train_data = tf.data.Dataset.zip((train_data_x, train_data_y)).batch(batch_size)\n",
    "\n",
    "valid_data_x = tf.data.Dataset.from_tensor_slices((nodes_train[train_size:], in_edges_train[train_size:]))\n",
    "valid_data_y = tf.data.Dataset.from_tensor_slices(out_labels[train_size:])\n",
    "\n",
    "valid_data = tf.data.Dataset.zip((valid_data_x, valid_data_y)).batch(batch_size)\n",
    "\n",
    "# Disable AutoShard.\n",
    "options = tf.data.Options()\n",
    "options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.OFF\n",
    "train_data = train_data.with_options(options)\n",
    "valid_data= valid_data.with_options(options)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/2\n",
      "Learning rate:  0.001\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 4/4 [00:00<00:00,  6.13it/s]\n"
     ]
    },
    {
     "ename": "ValueError",
     "evalue": "in user code:\n\n    /p/home/jusers/ketkaew1/juwels/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/keras/engine/training.py:805 train_function  *\n        return step_function(self, iterator)\n    /tmp/ipykernel_24066/1473428871.py:6 mse  *\n        preds = tf.boolean_mask(preds,  mask)\n    /p/home/jusers/ketkaew1/juwels/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/util/dispatch.py:201 wrapper  **\n        return target(*args, **kwargs)\n    /p/home/jusers/ketkaew1/juwels/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/ops/array_ops.py:1831 boolean_mask_v2\n        return boolean_mask(tensor, mask, name, axis)\n    /p/home/jusers/ketkaew1/juwels/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/util/dispatch.py:201 wrapper\n        return target(*args, **kwargs)\n    /p/home/jusers/ketkaew1/juwels/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/ops/array_ops.py:1751 boolean_mask\n        shape_tensor[axis:axis + ndims_mask].assert_is_compatible_with(shape_mask)\n    /p/home/jusers/ketkaew1/juwels/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/framework/tensor_shape.py:1134 assert_is_compatible_with\n        raise ValueError(\"Shapes %s and %s are incompatible\" % (self, other))\n\n    ValueError: Shapes (None, 841) and (None, 1) are incompatible\n",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[0;32m/tmp/ipykernel_24066/1024687623.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m mpnn.fit(train_data, \n\u001b[0m\u001b[1;32m      2\u001b[0m          \u001b[0mbatch_size\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mbatch_size\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m          \u001b[0mepochs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mepochs\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m          \u001b[0mcallbacks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mlrate\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstop_early\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m          \u001b[0muse_multiprocessing\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/keras/engine/training.py\u001b[0m in \u001b[0;36mfit\u001b[0;34m(self, x, y, batch_size, epochs, verbose, callbacks, validation_split, validation_data, shuffle, class_weight, sample_weight, initial_epoch, steps_per_epoch, validation_steps, validation_batch_size, validation_freq, max_queue_size, workers, use_multiprocessing)\u001b[0m\n\u001b[1;32m   1098\u001b[0m                 _r=1):\n\u001b[1;32m   1099\u001b[0m               \u001b[0mcallbacks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mon_train_batch_begin\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1100\u001b[0;31m               \u001b[0mtmp_logs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrain_function\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0miterator\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1101\u001b[0m               \u001b[0;32mif\u001b[0m \u001b[0mdata_handler\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshould_sync\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1102\u001b[0m                 \u001b[0mcontext\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0masync_wait\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *args, **kwds)\u001b[0m\n\u001b[1;32m    826\u001b[0m     \u001b[0mtracing_count\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexperimental_get_tracing_count\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    827\u001b[0m     \u001b[0;32mwith\u001b[0m \u001b[0mtrace\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mTrace\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_name\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mtm\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 828\u001b[0;31m       \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwds\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    829\u001b[0m       \u001b[0mcompiler\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m\"xla\"\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_experimental_compile\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0;34m\"nonXla\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    830\u001b[0m       \u001b[0mnew_tracing_count\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexperimental_get_tracing_count\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py\u001b[0m in \u001b[0;36m_call\u001b[0;34m(self, *args, **kwds)\u001b[0m\n\u001b[1;32m    860\u001b[0m       \u001b[0;31m# In this case we have not created variables on the first call. So we can\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    861\u001b[0m       \u001b[0;31m# run the first trace but we should fail if variables are created.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 862\u001b[0;31m       \u001b[0mresults\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_stateful_fn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwds\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    863\u001b[0m       \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_created_variables\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    864\u001b[0m         raise ValueError(\"Creating variables on a non-first call to a function\"\n",
      "\u001b[0;32m~/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/eager/function.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   2939\u001b[0m     \u001b[0;32mwith\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_lock\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2940\u001b[0m       (graph_function,\n\u001b[0;32m-> 2941\u001b[0;31m        filtered_flat_args) = self._maybe_define_function(args, kwargs)\n\u001b[0m\u001b[1;32m   2942\u001b[0m     return graph_function._call_flat(\n\u001b[1;32m   2943\u001b[0m         filtered_flat_args, captured_inputs=graph_function.captured_inputs)  # pylint: disable=protected-access\n",
      "\u001b[0;32m~/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/eager/function.py\u001b[0m in \u001b[0;36m_maybe_define_function\u001b[0;34m(self, args, kwargs)\u001b[0m\n\u001b[1;32m   3355\u001b[0m               \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minput_signature\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m \u001b[0;32mand\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3356\u001b[0m               call_context_key in self._function_cache.missed):\n\u001b[0;32m-> 3357\u001b[0;31m             return self._define_function_with_shape_relaxation(\n\u001b[0m\u001b[1;32m   3358\u001b[0m                 args, kwargs, flat_args, filtered_flat_args, cache_key_context)\n\u001b[1;32m   3359\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/eager/function.py\u001b[0m in \u001b[0;36m_define_function_with_shape_relaxation\u001b[0;34m(self, args, kwargs, flat_args, filtered_flat_args, cache_key_context)\u001b[0m\n\u001b[1;32m   3277\u001b[0m           expand_composites=True)\n\u001b[1;32m   3278\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 3279\u001b[0;31m     graph_function = self._create_graph_function(\n\u001b[0m\u001b[1;32m   3280\u001b[0m         args, kwargs, override_flat_arg_shapes=relaxed_arg_shapes)\n\u001b[1;32m   3281\u001b[0m     \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_function_cache\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0marg_relaxed\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mrank_only_cache_key\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mgraph_function\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/eager/function.py\u001b[0m in \u001b[0;36m_create_graph_function\u001b[0;34m(self, args, kwargs, override_flat_arg_shapes)\u001b[0m\n\u001b[1;32m   3194\u001b[0m     \u001b[0marg_names\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mbase_arg_names\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mmissing_arg_names\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3195\u001b[0m     graph_function = ConcreteFunction(\n\u001b[0;32m-> 3196\u001b[0;31m         func_graph_module.func_graph_from_py_func(\n\u001b[0m\u001b[1;32m   3197\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_name\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3198\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_python_function\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py\u001b[0m in \u001b[0;36mfunc_graph_from_py_func\u001b[0;34m(name, python_func, args, kwargs, signature, func_graph, autograph, autograph_options, add_control_dependencies, arg_names, op_return_value, collections, capture_by_value, override_flat_arg_shapes)\u001b[0m\n\u001b[1;32m    988\u001b[0m         \u001b[0m_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moriginal_func\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtf_decorator\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0munwrap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpython_func\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    989\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 990\u001b[0;31m       \u001b[0mfunc_outputs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpython_func\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mfunc_args\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mfunc_kwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    991\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    992\u001b[0m       \u001b[0;31m# invariant: `func_outputs` contains only Tensors, CompositeTensors,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py\u001b[0m in \u001b[0;36mwrapped_fn\u001b[0;34m(*args, **kwds)\u001b[0m\n\u001b[1;32m    632\u001b[0m             \u001b[0mxla_context\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mExit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    633\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 634\u001b[0;31m           \u001b[0mout\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mweak_wrapped_fn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__wrapped__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwds\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    635\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mout\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    636\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py\u001b[0m in \u001b[0;36mwrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    975\u001b[0m           \u001b[0;32mexcept\u001b[0m \u001b[0mException\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m  \u001b[0;31m# pylint:disable=broad-except\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    976\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mhasattr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0me\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"ag_error_metadata\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 977\u001b[0;31m               \u001b[0;32mraise\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mag_error_metadata\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto_exception\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0me\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    978\u001b[0m             \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    979\u001b[0m               \u001b[0;32mraise\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mValueError\u001b[0m: in user code:\n\n    /p/home/jusers/ketkaew1/juwels/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/keras/engine/training.py:805 train_function  *\n        return step_function(self, iterator)\n    /tmp/ipykernel_24066/1473428871.py:6 mse  *\n        preds = tf.boolean_mask(preds,  mask)\n    /p/home/jusers/ketkaew1/juwels/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/util/dispatch.py:201 wrapper  **\n        return target(*args, **kwargs)\n    /p/home/jusers/ketkaew1/juwels/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/ops/array_ops.py:1831 boolean_mask_v2\n        return boolean_mask(tensor, mask, name, axis)\n    /p/home/jusers/ketkaew1/juwels/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/util/dispatch.py:201 wrapper\n        return target(*args, **kwargs)\n    /p/home/jusers/ketkaew1/juwels/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/ops/array_ops.py:1751 boolean_mask\n        shape_tensor[axis:axis + ndims_mask].assert_is_compatible_with(shape_mask)\n    /p/home/jusers/ketkaew1/juwels/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/framework/tensor_shape.py:1134 assert_is_compatible_with\n        raise ValueError(\"Shapes %s and %s are incompatible\" % (self, other))\n\n    ValueError: Shapes (None, 841) and (None, 1) are incompatible\n"
     ]
    }
   ],
   "source": [
    "mpnn.fit(train_data, \n",
    "         batch_size = batch_size, \n",
    "         epochs = epochs, \n",
    "         callbacks = [lrate, stop_early], \n",
    "         use_multiprocessing = True, \n",
    "         initial_epoch = 0, \n",
    "         verbose = 1, \n",
    "         validation_data = valid_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/2\n",
      "Learning rate:  0.001\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-07-16 00:48:56.339532: I tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc:116] None of the MLIR optimization passes are enabled (registered 2)\n",
      "2021-07-16 00:48:56.353625: I tensorflow/core/platform/profile_utils/cpu_utils.cc:112] CPU Frequency: 2799890000 Hz\n",
      "100%|██████████| 4/4 [00:00<00:00,  6.27it/s]\n"
     ]
    },
    {
     "ename": "ValueError",
     "evalue": "in user code:\n\n    /p/home/jusers/ketkaew1/juwels/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/keras/engine/training.py:805 train_function  *\n        return step_function(self, iterator)\n    /tmp/ipykernel_24066/1473428871.py:6 mse  *\n        preds = tf.boolean_mask(preds,  mask)\n    /p/home/jusers/ketkaew1/juwels/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/util/dispatch.py:201 wrapper  **\n        return target(*args, **kwargs)\n    /p/home/jusers/ketkaew1/juwels/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/ops/array_ops.py:1831 boolean_mask_v2\n        return boolean_mask(tensor, mask, name, axis)\n    /p/home/jusers/ketkaew1/juwels/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/util/dispatch.py:201 wrapper\n        return target(*args, **kwargs)\n    /p/home/jusers/ketkaew1/juwels/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/ops/array_ops.py:1751 boolean_mask\n        shape_tensor[axis:axis + ndims_mask].assert_is_compatible_with(shape_mask)\n    /p/home/jusers/ketkaew1/juwels/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/framework/tensor_shape.py:1134 assert_is_compatible_with\n        raise ValueError(\"Shapes %s and %s are incompatible\" % (self, other))\n\n    ValueError: Shapes (16, 841) and (16, 1) are incompatible\n",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[0;32m/tmp/ipykernel_24066/2349956925.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m mpnn.fit([nodes_train[:train_size], in_edges_train[:train_size]], \n\u001b[0m\u001b[1;32m      2\u001b[0m          \u001b[0my\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mout_labels\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0mtrain_size\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m          \u001b[0mbatch_size\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mbatch_size\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m          \u001b[0mepochs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mepochs\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m          \u001b[0mcallbacks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mlrate\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstop_early\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/keras/engine/training.py\u001b[0m in \u001b[0;36mfit\u001b[0;34m(self, x, y, batch_size, epochs, verbose, callbacks, validation_split, validation_data, shuffle, class_weight, sample_weight, initial_epoch, steps_per_epoch, validation_steps, validation_batch_size, validation_freq, max_queue_size, workers, use_multiprocessing)\u001b[0m\n\u001b[1;32m   1098\u001b[0m                 _r=1):\n\u001b[1;32m   1099\u001b[0m               \u001b[0mcallbacks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mon_train_batch_begin\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1100\u001b[0;31m               \u001b[0mtmp_logs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrain_function\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0miterator\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1101\u001b[0m               \u001b[0;32mif\u001b[0m \u001b[0mdata_handler\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshould_sync\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1102\u001b[0m                 \u001b[0mcontext\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0masync_wait\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *args, **kwds)\u001b[0m\n\u001b[1;32m    826\u001b[0m     \u001b[0mtracing_count\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexperimental_get_tracing_count\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    827\u001b[0m     \u001b[0;32mwith\u001b[0m \u001b[0mtrace\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mTrace\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_name\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mtm\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 828\u001b[0;31m       \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwds\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    829\u001b[0m       \u001b[0mcompiler\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m\"xla\"\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_experimental_compile\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0;34m\"nonXla\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    830\u001b[0m       \u001b[0mnew_tracing_count\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexperimental_get_tracing_count\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py\u001b[0m in \u001b[0;36m_call\u001b[0;34m(self, *args, **kwds)\u001b[0m\n\u001b[1;32m    860\u001b[0m       \u001b[0;31m# In this case we have not created variables on the first call. So we can\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    861\u001b[0m       \u001b[0;31m# run the first trace but we should fail if variables are created.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 862\u001b[0;31m       \u001b[0mresults\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_stateful_fn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwds\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    863\u001b[0m       \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_created_variables\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    864\u001b[0m         raise ValueError(\"Creating variables on a non-first call to a function\"\n",
      "\u001b[0;32m~/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/eager/function.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   2939\u001b[0m     \u001b[0;32mwith\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_lock\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2940\u001b[0m       (graph_function,\n\u001b[0;32m-> 2941\u001b[0;31m        filtered_flat_args) = self._maybe_define_function(args, kwargs)\n\u001b[0m\u001b[1;32m   2942\u001b[0m     return graph_function._call_flat(\n\u001b[1;32m   2943\u001b[0m         filtered_flat_args, captured_inputs=graph_function.captured_inputs)  # pylint: disable=protected-access\n",
      "\u001b[0;32m~/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/eager/function.py\u001b[0m in \u001b[0;36m_maybe_define_function\u001b[0;34m(self, args, kwargs)\u001b[0m\n\u001b[1;32m   3355\u001b[0m               \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minput_signature\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m \u001b[0;32mand\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3356\u001b[0m               call_context_key in self._function_cache.missed):\n\u001b[0;32m-> 3357\u001b[0;31m             return self._define_function_with_shape_relaxation(\n\u001b[0m\u001b[1;32m   3358\u001b[0m                 args, kwargs, flat_args, filtered_flat_args, cache_key_context)\n\u001b[1;32m   3359\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/eager/function.py\u001b[0m in \u001b[0;36m_define_function_with_shape_relaxation\u001b[0;34m(self, args, kwargs, flat_args, filtered_flat_args, cache_key_context)\u001b[0m\n\u001b[1;32m   3277\u001b[0m           expand_composites=True)\n\u001b[1;32m   3278\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 3279\u001b[0;31m     graph_function = self._create_graph_function(\n\u001b[0m\u001b[1;32m   3280\u001b[0m         args, kwargs, override_flat_arg_shapes=relaxed_arg_shapes)\n\u001b[1;32m   3281\u001b[0m     \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_function_cache\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0marg_relaxed\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mrank_only_cache_key\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mgraph_function\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/eager/function.py\u001b[0m in \u001b[0;36m_create_graph_function\u001b[0;34m(self, args, kwargs, override_flat_arg_shapes)\u001b[0m\n\u001b[1;32m   3194\u001b[0m     \u001b[0marg_names\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mbase_arg_names\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mmissing_arg_names\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3195\u001b[0m     graph_function = ConcreteFunction(\n\u001b[0;32m-> 3196\u001b[0;31m         func_graph_module.func_graph_from_py_func(\n\u001b[0m\u001b[1;32m   3197\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_name\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3198\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_python_function\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py\u001b[0m in \u001b[0;36mfunc_graph_from_py_func\u001b[0;34m(name, python_func, args, kwargs, signature, func_graph, autograph, autograph_options, add_control_dependencies, arg_names, op_return_value, collections, capture_by_value, override_flat_arg_shapes)\u001b[0m\n\u001b[1;32m    988\u001b[0m         \u001b[0m_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moriginal_func\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtf_decorator\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0munwrap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpython_func\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    989\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 990\u001b[0;31m       \u001b[0mfunc_outputs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpython_func\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mfunc_args\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mfunc_kwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    991\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    992\u001b[0m       \u001b[0;31m# invariant: `func_outputs` contains only Tensors, CompositeTensors,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py\u001b[0m in \u001b[0;36mwrapped_fn\u001b[0;34m(*args, **kwds)\u001b[0m\n\u001b[1;32m    632\u001b[0m             \u001b[0mxla_context\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mExit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    633\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 634\u001b[0;31m           \u001b[0mout\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mweak_wrapped_fn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__wrapped__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwds\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    635\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mout\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    636\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py\u001b[0m in \u001b[0;36mwrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    975\u001b[0m           \u001b[0;32mexcept\u001b[0m \u001b[0mException\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m  \u001b[0;31m# pylint:disable=broad-except\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    976\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mhasattr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0me\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"ag_error_metadata\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 977\u001b[0;31m               \u001b[0;32mraise\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mag_error_metadata\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto_exception\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0me\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    978\u001b[0m             \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    979\u001b[0m               \u001b[0;32mraise\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mValueError\u001b[0m: in user code:\n\n    /p/home/jusers/ketkaew1/juwels/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/keras/engine/training.py:805 train_function  *\n        return step_function(self, iterator)\n    /tmp/ipykernel_24066/1473428871.py:6 mse  *\n        preds = tf.boolean_mask(preds,  mask)\n    /p/home/jusers/ketkaew1/juwels/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/util/dispatch.py:201 wrapper  **\n        return target(*args, **kwargs)\n    /p/home/jusers/ketkaew1/juwels/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/ops/array_ops.py:1831 boolean_mask_v2\n        return boolean_mask(tensor, mask, name, axis)\n    /p/home/jusers/ketkaew1/juwels/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/util/dispatch.py:201 wrapper\n        return target(*args, **kwargs)\n    /p/home/jusers/ketkaew1/juwels/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/ops/array_ops.py:1751 boolean_mask\n        shape_tensor[axis:axis + ndims_mask].assert_is_compatible_with(shape_mask)\n    /p/home/jusers/ketkaew1/juwels/jupyter/kernels/py38-tf/lib/python3.8/site-packages/tensorflow/python/framework/tensor_shape.py:1134 assert_is_compatible_with\n        raise ValueError(\"Shapes %s and %s are incompatible\" % (self, other))\n\n    ValueError: Shapes (16, 841) and (16, 1) are incompatible\n"
     ]
    }
   ],
   "source": [
    "# mpnn.fit([nodes_train[:train_size], in_edges_train[:train_size]], \n",
    "#          y = out_labels[:train_size], \n",
    "#          batch_size = batch_size, \n",
    "#          epochs = epochs, \n",
    "#          callbacks = [lrate, stop_early], \n",
    "#          use_multiprocessing = True, \n",
    "#          initial_epoch = 0, \n",
    "#          verbose = 1, \n",
    "#          validation_data = ([nodes_train[train_size:], in_edges_train[train_size:]],\n",
    "#                             out_labels[train_size:]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 4/4 [00:00<00:00,  6.37it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "16/16 [==============================] - 7s 367ms/step\n"
     ]
    }
   ],
   "source": [
    "preds = mpnn.predict([nodes_test, in_edges_test], use_multiprocessing=True, verbose=1)\n",
    "np.save(\"preds_kernel.npy\" , preds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(500, 841, 1)"
      ]
     },
     "execution_count": 69,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "preds.shape"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.10 ('py38')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "f2b0cf5112dfd2af73b780f4815bf58f6493434a6fee69d3bee328e26a09a438"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}