{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":58266,"databundleVersionId":6641124,"sourceType":"competition"},{"sourceId":144043388,"sourceType":"kernelVersion"},{"sourceId":144045966,"sourceType":"kernelVersion"},{"sourceId":144045983,"sourceType":"kernelVersion"}],"dockerImageVersionId":30626,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n        break\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-12-30T03:34:29.883725Z","iopub.execute_input":"2023-12-30T03:34:29.885183Z","iopub.status.idle":"2023-12-30T03:34:31.391930Z","shell.execute_reply.started":"2023-12-30T03:34:29.885125Z","shell.execute_reply":"2023-12-30T03:34:31.390776Z"},"collapsed":true,"jupyter":{"outputs_hidden":true},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install torch_geometric","metadata":{"execution":{"iopub.status.busy":"2023-12-18T06:37:11.172206Z","iopub.execute_input":"2023-12-18T06:37:11.173843Z","iopub.status.idle":"2023-12-18T06:37:21.673802Z","shell.execute_reply.started":"2023-12-18T06:37:11.173778Z","shell.execute_reply":"2023-12-18T06:37:21.671598Z"},"collapsed":true,"jupyter":{"outputs_hidden":true},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install tensorflow_gnn --pre\n!pip install tensorflow_ranking","metadata":{"execution":{"iopub.status.busy":"2023-12-18T06:37:21.678728Z","iopub.execute_input":"2023-12-18T06:37:21.680764Z","iopub.status.idle":"2023-12-18T06:37:42.737734Z","shell.execute_reply.started":"2023-12-18T06:37:21.680695Z","shell.execute_reply":"2023-12-18T06:37:42.736034Z"},"collapsed":true,"jupyter":{"outputs_hidden":true},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from tqdm import tqdm\nfrom sklearn.model_selection import train_test_split\nimport torch\nfrom torch import nn\nfrom torch import Tensor\nfrom torch_geometric.nn import GCNConv,SAGEConv\nfrom torch_geometric.datasets import Planetoid\nfrom torch.utils.data import DataLoader,Dataset\nfrom timm.scheduler import CosineLRScheduler\nimport matplotlib.pyplot as plt\ndevice=\"cuda\" if torch.cuda.is_available() else \"cpu\"\nimport tensorflow as tf\nimport tensorflow_gnn as tfgnn\nimport tensorflow_ranking as tfr","metadata":{"execution":{"iopub.status.busy":"2023-12-18T06:37:42.740851Z","iopub.execute_input":"2023-12-18T06:37:42.741325Z","iopub.status.idle":"2023-12-18T06:37:42.753093Z","shell.execute_reply.started":"2023-12-18T06:37:42.741283Z","shell.execute_reply":"2023-12-18T06:37:42.751209Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import tpugraphsv1_layout_data_py as layout_data\nimport tpugraphsv1_tile_data_py as title_data\nimport tpugraphsv1_implicit_py as implicit","metadata":{"execution":{"iopub.status.busy":"2023-12-18T06:38:10.607203Z","iopub.execute_input":"2023-12-18T06:38:10.607882Z","iopub.status.idle":"2023-12-18T06:38:10.615251Z","shell.execute_reply.started":"2023-12-18T06:38:10.607832Z","shell.execute_reply":"2023-12-18T06:38:10.613711Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def _mlp(dims, hidden_activation, l2reg=1e-4, use_bias=True):\n  \"\"\"Helper function for multi-layer perceptron (MLP).\"\"\"\n  layers = []\n  for i, dim in enumerate(dims):\n    if i > 0:\n      layers.append(tf.keras.layers.Activation(hidden_activation))\n    layers.append(tf.keras.layers.Dense(\n        dim, kernel_regularizer=tf.keras.regularizers.l2(l2reg),\n        use_bias=use_bias))\n  return tf.keras.Sequential(layers)\n\n\nclass _OpEmbedding(tf.keras.Model):\n  \"\"\"Embeds GraphTensor.node_sets['op']['op'] nodes into feature 'op_e'.\"\"\"\n\n  def __init__(self, num_ops: int, embed_d: int, l2reg: float = 1e-4):\n    super().__init__()\n    self.embedding_layer = tf.keras.layers.Embedding(\n        num_ops, embed_d, activity_regularizer=tf.keras.regularizers.l2(l2reg))\n\n  def call(\n      self, graph: tfgnn.GraphTensor,\n      training: bool = False) -> tfgnn.GraphTensor:\n    op_features = dict(graph.node_sets['op'].features)\n    op_features['op_e'] = self.embedding_layer(\n        tf.cast(graph.node_sets['op']['op'], tf.int32))\n    return graph.replace_features(node_sets={'op': op_features})","metadata":{"execution":{"iopub.status.busy":"2023-12-18T06:38:12.664578Z","iopub.execute_input":"2023-12-18T06:38:12.665042Z","iopub.status.idle":"2023-12-18T06:38:12.675925Z","shell.execute_reply.started":"2023-12-18T06:38:12.665009Z","shell.execute_reply":"2023-12-18T06:38:12.674766Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"LAYOUT_DATA_ROOT = '/kaggle/input/predict-ai-model-runtime/npz_all/npz/layout'\nSOURCE = 'xla'  # Can be \"xla\" or \"nlp\"\nSEARCH = 'random'  # Can be \"random\" or \"default\"\n\n# Batch size information.\nBATCH_SIZE = 16  # Number of graphs per batch.\nCONFIGS_PER_GRAPH = 5  # Number of configurations (features and target values) per graph.\nMAX_KEEP_NODES = 1000","metadata":{"execution":{"iopub.status.busy":"2023-12-18T06:38:15.541407Z","iopub.execute_input":"2023-12-18T06:38:15.542105Z","iopub.status.idle":"2023-12-18T06:38:15.549853Z","shell.execute_reply.started":"2023-12-18T06:38:15.542074Z","shell.execute_reply":"2023-12-18T06:38:15.547041Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nlayout_data_root_dir = os.path.join(\n      os.path.expanduser(LAYOUT_DATA_ROOT), SOURCE, SEARCH)\n\nlayout_npz_dataset = layout_data.get_npz_dataset(\n    layout_data_root_dir,\n    min_train_configs=CONFIGS_PER_GRAPH,\n    max_train_configs=500,  # If any graph has more than this configurations, it will be filtered [speeds up loading + training]\n    cache_dir='cache'\n)","metadata":{"execution":{"iopub.status.busy":"2023-12-18T06:38:17.993838Z","iopub.execute_input":"2023-12-18T06:38:17.994332Z","iopub.status.idle":"2023-12-18T06:38:36.888474Z","shell.execute_reply.started":"2023-12-18T06:38:17.994295Z","shell.execute_reply":"2023-12-18T06:38:36.886797Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def pair_layout_graph_with_label(graph: tfgnn.GraphTensor):\n    \n    label=tf.cast(graph.node_sets[\"g\"][\"runtimes\"],tf.float32)/1e7\n    return graph,label\n\n","metadata":{"execution":{"iopub.status.busy":"2023-12-18T06:38:50.508311Z","iopub.execute_input":"2023-12-18T06:38:50.508781Z","iopub.status.idle":"2023-12-18T06:38:50.516507Z","shell.execute_reply.started":"2023-12-18T06:38:50.508754Z","shell.execute_reply":"2023-12-18T06:38:50.514622Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"layout_train=(\n    layout_npz_dataset.train.get_graph_tensors_dataset(\n        CONFIGS_PER_GRAPH,max_nodes=MAX_KEEP_NODES)\n        .shuffle(100,reshuffle_each_iteration=True)\n        .batch(BATCH_SIZE,drop_remainder=False)\n        .map(tfgnn.GraphTensor.merge_batch_to_components)\n        .map(pair_layout_graph_with_label))\nlayout_valid=(\n    layout_npz_dataset.validation.get_graph_tensors_dataset(\n        CONFIGS_PER_GRAPH)\n        .batch(BATCH_SIZE,drop_remainder=False)\n        .map(tfgnn.GraphTensor.merge_batch_to_components)\n        .map(pair_layout_graph_with_label))","metadata":{"execution":{"iopub.status.busy":"2023-12-18T06:38:53.060238Z","iopub.execute_input":"2023-12-18T06:38:53.060758Z","iopub.status.idle":"2023-12-18T06:38:55.118736Z","shell.execute_reply.started":"2023-12-18T06:38:53.060722Z","shell.execute_reply":"2023-12-18T06:38:55.117737Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"graph_batch,config_runtimes=next(iter(layout_train.take(1)))\nprint(\"graph_batch=\")\nprint(graph_batch)\nprint(\"\\n\\n\")\nprint(\"config_runtimes=\")\nprint(config_runtimes)","metadata":{"execution":{"iopub.status.busy":"2023-12-18T06:38:55.120569Z","iopub.execute_input":"2023-12-18T06:38:55.120870Z","iopub.status.idle":"2023-12-18T06:38:57.948987Z","shell.execute_reply.started":"2023-12-18T06:38:55.120844Z","shell.execute_reply":"2023-12-18T06:38:57.947453Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print('graph_batch.context =', graph_batch.context)\n# Note: graph_batch.context.sizes must be equal to BATCH_SIZE.\n# Lets print-out all features for all nodesets.\n\nfor node_set_name in sorted(graph_batch.node_sets.keys()):\n    print(f'\\n\\n #####  NODE SET \"{node_set_name}\" #########')\n    print('** Has sizes: ', graph_batch.node_sets[node_set_name].sizes)\n    for feature_name in graph_batch.node_sets[node_set_name].features.keys():\n        print(f'\\n Feature \"{feature_name}\" has values')\n        print(graph_batch.node_sets[node_set_name][feature_name])","metadata":{"execution":{"iopub.status.busy":"2023-12-18T06:38:57.952353Z","iopub.execute_input":"2023-12-18T06:38:57.952790Z","iopub.status.idle":"2023-12-18T06:38:57.962357Z","shell.execute_reply.started":"2023-12-18T06:38:57.952759Z","shell.execute_reply":"2023-12-18T06:38:57.961505Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print('\\n config edge set: ', graph_batch.edge_sets['config'])  \nprint('\\n config source nodes: ', graph_batch.edge_sets['config'].adjacency.source)\nprint('\\n config target nodes: ', graph_batch.edge_sets['config'].adjacency.target)\nprint('\\n g_op edge set: ', graph_batch.edge_sets['g_op'])\nprint('\\n g_config edge set: ', graph_batch.edge_sets['g_config'])\n","metadata":{"execution":{"iopub.status.busy":"2023-12-18T06:39:05.211305Z","iopub.execute_input":"2023-12-18T06:39:05.211746Z","iopub.status.idle":"2023-12-18T06:39:05.221077Z","shell.execute_reply.started":"2023-12-18T06:39:05.211718Z","shell.execute_reply":"2023-12-18T06:39:05.220300Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(graph_batch.edge_sets['config'])\nprint(graph_batch.edge_sets['config'].adjacency.source)\nprint(graph_batch.edge_sets['config'].adjacency.target)","metadata":{"execution":{"iopub.status.busy":"2023-12-18T06:39:07.710532Z","iopub.execute_input":"2023-12-18T06:39:07.710973Z","iopub.status.idle":"2023-12-18T06:39:07.717811Z","shell.execute_reply.started":"2023-12-18T06:39:07.710944Z","shell.execute_reply":"2023-12-18T06:39:07.717091Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"ops=layout_npz_dataset.num_ops\nprint(f\"number of ops in the dataset:\",{ops})\nembedding_layer = _OpEmbedding(ops,16)\ngraph_batch_embedded_ops=embedding_layer(graph_batch)\nprint(f\"\\n\\n Before embedding,node-set op=\\n' ,{graph_batch.node_sets['op']}\")\nprint(f\"\\n\\n After embedding,node-set op=\\n' ,{graph_batch_embedded_ops.node_sets['op']}\")","metadata":{"execution":{"iopub.status.busy":"2023-12-18T06:39:09.363082Z","iopub.execute_input":"2023-12-18T06:39:09.363565Z","iopub.status.idle":"2023-12-18T06:39:09.433362Z","shell.execute_reply.started":"2023-12-18T06:39:09.363518Z","shell.execute_reply":"2023-12-18T06:39:09.430805Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"op_e=graph_batch_embedded_ops.node_sets['op']['op_e']\nconfig_features=graph_batch_embedded_ops.node_sets['nconfig']['feats']\nprint(f\"op_e.shape ==  {op_e.shape}\")\nprint(f\"config_feature.shape == {config_features.shape}\")","metadata":{"execution":{"iopub.status.busy":"2023-12-18T06:39:12.101591Z","iopub.execute_input":"2023-12-18T06:39:12.102009Z","iopub.status.idle":"2023-12-18T06:39:12.110726Z","shell.execute_reply.started":"2023-12-18T06:39:12.101980Z","shell.execute_reply":"2023-12-18T06:39:12.109093Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"config_adj = implicit.AdjacencyMultiplier(graph_batch_embedded_ops, 'config')\nprint('config_adj.shape =', config_adj.shape)\nresized_config_features = config_adj @ config_features\nprint('resized_config_features.shape =', resized_config_features.shape)","metadata":{"execution":{"iopub.status.busy":"2023-12-18T06:39:13.790744Z","iopub.execute_input":"2023-12-18T06:39:13.791177Z","iopub.status.idle":"2023-12-18T06:39:13.808113Z","shell.execute_reply.started":"2023-12-18T06:39:13.791149Z","shell.execute_reply":"2023-12-18T06:39:13.805791Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class ResModel(tf.keras.Model):\n    \"\"\"GNN with residual connections.\"\"\"\n\n    def __init__(\n        self, num_configs: int, num_ops: int, op_embed_dim: int = 32,\n        num_gnns: int = 2, mlp_layers: int = 2,\n        hidden_activation: str = 'leaky_relu',\n        hidden_dim: int = 32, reduction: str = 'sum'):\n        super().__init__()\n        self._num_configs = num_configs\n        self._num_ops = num_ops\n        self._op_embedding = _OpEmbedding(num_ops, op_embed_dim)\n        self._prenet = _mlp([hidden_dim] * mlp_layers, hidden_activation)\n        self._gc_layers = []\n        for _ in range(num_gnns):\n            self._gc_layers.append(_mlp([hidden_dim] * mlp_layers, hidden_activation))\n        self._postnet = _mlp([hidden_dim, 1], hidden_activation, use_bias=False)\n\n    def call(self, graph: tfgnn.GraphTensor, training: bool = False):\n        del training\n        return self.forward(graph, self._num_configs)\n\n    def _node_level_forward(\n        self, node_features: tf.Tensor,\n        config_features: tf.Tensor,\n        graph: tfgnn.GraphTensor, num_configs: int,\n        edgeset_prefix='') -> tf.Tensor:\n        adj_op_op = implicit.AdjacencyMultiplier(\n            graph, edgeset_prefix+'feed')  # op->op\n        adj_config = implicit.AdjacencyMultiplier(\n            graph, edgeset_prefix+'config')  # nconfig->op\n\n        adj_op_op_hat = (adj_op_op + adj_op_op.transpose()).add_eye()\n        adj_op_op_hat = adj_op_op_hat.normalize_symmetric()\n\n        x = node_features\n\n        x = tf.stack([x] * num_configs, axis=1)\n        config_features = 100 * (adj_config @ config_features)\n        x = tf.concat([config_features, x], axis=-1)\n        x = self._prenet(x)\n        x = tf.nn.leaky_relu(x)\n\n        for layer in self._gc_layers:\n            y = x\n            y = tf.concat([config_features, y], axis=-1)\n            y = tf.nn.leaky_relu(layer(adj_op_op_hat @ y))\n            x += y\n        return x\n\n    def forward(\n        self, graph: tfgnn.GraphTensor, num_configs: int,\n        backprop=True) -> tf.Tensor:\n        graph = self._op_embedding(graph)\n\n        config_features = graph.node_sets['nconfig']['feats']\n        node_features = tf.concat([\n            graph.node_sets['op']['feats'],\n            graph.node_sets['op']['op_e']\n        ], axis=-1)\n\n        x_full = self._node_level_forward(\n            node_features=tf.stop_gradient(node_features),\n            config_features=tf.stop_gradient(config_features),\n            graph=graph, num_configs=num_configs)\n\n        if backprop:\n            x_backprop = self._node_level_forward(\n                node_features=node_features,\n                config_features=config_features,\n                graph=graph, num_configs=num_configs, edgeset_prefix='sampled_')\n\n            is_selected = graph.node_sets['op']['selected']\n            # Need to expand twice as `is_selected` is a vector (num_nodes) but\n            # x_{backprop, full} are 3D tensors (num_nodes, num_configs, num_feats).\n            is_selected = tf.expand_dims(is_selected, axis=-1)\n            is_selected = tf.expand_dims(is_selected, axis=-1)\n            x = tf.where(is_selected, x_backprop, x_full)\n        else:\n            x = x_full\n\n        adj_config = implicit.AdjacencyMultiplier(graph, 'config')\n\n        # Features for configurable nodes.\n        config_feats = (adj_config.transpose() @ x)\n\n        # Global pooling\n        adj_pool_op_sum = implicit.AdjacencyMultiplier(graph, 'g_op').transpose()\n        adj_pool_op_mean = adj_pool_op_sum.normalize_right()\n        adj_pool_config_sum = implicit.AdjacencyMultiplier(\n            graph, 'g_config').transpose()\n        x = self._postnet(tf.concat([\n            # (A D^-1) @ Features\n            adj_pool_op_mean @ x,\n            # l2_normalize( A @ Features )\n            tf.nn.l2_normalize(adj_pool_op_sum @ x, axis=-1),\n            # l2_normalize( A @ Features )\n            tf.nn.l2_normalize(adj_pool_config_sum @ config_feats, axis=-1),\n        ], axis=-1))\n\n        x = tf.squeeze(x, -1)\n\n        return x","metadata":{"execution":{"iopub.status.busy":"2023-12-18T06:39:17.606301Z","iopub.execute_input":"2023-12-18T06:39:17.607017Z","iopub.status.idle":"2023-12-18T06:39:17.624587Z","shell.execute_reply.started":"2023-12-18T06:39:17.606986Z","shell.execute_reply":"2023-12-18T06:39:17.623207Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model=ResModel(CONFIGS_PER_GRAPH,layout_npz_dataset.num_ops)\nloss=tfr.keras.losses.ListMLELoss()\nopt=tf._optimizers.Adam(learning_rate=0.01,clipnorm=0.5)\nmodel.compile(loss=loss,optimizer=opt,metrics=[\n    tfr.keras.metrics.OPAMetric(name=\"opa_metric\"),\n])","metadata":{"execution":{"iopub.status.busy":"2023-12-18T06:39:23.841355Z","iopub.execute_input":"2023-12-18T06:39:23.841854Z","iopub.status.idle":"2023-12-18T06:39:23.895464Z","shell.execute_reply.started":"2023-12-18T06:39:23.841821Z","shell.execute_reply":"2023-12-18T06:39:23.894182Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"early_stopping=10\nbest_params=None\nbest_val_opa=-1\nbest_val_at_epoch=-1\nepochs=2\n\nfor i in range(epochs):\n    history = model.fit(\n        layout_train, epochs=1, verbose=1, validation_data=layout_valid,\n        validation_freq=1)\n\n    train_loss = history.history['loss'][-1]\n    train_opa = history.history['opa_metric'][-1]\n    val_loss = history.history['val_loss'][-1]\n    val_opa = history.history['val_opa_metric'][-1]\n    if val_opa > best_val_opa:\n        best_val_opa = val_opa\n        best_val_at_epoch = i\n        best_params = {v.ref: v + 0 for v in model.trainable_variables}\n        print(' * [@%i] Validation (NEW BEST): %s' % (i, str(val_opa)))\n    elif early_stopping > 0 and i - best_val_at_epoch >= early_stopping:\n      print('[@%i] Best accuracy was attained at epoch %i. Stopping.' % (i, best_val_at_epoch))\n      break\n\n# Restore best parameters.\nprint('Restoring parameters corresponding to the best validation OPA.')\nassert best_params is not None\nfor v in model.trainable_variables:\n    v.assign(best_params[v.ref])","metadata":{"execution":{"iopub.status.busy":"2023-12-18T06:39:26.123344Z","iopub.execute_input":"2023-12-18T06:39:26.124658Z","iopub.status.idle":"2023-12-18T06:46:00.039221Z","shell.execute_reply.started":"2023-12-18T06:39:26.124619Z","shell.execute_reply":"2023-12-18T06:46:00.038404Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from tensorflow.keras import layers\nclass CNConv(layers.Layer):\n    def __init__(self, num_filters, k, activation='relu', **kwargs):\n        super(CNConv, self).__init__(**kwargs)\n        self.num_filters = num_filters\n        self.k = k\n        self.activation = tf.keras.activations.get(activation)\n\n    def build(self, input_shape):\n        self.kernel = self.add_weight(\"kernel\",\n                                      (self.k, input_shape[-1], self.num_filters))\n\n    def cheb_polynomials(self, laplacian, k):\n        cheb_polys = [tf.eye(tf.shape(laplacian)[0])]  # T_0(x) = I\n        cheb_polys.append(laplacian)  # T_1(x) = L\n        for i in range(2, k):\n            cheb_polys.append(2 * laplacian * cheb_polys[-1] - cheb_polys[-2])\n        return cheb_polys\n\n    def call(self, inputs, laplacian):\n        laplacian = tf.convert_to_tensor(laplacian, tf.float32)\n        graph_signal = inputs\n\n        # Compute ChebNet polynomials\n        cheb_polys = self.cheb_polynomials(laplacian, self.k)\n\n        # ChebNet convolution\n        cheb_filter = tf.linalg.matrix_transpose(self.kernel)\n        cheb_conv = tf.matmul(tf.concat([tf.matmul(graph_signal, cheb_poly) for cheb_poly in cheb_polys], axis=-1), cheb_filter)\n        output = self.activation(cheb_conv)\n\n        return output\n\nclass SAGEConv(layers.Layer):\n    def __init__(self, num_neighbors, num_filters, activation='relu', **kwargs):\n        super(SAGEConv, self).__init__(**kwargs)\n        self.num_neighbors = num_neighbors\n        self.num_filters = num_filters\n        self.activation = tf.keras.activations.get(activation)\n\n    def build(self, input_shape):\n        self.kernel = self.add_weight(\"kernel\",\n                                      (input_shape[-1] * (self.num_neighbors + 1),\n                                       self.num_filters))\n\n    def call(self, inputs, adjacency_matrix):\n        adjacency_matrix = tf.cast(adjacency_matrix, tf.float32)\n        graph_signal = inputs\n\n        # Aggregate neighbors using mean\n        aggregated_neighbors = tf.matmul(adjacency_matrix, graph_signal) / (\n                tf.reduce_sum(adjacency_matrix, axis=-1, keepdims=True) + 1e-7)\n\n        # Concatenate aggregated neighbors and input features\n        concatenated_features = tf.concat([graph_signal, aggregated_neighbors], axis=-1)\n\n        # GraphSAGE convolution\n        output = tf.matmul(concatenated_features, self.kernel)\n        output = self.activation(output)\n\n        return output\nclass GraphNeuralNetwork(tf.keras.Model):\n    def __init__(self, ops, embed_dim):\n        super(GraphNeuralNetwork, self).__init__()\n\n        # Embedding layer for \"op\" node set\n        self.embedding_layer = _OpEmbedding(ops, embed_dim)\n\n        # ChebNet layer\n        self.chebnet_layer = CNConv(num_filters=64, k=2)\n\n        # SAGEConv layer\n        self.sageconv_layer = SAGEConv(num_neighbors=3, num_filters=32)\n\n        # Output layer (you may modify this based on your specific task)\n        self.output_layer = tf.keras.layers.Dense(units=1, activation='linear')\n\n    def call(self, graph, training=False):\n        # Embedding layer\n        graph_embedded_ops = self.embedding_layer(graph)\n\n        # ChebNet layer\n        config_adj = implicit.AdjacencyMultiplier(graph_embedded_ops, 'config')\n        chebnet_output = self.chebnet_layer(graph_embedded_ops.node_sets['nconfig']['feats'], config_adj)\n\n        # SAGEConv layer\n        sageconv_output = self.sageconv_layer(chebnet_output, config_adj)\n\n        # Global pooling (you may want to customize this based on your task)\n        global_pooled = tf.reduce_mean(sageconv_output, axis=1)\n\n        # Output layer\n        model_output = self.output_layer(global_pooled)\n\n        return model_output","metadata":{"execution":{"iopub.status.busy":"2023-12-18T06:34:32.910910Z","iopub.execute_input":"2023-12-18T06:34:32.911433Z","iopub.status.idle":"2023-12-18T06:34:32.932434Z","shell.execute_reply.started":"2023-12-18T06:34:32.911396Z","shell.execute_reply":"2023-12-18T06:34:32.930925Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"gnn_model = GraphNeuralNetwork(ops=ops, embed_dim=16)","metadata":{"execution":{"iopub.status.busy":"2023-12-18T06:34:36.416691Z","iopub.execute_input":"2023-12-18T06:34:36.418277Z","iopub.status.idle":"2023-12-18T06:34:36.436280Z","shell.execute_reply.started":"2023-12-18T06:34:36.418236Z","shell.execute_reply":"2023-12-18T06:34:36.435150Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"gnn_model.compile(optimizer='adam', loss='mse')","metadata":{"execution":{"iopub.status.busy":"2023-12-18T06:34:38.621949Z","iopub.execute_input":"2023-12-18T06:34:38.622476Z","iopub.status.idle":"2023-12-18T06:34:38.641799Z","shell.execute_reply.started":"2023-12-18T06:34:38.622438Z","shell.execute_reply":"2023-12-18T06:34:38.640581Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#gnn_model.fit(layout_train, validation_data=layout_valid, epochs=10)","metadata":{"execution":{"iopub.status.busy":"2023-12-18T06:35:49.767581Z","iopub.execute_input":"2023-12-18T06:35:49.768628Z","iopub.status.idle":"2023-12-18T06:35:49.772317Z","shell.execute_reply.started":"2023-12-18T06:35:49.768597Z","shell.execute_reply":"2023-12-18T06:35:49.771597Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import tqdm\n_INFERENCE_CONFIGS_BATCH_SIZE = 50\n\noutput_csv_filename = f'inference_layout_{SOURCE}_{SEARCH}.csv'\nprint('\\n\\n   Running inference on test set ...\\n\\n')\ntest_rankings = []\n\nassert layout_npz_dataset.test.graph_id is not None\nfor graph in tqdm.tqdm(layout_npz_dataset.test.iter_graph_tensors(),\n                       total=layout_npz_dataset.test.graph_id.shape[-1],\n                       desc='Inference'):\n    num_configs = graph.node_sets['g']['runtimes'].shape[-1]\n    all_scores = []\n    for i in tqdm.tqdm(range(0, num_configs, _INFERENCE_CONFIGS_BATCH_SIZE)):\n        end_i = min(i + _INFERENCE_CONFIGS_BATCH_SIZE, num_configs)\n        # Take a cut of the configs.\n        node_set_g = graph.node_sets['g']\n        subconfigs_graph = tfgnn.GraphTensor.from_pieces(\n            edge_sets=graph.edge_sets,\n            node_sets={\n                'op': graph.node_sets['op'],\n                'nconfig': tfgnn.NodeSet.from_fields(\n                    sizes=graph.node_sets['nconfig'].sizes,\n                    features={\n                        'feats': graph.node_sets['nconfig']['feats'][:, i:end_i],\n                    }),\n                'g': tfgnn.NodeSet.from_fields(\n                    sizes=tf.constant([1]),\n                    features={\n                        'graph_id': node_set_g['graph_id'],\n                        'runtimes': node_set_g['runtimes'][:, i:end_i],\n                        'kept_node_ratio': node_set_g['kept_node_ratio'],\n                    })\n            })\n        h = model.forward(subconfigs_graph, num_configs=(end_i - i),\n                          backprop=False)\n        all_scores.append(h[0])\n    all_scores = tf.concat(all_scores, axis=0)\n    graph_id = graph.node_sets['g']['graph_id'][0].numpy().decode()\n    sorted_indices = tf.strings.join(\n        tf.strings.as_string(tf.argsort(all_scores)), ';').numpy().decode()\n    test_rankings.append((graph_id, sorted_indices))\n\nwith tf.io.gfile.GFile(output_csv_filename, 'w') as fout:\n    fout.write('ID,TopConfigs\\n')\n    for graph_id, ranks in test_rankings:\n        fout.write(f'layout:{SOURCE}:{SEARCH}:{graph_id},{ranks}\\n')\nprint('\\n\\n   ***  Wrote', output_csv_filename, '\\n\\n')","metadata":{"execution":{"iopub.status.busy":"2023-12-18T06:36:20.872378Z","iopub.execute_input":"2023-12-18T06:36:20.872900Z","iopub.status.idle":"2023-12-18T06:36:21.078176Z","shell.execute_reply.started":"2023-12-18T06:36:20.872864Z","shell.execute_reply":"2023-12-18T06:36:21.076457Z"},"trusted":true},"execution_count":null,"outputs":[]}]}