).\n"
- ],
- "metadata": {}
+ ]
},
{
"cell_type": "code",
"execution_count": 12,
- "source": [
- "example_structs = val_structs[:10].tolist()\n",
- "example_targets = val_targets[:10].tolist()\n",
- "\n",
- "predicted, stddevs = prob_model.predict(example_structs)\n",
- "uncerts = 2 * stddevs\n"
- ],
+ "metadata": {},
"outputs": [
{
- "output_type": "stream",
"name": "stderr",
+ "output_type": "stream",
"text": [
"/home/awsm/anaconda3/envs/unlockNN/lib/python3.8/site-packages/tensorflow_probability/python/distributions/gaussian_process.py:363: UserWarning: Unable to detect statically whether the number of index_points is 1. As a result, defaulting to treating the marginal GP at `index_points` as a multivariate Gaussian. This makes some methods, like `cdf` unavailable.\n",
" warnings.warn(\n"
]
}
],
- "metadata": {}
+ "source": [
+ "example_structs = val_structs[:10].tolist()\n",
+ "example_targets = val_targets[:10].tolist()\n",
+ "\n",
+ "predicted, stddevs = prob_model.predict(example_structs)\n",
+ "uncerts = 2 * stddevs\n"
+ ]
},
{
"cell_type": "code",
"execution_count": 13,
- "source": [
- "pd.DataFrame(\n",
- " {\n",
- " \"Composition\": [struct.composition.reduced_formula for struct in example_structs],\n",
- " \"Formation energy per atom / eV\": example_targets,\n",
- " \"Predicted / eV\": [\n",
- " f\"{pred:.2f} ± {uncert:.2f}\" for pred, uncert in zip(predicted, uncerts)\n",
- " ],\n",
- " }\n",
- ")\n"
- ],
+ "metadata": {},
"outputs": [
{
- "output_type": "execute_result",
"data": {
"text/html": [
"\n",
@@ -839,36 +828,39 @@
"9 HoTl3 -0.215986 -0.19 ± 0.04"
]
},
+ "execution_count": 13,
"metadata": {},
- "execution_count": 13
+ "output_type": "execute_result"
}
],
- "metadata": {}
+ "source": [
+ "pd.DataFrame(\n",
+ " {\n",
+ " \"Composition\": [struct.composition.reduced_formula for struct in example_structs],\n",
+ " \"Formation energy per atom / eV\": example_targets,\n",
+ " \"Predicted / eV\": [\n",
+ " f\"{pred:.2f} ± {uncert:.2f}\" for pred, uncert in zip(predicted, uncerts)\n",
+ " ],\n",
+ " }\n",
+ ")\n"
+ ]
},
{
"cell_type": "code",
"execution_count": 14,
- "source": [
- "val_metrics = evaluate_uq_metrics(prob_model, val_structs, val_targets)\n",
- "train_metrics = evaluate_uq_metrics(prob_model, train_structs, train_targets)\n",
- "\n",
- "print(\"Validation metrics:\")\n",
- "pprint(val_metrics)\n",
- "print(\"Training metrics:\")\n",
- "pprint(train_metrics)\n"
- ],
+ "metadata": {},
"outputs": [
{
- "output_type": "stream",
"name": "stderr",
+ "output_type": "stream",
"text": [
"/home/awsm/anaconda3/envs/unlockNN/lib/python3.8/site-packages/tensorflow_probability/python/distributions/gaussian_process.py:363: UserWarning: Unable to detect statically whether the number of index_points is 1. As a result, defaulting to treating the marginal GP at `index_points` as a multivariate Gaussian. This makes some methods, like `cdf` unavailable.\n",
" warnings.warn(\n"
]
},
{
- "output_type": "stream",
"name": "stdout",
+ "output_type": "stream",
"text": [
"Validation metrics:\n",
"{'mae': 0.04725769517605442,\n",
@@ -887,7 +879,15 @@
]
}
],
- "metadata": {}
+ "source": [
+ "val_metrics = evaluate_uq_metrics(prob_model, val_structs, val_targets)\n",
+ "train_metrics = evaluate_uq_metrics(prob_model, train_structs, train_targets)\n",
+ "\n",
+ "print(\"Validation metrics:\")\n",
+ "pprint(val_metrics)\n",
+ "print(\"Training metrics:\")\n",
+ "pprint(train_metrics)\n"
+ ]
}
],
"metadata": {
@@ -914,4 +914,4 @@
},
"nbformat": 4,
"nbformat_minor": 2
-}
\ No newline at end of file
+}
diff --git a/paper/paper.bib b/paper/paper.bib
index 3a94081..ae90ec2 100644
--- a/paper/paper.bib
+++ b/paper/paper.bib
@@ -194,22 +194,20 @@ @article{schuttSchNetDeepLearning2018
file = {/home/awsm/Zotero/storage/DEJUW4IW/Schütt et al. - 2018 - SchNet – A deep learning architecture for molecule.pdf}
}
-@article{tranMethodsComparingUncertainty2020a,
+@online{tranMethodsComparingUncertainty2020,
+ ids = {tranMethodsComparingUncertainty2020a},
title = {Methods for Comparing Uncertainty Quantifications for Material Property Predictions},
author = {Tran, Kevin and Neiswanger, Willie and Yoon, Junwoong and Zhang, Qingyang and Xing, Eric and Ulissi, Zachary W.},
- date = {2020-05},
- journaltitle = {Machine Learning: Science and Technology},
- shortjournal = {Mach. Learn.: Sci. Technol.},
- volume = {1},
- number = {2},
- pages = {025006},
+ date = {2020-02-20},
+ eprint = {1912.10066},
+ eprinttype = {arxiv},
+ primaryclass = {cond-mat, physics:physics},
publisher = {{IOP Publishing}},
- issn = {2632-2153},
- doi = {10.1088/2632-2153/ab7e1a},
- url = {https://doi.org/10.1088/2632-2153/ab7e1a},
- urldate = {2021-09-04},
+ url = {http://arxiv.org/abs/1912.10066},
+ urldate = {2020-09-20},
abstract = {Data science and informatics tools have been proliferating recently within the computational materials science and catalysis fields. This proliferation has spurned the creation of various frameworks for automated materials screening, discovery, and design. Underpinning these frameworks are surrogate models with uncertainty estimates on their predictions. These uncertainty estimates are instrumental for determining which materials to screen next, but the computational catalysis field does not yet have a standard procedure for judging the quality of such uncertainty estimates. Here we present a suite of figures and performance metrics derived from the machine learning community that can be used to judge the quality of such uncertainty estimates. This suite probes the accuracy, calibration, and sharpness of a model quantitatively. We then show a case study where we judge various methods for predicting density-functional-theory-calculated adsorption energies. Of the methods studied here, we find that the best performer is a model where a convolutional neural network is used to supply features to a Gaussian process regressor, which then makes predictions of adsorption energies along with corresponding uncertainty estimates.},
- langid = {english},
+ archiveprefix = {arXiv},
+ keywords = {Condensed Matter - Materials Science,Physics - Computational Physics},
file = {/home/awsm/Zotero/storage/4I3UMCUN/Tran et al. - 2020 - Methods for comparing uncertainty quantifications .pdf}
}
diff --git a/paper/paper.md b/paper/paper.md
index 6e451be..302011a 100644
--- a/paper/paper.md
+++ b/paper/paper.md
@@ -54,7 +54,10 @@ UnlockNN provides an API to add uncertainty quantification to Keras-based models
and comes packaged with a specific implementation for compatibility with MEGNet
[@chenGraphNetworksUniversal2019], which is a graph NN implementation for
materials property prediction that has achieved state-of-the-art accuracy on
-many benchmark tasks [@dunnBenchmarkingMaterialsProperty2020].
+many benchmark tasks [@dunnBenchmarkingMaterialsProperty2020]. The package is
+designed for materials and chemistry researchers to improve their models'
+reliability and identify the domain(s) of materials on which the models can
+perform well.
This uncertainty quantification is achieved by supplanting the output layer of
the model with a variational Gaussian process (VGP)
@@ -75,12 +78,14 @@ the training iterations needed to train the base NN it is modifying.
The primary interface for unlockNN is the `model` module, which contains an
extensible `ProbNN` class for adding uncertainty quantification to arbitrary
Keras models. It also contains a `MEGNetProbModel` class for adding uncertainty
-quantification to MEGNet.
+quantification to MEGNet, and a download suite for accessing example models
+and data. The models can be configured with different kernels for the VGP,
+which are implemented in a `kernel_layers` module.
Neural network-fed VGPs share a similar principle to the convolution-fed
-Gaussian processes formulated by @tranMethodsComparingUncertainty2020a. UnlockNN
+Gaussian processes formulated by @tranMethodsComparingUncertainty2020. UnlockNN
also implements tools for calculating the performance metrics suggested by
-@tranMethodsComparingUncertainty2020a, including sharpness and calibration error,
+@tranMethodsComparingUncertainty2020, including sharpness and calibration error,
via its `metrics` module.
## Acknowledgements
diff --git a/pytest.ini b/pytest.ini
new file mode 100644
index 0000000..19f5fb5
--- /dev/null
+++ b/pytest.ini
@@ -0,0 +1,10 @@
+[pytest]
+addopts = --cov=unlocknn
+testpaths =
+ tests
+filterwarnings =
+ ignore::DeprecationWarning
+ ignore::PendingDeprecationWarning
+ ignore:Unable to detect statically whether the number of index_points is 1:UserWarning
+ ignore:Converting sparse .* to a dense Tensor of unknown shape:UserWarning
+ ignore:Custom mask layers require a config and must override get_config:Warning
diff --git a/tests/test_download.py b/tests/test_download.py
index ab724d8..daf4313 100644
--- a/tests/test_download.py
+++ b/tests/test_download.py
@@ -2,22 +2,35 @@
import os
from pathlib import Path
-os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" # see issue #152
+import pytest
+import requests
+from pytest_mock import MockerFixture
+
+os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+# See https://github.com/tensorflow/tensorflow/issues/152#issuecomment-273663277
os.environ["CUDA_VISIBLE_DEVICES"] = ""
from unlocknn.download import load_data, load_pretrained
-def test_load_model(tmp_path: Path):
+def test_load_model(tmp_path: Path, mocker: MockerFixture):
"""Test downloading formation energies model."""
- orig_model = load_pretrained("binary_e_form", save_dir=tmp_path)
- reload_model = load_pretrained("binary_e_form", save_dir=tmp_path)
+ get_spy = mocker.spy(requests, "get")
+ load_pretrained("binary_e_form", save_dir=tmp_path)
+ get_spy.assert_called_once()
+ load_pretrained("binary_e_form", save_dir=tmp_path)
+ # Second call should load from disk
+ get_spy.assert_called_once()
-def test_load_data(tmp_path: Path):
+def test_load_data(tmp_path: Path, mocker: MockerFixture):
"""Test downloading formation energies data."""
+ get_spy = mocker.spy(requests, "get")
orig_data = load_data("binary_e_form", save_dir=tmp_path)
+ get_spy.assert_called_once()
reload_data = load_data("binary_e_form", save_dir=tmp_path)
+ # Second call should load from disk
+ get_spy.assert_called_once()
- assert orig_data.loc[0, "formation_energy_per_atom"] == -0.7374389025000001
- assert reload_data.loc[0, "formation_energy_per_atom"] == -0.7374389025000001
+ assert orig_data.loc[0, "formation_energy_per_atom"] == pytest.approx(-0.7374389025, abs=1e-10)
+ assert reload_data.loc[0, "formation_energy_per_atom"] == pytest.approx(-0.7374389025, abs=1e-10)
diff --git a/tests/test_kernels.py b/tests/test_kernels.py
index 7db5acc..d17ad48 100644
--- a/tests/test_kernels.py
+++ b/tests/test_kernels.py
@@ -1,7 +1,7 @@
"""Test the `kernel_layers` package."""
import os
-os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" # see issue #152
+os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = ""
from pathlib import Path
@@ -9,20 +9,56 @@
import numpy as np
import pytest
+import tensorflow as tf
+import tensorflow_probability as tfp
from unlocknn.kernel_layers import *
+class ExampleKernel(KernelLayer):
+ """An example kernel for testing."""
+
+ def __init__(self, bias_variance: float=1.0, **kwargs):
+ """Initialize the bias_variance parameter."""
+ super().__init__(**kwargs)
+ dtype = kwargs.get("dtype", tf.float64)
+ self.bias_variance = self.add_weight(initializer=tf.constant_initializer(bias_variance), dtype=dtype, name="bias_variance")
+
+ def kernel(self) -> tfp.math.psd_kernels.PositiveSemidefiniteKernel:
+ """Get a linear kernel that's parameterised by a given bias variance."""
+ return tfp.math.psd_kernels.Linear(bias_variance=self.bias_variance)
+
+ @property
+ def config(self) -> dict:
+ """Return an empty dict: no kwargs are required."""
+ return dict()
+
+def test_custom_kernel_reloading(tmp_path: Path):
+ """Test saving and loading with a custom kernel."""
+ save_path = tmp_path / "kernel"
+ kernel = ExampleKernel(bias_variance=2.0)
+ kernel.save(save_path)
+ reload_kernel = load_kernel(save_path, ExampleKernel)
+ assert reload_kernel.get_weights()[0] == pytest.approx(2.0)
+
@pytest.mark.parametrize("kernel_type", [RBFKernelFn, MaternOneHalfFn])
def test_reload(tmp_path: Path, kernel_type: Type[AmpAndLengthScaleFn]):
- """Test saving and reloading a kernel to/from disk."""
- # Example layer weights: [_amplitude, _length_scale]
+ """Test saving and reloading a builtin kernel to/from disk."""
+ # Example layer weights: [_amplitude_basis, _length_scale_basis]
example_weights = [np.array(2.0), np.array(3.0)]
save_dir = tmp_path / "kernel"
orig_kernel = kernel_type()
orig_kernel.set_weights(example_weights)
orig_kernel.save(save_dir)
- loaded_kernel = load_kernel(save_dir)
+ orig_amplitude = orig_kernel.amplitude
+ orig_length_scale = orig_kernel.length_scale
+
+ loaded_kernel: AmpAndLengthScaleFn = load_kernel(save_dir)
assert isinstance(loaded_kernel, kernel_type)
loaded_weights = loaded_kernel.get_weights()
assert loaded_weights == example_weights
+
+ loaded_amp = loaded_kernel.amplitude.numpy()
+ loaded_ls = loaded_kernel.length_scale.numpy()
+ assert loaded_amp == pytest.approx(orig_amplitude)
+ assert loaded_ls == pytest.approx(orig_length_scale)
diff --git a/tests/test_meg_utils.py b/tests/test_meg_utils.py
new file mode 100644
index 0000000..024532f
--- /dev/null
+++ b/tests/test_meg_utils.py
@@ -0,0 +1,51 @@
+"""Test for the MEGNet utilities module"""
+import os
+
+
+os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = ""
+
+from pathlib import Path
+
+import numpy as np
+import pytest
+from megnet.models.megnet import MEGNetModel
+from megnet.utils.preprocessing import Scaler
+from unlocknn.megnet_utils import create_megnet_input
+
+from .utils import datadir, load_df_head
+
+
+class ExampleScaler(Scaler):
+ """An example scaler for testing.
+
+ Performs transforming by dividing by number of atoms,
+ and inverse transforming by multiplying.
+
+ """
+ def transform(self, target: np.ndarray, n: int = 1) -> np.ndarray:
+ return target / n
+
+ def inverse_transform(self, transformed_target: np.ndarray, n: int = 1) -> np.ndarray:
+ return transformed_target * n
+
+def test_input_with_scaler(datadir: Path):
+ """Test input generation."""
+ binary_dir = datadir / "mp_binary_on_hull.pkl"
+ binary_df = load_df_head(binary_dir)
+
+ meg_model = MEGNetModel.from_file(str(datadir / "formation_energy.hdf5"))
+ meg_model.target_scaler = ExampleScaler()
+
+ input_gen, _ = create_megnet_input(
+ meg_model,
+ binary_df["structure"],
+ binary_df["formation_energy_per_atom"],
+ batch_size=100, # We have just one batch
+ shuffle=False
+ )
+
+ # Get first batch (the whole input)
+ _, scaled_targets = input_gen.__getitem__(0)
+ # Check targets are scaled
+ scaled_targets == pytest.approx(binary_df["formation_energy_per_atom"] / binary_df["num_atoms"], rel=1e-6)
diff --git a/tests/test_meg_utils/formation_energy.hdf5 b/tests/test_meg_utils/formation_energy.hdf5
new file mode 100644
index 0000000..86b1cc4
Binary files /dev/null and b/tests/test_meg_utils/formation_energy.hdf5 differ
diff --git a/tests/test_meg_utils/formation_energy.hdf5.json b/tests/test_meg_utils/formation_energy.hdf5.json
new file mode 100644
index 0000000..5b2f386
--- /dev/null
+++ b/tests/test_meg_utils/formation_energy.hdf5.json
@@ -0,0 +1,8 @@
+{"graph_converter": {"@module": "megnet.data.crystal", "@class": "CrystalGraph", "atom_converter": {"@module": "megnet.data.graph", "@class": "DummyConverter"}, "bond_converter": {"@module": "megnet.data.graph", "@class": "GaussianDistance", "centers": {"@module": "numpy", "@class": "array", "dtype": "float64", "data": [0.0, 0.06060606060606061, 0.12121212121212122, 0.18181818181818182, 0.24242424242424243, 0.30303030303030304, 0.36363636363636365, 0.42424242424242425, 0.48484848484848486, 0.5454545454545454, 0.6060606060606061, 0.6666666666666667, 0.7272727272727273, 0.7878787878787878, 0.8484848484848485, 0.9090909090909092, 0.9696969696969697, 1.0303030303030303, 1.0909090909090908, 1.1515151515151516, 1.2121212121212122, 1.2727272727272727, 1.3333333333333335, 1.393939393939394, 1.4545454545454546, 1.5151515151515151, 1.5757575757575757, 1.6363636363636365, 1.696969696969697, 1.7575757575757576, 1.8181818181818183, 1.878787878787879, 1.9393939393939394, 2.0, 2.0606060606060606, 2.121212121212121, 2.1818181818181817, 2.2424242424242427, 2.303030303030303, 2.3636363636363638, 2.4242424242424243, 2.484848484848485, 2.5454545454545454, 2.606060606060606, 2.666666666666667, 2.7272727272727275, 2.787878787878788, 2.8484848484848486, 2.909090909090909, 2.9696969696969697, 3.0303030303030303, 3.090909090909091, 3.1515151515151514, 3.2121212121212124, 3.272727272727273, 3.3333333333333335, 3.393939393939394, 3.4545454545454546, 3.515151515151515, 3.5757575757575757, 3.6363636363636367, 3.6969696969696972, 3.757575757575758, 3.8181818181818183, 3.878787878787879, 3.9393939393939394, 4.0, 4.0606060606060606, 4.121212121212121, 4.181818181818182, 4.242424242424242, 4.303030303030303, 4.363636363636363, 4.424242424242425, 4.484848484848485, 4.545454545454546, 4.606060606060606, 4.666666666666667, 4.7272727272727275, 4.787878787878788, 4.848484848484849, 4.909090909090909, 4.96969696969697, 5.03030303030303, 5.090909090909091, 5.151515151515151, 5.212121212121212, 5.2727272727272725, 5.333333333333334, 5.3939393939393945, 5.454545454545455, 5.515151515151516, 5.575757575757576, 5.636363636363637, 5.696969696969697, 5.757575757575758, 5.818181818181818, 5.878787878787879, 5.9393939393939394, 6.0]}, "width": 0.5}, "cutoff": 5.0},
+ "metadata": {"name": "Formation energy",
+ "unit": "eV/atom",
+ "metric": "mae",
+ "metric value": 0.026,
+ "training dataset": "mp-2019.4.1",
+ "description": "This model was trained using structures and formation energy of Materials Project data base downloaded on April 1, 2019. The total data size is 133420. The model was trained by a 0.8-0.1-0.1, train, validation and test data ratio"}
+}
\ No newline at end of file
diff --git a/tests/test_meg_utils/mp_binary_on_hull.pkl b/tests/test_meg_utils/mp_binary_on_hull.pkl
new file mode 100644
index 0000000..18cb1b2
Binary files /dev/null and b/tests/test_meg_utils/mp_binary_on_hull.pkl differ
diff --git a/tests/test_metrics.py b/tests/test_metrics.py
new file mode 100644
index 0000000..2d3cee6
--- /dev/null
+++ b/tests/test_metrics.py
@@ -0,0 +1,58 @@
+
+"""Test metrics."""
+import os
+from typing import Callable, List, NamedTuple
+
+os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = ""
+
+import pytest
+from unlocknn.megnet_utils import Targets
+from unlocknn.metrics import neg_log_likelihood, MAE, MSE, RMSE, variation, sharpness
+
+class ToyData(NamedTuple):
+ """Container for some example data."""
+
+ predictions: Targets
+ stddevs: Targets
+ true_vals: Targets
+
+class ExpectedMetrics(NamedTuple):
+ """Container for expected values for metrics."""
+
+ nll: float
+ mae: float
+ mse: float
+ rmse: float
+ sharpness: float
+ variation: float
+
+# This list should match the order for ExpectedMetrics
+METRICS_TEST_ORDER: List[Callable[[Targets, Targets, Targets], float]] = [
+ neg_log_likelihood,
+ MAE,
+ MSE,
+ RMSE,
+ sharpness,
+ variation
+]
+
+TEST_DATA: List[ToyData] = [
+ ToyData([0.0, 1.0, 2.0], [1.0, 2.0, 3.0], [0.0, 1.0, 2.0]),
+ ToyData([0.0, 1.0, 2.0], [1.0, 2.0, 3.0], [-1.0, 1.0, 2.0]),
+ ToyData([0.0, 1.0, 2.0], [1.0, 1.0, 3.0], [1.0, 1.0, 2.0]),
+]
+
+# TODO: WIP Metrics values!
+TEST_EXPECTED: List[ExpectedMetrics] = [
+ ExpectedMetrics(4.5486, 0.0, 0.0, 0.0, 2.1602, 0.5),
+ ExpectedMetrics(5.0486, 1/3, 1/3, 0.57735, 2.1602, 0.5),
+ ExpectedMetrics(4.3554, 1/3, 1/3, 0.57735, 1.9149, 0.69282),
+]
+
+
+@pytest.mark.parametrize("toy_data,expected", zip(TEST_DATA, TEST_EXPECTED))
+def test_metrics(toy_data: ToyData, expected: ExpectedMetrics):
+ """Test that the metrics for the toy data match expected values."""
+ for metric, expected_value in zip(METRICS_TEST_ORDER, expected):
+ assert metric(*toy_data) == pytest.approx(expected_value, abs=1e-4)
diff --git a/tests/test_model.py b/tests/test_model.py
index f03f85b..d5cbafd 100644
--- a/tests/test_model.py
+++ b/tests/test_model.py
@@ -1,115 +1,169 @@
"""Test model features."""
import os
-os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" # see issue #152
+os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = ""
import random as python_random
-from distutils import dir_util
-from math import floor
+from collections import Counter
from pathlib import Path
-from typing import List, Tuple
+from typing import List, Set
import numpy as np
-import pandas as pd
import pytest
import tensorflow as tf
from megnet.models import MEGNetModel
+from pymatgen.core.structure import Structure
from unlocknn import MEGNetProbModel
from unlocknn.initializers import SampleInitializer
+from unlocknn.megnet_utils import MEGNetGraph
+
+from .utils import (SplitData, datadir, load_df_head, train_test_split,
+ weights_equal)
np.random.seed(123)
python_random.seed(123)
tf.random.set_seed(123)
-@pytest.fixture
-def datadir(tmpdir, request):
- """Access data directory.
-
- Fixture responsible for searching a folder with the same name of test
- module and, if available, moving all contents to a temporary directory so
- tests can use them freely.
+def structs_to_graphs(meg_model: MEGNetModel, structures: List[Structure]) -> List[MEGNetGraph]:
+ """Convert Structures to graphs."""
+ dummy_targets = [0.0] * len(structures)
+ return meg_model.get_all_graphs_targets(structures, dummy_targets)[0]
- Source: https://stackoverflow.com/a/29631801/
+def find_duplicate_weights(prob_model: MEGNetProbModel) -> Set[str]:
+ """Find duplicate weight names."""
+ count = Counter(weight.name for weight in prob_model.model.weights)
+ return set(pair[0] for pair in count.items() if pair[1] > 1)
+@pytest.fixture
+def split_data(datadir: Path) -> SplitData:
+ """Get some example data, split into training and test subsets.
+
+ Returns:
+ train_structs
+
+ train_targets
+
+ test_structs
+
+ test_targets
+
"""
- filename = request.module.__file__
- test_dir, _ = os.path.splitext(filename)
-
- if os.path.isdir(test_dir):
- dir_util.copy_tree(test_dir, str(tmpdir))
-
- return tmpdir
-
+ binary_dir = datadir / "mp_binary_on_hull.pkl"
+ binary_df = load_df_head(binary_dir)
+ structures = binary_df["structure"].tolist()
+ formation_energies = binary_df["formation_energy_per_atom"].tolist()
-def weights_equal(weights_a: List[np.ndarray], weights_b: List[np.ndarray]) -> bool:
- """Check equality between weights."""
- return all(
- (weight1 == weight2).all() for weight1, weight2 in zip(weights_a, weights_b)
+ return train_test_split(
+ structures, formation_energies
)
-def train_test_split(
- structures: list, targets: list, train_frac: float = 0.8
-) -> Tuple[Tuple[list, list], Tuple[list, list]]:
- """Split structures and targets into training and testing subsets."""
- num_train = floor(len(structures) * train_frac)
- return (
- (structures[:num_train], targets[:num_train]),
- (structures[num_train:], targets[num_train:]),
- )
+def test_sample_init(datadir: Path, split_data: SplitData):
+ """Test the SampleInitializer."""
+ megnet_e_form_model = MEGNetModel.from_file(str(datadir / "formation_energy.hdf5"))
+ (train_structs, _), (_, _) = split_data
+
+ initializer = SampleInitializer(train_structs, megnet_e_form_model, batch_size=32)
+ MEGNetProbModel(megnet_e_form_model, 10, index_initializer=initializer)
+ # If this works without any errors, we're doing OK
-def test_sample_init(datadir: Path):
- """Test the SampleInitializer."""
+@pytest.mark.parametrize("use_norm", [True, False])
+def test_model_reload(tmp_path: Path, datadir: Path, use_norm: bool):
+ """Test saving and reloading a model from disk.
+
+ Asserts weight equality.
+
+ """
+ save_dir = tmp_path / ("norm_model" if use_norm else "unnorm_model")
megnet_e_form_model = MEGNetModel.from_file(str(datadir / "formation_energy.hdf5"))
- binary_dir = datadir / "mp_binary_on_hull.pkl"
+ prob_model = MEGNetProbModel(megnet_e_form_model, 10, use_normalization=use_norm)
- try:
- binary_df = pd.read_pickle(binary_dir)[:100]
- except ValueError:
- # Older python version
- import pickle5 as pkl
+ prob_model.save(save_dir, ckpt_path=None)
+ loaded_model = MEGNetProbModel.load(save_dir, load_ckpt=False)
+ assert weights_equal(
+ prob_model.model.get_weights(), loaded_model.model.get_weights()
+ )
- with binary_dir.open("rb") as f:
- binary_df = pkl.load(f)
+@pytest.mark.parametrize("use_graphs", [True, False])
+@pytest.mark.parametrize("use_norm", [True, False])
+def test_model_prediction(datadir: Path, split_data: SplitData, use_norm: bool, use_graphs: bool):
+ """Test that model prediction has expected dimensions."""
+ model_name = "prob_e_form_" + ("" if use_norm else "un") + "norm"
+ prob_model = MEGNetProbModel.load(datadir / model_name, load_ckpt=False)
+ (train_structs, _), (_, _) = split_data
+ train_input = structs_to_graphs(prob_model.meg_model, train_structs) if use_graphs else train_structs
- structures = binary_df["structure"].tolist()
- formation_energies = binary_df["formation_energy_per_atom"].tolist()
+ prob_model.update_pred_model()
- (train_structs, _), (_, _) = train_test_split(structures, formation_energies)
- initializer = SampleInitializer(train_structs, megnet_e_form_model, batch_size=32)
- MEGNetProbModel(megnet_e_form_model, 10, index_initializer=initializer)
- # If this works without any errors, we're doing OK
+ prediction, stddev = prob_model.predict(train_input, batch_size=16)
+ # Arrays should have the same shape: flat, with one entry per structure
+ expected_shape = (len(train_structs),)
+ assert prediction.shape == expected_shape
+ assert stddev.shape == expected_shape
+
+@pytest.mark.parametrize("use_graphs", [True, False])
@pytest.mark.parametrize("use_norm", [True, False])
-def test_meg_prob(tmp_path: Path, datadir: Path, use_norm: bool):
- """Test creation, training and I/O of a `MEGNetProbModel`."""
+def test_model_training(tmp_path: Path, datadir: Path, split_data: SplitData, use_norm: bool, use_graphs: bool):
+ """Test training then saving a model.
+
+ Check if expected weights update/are frozen. Saves model after training to
+ check that checkpoints are handled correctly.
+
+ """
save_dir = tmp_path / ("norm_model" if use_norm else "unnorm_model")
ckpt_path = tmp_path / "checkpoint.h5"
- megnet_e_form_model = MEGNetModel.from_file(str(datadir / "formation_energy.hdf5"))
- binary_dir = datadir / "mp_binary_on_hull.pkl"
+ model_name = "prob_e_form_{}norm".format("" if use_norm else "un")
+ (train_structs, train_targets), (test_structs, test_targets) = split_data
+ last_nn_idx = -2 if use_norm else -1
- try:
- binary_df = pd.read_pickle(binary_dir)[:100]
- except ValueError:
- # Older python version
- import pickle5 as pkl
+ prob_model = MEGNetProbModel.load(datadir / model_name, load_ckpt=False)
+ train_input = structs_to_graphs(prob_model.meg_model, train_structs) if use_graphs else train_structs
+ test_input = structs_to_graphs(prob_model.meg_model, test_structs) if use_graphs else test_structs
- with binary_dir.open("rb") as f:
- binary_df = pkl.load(f)
+ for layer in prob_model.model.layers:
+ print(layer.name)
- structures = binary_df["structure"].tolist()
- formation_energies = binary_df["formation_energy_per_atom"].tolist()
+ init_weights = [layer.get_weights() for layer in prob_model.model.layers]
+ init_nn_weights = init_weights[:last_nn_idx]
+ init_uq_weights = init_weights[last_nn_idx:]
- (train_structs, train_targets), (test_structs, test_targets) = train_test_split(
- structures, formation_energies
+ # For initial training, we expect the weights before the
+ # `Norm`/`VGP` layers to be frozen initially, so these shouldn't change.
+ # The rest should, as the model isn't optimised.
+ prob_model.train(
+ train_input, train_targets, 1, test_input, test_targets, batch_size=32, ckpt_path=ckpt_path
)
+
+ final_weights = [layer.get_weights() for layer in prob_model.model.layers]
+ final_nn_weights = final_weights[:last_nn_idx]
+ final_uq_weights = final_weights[last_nn_idx:]
+
+ for init_nn_weight, final_nn_weight in zip(init_nn_weights, final_nn_weights):
+ assert weights_equal(init_nn_weight, final_nn_weight)
+
+ for init_uq_weight, final_uq_weight in zip(init_uq_weights, final_uq_weights):
+ assert not weights_equal(init_uq_weight, final_uq_weight)
+
+ # * Test saving the model
+ # If this runs without errors, we're OK.
+ prob_model.save(save_dir, ckpt_path=ckpt_path)
+
+@pytest.mark.parametrize("use_norm", [True, False])
+def test_meg_weights_preserved(datadir: Path, use_norm: bool):
+ """Test that the MEGNet weights are correctly transferred to a MEGNetProbModel."""
+ megnet_e_form_model = MEGNetModel.from_file(str(datadir / "formation_energy.hdf5"))
+
+ # * Initialize model
prob_model = MEGNetProbModel(megnet_e_form_model, 10, use_normalization=use_norm)
- # Test weights equality
+ # * Test weights equality
+ # Ensure that the weights of the model that we copied over are unchanged
+ # and in the correct layers.
last_nn_idx = -2 if use_norm else -1
meg_nn_weights = [
layer.get_weights() for layer in megnet_e_form_model.model.layers[:-1]
@@ -120,30 +174,50 @@ def test_meg_prob(tmp_path: Path, datadir: Path, use_norm: bool):
for meg_layer, prob_layer in zip(meg_nn_weights, prob_model_nn_weights):
assert weights_equal(meg_layer, prob_layer)
- init_performance = prob_model.evaluate(test_structs, test_targets)
- init_loss = init_performance["loss"]
-
- # Test training without validation
- prob_model.train(
- train_structs, train_targets, 1, batch_size=32, ckpt_path=ckpt_path
- )
- # Test training with validation
- prob_model.train(
- train_structs,
- train_targets,
- 1,
- test_structs,
- test_targets,
- batch_size=32,
- ckpt_path=ckpt_path,
- )
-
- # Save and reload model from disk
- prob_model.save(save_dir, ckpt_path=ckpt_path)
- loaded_model = MEGNetProbModel.load(save_dir, load_ckpt=False)
- assert weights_equal(
- prob_model.model.get_weights(), loaded_model.model.get_weights()
- )
-
- # Test prediction
- prob_model.predict(train_structs, batch_size=32)
+@pytest.mark.parametrize("use_norm", [True, False])
+def test_freezing(datadir: Path, use_norm: bool):
+ """Test that model layers can be frozen and thawed.
+
+ Also ensure that no weight names are duplicated during this process.
+
+ """
+ model_name = "prob_e_form_" + ("" if use_norm else "un") + "norm"
+ prob_model = MEGNetProbModel.load(datadir / model_name, load_ckpt=False)
+
+ # Loaded model should have NN frozen and VGP and Norm unfrozen
+ assert prob_model.nn_frozen
+ assert not prob_model.vgp_frozen
+ if use_norm:
+ assert prob_model.norm_frozen is False
+ else:
+ assert prob_model.norm_frozen is None
+
+ # Try freezing/thawing
+ prob_model.set_frozen("VGP")
+ assert prob_model.vgp_frozen
+
+ # Test thawing with a single item list
+ prob_model.set_frozen(["VGP"], False)
+ assert not prob_model.vgp_frozen
+
+ # Test error handling
+ if use_norm:
+ prob_model.set_frozen("Norm")
+ assert prob_model.norm_frozen
+ else:
+ with pytest.raises(ValueError):
+ prob_model.set_frozen("Norm")
+
+ # Test thawing with a list
+ prob_model.set_frozen(["NN", "VGP"], False)
+ assert not prob_model.nn_frozen
+ assert not prob_model.vgp_frozen
+
+ # Now in two steps
+ prob_model.set_frozen(["VGP"], recompile=False)
+ prob_model.set_frozen(["NN"], freeze=False)
+ assert not prob_model.nn_frozen
+ assert prob_model.vgp_frozen
+
+ duplicate_names = find_duplicate_weights(prob_model)
+ assert len(duplicate_names) == 0
\ No newline at end of file
diff --git a/tests/test_model/prob_e_form_norm/checkpoint b/tests/test_model/prob_e_form_norm/checkpoint
new file mode 100644
index 0000000..efd0621
--- /dev/null
+++ b/tests/test_model/prob_e_form_norm/checkpoint
@@ -0,0 +1,2 @@
+model_checkpoint_path: "weights"
+all_model_checkpoint_paths: "weights"
diff --git a/tests/test_model/prob_e_form_norm/config.json b/tests/test_model/prob_e_form_norm/config.json
new file mode 100644
index 0000000..e237cb9
--- /dev/null
+++ b/tests/test_model/prob_e_form_norm/config.json
@@ -0,0 +1 @@
+{"num_inducing_points": 10, "metrics": [], "kl_weight": 1.0, "latent_layer": -2, "target_shape": [1], "use_normalization": true}
\ No newline at end of file
diff --git a/tests/test_model/prob_e_form_norm/kernel/config.json b/tests/test_model/prob_e_form_norm/kernel/config.json
new file mode 100644
index 0000000..b8a0da1
--- /dev/null
+++ b/tests/test_model/prob_e_form_norm/kernel/config.json
@@ -0,0 +1 @@
+{"type": "rbf"}
\ No newline at end of file
diff --git a/tests/test_model/prob_e_form_norm/kernel/params.npy b/tests/test_model/prob_e_form_norm/kernel/params.npy
new file mode 100644
index 0000000..6d8806c
Binary files /dev/null and b/tests/test_model/prob_e_form_norm/kernel/params.npy differ
diff --git a/tests/test_model/prob_e_form_norm/megnet.json b/tests/test_model/prob_e_form_norm/megnet.json
new file mode 100644
index 0000000..c931337
--- /dev/null
+++ b/tests/test_model/prob_e_form_norm/megnet.json
@@ -0,0 +1 @@
+{"graph_converter": {"@module": "megnet.data.crystal", "@class": "CrystalGraph", "@version": "1.2.8", "atom_converter": {"@module": "megnet.data.graph", "@class": "DummyConverter", "@version": "1.2.8"}, "bond_converter": {"@module": "megnet.data.graph", "@class": "GaussianDistance", "@version": "1.2.8", "centers": {"@module": "numpy", "@class": "array", "dtype": "float64", "data": [0.0, 0.06060606060606061, 0.12121212121212122, 0.18181818181818182, 0.24242424242424243, 0.30303030303030304, 0.36363636363636365, 0.42424242424242425, 0.48484848484848486, 0.5454545454545454, 0.6060606060606061, 0.6666666666666667, 0.7272727272727273, 0.7878787878787878, 0.8484848484848485, 0.9090909090909092, 0.9696969696969697, 1.0303030303030303, 1.0909090909090908, 1.1515151515151516, 1.2121212121212122, 1.2727272727272727, 1.3333333333333335, 1.393939393939394, 1.4545454545454546, 1.5151515151515151, 1.5757575757575757, 1.6363636363636365, 1.696969696969697, 1.7575757575757576, 1.8181818181818183, 1.878787878787879, 1.9393939393939394, 2.0, 2.0606060606060606, 2.121212121212121, 2.1818181818181817, 2.2424242424242427, 2.303030303030303, 2.3636363636363638, 2.4242424242424243, 2.484848484848485, 2.5454545454545454, 2.606060606060606, 2.666666666666667, 2.7272727272727275, 2.787878787878788, 2.8484848484848486, 2.909090909090909, 2.9696969696969697, 3.0303030303030303, 3.090909090909091, 3.1515151515151514, 3.2121212121212124, 3.272727272727273, 3.3333333333333335, 3.393939393939394, 3.4545454545454546, 3.515151515151515, 3.5757575757575757, 3.6363636363636367, 3.6969696969696972, 3.757575757575758, 3.8181818181818183, 3.878787878787879, 3.9393939393939394, 4.0, 4.0606060606060606, 4.121212121212121, 4.181818181818182, 4.242424242424242, 4.303030303030303, 4.363636363636363, 4.424242424242425, 4.484848484848485, 4.545454545454546, 4.606060606060606, 4.666666666666667, 4.7272727272727275, 4.787878787878788, 4.848484848484849, 4.909090909090909, 4.96969696969697, 5.03030303030303, 5.090909090909091, 5.151515151515151, 5.212121212121212, 5.2727272727272725, 5.333333333333334, 5.3939393939393945, 5.454545454545455, 5.515151515151516, 5.575757575757576, 5.636363636363637, 5.696969696969697, 5.757575757575758, 5.818181818181818, 5.878787878787879, 5.9393939393939394, 6.0]}, "width": 0.5}, "cutoff": 5.0, "nn_strategy": {"@module": "megnet.data.local_env", "@class": "MinimumDistanceNNAll", "cutoff": 5.0}}, "target_scaler": {"@module": "megnet.utils.preprocessing", "@class": "DummyScaler", "@version": "1.2.8"}, "metadata": {"name": "Formation energy", "unit": "eV/atom", "metric": "mae", "metric value": 0.026, "training dataset": "mp-2019.4.1", "description": "This model was trained using structures and formation energy of Materials Project data base downloaded on April 1, 2019. The total data size is 133420. The model was trained by a 0.8-0.1-0.1, train, validation and test data ratio"}}
\ No newline at end of file
diff --git a/tests/test_model/prob_e_form_norm/megnet/keras_metadata.pb b/tests/test_model/prob_e_form_norm/megnet/keras_metadata.pb
new file mode 100644
index 0000000..003327a
--- /dev/null
+++ b/tests/test_model/prob_e_form_norm/megnet/keras_metadata.pb
@@ -0,0 +1,48 @@
+
+root"_tf_keras_network*{"name": "model_1", "trainable": true, "expects_training_arg": true, "dtype": "float32", "batch_input_shape": null, "must_restore_from_config": false, "class_name": "Functional", "config": {"name": "model_1", "layers": [{"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_1"}, "name": "input_1", "inbound_nodes": []}, {"class_name": "Embedding", "config": {"name": "embedding_1", "trainable": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "float32", "input_dim": 95, "output_dim": 16, "embeddings_initializer": {"class_name": "RandomUniform", "config": {"minval": -0.05, "maxval": 0.05, "seed": null}}, "embeddings_regularizer": null, "activity_regularizer": null, "embeddings_constraint": null, "mask_zero": false, "input_length": null}, "name": "embedding_1", "inbound_nodes": [[["input_1", 0, 0, {}]]]}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null, 100]}, "dtype": "float32", "sparse": false, "ragged": false, "name": "input_2"}, "name": "input_2", "inbound_nodes": []}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null, 2]}, "dtype": "float32", "sparse": false, "ragged": false, "name": "input_3"}, "name": "input_3", "inbound_nodes": []}, {"class_name": "Dense", "config": {"name": "dense_1", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_1", "inbound_nodes": [[["embedding_1", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_3", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_3", "inbound_nodes": [[["input_2", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_5", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_5", "inbound_nodes": [[["input_3", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_2", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_2", "inbound_nodes": [[["dense_1", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_4", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_4", "inbound_nodes": [[["dense_3", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_6", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_6", "inbound_nodes": [[["dense_5", 0, 0, {}]]]}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_4"}, "name": "input_4", "inbound_nodes": []}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_5"}, "name": "input_5", "inbound_nodes": []}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_6"}, "name": "input_6", "inbound_nodes": []}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_7"}, "name": "input_7", "inbound_nodes": []}, {"class_name": "MEGNetLayer", "config": {"name": "meg_net_layer_1", "trainable": false, "dtype": "float32", "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "units_e": [64, 64, 32], "units_v": [64, 64, 32], "units_u": [64, 64, 32], "pool_method": "mean"}, "name": "meg_net_layer_1", "inbound_nodes": [[["dense_2", 0, 0, {}], ["dense_4", 0, 0, {}], ["dense_6", 0, 0, {}], ["input_4", 0, 0, {}], ["input_5", 0, 0, {}], ["input_6", 0, 0, {}], ["input_7", 0, 0, {}]]]}, {"class_name": "Add", "config": {"name": "add_1", "trainable": false, "dtype": "float32"}, "name": "add_1", "inbound_nodes": [[["dense_2", 0, 0, {}], ["meg_net_layer_1", 0, 0, {}]]]}, {"class_name": "Add", "config": {"name": "add_2", "trainable": false, "dtype": "float32"}, "name": "add_2", "inbound_nodes": [[["dense_4", 0, 0, {}], ["meg_net_layer_1", 0, 1, {}]]]}, {"class_name": "Add", "config": {"name": "add_3", "trainable": false, "dtype": "float32"}, "name": "add_3", "inbound_nodes": [[["dense_6", 0, 0, {}], ["meg_net_layer_1", 0, 2, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_7", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_7", "inbound_nodes": [[["add_1", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_9", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_9", "inbound_nodes": [[["add_2", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_11", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_11", "inbound_nodes": [[["add_3", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_8", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_8", "inbound_nodes": [[["dense_7", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_10", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_10", "inbound_nodes": [[["dense_9", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_12", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_12", "inbound_nodes": [[["dense_11", 0, 0, {}]]]}, {"class_name": "MEGNetLayer", "config": {"name": "meg_net_layer_2", "trainable": false, "dtype": "float32", "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "units_e": [64, 64, 32], "units_v": [64, 64, 32], "units_u": [64, 64, 32], "pool_method": "mean"}, "name": "meg_net_layer_2", "inbound_nodes": [[["dense_8", 0, 0, {}], ["dense_10", 0, 0, {}], ["dense_12", 0, 0, {}], ["input_4", 0, 0, {}], ["input_5", 0, 0, {}], ["input_6", 0, 0, {}], ["input_7", 0, 0, {}]]]}, {"class_name": "Add", "config": {"name": "add_4", "trainable": false, "dtype": "float32"}, "name": "add_4", "inbound_nodes": [[["add_1", 0, 0, {}], ["meg_net_layer_2", 0, 0, {}]]]}, {"class_name": "Add", "config": {"name": "add_5", "trainable": false, "dtype": "float32"}, "name": "add_5", "inbound_nodes": [[["add_2", 0, 0, {}], ["meg_net_layer_2", 0, 1, {}]]]}, {"class_name": "Add", "config": {"name": "add_6", "trainable": false, "dtype": "float32"}, "name": "add_6", "inbound_nodes": [[["add_3", 0, 0, {}], ["meg_net_layer_2", 0, 2, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_13", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_13", "inbound_nodes": [[["add_4", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_15", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_15", "inbound_nodes": [[["add_5", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_17", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_17", "inbound_nodes": [[["add_6", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_14", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_14", "inbound_nodes": [[["dense_13", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_16", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_16", "inbound_nodes": [[["dense_15", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_18", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_18", "inbound_nodes": [[["dense_17", 0, 0, {}]]]}, {"class_name": "MEGNetLayer", "config": {"name": "meg_net_layer_3", "trainable": false, "dtype": "float32", "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "units_e": [64, 64, 32], "units_v": [64, 64, 32], "units_u": [64, 64, 32], "pool_method": "mean"}, "name": "meg_net_layer_3", "inbound_nodes": [[["dense_14", 0, 0, {}], ["dense_16", 0, 0, {}], ["dense_18", 0, 0, {}], ["input_4", 0, 0, {}], ["input_5", 0, 0, {}], ["input_6", 0, 0, {}], ["input_7", 0, 0, {}]]]}, {"class_name": "Add", "config": {"name": "add_7", "trainable": false, "dtype": "float32"}, "name": "add_7", "inbound_nodes": [[["add_4", 0, 0, {}], ["meg_net_layer_3", 0, 0, {}]]]}, {"class_name": "Add", "config": {"name": "add_8", "trainable": false, "dtype": "float32"}, "name": "add_8", "inbound_nodes": [[["add_5", 0, 0, {}], ["meg_net_layer_3", 0, 1, {}]]]}, {"class_name": "Set2Set", "config": {"name": "set2_set_1", "trainable": false, "dtype": "float32", "T": 2, "n_hidden": 16, "activation": "linear", "activation_lstm": "tanh", "recurrent_activation": "hard_sigmoid", "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "recurrent_initializer": {"class_name": "Orthogonal", "config": {"gain": 1.0, "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "use_bias": true, "unit_forget_bias": true, "kernel_regularizer": null, "recurrent_regularizer": null, "bias_regularizer": null, "kernel_constraint": null, "recurrent_constraint": null, "bias_constraint": null}, "name": "set2_set_1", "inbound_nodes": [[["add_7", 0, 0, {}], ["input_6", 0, 0, {}]]]}, {"class_name": "Set2Set", "config": {"name": "set2_set_2", "trainable": false, "dtype": "float32", "T": 2, "n_hidden": 16, "activation": "linear", "activation_lstm": "tanh", "recurrent_activation": "hard_sigmoid", "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "recurrent_initializer": {"class_name": "Orthogonal", "config": {"gain": 1.0, "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "use_bias": true, "unit_forget_bias": true, "kernel_regularizer": null, "recurrent_regularizer": null, "bias_regularizer": null, "kernel_constraint": null, "recurrent_constraint": null, "bias_constraint": null}, "name": "set2_set_2", "inbound_nodes": [[["add_8", 0, 0, {}], ["input_7", 0, 0, {}]]]}, {"class_name": "Add", "config": {"name": "add_9", "trainable": false, "dtype": "float32"}, "name": "add_9", "inbound_nodes": [[["add_6", 0, 0, {}], ["meg_net_layer_3", 0, 2, {}]]]}, {"class_name": "Concatenate", "config": {"name": "concatenate_1", "trainable": false, "dtype": "float32", "axis": -1}, "name": "concatenate_1", "inbound_nodes": [[["set2_set_1", 0, 0, {}], ["set2_set_2", 0, 0, {}], ["add_9", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_19", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_19", "inbound_nodes": [[["concatenate_1", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_20", "trainable": false, "dtype": "float32", "units": 16, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_20", "inbound_nodes": [[["dense_19", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_21", "trainable": true, "dtype": "float32", "units": 1, "activation": "linear", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_21", "inbound_nodes": [[["dense_20", 0, 0, {}]]]}], "input_layers": [["input_1", 0, 0], ["input_2", 0, 0], ["input_3", 0, 0], ["input_4", 0, 0], ["input_5", 0, 0], ["input_6", 0, 0], ["input_7", 0, 0]], "output_layers": [["dense_21", 0, 0]]}, "shared_object_id": 99, "input_spec": [{"class_name": "InputSpec", "config": {"dtype": null, "shape": {"class_name": "__tuple__", "items": [null, null]}, "ndim": 2, "max_ndim": null, "min_ndim": null, "axes": {}}}, {"class_name": "InputSpec", "config": {"dtype": null, "shape": {"class_name": "__tuple__", "items": [null, null, 100]}, "ndim": 3, "max_ndim": null, "min_ndim": null, "axes": {}}}, {"class_name": "InputSpec", "config": {"dtype": null, "shape": {"class_name": "__tuple__", "items": [null, null, 2]}, "ndim": 3, "max_ndim": null, "min_ndim": null, "axes": {}}}, {"class_name": "InputSpec", "config": {"dtype": null, "shape": {"class_name": "__tuple__", "items": [null, null]}, "ndim": 2, "max_ndim": null, "min_ndim": null, "axes": {}}}, {"class_name": "InputSpec", "config": {"dtype": null, "shape": {"class_name": "__tuple__", "items": [null, null]}, "ndim": 2, "max_ndim": null, "min_ndim": null, "axes": {}}}, {"class_name": "InputSpec", "config": {"dtype": null, "shape": {"class_name": "__tuple__", "items": [null, null]}, "ndim": 2, "max_ndim": null, "min_ndim": null, "axes": {}}}, {"class_name": "InputSpec", "config": {"dtype": null, "shape": {"class_name": "__tuple__", "items": [null, null]}, "ndim": 2, "max_ndim": null, "min_ndim": null, "axes": {}}}], "build_input_shape": [{"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null, 100]}, {"class_name": "TensorShape", "items": [null, null, 2]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}], "is_graph_network": true, "full_save_spec": {"class_name": "__tuple__", "items": [[[{"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_1"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null, 100]}, "float32", "input_2"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null, 2]}, "float32", "input_3"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_4"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_5"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_6"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_7"]}]], {}]}, "save_spec": [{"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_1"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null, 100]}, "float32", "input_2"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null, 2]}, "float32", "input_3"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_4"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_5"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_6"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_7"]}], "keras_version": "2.6.0", "backend": "tensorflow", "model_config": {"class_name": "Functional", "config": {"name": "model_1", "layers": [{"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_1"}, "name": "input_1", "inbound_nodes": [], "shared_object_id": 0}, {"class_name": "Embedding", "config": {"name": "embedding_1", "trainable": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "float32", "input_dim": 95, "output_dim": 16, "embeddings_initializer": {"class_name": "RandomUniform", "config": {"minval": -0.05, "maxval": 0.05, "seed": null}, "shared_object_id": 1}, "embeddings_regularizer": null, "activity_regularizer": null, "embeddings_constraint": null, "mask_zero": false, "input_length": null}, "name": "embedding_1", "inbound_nodes": [[["input_1", 0, 0, {}]]], "shared_object_id": 2}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null, 100]}, "dtype": "float32", "sparse": false, "ragged": false, "name": "input_2"}, "name": "input_2", "inbound_nodes": [], "shared_object_id": 3}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null, 2]}, "dtype": "float32", "sparse": false, "ragged": false, "name": "input_3"}, "name": "input_3", "inbound_nodes": [], "shared_object_id": 4}, {"class_name": "Dense", "config": {"name": "dense_1", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 5}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 6}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_1", "inbound_nodes": [[["embedding_1", 0, 0, {}]]], "shared_object_id": 7}, {"class_name": "Dense", "config": {"name": "dense_3", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 8}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 9}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_3", "inbound_nodes": [[["input_2", 0, 0, {}]]], "shared_object_id": 10}, {"class_name": "Dense", "config": {"name": "dense_5", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 11}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 12}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_5", "inbound_nodes": [[["input_3", 0, 0, {}]]], "shared_object_id": 13}, {"class_name": "Dense", "config": {"name": "dense_2", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 14}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 15}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_2", "inbound_nodes": [[["dense_1", 0, 0, {}]]], "shared_object_id": 16}, {"class_name": "Dense", "config": {"name": "dense_4", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 17}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 18}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_4", "inbound_nodes": [[["dense_3", 0, 0, {}]]], "shared_object_id": 19}, {"class_name": "Dense", "config": {"name": "dense_6", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 20}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 21}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_6", "inbound_nodes": [[["dense_5", 0, 0, {}]]], "shared_object_id": 22}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_4"}, "name": "input_4", "inbound_nodes": [], "shared_object_id": 23}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_5"}, "name": "input_5", "inbound_nodes": [], "shared_object_id": 24}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_6"}, "name": "input_6", "inbound_nodes": [], "shared_object_id": 25}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_7"}, "name": "input_7", "inbound_nodes": [], "shared_object_id": 26}, {"class_name": "MEGNetLayer", "config": {"name": "meg_net_layer_1", "trainable": false, "dtype": "float32", "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 27}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 28}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "units_e": [64, 64, 32], "units_v": [64, 64, 32], "units_u": [64, 64, 32], "pool_method": "mean"}, "name": "meg_net_layer_1", "inbound_nodes": [[["dense_2", 0, 0, {}], ["dense_4", 0, 0, {}], ["dense_6", 0, 0, {}], ["input_4", 0, 0, {}], ["input_5", 0, 0, {}], ["input_6", 0, 0, {}], ["input_7", 0, 0, {}]]], "shared_object_id": 29}, {"class_name": "Add", "config": {"name": "add_1", "trainable": false, "dtype": "float32"}, "name": "add_1", "inbound_nodes": [[["dense_2", 0, 0, {}], ["meg_net_layer_1", 0, 0, {}]]], "shared_object_id": 30}, {"class_name": "Add", "config": {"name": "add_2", "trainable": false, "dtype": "float32"}, "name": "add_2", "inbound_nodes": [[["dense_4", 0, 0, {}], ["meg_net_layer_1", 0, 1, {}]]], "shared_object_id": 31}, {"class_name": "Add", "config": {"name": "add_3", "trainable": false, "dtype": "float32"}, "name": "add_3", "inbound_nodes": [[["dense_6", 0, 0, {}], ["meg_net_layer_1", 0, 2, {}]]], "shared_object_id": 32}, {"class_name": "Dense", "config": {"name": "dense_7", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 33}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 34}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_7", "inbound_nodes": [[["add_1", 0, 0, {}]]], "shared_object_id": 35}, {"class_name": "Dense", "config": {"name": "dense_9", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 36}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 37}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_9", "inbound_nodes": [[["add_2", 0, 0, {}]]], "shared_object_id": 38}, {"class_name": "Dense", "config": {"name": "dense_11", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 39}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 40}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_11", "inbound_nodes": [[["add_3", 0, 0, {}]]], "shared_object_id": 41}, {"class_name": "Dense", "config": {"name": "dense_8", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 42}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 43}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_8", "inbound_nodes": [[["dense_7", 0, 0, {}]]], "shared_object_id": 44}, {"class_name": "Dense", "config": {"name": "dense_10", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 45}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 46}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_10", "inbound_nodes": [[["dense_9", 0, 0, {}]]], "shared_object_id": 47}, {"class_name": "Dense", "config": {"name": "dense_12", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 48}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 49}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_12", "inbound_nodes": [[["dense_11", 0, 0, {}]]], "shared_object_id": 50}, {"class_name": "MEGNetLayer", "config": {"name": "meg_net_layer_2", "trainable": false, "dtype": "float32", "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 51}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 52}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "units_e": [64, 64, 32], "units_v": [64, 64, 32], "units_u": [64, 64, 32], "pool_method": "mean"}, "name": "meg_net_layer_2", "inbound_nodes": [[["dense_8", 0, 0, {}], ["dense_10", 0, 0, {}], ["dense_12", 0, 0, {}], ["input_4", 0, 0, {}], ["input_5", 0, 0, {}], ["input_6", 0, 0, {}], ["input_7", 0, 0, {}]]], "shared_object_id": 53}, {"class_name": "Add", "config": {"name": "add_4", "trainable": false, "dtype": "float32"}, "name": "add_4", "inbound_nodes": [[["add_1", 0, 0, {}], ["meg_net_layer_2", 0, 0, {}]]], "shared_object_id": 54}, {"class_name": "Add", "config": {"name": "add_5", "trainable": false, "dtype": "float32"}, "name": "add_5", "inbound_nodes": [[["add_2", 0, 0, {}], ["meg_net_layer_2", 0, 1, {}]]], "shared_object_id": 55}, {"class_name": "Add", "config": {"name": "add_6", "trainable": false, "dtype": "float32"}, "name": "add_6", "inbound_nodes": [[["add_3", 0, 0, {}], ["meg_net_layer_2", 0, 2, {}]]], "shared_object_id": 56}, {"class_name": "Dense", "config": {"name": "dense_13", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 57}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 58}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_13", "inbound_nodes": [[["add_4", 0, 0, {}]]], "shared_object_id": 59}, {"class_name": "Dense", "config": {"name": "dense_15", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 60}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 61}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_15", "inbound_nodes": [[["add_5", 0, 0, {}]]], "shared_object_id": 62}, {"class_name": "Dense", "config": {"name": "dense_17", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 63}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 64}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_17", "inbound_nodes": [[["add_6", 0, 0, {}]]], "shared_object_id": 65}, {"class_name": "Dense", "config": {"name": "dense_14", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 66}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 67}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_14", "inbound_nodes": [[["dense_13", 0, 0, {}]]], "shared_object_id": 68}, {"class_name": "Dense", "config": {"name": "dense_16", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 69}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 70}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_16", "inbound_nodes": [[["dense_15", 0, 0, {}]]], "shared_object_id": 71}, {"class_name": "Dense", "config": {"name": "dense_18", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 72}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 73}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_18", "inbound_nodes": [[["dense_17", 0, 0, {}]]], "shared_object_id": 74}, {"class_name": "MEGNetLayer", "config": {"name": "meg_net_layer_3", "trainable": false, "dtype": "float32", "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 75}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 76}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "units_e": [64, 64, 32], "units_v": [64, 64, 32], "units_u": [64, 64, 32], "pool_method": "mean"}, "name": "meg_net_layer_3", "inbound_nodes": [[["dense_14", 0, 0, {}], ["dense_16", 0, 0, {}], ["dense_18", 0, 0, {}], ["input_4", 0, 0, {}], ["input_5", 0, 0, {}], ["input_6", 0, 0, {}], ["input_7", 0, 0, {}]]], "shared_object_id": 77}, {"class_name": "Add", "config": {"name": "add_7", "trainable": false, "dtype": "float32"}, "name": "add_7", "inbound_nodes": [[["add_4", 0, 0, {}], ["meg_net_layer_3", 0, 0, {}]]], "shared_object_id": 78}, {"class_name": "Add", "config": {"name": "add_8", "trainable": false, "dtype": "float32"}, "name": "add_8", "inbound_nodes": [[["add_5", 0, 0, {}], ["meg_net_layer_3", 0, 1, {}]]], "shared_object_id": 79}, {"class_name": "Set2Set", "config": {"name": "set2_set_1", "trainable": false, "dtype": "float32", "T": 2, "n_hidden": 16, "activation": "linear", "activation_lstm": "tanh", "recurrent_activation": "hard_sigmoid", "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 80}, "recurrent_initializer": {"class_name": "Orthogonal", "config": {"gain": 1.0, "seed": null}, "shared_object_id": 81}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 82}, "use_bias": true, "unit_forget_bias": true, "kernel_regularizer": null, "recurrent_regularizer": null, "bias_regularizer": null, "kernel_constraint": null, "recurrent_constraint": null, "bias_constraint": null}, "name": "set2_set_1", "inbound_nodes": [[["add_7", 0, 0, {}], ["input_6", 0, 0, {}]]], "shared_object_id": 83}, {"class_name": "Set2Set", "config": {"name": "set2_set_2", "trainable": false, "dtype": "float32", "T": 2, "n_hidden": 16, "activation": "linear", "activation_lstm": "tanh", "recurrent_activation": "hard_sigmoid", "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 84}, "recurrent_initializer": {"class_name": "Orthogonal", "config": {"gain": 1.0, "seed": null}, "shared_object_id": 85}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 86}, "use_bias": true, "unit_forget_bias": true, "kernel_regularizer": null, "recurrent_regularizer": null, "bias_regularizer": null, "kernel_constraint": null, "recurrent_constraint": null, "bias_constraint": null}, "name": "set2_set_2", "inbound_nodes": [[["add_8", 0, 0, {}], ["input_7", 0, 0, {}]]], "shared_object_id": 87}, {"class_name": "Add", "config": {"name": "add_9", "trainable": false, "dtype": "float32"}, "name": "add_9", "inbound_nodes": [[["add_6", 0, 0, {}], ["meg_net_layer_3", 0, 2, {}]]], "shared_object_id": 88}, {"class_name": "Concatenate", "config": {"name": "concatenate_1", "trainable": false, "dtype": "float32", "axis": -1}, "name": "concatenate_1", "inbound_nodes": [[["set2_set_1", 0, 0, {}], ["set2_set_2", 0, 0, {}], ["add_9", 0, 0, {}]]], "shared_object_id": 89}, {"class_name": "Dense", "config": {"name": "dense_19", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 90}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 91}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_19", "inbound_nodes": [[["concatenate_1", 0, 0, {}]]], "shared_object_id": 92}, {"class_name": "Dense", "config": {"name": "dense_20", "trainable": false, "dtype": "float32", "units": 16, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 93}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 94}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_20", "inbound_nodes": [[["dense_19", 0, 0, {}]]], "shared_object_id": 95}, {"class_name": "Dense", "config": {"name": "dense_21", "trainable": true, "dtype": "float32", "units": 1, "activation": "linear", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 96}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 97}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_21", "inbound_nodes": [[["dense_20", 0, 0, {}]]], "shared_object_id": 98}], "input_layers": [["input_1", 0, 0], ["input_2", 0, 0], ["input_3", 0, 0], ["input_4", 0, 0], ["input_5", 0, 0], ["input_6", 0, 0], ["input_7", 0, 0]], "output_layers": [["dense_21", 0, 0]]}}, "training_config": {"loss": "mean_squared_error", "metrics": [[]], "weighted_metrics": null, "loss_weights": null, "optimizer_config": {"class_name": "Adam", "config": {"name": "Adam", "learning_rate": 0.0010000000474974513, "decay": 0.0, "beta_1": 0.8999999761581421, "beta_2": 0.9990000128746033, "epsilon": 1e-07, "amsgrad": false}}}}2
+ root.layer-0"_tf_keras_input_layer*{"class_name": "InputLayer", "name": "input_1", "dtype": "int32", "sparse": false, "ragged": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_1"}}2
+root.layer_with_weights-0"_tf_keras_layer*{"name": "embedding_1", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "stateful": false, "must_restore_from_config": false, "class_name": "Embedding", "config": {"name": "embedding_1", "trainable": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "float32", "input_dim": 95, "output_dim": 16, "embeddings_initializer": {"class_name": "RandomUniform", "config": {"minval": -0.05, "maxval": 0.05, "seed": null}, "shared_object_id": 1}, "embeddings_regularizer": null, "activity_regularizer": null, "embeddings_constraint": null, "mask_zero": false, "input_length": null}, "inbound_nodes": [[["input_1", 0, 0, {}]]], "shared_object_id": 2, "build_input_shape": {"class_name": "TensorShape", "items": [null, null]}}2
+ root.layer-2"_tf_keras_input_layer*{"class_name": "InputLayer", "name": "input_2", "dtype": "float32", "sparse": false, "ragged": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null, 100]}, "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null, 100]}, "dtype": "float32", "sparse": false, "ragged": false, "name": "input_2"}}2
+ root.layer-3"_tf_keras_input_layer*{"class_name": "InputLayer", "name": "input_3", "dtype": "float32", "sparse": false, "ragged": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null, 2]}, "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null, 2]}, "dtype": "float32", "sparse": false, "ragged": false, "name": "input_3"}}2
+root.layer_with_weights-1"_tf_keras_layer*{"name": "dense_1", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_1", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 5}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 6}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["embedding_1", 0, 0, {}]]], "shared_object_id": 7, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 16}}, "shared_object_id": 107}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 16]}}2
+root.layer_with_weights-2"_tf_keras_layer*{"name": "dense_3", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_3", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 8}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 9}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["input_2", 0, 0, {}]]], "shared_object_id": 10, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 100}}, "shared_object_id": 108}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 100]}}2
+root.layer_with_weights-3"_tf_keras_layer*{"name": "dense_5", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_5", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 11}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 12}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["input_3", 0, 0, {}]]], "shared_object_id": 13, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 2}}, "shared_object_id": 109}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 2]}}2
+root.layer_with_weights-4"_tf_keras_layer*{"name": "dense_2", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_2", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 14}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 15}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_1", 0, 0, {}]]], "shared_object_id": 16, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 64}}, "shared_object_id": 110}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 64]}}2
+ root.layer_with_weights-5"_tf_keras_layer*{"name": "dense_4", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_4", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 17}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 18}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_3", 0, 0, {}]]], "shared_object_id": 19, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 64}}, "shared_object_id": 111}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 64]}}2
+
+root.layer_with_weights-6"_tf_keras_layer*{"name": "dense_6", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_6", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 20}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 21}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_5", 0, 0, {}]]], "shared_object_id": 22, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 64}}, "shared_object_id": 112}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 64]}}2
+
root.layer-10"_tf_keras_input_layer*{"class_name": "InputLayer", "name": "input_4", "dtype": "int32", "sparse": false, "ragged": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_4"}}2
+
root.layer-11"_tf_keras_input_layer*{"class_name": "InputLayer", "name": "input_5", "dtype": "int32", "sparse": false, "ragged": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_5"}}2
+
root.layer-12"_tf_keras_input_layer*{"class_name": "InputLayer", "name": "input_6", "dtype": "int32", "sparse": false, "ragged": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_6"}}2
+
root.layer-13"_tf_keras_input_layer*{"class_name": "InputLayer", "name": "input_7", "dtype": "int32", "sparse": false, "ragged": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_7"}}2
+ root.layer_with_weights-7"_tf_keras_layer* {"name": "meg_net_layer_1", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "MEGNetLayer", "config": {"name": "meg_net_layer_1", "trainable": false, "dtype": "float32", "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 27}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 28}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "units_e": [64, 64, 32], "units_v": [64, 64, 32], "units_u": [64, 64, 32], "pool_method": "mean"}, "inbound_nodes": [[["dense_2", 0, 0, {}], ["dense_4", 0, 0, {}], ["dense_6", 0, 0, {}], ["input_4", 0, 0, {}], ["input_5", 0, 0, {}], ["input_6", 0, 0, {}], ["input_7", 0, 0, {}]]], "shared_object_id": 29, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}]}2
+
root.layer-15"_tf_keras_layer*{"name": "add_1", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Add", "config": {"name": "add_1", "trainable": false, "dtype": "float32"}, "inbound_nodes": [[["dense_2", 0, 0, {}], ["meg_net_layer_1", 0, 0, {}]]], "shared_object_id": 30, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}]}2
+
root.layer-16"_tf_keras_layer*{"name": "add_2", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Add", "config": {"name": "add_2", "trainable": false, "dtype": "float32"}, "inbound_nodes": [[["dense_4", 0, 0, {}], ["meg_net_layer_1", 0, 1, {}]]], "shared_object_id": 31, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}]}2
+
root.layer-17"_tf_keras_layer*{"name": "add_3", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Add", "config": {"name": "add_3", "trainable": false, "dtype": "float32"}, "inbound_nodes": [[["dense_6", 0, 0, {}], ["meg_net_layer_1", 0, 2, {}]]], "shared_object_id": 32, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [1, null, 32]}]}2
+root.layer_with_weights-8"_tf_keras_layer*{"name": "dense_7", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_7", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 33}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 34}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["add_1", 0, 0, {}]]], "shared_object_id": 35, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 32}}, "shared_object_id": 113}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 32]}}2
+root.layer_with_weights-9"_tf_keras_layer*{"name": "dense_9", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_9", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 36}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 37}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["add_2", 0, 0, {}]]], "shared_object_id": 38, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 32}}, "shared_object_id": 114}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 32]}}2
+root.layer_with_weights-10"_tf_keras_layer*{"name": "dense_11", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_11", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 39}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 40}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["add_3", 0, 0, {}]]], "shared_object_id": 41, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 32}}, "shared_object_id": 115}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 32]}}2
+root.layer_with_weights-11"_tf_keras_layer*{"name": "dense_8", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_8", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 42}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 43}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_7", 0, 0, {}]]], "shared_object_id": 44, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 64}}, "shared_object_id": 116}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 64]}}2
+root.layer_with_weights-12"_tf_keras_layer*{"name": "dense_10", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_10", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 45}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 46}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_9", 0, 0, {}]]], "shared_object_id": 47, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 64}}, "shared_object_id": 117}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 64]}}2
+root.layer_with_weights-13"_tf_keras_layer*{"name": "dense_12", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_12", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 48}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 49}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_11", 0, 0, {}]]], "shared_object_id": 50, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 64}}, "shared_object_id": 118}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 64]}}2
+ root.layer_with_weights-14"_tf_keras_layer* {"name": "meg_net_layer_2", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "MEGNetLayer", "config": {"name": "meg_net_layer_2", "trainable": false, "dtype": "float32", "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 51}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 52}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "units_e": [64, 64, 32], "units_v": [64, 64, 32], "units_u": [64, 64, 32], "pool_method": "mean"}, "inbound_nodes": [[["dense_8", 0, 0, {}], ["dense_10", 0, 0, {}], ["dense_12", 0, 0, {}], ["input_4", 0, 0, {}], ["input_5", 0, 0, {}], ["input_6", 0, 0, {}], ["input_7", 0, 0, {}]]], "shared_object_id": 53, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}]}2
+
root.layer-25"_tf_keras_layer*{"name": "add_4", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Add", "config": {"name": "add_4", "trainable": false, "dtype": "float32"}, "inbound_nodes": [[["add_1", 0, 0, {}], ["meg_net_layer_2", 0, 0, {}]]], "shared_object_id": 54, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}]}2
+
root.layer-26"_tf_keras_layer*{"name": "add_5", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Add", "config": {"name": "add_5", "trainable": false, "dtype": "float32"}, "inbound_nodes": [[["add_2", 0, 0, {}], ["meg_net_layer_2", 0, 1, {}]]], "shared_object_id": 55, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}]}2
+
root.layer-27"_tf_keras_layer*{"name": "add_6", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Add", "config": {"name": "add_6", "trainable": false, "dtype": "float32"}, "inbound_nodes": [[["add_3", 0, 0, {}], ["meg_net_layer_2", 0, 2, {}]]], "shared_object_id": 56, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [1, null, 32]}]}2
+ root.layer_with_weights-15"_tf_keras_layer*{"name": "dense_13", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_13", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 57}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 58}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["add_4", 0, 0, {}]]], "shared_object_id": 59, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 32}}, "shared_object_id": 119}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 32]}}2
+ root.layer_with_weights-16"_tf_keras_layer*{"name": "dense_15", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_15", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 60}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 61}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["add_5", 0, 0, {}]]], "shared_object_id": 62, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 32}}, "shared_object_id": 120}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 32]}}2
+ root.layer_with_weights-17"_tf_keras_layer*{"name": "dense_17", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_17", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 63}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 64}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["add_6", 0, 0, {}]]], "shared_object_id": 65, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 32}}, "shared_object_id": 121}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 32]}}2
+ root.layer_with_weights-18"_tf_keras_layer*{"name": "dense_14", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_14", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 66}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 67}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_13", 0, 0, {}]]], "shared_object_id": 68, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 64}}, "shared_object_id": 122}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 64]}}2
+!root.layer_with_weights-19"_tf_keras_layer*{"name": "dense_16", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_16", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 69}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 70}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_15", 0, 0, {}]]], "shared_object_id": 71, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 64}}, "shared_object_id": 123}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 64]}}2
+"root.layer_with_weights-20"_tf_keras_layer*{"name": "dense_18", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_18", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 72}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 73}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_17", 0, 0, {}]]], "shared_object_id": 74, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 64}}, "shared_object_id": 124}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 64]}}2
+ #root.layer_with_weights-21"_tf_keras_layer* {"name": "meg_net_layer_3", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "MEGNetLayer", "config": {"name": "meg_net_layer_3", "trainable": false, "dtype": "float32", "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 75}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 76}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "units_e": [64, 64, 32], "units_v": [64, 64, 32], "units_u": [64, 64, 32], "pool_method": "mean"}, "inbound_nodes": [[["dense_14", 0, 0, {}], ["dense_16", 0, 0, {}], ["dense_18", 0, 0, {}], ["input_4", 0, 0, {}], ["input_5", 0, 0, {}], ["input_6", 0, 0, {}], ["input_7", 0, 0, {}]]], "shared_object_id": 77, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}]}2
+$
root.layer-35"_tf_keras_layer*{"name": "add_7", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Add", "config": {"name": "add_7", "trainable": false, "dtype": "float32"}, "inbound_nodes": [[["add_4", 0, 0, {}], ["meg_net_layer_3", 0, 0, {}]]], "shared_object_id": 78, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}]}2
+%
root.layer-36"_tf_keras_layer*{"name": "add_8", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Add", "config": {"name": "add_8", "trainable": false, "dtype": "float32"}, "inbound_nodes": [[["add_5", 0, 0, {}], ["meg_net_layer_3", 0, 1, {}]]], "shared_object_id": 79, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}]}2
+ &root.layer_with_weights-22"_tf_keras_layer* {"name": "set2_set_1", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Set2Set", "config": {"name": "set2_set_1", "trainable": false, "dtype": "float32", "T": 2, "n_hidden": 16, "activation": "linear", "activation_lstm": "tanh", "recurrent_activation": "hard_sigmoid", "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 80}, "recurrent_initializer": {"class_name": "Orthogonal", "config": {"gain": 1.0, "seed": null}, "shared_object_id": 81}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 82}, "use_bias": true, "unit_forget_bias": true, "kernel_regularizer": null, "recurrent_regularizer": null, "bias_regularizer": null, "kernel_constraint": null, "recurrent_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["add_7", 0, 0, {}], ["input_6", 0, 0, {}]]], "shared_object_id": 83, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null]}]}2
+ 'root.layer_with_weights-23"_tf_keras_layer* {"name": "set2_set_2", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Set2Set", "config": {"name": "set2_set_2", "trainable": false, "dtype": "float32", "T": 2, "n_hidden": 16, "activation": "linear", "activation_lstm": "tanh", "recurrent_activation": "hard_sigmoid", "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 84}, "recurrent_initializer": {"class_name": "Orthogonal", "config": {"gain": 1.0, "seed": null}, "shared_object_id": 85}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 86}, "use_bias": true, "unit_forget_bias": true, "kernel_regularizer": null, "recurrent_regularizer": null, "bias_regularizer": null, "kernel_constraint": null, "recurrent_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["add_8", 0, 0, {}], ["input_7", 0, 0, {}]]], "shared_object_id": 87, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null]}]}2
+(
root.layer-39"_tf_keras_layer*{"name": "add_9", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Add", "config": {"name": "add_9", "trainable": false, "dtype": "float32"}, "inbound_nodes": [[["add_6", 0, 0, {}], ["meg_net_layer_3", 0, 2, {}]]], "shared_object_id": 88, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [1, null, 32]}]}2
+)
root.layer-40"_tf_keras_layer*{"name": "concatenate_1", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Concatenate", "config": {"name": "concatenate_1", "trainable": false, "dtype": "float32", "axis": -1}, "inbound_nodes": [[["set2_set_1", 0, 0, {}], ["set2_set_2", 0, 0, {}], ["add_9", 0, 0, {}]]], "shared_object_id": 89, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}]}2
+*root.layer_with_weights-24"_tf_keras_layer*{"name": "dense_19", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_19", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 90}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 91}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["concatenate_1", 0, 0, {}]]], "shared_object_id": 92, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 96}}, "shared_object_id": 125}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 96]}}2
++root.layer_with_weights-25"_tf_keras_layer*{"name": "dense_20", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_20", "trainable": false, "dtype": "float32", "units": 16, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 93}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 94}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_19", 0, 0, {}]]], "shared_object_id": 95, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 32}}, "shared_object_id": 126}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 32]}}2
+,root.layer_with_weights-26"_tf_keras_layer*{"name": "dense_21", "trainable": true, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_21", "trainable": true, "dtype": "float32", "units": 1, "activation": "linear", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 96}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 97}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_20", 0, 0, {}]]], "shared_object_id": 98, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 16}}, "shared_object_id": 127}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 16]}}2
+root.keras_api.metrics.0"_tf_keras_metric*{"class_name": "Mean", "name": "loss", "dtype": "float32", "config": {"name": "loss", "dtype": "float32"}, "shared_object_id": 128}2
\ No newline at end of file
diff --git a/tests/test_model/prob_e_form_norm/megnet/saved_model.pb b/tests/test_model/prob_e_form_norm/megnet/saved_model.pb
new file mode 100644
index 0000000..fcb837e
Binary files /dev/null and b/tests/test_model/prob_e_form_norm/megnet/saved_model.pb differ
diff --git a/tests/test_model/prob_e_form_norm/megnet/variables/variables.data-00000-of-00001 b/tests/test_model/prob_e_form_norm/megnet/variables/variables.data-00000-of-00001
new file mode 100644
index 0000000..b28b286
Binary files /dev/null and b/tests/test_model/prob_e_form_norm/megnet/variables/variables.data-00000-of-00001 differ
diff --git a/tests/test_model/prob_e_form_norm/megnet/variables/variables.index b/tests/test_model/prob_e_form_norm/megnet/variables/variables.index
new file mode 100644
index 0000000..78fdd32
Binary files /dev/null and b/tests/test_model/prob_e_form_norm/megnet/variables/variables.index differ
diff --git a/tests/test_model/prob_e_form_norm/nn/keras_metadata.pb b/tests/test_model/prob_e_form_norm/nn/keras_metadata.pb
new file mode 100644
index 0000000..8b9fcf7
--- /dev/null
+++ b/tests/test_model/prob_e_form_norm/nn/keras_metadata.pb
@@ -0,0 +1,48 @@
+
+root"_tf_keras_network*{"name": "model_1", "trainable": true, "expects_training_arg": true, "dtype": "float32", "batch_input_shape": null, "must_restore_from_config": false, "class_name": "Functional", "config": {"name": "model_1", "layers": [{"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_1"}, "name": "input_1", "inbound_nodes": []}, {"class_name": "Embedding", "config": {"name": "embedding_1", "trainable": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "float32", "input_dim": 95, "output_dim": 16, "embeddings_initializer": {"class_name": "RandomUniform", "config": {"minval": -0.05, "maxval": 0.05, "seed": null}}, "embeddings_regularizer": null, "activity_regularizer": null, "embeddings_constraint": null, "mask_zero": false, "input_length": null}, "name": "embedding_1", "inbound_nodes": [[["input_1", 0, 0, {}]]]}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null, 100]}, "dtype": "float32", "sparse": false, "ragged": false, "name": "input_2"}, "name": "input_2", "inbound_nodes": []}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null, 2]}, "dtype": "float32", "sparse": false, "ragged": false, "name": "input_3"}, "name": "input_3", "inbound_nodes": []}, {"class_name": "Dense", "config": {"name": "dense_1", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_1", "inbound_nodes": [[["embedding_1", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_3", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_3", "inbound_nodes": [[["input_2", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_5", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_5", "inbound_nodes": [[["input_3", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_2", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_2", "inbound_nodes": [[["dense_1", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_4", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_4", "inbound_nodes": [[["dense_3", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_6", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_6", "inbound_nodes": [[["dense_5", 0, 0, {}]]]}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_4"}, "name": "input_4", "inbound_nodes": []}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_5"}, "name": "input_5", "inbound_nodes": []}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_6"}, "name": "input_6", "inbound_nodes": []}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_7"}, "name": "input_7", "inbound_nodes": []}, {"class_name": "MEGNetLayer", "config": {"name": "meg_net_layer_1", "trainable": false, "dtype": "float32", "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "units_e": [64, 64, 32], "units_v": [64, 64, 32], "units_u": [64, 64, 32], "pool_method": "mean"}, "name": "meg_net_layer_1", "inbound_nodes": [[["dense_2", 0, 0, {}], ["dense_4", 0, 0, {}], ["dense_6", 0, 0, {}], ["input_4", 0, 0, {}], ["input_5", 0, 0, {}], ["input_6", 0, 0, {}], ["input_7", 0, 0, {}]]]}, {"class_name": "Add", "config": {"name": "add_1", "trainable": false, "dtype": "float32"}, "name": "add_1", "inbound_nodes": [[["dense_2", 0, 0, {}], ["meg_net_layer_1", 0, 0, {}]]]}, {"class_name": "Add", "config": {"name": "add_2", "trainable": false, "dtype": "float32"}, "name": "add_2", "inbound_nodes": [[["dense_4", 0, 0, {}], ["meg_net_layer_1", 0, 1, {}]]]}, {"class_name": "Add", "config": {"name": "add_3", "trainable": false, "dtype": "float32"}, "name": "add_3", "inbound_nodes": [[["dense_6", 0, 0, {}], ["meg_net_layer_1", 0, 2, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_7", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_7", "inbound_nodes": [[["add_1", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_9", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_9", "inbound_nodes": [[["add_2", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_11", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_11", "inbound_nodes": [[["add_3", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_8", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_8", "inbound_nodes": [[["dense_7", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_10", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_10", "inbound_nodes": [[["dense_9", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_12", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_12", "inbound_nodes": [[["dense_11", 0, 0, {}]]]}, {"class_name": "MEGNetLayer", "config": {"name": "meg_net_layer_2", "trainable": false, "dtype": "float32", "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "units_e": [64, 64, 32], "units_v": [64, 64, 32], "units_u": [64, 64, 32], "pool_method": "mean"}, "name": "meg_net_layer_2", "inbound_nodes": [[["dense_8", 0, 0, {}], ["dense_10", 0, 0, {}], ["dense_12", 0, 0, {}], ["input_4", 0, 0, {}], ["input_5", 0, 0, {}], ["input_6", 0, 0, {}], ["input_7", 0, 0, {}]]]}, {"class_name": "Add", "config": {"name": "add_4", "trainable": false, "dtype": "float32"}, "name": "add_4", "inbound_nodes": [[["add_1", 0, 0, {}], ["meg_net_layer_2", 0, 0, {}]]]}, {"class_name": "Add", "config": {"name": "add_5", "trainable": false, "dtype": "float32"}, "name": "add_5", "inbound_nodes": [[["add_2", 0, 0, {}], ["meg_net_layer_2", 0, 1, {}]]]}, {"class_name": "Add", "config": {"name": "add_6", "trainable": false, "dtype": "float32"}, "name": "add_6", "inbound_nodes": [[["add_3", 0, 0, {}], ["meg_net_layer_2", 0, 2, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_13", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_13", "inbound_nodes": [[["add_4", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_15", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_15", "inbound_nodes": [[["add_5", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_17", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_17", "inbound_nodes": [[["add_6", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_14", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_14", "inbound_nodes": [[["dense_13", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_16", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_16", "inbound_nodes": [[["dense_15", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_18", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_18", "inbound_nodes": [[["dense_17", 0, 0, {}]]]}, {"class_name": "MEGNetLayer", "config": {"name": "meg_net_layer_3", "trainable": false, "dtype": "float32", "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "units_e": [64, 64, 32], "units_v": [64, 64, 32], "units_u": [64, 64, 32], "pool_method": "mean"}, "name": "meg_net_layer_3", "inbound_nodes": [[["dense_14", 0, 0, {}], ["dense_16", 0, 0, {}], ["dense_18", 0, 0, {}], ["input_4", 0, 0, {}], ["input_5", 0, 0, {}], ["input_6", 0, 0, {}], ["input_7", 0, 0, {}]]]}, {"class_name": "Add", "config": {"name": "add_7", "trainable": false, "dtype": "float32"}, "name": "add_7", "inbound_nodes": [[["add_4", 0, 0, {}], ["meg_net_layer_3", 0, 0, {}]]]}, {"class_name": "Add", "config": {"name": "add_8", "trainable": false, "dtype": "float32"}, "name": "add_8", "inbound_nodes": [[["add_5", 0, 0, {}], ["meg_net_layer_3", 0, 1, {}]]]}, {"class_name": "Set2Set", "config": {"name": "set2_set_1", "trainable": false, "dtype": "float32", "T": 2, "n_hidden": 16, "activation": "linear", "activation_lstm": "tanh", "recurrent_activation": "hard_sigmoid", "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "recurrent_initializer": {"class_name": "Orthogonal", "config": {"gain": 1.0, "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "use_bias": true, "unit_forget_bias": true, "kernel_regularizer": null, "recurrent_regularizer": null, "bias_regularizer": null, "kernel_constraint": null, "recurrent_constraint": null, "bias_constraint": null}, "name": "set2_set_1", "inbound_nodes": [[["add_7", 0, 0, {}], ["input_6", 0, 0, {}]]]}, {"class_name": "Set2Set", "config": {"name": "set2_set_2", "trainable": false, "dtype": "float32", "T": 2, "n_hidden": 16, "activation": "linear", "activation_lstm": "tanh", "recurrent_activation": "hard_sigmoid", "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "recurrent_initializer": {"class_name": "Orthogonal", "config": {"gain": 1.0, "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "use_bias": true, "unit_forget_bias": true, "kernel_regularizer": null, "recurrent_regularizer": null, "bias_regularizer": null, "kernel_constraint": null, "recurrent_constraint": null, "bias_constraint": null}, "name": "set2_set_2", "inbound_nodes": [[["add_8", 0, 0, {}], ["input_7", 0, 0, {}]]]}, {"class_name": "Add", "config": {"name": "add_9", "trainable": false, "dtype": "float32"}, "name": "add_9", "inbound_nodes": [[["add_6", 0, 0, {}], ["meg_net_layer_3", 0, 2, {}]]]}, {"class_name": "Concatenate", "config": {"name": "concatenate_1", "trainable": false, "dtype": "float32", "axis": -1}, "name": "concatenate_1", "inbound_nodes": [[["set2_set_1", 0, 0, {}], ["set2_set_2", 0, 0, {}], ["add_9", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_19", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_19", "inbound_nodes": [[["concatenate_1", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_20", "trainable": false, "dtype": "float32", "units": 16, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_20", "inbound_nodes": [[["dense_19", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_21", "trainable": true, "dtype": "float32", "units": 1, "activation": "linear", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_21", "inbound_nodes": [[["dense_20", 0, 0, {}]]]}], "input_layers": [["input_1", 0, 0], ["input_2", 0, 0], ["input_3", 0, 0], ["input_4", 0, 0], ["input_5", 0, 0], ["input_6", 0, 0], ["input_7", 0, 0]], "output_layers": [["dense_21", 0, 0]]}, "shared_object_id": 99, "input_spec": [{"class_name": "InputSpec", "config": {"dtype": null, "shape": {"class_name": "__tuple__", "items": [null, null]}, "ndim": 2, "max_ndim": null, "min_ndim": null, "axes": {}}}, {"class_name": "InputSpec", "config": {"dtype": null, "shape": {"class_name": "__tuple__", "items": [null, null, 100]}, "ndim": 3, "max_ndim": null, "min_ndim": null, "axes": {}}}, {"class_name": "InputSpec", "config": {"dtype": null, "shape": {"class_name": "__tuple__", "items": [null, null, 2]}, "ndim": 3, "max_ndim": null, "min_ndim": null, "axes": {}}}, {"class_name": "InputSpec", "config": {"dtype": null, "shape": {"class_name": "__tuple__", "items": [null, null]}, "ndim": 2, "max_ndim": null, "min_ndim": null, "axes": {}}}, {"class_name": "InputSpec", "config": {"dtype": null, "shape": {"class_name": "__tuple__", "items": [null, null]}, "ndim": 2, "max_ndim": null, "min_ndim": null, "axes": {}}}, {"class_name": "InputSpec", "config": {"dtype": null, "shape": {"class_name": "__tuple__", "items": [null, null]}, "ndim": 2, "max_ndim": null, "min_ndim": null, "axes": {}}}, {"class_name": "InputSpec", "config": {"dtype": null, "shape": {"class_name": "__tuple__", "items": [null, null]}, "ndim": 2, "max_ndim": null, "min_ndim": null, "axes": {}}}], "build_input_shape": [{"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null, 100]}, {"class_name": "TensorShape", "items": [null, null, 2]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}], "is_graph_network": true, "full_save_spec": {"class_name": "__tuple__", "items": [[[{"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_1"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null, 100]}, "float32", "input_2"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null, 2]}, "float32", "input_3"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_4"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_5"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_6"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_7"]}]], {}]}, "save_spec": [{"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_1"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null, 100]}, "float32", "input_2"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null, 2]}, "float32", "input_3"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_4"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_5"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_6"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_7"]}], "keras_version": "2.6.0", "backend": "tensorflow", "model_config": {"class_name": "Functional", "config": {"name": "model_1", "layers": [{"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_1"}, "name": "input_1", "inbound_nodes": [], "shared_object_id": 0}, {"class_name": "Embedding", "config": {"name": "embedding_1", "trainable": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "float32", "input_dim": 95, "output_dim": 16, "embeddings_initializer": {"class_name": "RandomUniform", "config": {"minval": -0.05, "maxval": 0.05, "seed": null}, "shared_object_id": 1}, "embeddings_regularizer": null, "activity_regularizer": null, "embeddings_constraint": null, "mask_zero": false, "input_length": null}, "name": "embedding_1", "inbound_nodes": [[["input_1", 0, 0, {}]]], "shared_object_id": 2}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null, 100]}, "dtype": "float32", "sparse": false, "ragged": false, "name": "input_2"}, "name": "input_2", "inbound_nodes": [], "shared_object_id": 3}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null, 2]}, "dtype": "float32", "sparse": false, "ragged": false, "name": "input_3"}, "name": "input_3", "inbound_nodes": [], "shared_object_id": 4}, {"class_name": "Dense", "config": {"name": "dense_1", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 5}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 6}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_1", "inbound_nodes": [[["embedding_1", 0, 0, {}]]], "shared_object_id": 7}, {"class_name": "Dense", "config": {"name": "dense_3", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 8}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 9}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_3", "inbound_nodes": [[["input_2", 0, 0, {}]]], "shared_object_id": 10}, {"class_name": "Dense", "config": {"name": "dense_5", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 11}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 12}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_5", "inbound_nodes": [[["input_3", 0, 0, {}]]], "shared_object_id": 13}, {"class_name": "Dense", "config": {"name": "dense_2", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 14}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 15}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_2", "inbound_nodes": [[["dense_1", 0, 0, {}]]], "shared_object_id": 16}, {"class_name": "Dense", "config": {"name": "dense_4", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 17}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 18}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_4", "inbound_nodes": [[["dense_3", 0, 0, {}]]], "shared_object_id": 19}, {"class_name": "Dense", "config": {"name": "dense_6", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 20}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 21}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_6", "inbound_nodes": [[["dense_5", 0, 0, {}]]], "shared_object_id": 22}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_4"}, "name": "input_4", "inbound_nodes": [], "shared_object_id": 23}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_5"}, "name": "input_5", "inbound_nodes": [], "shared_object_id": 24}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_6"}, "name": "input_6", "inbound_nodes": [], "shared_object_id": 25}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_7"}, "name": "input_7", "inbound_nodes": [], "shared_object_id": 26}, {"class_name": "MEGNetLayer", "config": {"name": "meg_net_layer_1", "trainable": false, "dtype": "float32", "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 27}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 28}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "units_e": [64, 64, 32], "units_v": [64, 64, 32], "units_u": [64, 64, 32], "pool_method": "mean"}, "name": "meg_net_layer_1", "inbound_nodes": [[["dense_2", 0, 0, {}], ["dense_4", 0, 0, {}], ["dense_6", 0, 0, {}], ["input_4", 0, 0, {}], ["input_5", 0, 0, {}], ["input_6", 0, 0, {}], ["input_7", 0, 0, {}]]], "shared_object_id": 29}, {"class_name": "Add", "config": {"name": "add_1", "trainable": false, "dtype": "float32"}, "name": "add_1", "inbound_nodes": [[["dense_2", 0, 0, {}], ["meg_net_layer_1", 0, 0, {}]]], "shared_object_id": 30}, {"class_name": "Add", "config": {"name": "add_2", "trainable": false, "dtype": "float32"}, "name": "add_2", "inbound_nodes": [[["dense_4", 0, 0, {}], ["meg_net_layer_1", 0, 1, {}]]], "shared_object_id": 31}, {"class_name": "Add", "config": {"name": "add_3", "trainable": false, "dtype": "float32"}, "name": "add_3", "inbound_nodes": [[["dense_6", 0, 0, {}], ["meg_net_layer_1", 0, 2, {}]]], "shared_object_id": 32}, {"class_name": "Dense", "config": {"name": "dense_7", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 33}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 34}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_7", "inbound_nodes": [[["add_1", 0, 0, {}]]], "shared_object_id": 35}, {"class_name": "Dense", "config": {"name": "dense_9", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 36}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 37}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_9", "inbound_nodes": [[["add_2", 0, 0, {}]]], "shared_object_id": 38}, {"class_name": "Dense", "config": {"name": "dense_11", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 39}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 40}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_11", "inbound_nodes": [[["add_3", 0, 0, {}]]], "shared_object_id": 41}, {"class_name": "Dense", "config": {"name": "dense_8", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 42}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 43}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_8", "inbound_nodes": [[["dense_7", 0, 0, {}]]], "shared_object_id": 44}, {"class_name": "Dense", "config": {"name": "dense_10", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 45}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 46}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_10", "inbound_nodes": [[["dense_9", 0, 0, {}]]], "shared_object_id": 47}, {"class_name": "Dense", "config": {"name": "dense_12", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 48}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 49}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_12", "inbound_nodes": [[["dense_11", 0, 0, {}]]], "shared_object_id": 50}, {"class_name": "MEGNetLayer", "config": {"name": "meg_net_layer_2", "trainable": false, "dtype": "float32", "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 51}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 52}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "units_e": [64, 64, 32], "units_v": [64, 64, 32], "units_u": [64, 64, 32], "pool_method": "mean"}, "name": "meg_net_layer_2", "inbound_nodes": [[["dense_8", 0, 0, {}], ["dense_10", 0, 0, {}], ["dense_12", 0, 0, {}], ["input_4", 0, 0, {}], ["input_5", 0, 0, {}], ["input_6", 0, 0, {}], ["input_7", 0, 0, {}]]], "shared_object_id": 53}, {"class_name": "Add", "config": {"name": "add_4", "trainable": false, "dtype": "float32"}, "name": "add_4", "inbound_nodes": [[["add_1", 0, 0, {}], ["meg_net_layer_2", 0, 0, {}]]], "shared_object_id": 54}, {"class_name": "Add", "config": {"name": "add_5", "trainable": false, "dtype": "float32"}, "name": "add_5", "inbound_nodes": [[["add_2", 0, 0, {}], ["meg_net_layer_2", 0, 1, {}]]], "shared_object_id": 55}, {"class_name": "Add", "config": {"name": "add_6", "trainable": false, "dtype": "float32"}, "name": "add_6", "inbound_nodes": [[["add_3", 0, 0, {}], ["meg_net_layer_2", 0, 2, {}]]], "shared_object_id": 56}, {"class_name": "Dense", "config": {"name": "dense_13", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 57}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 58}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_13", "inbound_nodes": [[["add_4", 0, 0, {}]]], "shared_object_id": 59}, {"class_name": "Dense", "config": {"name": "dense_15", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 60}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 61}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_15", "inbound_nodes": [[["add_5", 0, 0, {}]]], "shared_object_id": 62}, {"class_name": "Dense", "config": {"name": "dense_17", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 63}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 64}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_17", "inbound_nodes": [[["add_6", 0, 0, {}]]], "shared_object_id": 65}, {"class_name": "Dense", "config": {"name": "dense_14", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 66}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 67}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_14", "inbound_nodes": [[["dense_13", 0, 0, {}]]], "shared_object_id": 68}, {"class_name": "Dense", "config": {"name": "dense_16", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 69}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 70}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_16", "inbound_nodes": [[["dense_15", 0, 0, {}]]], "shared_object_id": 71}, {"class_name": "Dense", "config": {"name": "dense_18", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 72}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 73}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_18", "inbound_nodes": [[["dense_17", 0, 0, {}]]], "shared_object_id": 74}, {"class_name": "MEGNetLayer", "config": {"name": "meg_net_layer_3", "trainable": false, "dtype": "float32", "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 75}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 76}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "units_e": [64, 64, 32], "units_v": [64, 64, 32], "units_u": [64, 64, 32], "pool_method": "mean"}, "name": "meg_net_layer_3", "inbound_nodes": [[["dense_14", 0, 0, {}], ["dense_16", 0, 0, {}], ["dense_18", 0, 0, {}], ["input_4", 0, 0, {}], ["input_5", 0, 0, {}], ["input_6", 0, 0, {}], ["input_7", 0, 0, {}]]], "shared_object_id": 77}, {"class_name": "Add", "config": {"name": "add_7", "trainable": false, "dtype": "float32"}, "name": "add_7", "inbound_nodes": [[["add_4", 0, 0, {}], ["meg_net_layer_3", 0, 0, {}]]], "shared_object_id": 78}, {"class_name": "Add", "config": {"name": "add_8", "trainable": false, "dtype": "float32"}, "name": "add_8", "inbound_nodes": [[["add_5", 0, 0, {}], ["meg_net_layer_3", 0, 1, {}]]], "shared_object_id": 79}, {"class_name": "Set2Set", "config": {"name": "set2_set_1", "trainable": false, "dtype": "float32", "T": 2, "n_hidden": 16, "activation": "linear", "activation_lstm": "tanh", "recurrent_activation": "hard_sigmoid", "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 80}, "recurrent_initializer": {"class_name": "Orthogonal", "config": {"gain": 1.0, "seed": null}, "shared_object_id": 81}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 82}, "use_bias": true, "unit_forget_bias": true, "kernel_regularizer": null, "recurrent_regularizer": null, "bias_regularizer": null, "kernel_constraint": null, "recurrent_constraint": null, "bias_constraint": null}, "name": "set2_set_1", "inbound_nodes": [[["add_7", 0, 0, {}], ["input_6", 0, 0, {}]]], "shared_object_id": 83}, {"class_name": "Set2Set", "config": {"name": "set2_set_2", "trainable": false, "dtype": "float32", "T": 2, "n_hidden": 16, "activation": "linear", "activation_lstm": "tanh", "recurrent_activation": "hard_sigmoid", "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 84}, "recurrent_initializer": {"class_name": "Orthogonal", "config": {"gain": 1.0, "seed": null}, "shared_object_id": 85}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 86}, "use_bias": true, "unit_forget_bias": true, "kernel_regularizer": null, "recurrent_regularizer": null, "bias_regularizer": null, "kernel_constraint": null, "recurrent_constraint": null, "bias_constraint": null}, "name": "set2_set_2", "inbound_nodes": [[["add_8", 0, 0, {}], ["input_7", 0, 0, {}]]], "shared_object_id": 87}, {"class_name": "Add", "config": {"name": "add_9", "trainable": false, "dtype": "float32"}, "name": "add_9", "inbound_nodes": [[["add_6", 0, 0, {}], ["meg_net_layer_3", 0, 2, {}]]], "shared_object_id": 88}, {"class_name": "Concatenate", "config": {"name": "concatenate_1", "trainable": false, "dtype": "float32", "axis": -1}, "name": "concatenate_1", "inbound_nodes": [[["set2_set_1", 0, 0, {}], ["set2_set_2", 0, 0, {}], ["add_9", 0, 0, {}]]], "shared_object_id": 89}, {"class_name": "Dense", "config": {"name": "dense_19", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 90}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 91}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_19", "inbound_nodes": [[["concatenate_1", 0, 0, {}]]], "shared_object_id": 92}, {"class_name": "Dense", "config": {"name": "dense_20", "trainable": false, "dtype": "float32", "units": 16, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 93}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 94}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_20", "inbound_nodes": [[["dense_19", 0, 0, {}]]], "shared_object_id": 95}, {"class_name": "Dense", "config": {"name": "dense_21", "trainable": true, "dtype": "float32", "units": 1, "activation": "linear", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 96}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 97}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_21", "inbound_nodes": [[["dense_20", 0, 0, {}]]], "shared_object_id": 98}], "input_layers": [["input_1", 0, 0], ["input_2", 0, 0], ["input_3", 0, 0], ["input_4", 0, 0], ["input_5", 0, 0], ["input_6", 0, 0], ["input_7", 0, 0]], "output_layers": [["dense_21", 0, 0]]}}}2
+ root.layer-0"_tf_keras_input_layer*{"class_name": "InputLayer", "name": "input_1", "dtype": "int32", "sparse": false, "ragged": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_1"}}2
+root.layer_with_weights-0"_tf_keras_layer*{"name": "embedding_1", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "stateful": false, "must_restore_from_config": false, "class_name": "Embedding", "config": {"name": "embedding_1", "trainable": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "float32", "input_dim": 95, "output_dim": 16, "embeddings_initializer": {"class_name": "RandomUniform", "config": {"minval": -0.05, "maxval": 0.05, "seed": null}, "shared_object_id": 1}, "embeddings_regularizer": null, "activity_regularizer": null, "embeddings_constraint": null, "mask_zero": false, "input_length": null}, "inbound_nodes": [[["input_1", 0, 0, {}]]], "shared_object_id": 2, "build_input_shape": {"class_name": "TensorShape", "items": [null, null]}}2
+ root.layer-2"_tf_keras_input_layer*{"class_name": "InputLayer", "name": "input_2", "dtype": "float32", "sparse": false, "ragged": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null, 100]}, "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null, 100]}, "dtype": "float32", "sparse": false, "ragged": false, "name": "input_2"}}2
+ root.layer-3"_tf_keras_input_layer*{"class_name": "InputLayer", "name": "input_3", "dtype": "float32", "sparse": false, "ragged": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null, 2]}, "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null, 2]}, "dtype": "float32", "sparse": false, "ragged": false, "name": "input_3"}}2
+root.layer_with_weights-1"_tf_keras_layer*{"name": "dense_1", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_1", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 5}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 6}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["embedding_1", 0, 0, {}]]], "shared_object_id": 7, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 16}}, "shared_object_id": 107}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 16]}}2
+root.layer_with_weights-2"_tf_keras_layer*{"name": "dense_3", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_3", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 8}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 9}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["input_2", 0, 0, {}]]], "shared_object_id": 10, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 100}}, "shared_object_id": 108}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 100]}}2
+root.layer_with_weights-3"_tf_keras_layer*{"name": "dense_5", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_5", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 11}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 12}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["input_3", 0, 0, {}]]], "shared_object_id": 13, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 2}}, "shared_object_id": 109}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 2]}}2
+root.layer_with_weights-4"_tf_keras_layer*{"name": "dense_2", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_2", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 14}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 15}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_1", 0, 0, {}]]], "shared_object_id": 16, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 64}}, "shared_object_id": 110}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 64]}}2
+ root.layer_with_weights-5"_tf_keras_layer*{"name": "dense_4", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_4", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 17}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 18}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_3", 0, 0, {}]]], "shared_object_id": 19, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 64}}, "shared_object_id": 111}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 64]}}2
+
+root.layer_with_weights-6"_tf_keras_layer*{"name": "dense_6", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_6", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 20}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 21}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_5", 0, 0, {}]]], "shared_object_id": 22, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 64}}, "shared_object_id": 112}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 64]}}2
+
root.layer-10"_tf_keras_input_layer*{"class_name": "InputLayer", "name": "input_4", "dtype": "int32", "sparse": false, "ragged": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_4"}}2
+
root.layer-11"_tf_keras_input_layer*{"class_name": "InputLayer", "name": "input_5", "dtype": "int32", "sparse": false, "ragged": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_5"}}2
+
root.layer-12"_tf_keras_input_layer*{"class_name": "InputLayer", "name": "input_6", "dtype": "int32", "sparse": false, "ragged": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_6"}}2
+
root.layer-13"_tf_keras_input_layer*{"class_name": "InputLayer", "name": "input_7", "dtype": "int32", "sparse": false, "ragged": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_7"}}2
+ root.layer_with_weights-7"_tf_keras_layer* {"name": "meg_net_layer_1", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "MEGNetLayer", "config": {"name": "meg_net_layer_1", "trainable": false, "dtype": "float32", "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 27}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 28}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "units_e": [64, 64, 32], "units_v": [64, 64, 32], "units_u": [64, 64, 32], "pool_method": "mean"}, "inbound_nodes": [[["dense_2", 0, 0, {}], ["dense_4", 0, 0, {}], ["dense_6", 0, 0, {}], ["input_4", 0, 0, {}], ["input_5", 0, 0, {}], ["input_6", 0, 0, {}], ["input_7", 0, 0, {}]]], "shared_object_id": 29, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}]}2
+
root.layer-15"_tf_keras_layer*{"name": "add_1", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Add", "config": {"name": "add_1", "trainable": false, "dtype": "float32"}, "inbound_nodes": [[["dense_2", 0, 0, {}], ["meg_net_layer_1", 0, 0, {}]]], "shared_object_id": 30, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}]}2
+
root.layer-16"_tf_keras_layer*{"name": "add_2", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Add", "config": {"name": "add_2", "trainable": false, "dtype": "float32"}, "inbound_nodes": [[["dense_4", 0, 0, {}], ["meg_net_layer_1", 0, 1, {}]]], "shared_object_id": 31, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}]}2
+
root.layer-17"_tf_keras_layer*{"name": "add_3", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Add", "config": {"name": "add_3", "trainable": false, "dtype": "float32"}, "inbound_nodes": [[["dense_6", 0, 0, {}], ["meg_net_layer_1", 0, 2, {}]]], "shared_object_id": 32, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [1, null, 32]}]}2
+root.layer_with_weights-8"_tf_keras_layer*{"name": "dense_7", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_7", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 33}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 34}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["add_1", 0, 0, {}]]], "shared_object_id": 35, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 32}}, "shared_object_id": 113}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 32]}}2
+root.layer_with_weights-9"_tf_keras_layer*{"name": "dense_9", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_9", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 36}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 37}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["add_2", 0, 0, {}]]], "shared_object_id": 38, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 32}}, "shared_object_id": 114}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 32]}}2
+root.layer_with_weights-10"_tf_keras_layer*{"name": "dense_11", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_11", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 39}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 40}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["add_3", 0, 0, {}]]], "shared_object_id": 41, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 32}}, "shared_object_id": 115}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 32]}}2
+root.layer_with_weights-11"_tf_keras_layer*{"name": "dense_8", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_8", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 42}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 43}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_7", 0, 0, {}]]], "shared_object_id": 44, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 64}}, "shared_object_id": 116}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 64]}}2
+root.layer_with_weights-12"_tf_keras_layer*{"name": "dense_10", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_10", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 45}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 46}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_9", 0, 0, {}]]], "shared_object_id": 47, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 64}}, "shared_object_id": 117}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 64]}}2
+root.layer_with_weights-13"_tf_keras_layer*{"name": "dense_12", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_12", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 48}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 49}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_11", 0, 0, {}]]], "shared_object_id": 50, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 64}}, "shared_object_id": 118}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 64]}}2
+ root.layer_with_weights-14"_tf_keras_layer* {"name": "meg_net_layer_2", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "MEGNetLayer", "config": {"name": "meg_net_layer_2", "trainable": false, "dtype": "float32", "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 51}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 52}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "units_e": [64, 64, 32], "units_v": [64, 64, 32], "units_u": [64, 64, 32], "pool_method": "mean"}, "inbound_nodes": [[["dense_8", 0, 0, {}], ["dense_10", 0, 0, {}], ["dense_12", 0, 0, {}], ["input_4", 0, 0, {}], ["input_5", 0, 0, {}], ["input_6", 0, 0, {}], ["input_7", 0, 0, {}]]], "shared_object_id": 53, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}]}2
+
root.layer-25"_tf_keras_layer*{"name": "add_4", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Add", "config": {"name": "add_4", "trainable": false, "dtype": "float32"}, "inbound_nodes": [[["add_1", 0, 0, {}], ["meg_net_layer_2", 0, 0, {}]]], "shared_object_id": 54, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}]}2
+
root.layer-26"_tf_keras_layer*{"name": "add_5", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Add", "config": {"name": "add_5", "trainable": false, "dtype": "float32"}, "inbound_nodes": [[["add_2", 0, 0, {}], ["meg_net_layer_2", 0, 1, {}]]], "shared_object_id": 55, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}]}2
+
root.layer-27"_tf_keras_layer*{"name": "add_6", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Add", "config": {"name": "add_6", "trainable": false, "dtype": "float32"}, "inbound_nodes": [[["add_3", 0, 0, {}], ["meg_net_layer_2", 0, 2, {}]]], "shared_object_id": 56, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [1, null, 32]}]}2
+ root.layer_with_weights-15"_tf_keras_layer*{"name": "dense_13", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_13", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 57}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 58}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["add_4", 0, 0, {}]]], "shared_object_id": 59, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 32}}, "shared_object_id": 119}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 32]}}2
+ root.layer_with_weights-16"_tf_keras_layer*{"name": "dense_15", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_15", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 60}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 61}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["add_5", 0, 0, {}]]], "shared_object_id": 62, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 32}}, "shared_object_id": 120}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 32]}}2
+ root.layer_with_weights-17"_tf_keras_layer*{"name": "dense_17", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_17", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 63}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 64}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["add_6", 0, 0, {}]]], "shared_object_id": 65, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 32}}, "shared_object_id": 121}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 32]}}2
+ root.layer_with_weights-18"_tf_keras_layer*{"name": "dense_14", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_14", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 66}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 67}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_13", 0, 0, {}]]], "shared_object_id": 68, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 64}}, "shared_object_id": 122}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 64]}}2
+!root.layer_with_weights-19"_tf_keras_layer*{"name": "dense_16", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_16", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 69}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 70}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_15", 0, 0, {}]]], "shared_object_id": 71, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 64}}, "shared_object_id": 123}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 64]}}2
+"root.layer_with_weights-20"_tf_keras_layer*{"name": "dense_18", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_18", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 72}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 73}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_17", 0, 0, {}]]], "shared_object_id": 74, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 64}}, "shared_object_id": 124}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 64]}}2
+ #root.layer_with_weights-21"_tf_keras_layer* {"name": "meg_net_layer_3", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "MEGNetLayer", "config": {"name": "meg_net_layer_3", "trainable": false, "dtype": "float32", "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 75}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 76}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "units_e": [64, 64, 32], "units_v": [64, 64, 32], "units_u": [64, 64, 32], "pool_method": "mean"}, "inbound_nodes": [[["dense_14", 0, 0, {}], ["dense_16", 0, 0, {}], ["dense_18", 0, 0, {}], ["input_4", 0, 0, {}], ["input_5", 0, 0, {}], ["input_6", 0, 0, {}], ["input_7", 0, 0, {}]]], "shared_object_id": 77, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}]}2
+$
root.layer-35"_tf_keras_layer*{"name": "add_7", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Add", "config": {"name": "add_7", "trainable": false, "dtype": "float32"}, "inbound_nodes": [[["add_4", 0, 0, {}], ["meg_net_layer_3", 0, 0, {}]]], "shared_object_id": 78, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}]}2
+%
root.layer-36"_tf_keras_layer*{"name": "add_8", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Add", "config": {"name": "add_8", "trainable": false, "dtype": "float32"}, "inbound_nodes": [[["add_5", 0, 0, {}], ["meg_net_layer_3", 0, 1, {}]]], "shared_object_id": 79, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}]}2
+ &root.layer_with_weights-22"_tf_keras_layer* {"name": "set2_set_1", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Set2Set", "config": {"name": "set2_set_1", "trainable": false, "dtype": "float32", "T": 2, "n_hidden": 16, "activation": "linear", "activation_lstm": "tanh", "recurrent_activation": "hard_sigmoid", "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 80}, "recurrent_initializer": {"class_name": "Orthogonal", "config": {"gain": 1.0, "seed": null}, "shared_object_id": 81}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 82}, "use_bias": true, "unit_forget_bias": true, "kernel_regularizer": null, "recurrent_regularizer": null, "bias_regularizer": null, "kernel_constraint": null, "recurrent_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["add_7", 0, 0, {}], ["input_6", 0, 0, {}]]], "shared_object_id": 83, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null]}]}2
+ 'root.layer_with_weights-23"_tf_keras_layer* {"name": "set2_set_2", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Set2Set", "config": {"name": "set2_set_2", "trainable": false, "dtype": "float32", "T": 2, "n_hidden": 16, "activation": "linear", "activation_lstm": "tanh", "recurrent_activation": "hard_sigmoid", "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 84}, "recurrent_initializer": {"class_name": "Orthogonal", "config": {"gain": 1.0, "seed": null}, "shared_object_id": 85}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 86}, "use_bias": true, "unit_forget_bias": true, "kernel_regularizer": null, "recurrent_regularizer": null, "bias_regularizer": null, "kernel_constraint": null, "recurrent_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["add_8", 0, 0, {}], ["input_7", 0, 0, {}]]], "shared_object_id": 87, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null]}]}2
+(
root.layer-39"_tf_keras_layer*{"name": "add_9", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Add", "config": {"name": "add_9", "trainable": false, "dtype": "float32"}, "inbound_nodes": [[["add_6", 0, 0, {}], ["meg_net_layer_3", 0, 2, {}]]], "shared_object_id": 88, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [1, null, 32]}]}2
+)
root.layer-40"_tf_keras_layer*{"name": "concatenate_1", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Concatenate", "config": {"name": "concatenate_1", "trainable": false, "dtype": "float32", "axis": -1}, "inbound_nodes": [[["set2_set_1", 0, 0, {}], ["set2_set_2", 0, 0, {}], ["add_9", 0, 0, {}]]], "shared_object_id": 89, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}]}2
+*root.layer_with_weights-24"_tf_keras_layer*{"name": "dense_19", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_19", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 90}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 91}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["concatenate_1", 0, 0, {}]]], "shared_object_id": 92, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 96}}, "shared_object_id": 125}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 96]}}2
++root.layer_with_weights-25"_tf_keras_layer*{"name": "dense_20", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_20", "trainable": false, "dtype": "float32", "units": 16, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 93}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 94}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_19", 0, 0, {}]]], "shared_object_id": 95, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 32}}, "shared_object_id": 126}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 32]}}2
+,root.layer_with_weights-26"_tf_keras_layer*{"name": "dense_21", "trainable": true, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_21", "trainable": true, "dtype": "float32", "units": 1, "activation": "linear", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 96}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 97}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_20", 0, 0, {}]]], "shared_object_id": 98, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 16}}, "shared_object_id": 127}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 16]}}2
+root.keras_api.metrics.0"_tf_keras_metric*{"class_name": "Mean", "name": "loss", "dtype": "float32", "config": {"name": "loss", "dtype": "float32"}, "shared_object_id": 128}2
\ No newline at end of file
diff --git a/tests/test_model/prob_e_form_norm/nn/saved_model.pb b/tests/test_model/prob_e_form_norm/nn/saved_model.pb
new file mode 100644
index 0000000..dc1bb64
Binary files /dev/null and b/tests/test_model/prob_e_form_norm/nn/saved_model.pb differ
diff --git a/tests/test_model/prob_e_form_norm/nn/variables/variables.data-00000-of-00001 b/tests/test_model/prob_e_form_norm/nn/variables/variables.data-00000-of-00001
new file mode 100644
index 0000000..4b13792
Binary files /dev/null and b/tests/test_model/prob_e_form_norm/nn/variables/variables.data-00000-of-00001 differ
diff --git a/tests/test_model/prob_e_form_norm/nn/variables/variables.index b/tests/test_model/prob_e_form_norm/nn/variables/variables.index
new file mode 100644
index 0000000..a42bc06
Binary files /dev/null and b/tests/test_model/prob_e_form_norm/nn/variables/variables.index differ
diff --git a/tests/test_model/prob_e_form_norm/weights.data-00000-of-00001 b/tests/test_model/prob_e_form_norm/weights.data-00000-of-00001
new file mode 100644
index 0000000..953da0f
Binary files /dev/null and b/tests/test_model/prob_e_form_norm/weights.data-00000-of-00001 differ
diff --git a/tests/test_model/prob_e_form_norm/weights.index b/tests/test_model/prob_e_form_norm/weights.index
new file mode 100644
index 0000000..bc652c2
Binary files /dev/null and b/tests/test_model/prob_e_form_norm/weights.index differ
diff --git a/tests/test_model/prob_e_form_unnorm/checkpoint b/tests/test_model/prob_e_form_unnorm/checkpoint
new file mode 100644
index 0000000..efd0621
--- /dev/null
+++ b/tests/test_model/prob_e_form_unnorm/checkpoint
@@ -0,0 +1,2 @@
+model_checkpoint_path: "weights"
+all_model_checkpoint_paths: "weights"
diff --git a/tests/test_model/prob_e_form_unnorm/config.json b/tests/test_model/prob_e_form_unnorm/config.json
new file mode 100644
index 0000000..c00977d
--- /dev/null
+++ b/tests/test_model/prob_e_form_unnorm/config.json
@@ -0,0 +1 @@
+{"num_inducing_points": 10, "metrics": [], "kl_weight": 1.0, "latent_layer": -2, "target_shape": [1], "use_normalization": false}
\ No newline at end of file
diff --git a/tests/test_model/prob_e_form_unnorm/kernel/config.json b/tests/test_model/prob_e_form_unnorm/kernel/config.json
new file mode 100644
index 0000000..b8a0da1
--- /dev/null
+++ b/tests/test_model/prob_e_form_unnorm/kernel/config.json
@@ -0,0 +1 @@
+{"type": "rbf"}
\ No newline at end of file
diff --git a/tests/test_model/prob_e_form_unnorm/kernel/params.npy b/tests/test_model/prob_e_form_unnorm/kernel/params.npy
new file mode 100644
index 0000000..6d8806c
Binary files /dev/null and b/tests/test_model/prob_e_form_unnorm/kernel/params.npy differ
diff --git a/tests/test_model/prob_e_form_unnorm/megnet.json b/tests/test_model/prob_e_form_unnorm/megnet.json
new file mode 100644
index 0000000..c931337
--- /dev/null
+++ b/tests/test_model/prob_e_form_unnorm/megnet.json
@@ -0,0 +1 @@
+{"graph_converter": {"@module": "megnet.data.crystal", "@class": "CrystalGraph", "@version": "1.2.8", "atom_converter": {"@module": "megnet.data.graph", "@class": "DummyConverter", "@version": "1.2.8"}, "bond_converter": {"@module": "megnet.data.graph", "@class": "GaussianDistance", "@version": "1.2.8", "centers": {"@module": "numpy", "@class": "array", "dtype": "float64", "data": [0.0, 0.06060606060606061, 0.12121212121212122, 0.18181818181818182, 0.24242424242424243, 0.30303030303030304, 0.36363636363636365, 0.42424242424242425, 0.48484848484848486, 0.5454545454545454, 0.6060606060606061, 0.6666666666666667, 0.7272727272727273, 0.7878787878787878, 0.8484848484848485, 0.9090909090909092, 0.9696969696969697, 1.0303030303030303, 1.0909090909090908, 1.1515151515151516, 1.2121212121212122, 1.2727272727272727, 1.3333333333333335, 1.393939393939394, 1.4545454545454546, 1.5151515151515151, 1.5757575757575757, 1.6363636363636365, 1.696969696969697, 1.7575757575757576, 1.8181818181818183, 1.878787878787879, 1.9393939393939394, 2.0, 2.0606060606060606, 2.121212121212121, 2.1818181818181817, 2.2424242424242427, 2.303030303030303, 2.3636363636363638, 2.4242424242424243, 2.484848484848485, 2.5454545454545454, 2.606060606060606, 2.666666666666667, 2.7272727272727275, 2.787878787878788, 2.8484848484848486, 2.909090909090909, 2.9696969696969697, 3.0303030303030303, 3.090909090909091, 3.1515151515151514, 3.2121212121212124, 3.272727272727273, 3.3333333333333335, 3.393939393939394, 3.4545454545454546, 3.515151515151515, 3.5757575757575757, 3.6363636363636367, 3.6969696969696972, 3.757575757575758, 3.8181818181818183, 3.878787878787879, 3.9393939393939394, 4.0, 4.0606060606060606, 4.121212121212121, 4.181818181818182, 4.242424242424242, 4.303030303030303, 4.363636363636363, 4.424242424242425, 4.484848484848485, 4.545454545454546, 4.606060606060606, 4.666666666666667, 4.7272727272727275, 4.787878787878788, 4.848484848484849, 4.909090909090909, 4.96969696969697, 5.03030303030303, 5.090909090909091, 5.151515151515151, 5.212121212121212, 5.2727272727272725, 5.333333333333334, 5.3939393939393945, 5.454545454545455, 5.515151515151516, 5.575757575757576, 5.636363636363637, 5.696969696969697, 5.757575757575758, 5.818181818181818, 5.878787878787879, 5.9393939393939394, 6.0]}, "width": 0.5}, "cutoff": 5.0, "nn_strategy": {"@module": "megnet.data.local_env", "@class": "MinimumDistanceNNAll", "cutoff": 5.0}}, "target_scaler": {"@module": "megnet.utils.preprocessing", "@class": "DummyScaler", "@version": "1.2.8"}, "metadata": {"name": "Formation energy", "unit": "eV/atom", "metric": "mae", "metric value": 0.026, "training dataset": "mp-2019.4.1", "description": "This model was trained using structures and formation energy of Materials Project data base downloaded on April 1, 2019. The total data size is 133420. The model was trained by a 0.8-0.1-0.1, train, validation and test data ratio"}}
\ No newline at end of file
diff --git a/tests/test_model/prob_e_form_unnorm/megnet/keras_metadata.pb b/tests/test_model/prob_e_form_unnorm/megnet/keras_metadata.pb
new file mode 100644
index 0000000..003327a
--- /dev/null
+++ b/tests/test_model/prob_e_form_unnorm/megnet/keras_metadata.pb
@@ -0,0 +1,48 @@
+
+root"_tf_keras_network*{"name": "model_1", "trainable": true, "expects_training_arg": true, "dtype": "float32", "batch_input_shape": null, "must_restore_from_config": false, "class_name": "Functional", "config": {"name": "model_1", "layers": [{"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_1"}, "name": "input_1", "inbound_nodes": []}, {"class_name": "Embedding", "config": {"name": "embedding_1", "trainable": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "float32", "input_dim": 95, "output_dim": 16, "embeddings_initializer": {"class_name": "RandomUniform", "config": {"minval": -0.05, "maxval": 0.05, "seed": null}}, "embeddings_regularizer": null, "activity_regularizer": null, "embeddings_constraint": null, "mask_zero": false, "input_length": null}, "name": "embedding_1", "inbound_nodes": [[["input_1", 0, 0, {}]]]}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null, 100]}, "dtype": "float32", "sparse": false, "ragged": false, "name": "input_2"}, "name": "input_2", "inbound_nodes": []}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null, 2]}, "dtype": "float32", "sparse": false, "ragged": false, "name": "input_3"}, "name": "input_3", "inbound_nodes": []}, {"class_name": "Dense", "config": {"name": "dense_1", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_1", "inbound_nodes": [[["embedding_1", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_3", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_3", "inbound_nodes": [[["input_2", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_5", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_5", "inbound_nodes": [[["input_3", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_2", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_2", "inbound_nodes": [[["dense_1", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_4", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_4", "inbound_nodes": [[["dense_3", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_6", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_6", "inbound_nodes": [[["dense_5", 0, 0, {}]]]}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_4"}, "name": "input_4", "inbound_nodes": []}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_5"}, "name": "input_5", "inbound_nodes": []}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_6"}, "name": "input_6", "inbound_nodes": []}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_7"}, "name": "input_7", "inbound_nodes": []}, {"class_name": "MEGNetLayer", "config": {"name": "meg_net_layer_1", "trainable": false, "dtype": "float32", "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "units_e": [64, 64, 32], "units_v": [64, 64, 32], "units_u": [64, 64, 32], "pool_method": "mean"}, "name": "meg_net_layer_1", "inbound_nodes": [[["dense_2", 0, 0, {}], ["dense_4", 0, 0, {}], ["dense_6", 0, 0, {}], ["input_4", 0, 0, {}], ["input_5", 0, 0, {}], ["input_6", 0, 0, {}], ["input_7", 0, 0, {}]]]}, {"class_name": "Add", "config": {"name": "add_1", "trainable": false, "dtype": "float32"}, "name": "add_1", "inbound_nodes": [[["dense_2", 0, 0, {}], ["meg_net_layer_1", 0, 0, {}]]]}, {"class_name": "Add", "config": {"name": "add_2", "trainable": false, "dtype": "float32"}, "name": "add_2", "inbound_nodes": [[["dense_4", 0, 0, {}], ["meg_net_layer_1", 0, 1, {}]]]}, {"class_name": "Add", "config": {"name": "add_3", "trainable": false, "dtype": "float32"}, "name": "add_3", "inbound_nodes": [[["dense_6", 0, 0, {}], ["meg_net_layer_1", 0, 2, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_7", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_7", "inbound_nodes": [[["add_1", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_9", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_9", "inbound_nodes": [[["add_2", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_11", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_11", "inbound_nodes": [[["add_3", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_8", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_8", "inbound_nodes": [[["dense_7", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_10", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_10", "inbound_nodes": [[["dense_9", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_12", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_12", "inbound_nodes": [[["dense_11", 0, 0, {}]]]}, {"class_name": "MEGNetLayer", "config": {"name": "meg_net_layer_2", "trainable": false, "dtype": "float32", "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "units_e": [64, 64, 32], "units_v": [64, 64, 32], "units_u": [64, 64, 32], "pool_method": "mean"}, "name": "meg_net_layer_2", "inbound_nodes": [[["dense_8", 0, 0, {}], ["dense_10", 0, 0, {}], ["dense_12", 0, 0, {}], ["input_4", 0, 0, {}], ["input_5", 0, 0, {}], ["input_6", 0, 0, {}], ["input_7", 0, 0, {}]]]}, {"class_name": "Add", "config": {"name": "add_4", "trainable": false, "dtype": "float32"}, "name": "add_4", "inbound_nodes": [[["add_1", 0, 0, {}], ["meg_net_layer_2", 0, 0, {}]]]}, {"class_name": "Add", "config": {"name": "add_5", "trainable": false, "dtype": "float32"}, "name": "add_5", "inbound_nodes": [[["add_2", 0, 0, {}], ["meg_net_layer_2", 0, 1, {}]]]}, {"class_name": "Add", "config": {"name": "add_6", "trainable": false, "dtype": "float32"}, "name": "add_6", "inbound_nodes": [[["add_3", 0, 0, {}], ["meg_net_layer_2", 0, 2, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_13", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_13", "inbound_nodes": [[["add_4", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_15", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_15", "inbound_nodes": [[["add_5", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_17", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_17", "inbound_nodes": [[["add_6", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_14", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_14", "inbound_nodes": [[["dense_13", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_16", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_16", "inbound_nodes": [[["dense_15", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_18", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_18", "inbound_nodes": [[["dense_17", 0, 0, {}]]]}, {"class_name": "MEGNetLayer", "config": {"name": "meg_net_layer_3", "trainable": false, "dtype": "float32", "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "units_e": [64, 64, 32], "units_v": [64, 64, 32], "units_u": [64, 64, 32], "pool_method": "mean"}, "name": "meg_net_layer_3", "inbound_nodes": [[["dense_14", 0, 0, {}], ["dense_16", 0, 0, {}], ["dense_18", 0, 0, {}], ["input_4", 0, 0, {}], ["input_5", 0, 0, {}], ["input_6", 0, 0, {}], ["input_7", 0, 0, {}]]]}, {"class_name": "Add", "config": {"name": "add_7", "trainable": false, "dtype": "float32"}, "name": "add_7", "inbound_nodes": [[["add_4", 0, 0, {}], ["meg_net_layer_3", 0, 0, {}]]]}, {"class_name": "Add", "config": {"name": "add_8", "trainable": false, "dtype": "float32"}, "name": "add_8", "inbound_nodes": [[["add_5", 0, 0, {}], ["meg_net_layer_3", 0, 1, {}]]]}, {"class_name": "Set2Set", "config": {"name": "set2_set_1", "trainable": false, "dtype": "float32", "T": 2, "n_hidden": 16, "activation": "linear", "activation_lstm": "tanh", "recurrent_activation": "hard_sigmoid", "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "recurrent_initializer": {"class_name": "Orthogonal", "config": {"gain": 1.0, "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "use_bias": true, "unit_forget_bias": true, "kernel_regularizer": null, "recurrent_regularizer": null, "bias_regularizer": null, "kernel_constraint": null, "recurrent_constraint": null, "bias_constraint": null}, "name": "set2_set_1", "inbound_nodes": [[["add_7", 0, 0, {}], ["input_6", 0, 0, {}]]]}, {"class_name": "Set2Set", "config": {"name": "set2_set_2", "trainable": false, "dtype": "float32", "T": 2, "n_hidden": 16, "activation": "linear", "activation_lstm": "tanh", "recurrent_activation": "hard_sigmoid", "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "recurrent_initializer": {"class_name": "Orthogonal", "config": {"gain": 1.0, "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "use_bias": true, "unit_forget_bias": true, "kernel_regularizer": null, "recurrent_regularizer": null, "bias_regularizer": null, "kernel_constraint": null, "recurrent_constraint": null, "bias_constraint": null}, "name": "set2_set_2", "inbound_nodes": [[["add_8", 0, 0, {}], ["input_7", 0, 0, {}]]]}, {"class_name": "Add", "config": {"name": "add_9", "trainable": false, "dtype": "float32"}, "name": "add_9", "inbound_nodes": [[["add_6", 0, 0, {}], ["meg_net_layer_3", 0, 2, {}]]]}, {"class_name": "Concatenate", "config": {"name": "concatenate_1", "trainable": false, "dtype": "float32", "axis": -1}, "name": "concatenate_1", "inbound_nodes": [[["set2_set_1", 0, 0, {}], ["set2_set_2", 0, 0, {}], ["add_9", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_19", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_19", "inbound_nodes": [[["concatenate_1", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_20", "trainable": false, "dtype": "float32", "units": 16, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_20", "inbound_nodes": [[["dense_19", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_21", "trainable": true, "dtype": "float32", "units": 1, "activation": "linear", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_21", "inbound_nodes": [[["dense_20", 0, 0, {}]]]}], "input_layers": [["input_1", 0, 0], ["input_2", 0, 0], ["input_3", 0, 0], ["input_4", 0, 0], ["input_5", 0, 0], ["input_6", 0, 0], ["input_7", 0, 0]], "output_layers": [["dense_21", 0, 0]]}, "shared_object_id": 99, "input_spec": [{"class_name": "InputSpec", "config": {"dtype": null, "shape": {"class_name": "__tuple__", "items": [null, null]}, "ndim": 2, "max_ndim": null, "min_ndim": null, "axes": {}}}, {"class_name": "InputSpec", "config": {"dtype": null, "shape": {"class_name": "__tuple__", "items": [null, null, 100]}, "ndim": 3, "max_ndim": null, "min_ndim": null, "axes": {}}}, {"class_name": "InputSpec", "config": {"dtype": null, "shape": {"class_name": "__tuple__", "items": [null, null, 2]}, "ndim": 3, "max_ndim": null, "min_ndim": null, "axes": {}}}, {"class_name": "InputSpec", "config": {"dtype": null, "shape": {"class_name": "__tuple__", "items": [null, null]}, "ndim": 2, "max_ndim": null, "min_ndim": null, "axes": {}}}, {"class_name": "InputSpec", "config": {"dtype": null, "shape": {"class_name": "__tuple__", "items": [null, null]}, "ndim": 2, "max_ndim": null, "min_ndim": null, "axes": {}}}, {"class_name": "InputSpec", "config": {"dtype": null, "shape": {"class_name": "__tuple__", "items": [null, null]}, "ndim": 2, "max_ndim": null, "min_ndim": null, "axes": {}}}, {"class_name": "InputSpec", "config": {"dtype": null, "shape": {"class_name": "__tuple__", "items": [null, null]}, "ndim": 2, "max_ndim": null, "min_ndim": null, "axes": {}}}], "build_input_shape": [{"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null, 100]}, {"class_name": "TensorShape", "items": [null, null, 2]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}], "is_graph_network": true, "full_save_spec": {"class_name": "__tuple__", "items": [[[{"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_1"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null, 100]}, "float32", "input_2"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null, 2]}, "float32", "input_3"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_4"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_5"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_6"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_7"]}]], {}]}, "save_spec": [{"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_1"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null, 100]}, "float32", "input_2"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null, 2]}, "float32", "input_3"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_4"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_5"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_6"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_7"]}], "keras_version": "2.6.0", "backend": "tensorflow", "model_config": {"class_name": "Functional", "config": {"name": "model_1", "layers": [{"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_1"}, "name": "input_1", "inbound_nodes": [], "shared_object_id": 0}, {"class_name": "Embedding", "config": {"name": "embedding_1", "trainable": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "float32", "input_dim": 95, "output_dim": 16, "embeddings_initializer": {"class_name": "RandomUniform", "config": {"minval": -0.05, "maxval": 0.05, "seed": null}, "shared_object_id": 1}, "embeddings_regularizer": null, "activity_regularizer": null, "embeddings_constraint": null, "mask_zero": false, "input_length": null}, "name": "embedding_1", "inbound_nodes": [[["input_1", 0, 0, {}]]], "shared_object_id": 2}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null, 100]}, "dtype": "float32", "sparse": false, "ragged": false, "name": "input_2"}, "name": "input_2", "inbound_nodes": [], "shared_object_id": 3}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null, 2]}, "dtype": "float32", "sparse": false, "ragged": false, "name": "input_3"}, "name": "input_3", "inbound_nodes": [], "shared_object_id": 4}, {"class_name": "Dense", "config": {"name": "dense_1", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 5}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 6}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_1", "inbound_nodes": [[["embedding_1", 0, 0, {}]]], "shared_object_id": 7}, {"class_name": "Dense", "config": {"name": "dense_3", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 8}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 9}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_3", "inbound_nodes": [[["input_2", 0, 0, {}]]], "shared_object_id": 10}, {"class_name": "Dense", "config": {"name": "dense_5", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 11}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 12}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_5", "inbound_nodes": [[["input_3", 0, 0, {}]]], "shared_object_id": 13}, {"class_name": "Dense", "config": {"name": "dense_2", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 14}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 15}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_2", "inbound_nodes": [[["dense_1", 0, 0, {}]]], "shared_object_id": 16}, {"class_name": "Dense", "config": {"name": "dense_4", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 17}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 18}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_4", "inbound_nodes": [[["dense_3", 0, 0, {}]]], "shared_object_id": 19}, {"class_name": "Dense", "config": {"name": "dense_6", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 20}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 21}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_6", "inbound_nodes": [[["dense_5", 0, 0, {}]]], "shared_object_id": 22}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_4"}, "name": "input_4", "inbound_nodes": [], "shared_object_id": 23}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_5"}, "name": "input_5", "inbound_nodes": [], "shared_object_id": 24}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_6"}, "name": "input_6", "inbound_nodes": [], "shared_object_id": 25}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_7"}, "name": "input_7", "inbound_nodes": [], "shared_object_id": 26}, {"class_name": "MEGNetLayer", "config": {"name": "meg_net_layer_1", "trainable": false, "dtype": "float32", "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 27}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 28}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "units_e": [64, 64, 32], "units_v": [64, 64, 32], "units_u": [64, 64, 32], "pool_method": "mean"}, "name": "meg_net_layer_1", "inbound_nodes": [[["dense_2", 0, 0, {}], ["dense_4", 0, 0, {}], ["dense_6", 0, 0, {}], ["input_4", 0, 0, {}], ["input_5", 0, 0, {}], ["input_6", 0, 0, {}], ["input_7", 0, 0, {}]]], "shared_object_id": 29}, {"class_name": "Add", "config": {"name": "add_1", "trainable": false, "dtype": "float32"}, "name": "add_1", "inbound_nodes": [[["dense_2", 0, 0, {}], ["meg_net_layer_1", 0, 0, {}]]], "shared_object_id": 30}, {"class_name": "Add", "config": {"name": "add_2", "trainable": false, "dtype": "float32"}, "name": "add_2", "inbound_nodes": [[["dense_4", 0, 0, {}], ["meg_net_layer_1", 0, 1, {}]]], "shared_object_id": 31}, {"class_name": "Add", "config": {"name": "add_3", "trainable": false, "dtype": "float32"}, "name": "add_3", "inbound_nodes": [[["dense_6", 0, 0, {}], ["meg_net_layer_1", 0, 2, {}]]], "shared_object_id": 32}, {"class_name": "Dense", "config": {"name": "dense_7", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 33}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 34}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_7", "inbound_nodes": [[["add_1", 0, 0, {}]]], "shared_object_id": 35}, {"class_name": "Dense", "config": {"name": "dense_9", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 36}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 37}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_9", "inbound_nodes": [[["add_2", 0, 0, {}]]], "shared_object_id": 38}, {"class_name": "Dense", "config": {"name": "dense_11", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 39}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 40}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_11", "inbound_nodes": [[["add_3", 0, 0, {}]]], "shared_object_id": 41}, {"class_name": "Dense", "config": {"name": "dense_8", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 42}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 43}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_8", "inbound_nodes": [[["dense_7", 0, 0, {}]]], "shared_object_id": 44}, {"class_name": "Dense", "config": {"name": "dense_10", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 45}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 46}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_10", "inbound_nodes": [[["dense_9", 0, 0, {}]]], "shared_object_id": 47}, {"class_name": "Dense", "config": {"name": "dense_12", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 48}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 49}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_12", "inbound_nodes": [[["dense_11", 0, 0, {}]]], "shared_object_id": 50}, {"class_name": "MEGNetLayer", "config": {"name": "meg_net_layer_2", "trainable": false, "dtype": "float32", "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 51}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 52}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "units_e": [64, 64, 32], "units_v": [64, 64, 32], "units_u": [64, 64, 32], "pool_method": "mean"}, "name": "meg_net_layer_2", "inbound_nodes": [[["dense_8", 0, 0, {}], ["dense_10", 0, 0, {}], ["dense_12", 0, 0, {}], ["input_4", 0, 0, {}], ["input_5", 0, 0, {}], ["input_6", 0, 0, {}], ["input_7", 0, 0, {}]]], "shared_object_id": 53}, {"class_name": "Add", "config": {"name": "add_4", "trainable": false, "dtype": "float32"}, "name": "add_4", "inbound_nodes": [[["add_1", 0, 0, {}], ["meg_net_layer_2", 0, 0, {}]]], "shared_object_id": 54}, {"class_name": "Add", "config": {"name": "add_5", "trainable": false, "dtype": "float32"}, "name": "add_5", "inbound_nodes": [[["add_2", 0, 0, {}], ["meg_net_layer_2", 0, 1, {}]]], "shared_object_id": 55}, {"class_name": "Add", "config": {"name": "add_6", "trainable": false, "dtype": "float32"}, "name": "add_6", "inbound_nodes": [[["add_3", 0, 0, {}], ["meg_net_layer_2", 0, 2, {}]]], "shared_object_id": 56}, {"class_name": "Dense", "config": {"name": "dense_13", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 57}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 58}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_13", "inbound_nodes": [[["add_4", 0, 0, {}]]], "shared_object_id": 59}, {"class_name": "Dense", "config": {"name": "dense_15", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 60}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 61}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_15", "inbound_nodes": [[["add_5", 0, 0, {}]]], "shared_object_id": 62}, {"class_name": "Dense", "config": {"name": "dense_17", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 63}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 64}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_17", "inbound_nodes": [[["add_6", 0, 0, {}]]], "shared_object_id": 65}, {"class_name": "Dense", "config": {"name": "dense_14", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 66}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 67}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_14", "inbound_nodes": [[["dense_13", 0, 0, {}]]], "shared_object_id": 68}, {"class_name": "Dense", "config": {"name": "dense_16", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 69}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 70}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_16", "inbound_nodes": [[["dense_15", 0, 0, {}]]], "shared_object_id": 71}, {"class_name": "Dense", "config": {"name": "dense_18", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 72}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 73}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_18", "inbound_nodes": [[["dense_17", 0, 0, {}]]], "shared_object_id": 74}, {"class_name": "MEGNetLayer", "config": {"name": "meg_net_layer_3", "trainable": false, "dtype": "float32", "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 75}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 76}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "units_e": [64, 64, 32], "units_v": [64, 64, 32], "units_u": [64, 64, 32], "pool_method": "mean"}, "name": "meg_net_layer_3", "inbound_nodes": [[["dense_14", 0, 0, {}], ["dense_16", 0, 0, {}], ["dense_18", 0, 0, {}], ["input_4", 0, 0, {}], ["input_5", 0, 0, {}], ["input_6", 0, 0, {}], ["input_7", 0, 0, {}]]], "shared_object_id": 77}, {"class_name": "Add", "config": {"name": "add_7", "trainable": false, "dtype": "float32"}, "name": "add_7", "inbound_nodes": [[["add_4", 0, 0, {}], ["meg_net_layer_3", 0, 0, {}]]], "shared_object_id": 78}, {"class_name": "Add", "config": {"name": "add_8", "trainable": false, "dtype": "float32"}, "name": "add_8", "inbound_nodes": [[["add_5", 0, 0, {}], ["meg_net_layer_3", 0, 1, {}]]], "shared_object_id": 79}, {"class_name": "Set2Set", "config": {"name": "set2_set_1", "trainable": false, "dtype": "float32", "T": 2, "n_hidden": 16, "activation": "linear", "activation_lstm": "tanh", "recurrent_activation": "hard_sigmoid", "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 80}, "recurrent_initializer": {"class_name": "Orthogonal", "config": {"gain": 1.0, "seed": null}, "shared_object_id": 81}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 82}, "use_bias": true, "unit_forget_bias": true, "kernel_regularizer": null, "recurrent_regularizer": null, "bias_regularizer": null, "kernel_constraint": null, "recurrent_constraint": null, "bias_constraint": null}, "name": "set2_set_1", "inbound_nodes": [[["add_7", 0, 0, {}], ["input_6", 0, 0, {}]]], "shared_object_id": 83}, {"class_name": "Set2Set", "config": {"name": "set2_set_2", "trainable": false, "dtype": "float32", "T": 2, "n_hidden": 16, "activation": "linear", "activation_lstm": "tanh", "recurrent_activation": "hard_sigmoid", "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 84}, "recurrent_initializer": {"class_name": "Orthogonal", "config": {"gain": 1.0, "seed": null}, "shared_object_id": 85}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 86}, "use_bias": true, "unit_forget_bias": true, "kernel_regularizer": null, "recurrent_regularizer": null, "bias_regularizer": null, "kernel_constraint": null, "recurrent_constraint": null, "bias_constraint": null}, "name": "set2_set_2", "inbound_nodes": [[["add_8", 0, 0, {}], ["input_7", 0, 0, {}]]], "shared_object_id": 87}, {"class_name": "Add", "config": {"name": "add_9", "trainable": false, "dtype": "float32"}, "name": "add_9", "inbound_nodes": [[["add_6", 0, 0, {}], ["meg_net_layer_3", 0, 2, {}]]], "shared_object_id": 88}, {"class_name": "Concatenate", "config": {"name": "concatenate_1", "trainable": false, "dtype": "float32", "axis": -1}, "name": "concatenate_1", "inbound_nodes": [[["set2_set_1", 0, 0, {}], ["set2_set_2", 0, 0, {}], ["add_9", 0, 0, {}]]], "shared_object_id": 89}, {"class_name": "Dense", "config": {"name": "dense_19", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 90}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 91}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_19", "inbound_nodes": [[["concatenate_1", 0, 0, {}]]], "shared_object_id": 92}, {"class_name": "Dense", "config": {"name": "dense_20", "trainable": false, "dtype": "float32", "units": 16, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 93}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 94}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_20", "inbound_nodes": [[["dense_19", 0, 0, {}]]], "shared_object_id": 95}, {"class_name": "Dense", "config": {"name": "dense_21", "trainable": true, "dtype": "float32", "units": 1, "activation": "linear", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 96}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 97}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_21", "inbound_nodes": [[["dense_20", 0, 0, {}]]], "shared_object_id": 98}], "input_layers": [["input_1", 0, 0], ["input_2", 0, 0], ["input_3", 0, 0], ["input_4", 0, 0], ["input_5", 0, 0], ["input_6", 0, 0], ["input_7", 0, 0]], "output_layers": [["dense_21", 0, 0]]}}, "training_config": {"loss": "mean_squared_error", "metrics": [[]], "weighted_metrics": null, "loss_weights": null, "optimizer_config": {"class_name": "Adam", "config": {"name": "Adam", "learning_rate": 0.0010000000474974513, "decay": 0.0, "beta_1": 0.8999999761581421, "beta_2": 0.9990000128746033, "epsilon": 1e-07, "amsgrad": false}}}}2
+ root.layer-0"_tf_keras_input_layer*{"class_name": "InputLayer", "name": "input_1", "dtype": "int32", "sparse": false, "ragged": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_1"}}2
+root.layer_with_weights-0"_tf_keras_layer*{"name": "embedding_1", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "stateful": false, "must_restore_from_config": false, "class_name": "Embedding", "config": {"name": "embedding_1", "trainable": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "float32", "input_dim": 95, "output_dim": 16, "embeddings_initializer": {"class_name": "RandomUniform", "config": {"minval": -0.05, "maxval": 0.05, "seed": null}, "shared_object_id": 1}, "embeddings_regularizer": null, "activity_regularizer": null, "embeddings_constraint": null, "mask_zero": false, "input_length": null}, "inbound_nodes": [[["input_1", 0, 0, {}]]], "shared_object_id": 2, "build_input_shape": {"class_name": "TensorShape", "items": [null, null]}}2
+ root.layer-2"_tf_keras_input_layer*{"class_name": "InputLayer", "name": "input_2", "dtype": "float32", "sparse": false, "ragged": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null, 100]}, "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null, 100]}, "dtype": "float32", "sparse": false, "ragged": false, "name": "input_2"}}2
+ root.layer-3"_tf_keras_input_layer*{"class_name": "InputLayer", "name": "input_3", "dtype": "float32", "sparse": false, "ragged": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null, 2]}, "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null, 2]}, "dtype": "float32", "sparse": false, "ragged": false, "name": "input_3"}}2
+root.layer_with_weights-1"_tf_keras_layer*{"name": "dense_1", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_1", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 5}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 6}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["embedding_1", 0, 0, {}]]], "shared_object_id": 7, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 16}}, "shared_object_id": 107}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 16]}}2
+root.layer_with_weights-2"_tf_keras_layer*{"name": "dense_3", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_3", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 8}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 9}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["input_2", 0, 0, {}]]], "shared_object_id": 10, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 100}}, "shared_object_id": 108}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 100]}}2
+root.layer_with_weights-3"_tf_keras_layer*{"name": "dense_5", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_5", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 11}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 12}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["input_3", 0, 0, {}]]], "shared_object_id": 13, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 2}}, "shared_object_id": 109}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 2]}}2
+root.layer_with_weights-4"_tf_keras_layer*{"name": "dense_2", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_2", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 14}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 15}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_1", 0, 0, {}]]], "shared_object_id": 16, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 64}}, "shared_object_id": 110}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 64]}}2
+ root.layer_with_weights-5"_tf_keras_layer*{"name": "dense_4", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_4", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 17}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 18}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_3", 0, 0, {}]]], "shared_object_id": 19, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 64}}, "shared_object_id": 111}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 64]}}2
+
+root.layer_with_weights-6"_tf_keras_layer*{"name": "dense_6", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_6", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 20}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 21}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_5", 0, 0, {}]]], "shared_object_id": 22, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 64}}, "shared_object_id": 112}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 64]}}2
+
root.layer-10"_tf_keras_input_layer*{"class_name": "InputLayer", "name": "input_4", "dtype": "int32", "sparse": false, "ragged": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_4"}}2
+
root.layer-11"_tf_keras_input_layer*{"class_name": "InputLayer", "name": "input_5", "dtype": "int32", "sparse": false, "ragged": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_5"}}2
+
root.layer-12"_tf_keras_input_layer*{"class_name": "InputLayer", "name": "input_6", "dtype": "int32", "sparse": false, "ragged": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_6"}}2
+
root.layer-13"_tf_keras_input_layer*{"class_name": "InputLayer", "name": "input_7", "dtype": "int32", "sparse": false, "ragged": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_7"}}2
+ root.layer_with_weights-7"_tf_keras_layer* {"name": "meg_net_layer_1", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "MEGNetLayer", "config": {"name": "meg_net_layer_1", "trainable": false, "dtype": "float32", "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 27}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 28}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "units_e": [64, 64, 32], "units_v": [64, 64, 32], "units_u": [64, 64, 32], "pool_method": "mean"}, "inbound_nodes": [[["dense_2", 0, 0, {}], ["dense_4", 0, 0, {}], ["dense_6", 0, 0, {}], ["input_4", 0, 0, {}], ["input_5", 0, 0, {}], ["input_6", 0, 0, {}], ["input_7", 0, 0, {}]]], "shared_object_id": 29, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}]}2
+
root.layer-15"_tf_keras_layer*{"name": "add_1", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Add", "config": {"name": "add_1", "trainable": false, "dtype": "float32"}, "inbound_nodes": [[["dense_2", 0, 0, {}], ["meg_net_layer_1", 0, 0, {}]]], "shared_object_id": 30, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}]}2
+
root.layer-16"_tf_keras_layer*{"name": "add_2", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Add", "config": {"name": "add_2", "trainable": false, "dtype": "float32"}, "inbound_nodes": [[["dense_4", 0, 0, {}], ["meg_net_layer_1", 0, 1, {}]]], "shared_object_id": 31, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}]}2
+
root.layer-17"_tf_keras_layer*{"name": "add_3", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Add", "config": {"name": "add_3", "trainable": false, "dtype": "float32"}, "inbound_nodes": [[["dense_6", 0, 0, {}], ["meg_net_layer_1", 0, 2, {}]]], "shared_object_id": 32, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [1, null, 32]}]}2
+root.layer_with_weights-8"_tf_keras_layer*{"name": "dense_7", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_7", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 33}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 34}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["add_1", 0, 0, {}]]], "shared_object_id": 35, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 32}}, "shared_object_id": 113}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 32]}}2
+root.layer_with_weights-9"_tf_keras_layer*{"name": "dense_9", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_9", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 36}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 37}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["add_2", 0, 0, {}]]], "shared_object_id": 38, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 32}}, "shared_object_id": 114}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 32]}}2
+root.layer_with_weights-10"_tf_keras_layer*{"name": "dense_11", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_11", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 39}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 40}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["add_3", 0, 0, {}]]], "shared_object_id": 41, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 32}}, "shared_object_id": 115}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 32]}}2
+root.layer_with_weights-11"_tf_keras_layer*{"name": "dense_8", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_8", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 42}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 43}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_7", 0, 0, {}]]], "shared_object_id": 44, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 64}}, "shared_object_id": 116}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 64]}}2
+root.layer_with_weights-12"_tf_keras_layer*{"name": "dense_10", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_10", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 45}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 46}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_9", 0, 0, {}]]], "shared_object_id": 47, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 64}}, "shared_object_id": 117}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 64]}}2
+root.layer_with_weights-13"_tf_keras_layer*{"name": "dense_12", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_12", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 48}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 49}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_11", 0, 0, {}]]], "shared_object_id": 50, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 64}}, "shared_object_id": 118}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 64]}}2
+ root.layer_with_weights-14"_tf_keras_layer* {"name": "meg_net_layer_2", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "MEGNetLayer", "config": {"name": "meg_net_layer_2", "trainable": false, "dtype": "float32", "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 51}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 52}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "units_e": [64, 64, 32], "units_v": [64, 64, 32], "units_u": [64, 64, 32], "pool_method": "mean"}, "inbound_nodes": [[["dense_8", 0, 0, {}], ["dense_10", 0, 0, {}], ["dense_12", 0, 0, {}], ["input_4", 0, 0, {}], ["input_5", 0, 0, {}], ["input_6", 0, 0, {}], ["input_7", 0, 0, {}]]], "shared_object_id": 53, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}]}2
+
root.layer-25"_tf_keras_layer*{"name": "add_4", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Add", "config": {"name": "add_4", "trainable": false, "dtype": "float32"}, "inbound_nodes": [[["add_1", 0, 0, {}], ["meg_net_layer_2", 0, 0, {}]]], "shared_object_id": 54, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}]}2
+
root.layer-26"_tf_keras_layer*{"name": "add_5", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Add", "config": {"name": "add_5", "trainable": false, "dtype": "float32"}, "inbound_nodes": [[["add_2", 0, 0, {}], ["meg_net_layer_2", 0, 1, {}]]], "shared_object_id": 55, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}]}2
+
root.layer-27"_tf_keras_layer*{"name": "add_6", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Add", "config": {"name": "add_6", "trainable": false, "dtype": "float32"}, "inbound_nodes": [[["add_3", 0, 0, {}], ["meg_net_layer_2", 0, 2, {}]]], "shared_object_id": 56, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [1, null, 32]}]}2
+ root.layer_with_weights-15"_tf_keras_layer*{"name": "dense_13", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_13", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 57}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 58}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["add_4", 0, 0, {}]]], "shared_object_id": 59, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 32}}, "shared_object_id": 119}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 32]}}2
+ root.layer_with_weights-16"_tf_keras_layer*{"name": "dense_15", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_15", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 60}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 61}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["add_5", 0, 0, {}]]], "shared_object_id": 62, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 32}}, "shared_object_id": 120}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 32]}}2
+ root.layer_with_weights-17"_tf_keras_layer*{"name": "dense_17", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_17", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 63}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 64}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["add_6", 0, 0, {}]]], "shared_object_id": 65, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 32}}, "shared_object_id": 121}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 32]}}2
+ root.layer_with_weights-18"_tf_keras_layer*{"name": "dense_14", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_14", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 66}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 67}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_13", 0, 0, {}]]], "shared_object_id": 68, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 64}}, "shared_object_id": 122}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 64]}}2
+!root.layer_with_weights-19"_tf_keras_layer*{"name": "dense_16", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_16", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 69}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 70}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_15", 0, 0, {}]]], "shared_object_id": 71, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 64}}, "shared_object_id": 123}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 64]}}2
+"root.layer_with_weights-20"_tf_keras_layer*{"name": "dense_18", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_18", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 72}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 73}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_17", 0, 0, {}]]], "shared_object_id": 74, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 64}}, "shared_object_id": 124}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 64]}}2
+ #root.layer_with_weights-21"_tf_keras_layer* {"name": "meg_net_layer_3", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "MEGNetLayer", "config": {"name": "meg_net_layer_3", "trainable": false, "dtype": "float32", "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 75}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 76}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "units_e": [64, 64, 32], "units_v": [64, 64, 32], "units_u": [64, 64, 32], "pool_method": "mean"}, "inbound_nodes": [[["dense_14", 0, 0, {}], ["dense_16", 0, 0, {}], ["dense_18", 0, 0, {}], ["input_4", 0, 0, {}], ["input_5", 0, 0, {}], ["input_6", 0, 0, {}], ["input_7", 0, 0, {}]]], "shared_object_id": 77, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}]}2
+$
root.layer-35"_tf_keras_layer*{"name": "add_7", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Add", "config": {"name": "add_7", "trainable": false, "dtype": "float32"}, "inbound_nodes": [[["add_4", 0, 0, {}], ["meg_net_layer_3", 0, 0, {}]]], "shared_object_id": 78, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}]}2
+%
root.layer-36"_tf_keras_layer*{"name": "add_8", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Add", "config": {"name": "add_8", "trainable": false, "dtype": "float32"}, "inbound_nodes": [[["add_5", 0, 0, {}], ["meg_net_layer_3", 0, 1, {}]]], "shared_object_id": 79, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}]}2
+ &root.layer_with_weights-22"_tf_keras_layer* {"name": "set2_set_1", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Set2Set", "config": {"name": "set2_set_1", "trainable": false, "dtype": "float32", "T": 2, "n_hidden": 16, "activation": "linear", "activation_lstm": "tanh", "recurrent_activation": "hard_sigmoid", "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 80}, "recurrent_initializer": {"class_name": "Orthogonal", "config": {"gain": 1.0, "seed": null}, "shared_object_id": 81}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 82}, "use_bias": true, "unit_forget_bias": true, "kernel_regularizer": null, "recurrent_regularizer": null, "bias_regularizer": null, "kernel_constraint": null, "recurrent_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["add_7", 0, 0, {}], ["input_6", 0, 0, {}]]], "shared_object_id": 83, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null]}]}2
+ 'root.layer_with_weights-23"_tf_keras_layer* {"name": "set2_set_2", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Set2Set", "config": {"name": "set2_set_2", "trainable": false, "dtype": "float32", "T": 2, "n_hidden": 16, "activation": "linear", "activation_lstm": "tanh", "recurrent_activation": "hard_sigmoid", "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 84}, "recurrent_initializer": {"class_name": "Orthogonal", "config": {"gain": 1.0, "seed": null}, "shared_object_id": 85}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 86}, "use_bias": true, "unit_forget_bias": true, "kernel_regularizer": null, "recurrent_regularizer": null, "bias_regularizer": null, "kernel_constraint": null, "recurrent_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["add_8", 0, 0, {}], ["input_7", 0, 0, {}]]], "shared_object_id": 87, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null]}]}2
+(
root.layer-39"_tf_keras_layer*{"name": "add_9", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Add", "config": {"name": "add_9", "trainable": false, "dtype": "float32"}, "inbound_nodes": [[["add_6", 0, 0, {}], ["meg_net_layer_3", 0, 2, {}]]], "shared_object_id": 88, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [1, null, 32]}]}2
+)
root.layer-40"_tf_keras_layer*{"name": "concatenate_1", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Concatenate", "config": {"name": "concatenate_1", "trainable": false, "dtype": "float32", "axis": -1}, "inbound_nodes": [[["set2_set_1", 0, 0, {}], ["set2_set_2", 0, 0, {}], ["add_9", 0, 0, {}]]], "shared_object_id": 89, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}]}2
+*root.layer_with_weights-24"_tf_keras_layer*{"name": "dense_19", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_19", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 90}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 91}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["concatenate_1", 0, 0, {}]]], "shared_object_id": 92, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 96}}, "shared_object_id": 125}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 96]}}2
++root.layer_with_weights-25"_tf_keras_layer*{"name": "dense_20", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_20", "trainable": false, "dtype": "float32", "units": 16, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 93}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 94}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_19", 0, 0, {}]]], "shared_object_id": 95, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 32}}, "shared_object_id": 126}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 32]}}2
+,root.layer_with_weights-26"_tf_keras_layer*{"name": "dense_21", "trainable": true, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_21", "trainable": true, "dtype": "float32", "units": 1, "activation": "linear", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 96}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 97}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_20", 0, 0, {}]]], "shared_object_id": 98, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 16}}, "shared_object_id": 127}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 16]}}2
+root.keras_api.metrics.0"_tf_keras_metric*{"class_name": "Mean", "name": "loss", "dtype": "float32", "config": {"name": "loss", "dtype": "float32"}, "shared_object_id": 128}2
\ No newline at end of file
diff --git a/tests/test_model/prob_e_form_unnorm/megnet/saved_model.pb b/tests/test_model/prob_e_form_unnorm/megnet/saved_model.pb
new file mode 100644
index 0000000..1126752
Binary files /dev/null and b/tests/test_model/prob_e_form_unnorm/megnet/saved_model.pb differ
diff --git a/tests/test_model/prob_e_form_unnorm/megnet/variables/variables.data-00000-of-00001 b/tests/test_model/prob_e_form_unnorm/megnet/variables/variables.data-00000-of-00001
new file mode 100644
index 0000000..b28b286
Binary files /dev/null and b/tests/test_model/prob_e_form_unnorm/megnet/variables/variables.data-00000-of-00001 differ
diff --git a/tests/test_model/prob_e_form_unnorm/megnet/variables/variables.index b/tests/test_model/prob_e_form_unnorm/megnet/variables/variables.index
new file mode 100644
index 0000000..78fdd32
Binary files /dev/null and b/tests/test_model/prob_e_form_unnorm/megnet/variables/variables.index differ
diff --git a/tests/test_model/prob_e_form_unnorm/nn/keras_metadata.pb b/tests/test_model/prob_e_form_unnorm/nn/keras_metadata.pb
new file mode 100644
index 0000000..8b9fcf7
--- /dev/null
+++ b/tests/test_model/prob_e_form_unnorm/nn/keras_metadata.pb
@@ -0,0 +1,48 @@
+
+root"_tf_keras_network*{"name": "model_1", "trainable": true, "expects_training_arg": true, "dtype": "float32", "batch_input_shape": null, "must_restore_from_config": false, "class_name": "Functional", "config": {"name": "model_1", "layers": [{"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_1"}, "name": "input_1", "inbound_nodes": []}, {"class_name": "Embedding", "config": {"name": "embedding_1", "trainable": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "float32", "input_dim": 95, "output_dim": 16, "embeddings_initializer": {"class_name": "RandomUniform", "config": {"minval": -0.05, "maxval": 0.05, "seed": null}}, "embeddings_regularizer": null, "activity_regularizer": null, "embeddings_constraint": null, "mask_zero": false, "input_length": null}, "name": "embedding_1", "inbound_nodes": [[["input_1", 0, 0, {}]]]}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null, 100]}, "dtype": "float32", "sparse": false, "ragged": false, "name": "input_2"}, "name": "input_2", "inbound_nodes": []}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null, 2]}, "dtype": "float32", "sparse": false, "ragged": false, "name": "input_3"}, "name": "input_3", "inbound_nodes": []}, {"class_name": "Dense", "config": {"name": "dense_1", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_1", "inbound_nodes": [[["embedding_1", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_3", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_3", "inbound_nodes": [[["input_2", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_5", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_5", "inbound_nodes": [[["input_3", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_2", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_2", "inbound_nodes": [[["dense_1", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_4", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_4", "inbound_nodes": [[["dense_3", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_6", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_6", "inbound_nodes": [[["dense_5", 0, 0, {}]]]}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_4"}, "name": "input_4", "inbound_nodes": []}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_5"}, "name": "input_5", "inbound_nodes": []}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_6"}, "name": "input_6", "inbound_nodes": []}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_7"}, "name": "input_7", "inbound_nodes": []}, {"class_name": "MEGNetLayer", "config": {"name": "meg_net_layer_1", "trainable": false, "dtype": "float32", "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "units_e": [64, 64, 32], "units_v": [64, 64, 32], "units_u": [64, 64, 32], "pool_method": "mean"}, "name": "meg_net_layer_1", "inbound_nodes": [[["dense_2", 0, 0, {}], ["dense_4", 0, 0, {}], ["dense_6", 0, 0, {}], ["input_4", 0, 0, {}], ["input_5", 0, 0, {}], ["input_6", 0, 0, {}], ["input_7", 0, 0, {}]]]}, {"class_name": "Add", "config": {"name": "add_1", "trainable": false, "dtype": "float32"}, "name": "add_1", "inbound_nodes": [[["dense_2", 0, 0, {}], ["meg_net_layer_1", 0, 0, {}]]]}, {"class_name": "Add", "config": {"name": "add_2", "trainable": false, "dtype": "float32"}, "name": "add_2", "inbound_nodes": [[["dense_4", 0, 0, {}], ["meg_net_layer_1", 0, 1, {}]]]}, {"class_name": "Add", "config": {"name": "add_3", "trainable": false, "dtype": "float32"}, "name": "add_3", "inbound_nodes": [[["dense_6", 0, 0, {}], ["meg_net_layer_1", 0, 2, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_7", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_7", "inbound_nodes": [[["add_1", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_9", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_9", "inbound_nodes": [[["add_2", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_11", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_11", "inbound_nodes": [[["add_3", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_8", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_8", "inbound_nodes": [[["dense_7", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_10", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_10", "inbound_nodes": [[["dense_9", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_12", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_12", "inbound_nodes": [[["dense_11", 0, 0, {}]]]}, {"class_name": "MEGNetLayer", "config": {"name": "meg_net_layer_2", "trainable": false, "dtype": "float32", "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "units_e": [64, 64, 32], "units_v": [64, 64, 32], "units_u": [64, 64, 32], "pool_method": "mean"}, "name": "meg_net_layer_2", "inbound_nodes": [[["dense_8", 0, 0, {}], ["dense_10", 0, 0, {}], ["dense_12", 0, 0, {}], ["input_4", 0, 0, {}], ["input_5", 0, 0, {}], ["input_6", 0, 0, {}], ["input_7", 0, 0, {}]]]}, {"class_name": "Add", "config": {"name": "add_4", "trainable": false, "dtype": "float32"}, "name": "add_4", "inbound_nodes": [[["add_1", 0, 0, {}], ["meg_net_layer_2", 0, 0, {}]]]}, {"class_name": "Add", "config": {"name": "add_5", "trainable": false, "dtype": "float32"}, "name": "add_5", "inbound_nodes": [[["add_2", 0, 0, {}], ["meg_net_layer_2", 0, 1, {}]]]}, {"class_name": "Add", "config": {"name": "add_6", "trainable": false, "dtype": "float32"}, "name": "add_6", "inbound_nodes": [[["add_3", 0, 0, {}], ["meg_net_layer_2", 0, 2, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_13", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_13", "inbound_nodes": [[["add_4", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_15", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_15", "inbound_nodes": [[["add_5", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_17", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_17", "inbound_nodes": [[["add_6", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_14", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_14", "inbound_nodes": [[["dense_13", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_16", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_16", "inbound_nodes": [[["dense_15", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_18", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_18", "inbound_nodes": [[["dense_17", 0, 0, {}]]]}, {"class_name": "MEGNetLayer", "config": {"name": "meg_net_layer_3", "trainable": false, "dtype": "float32", "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "units_e": [64, 64, 32], "units_v": [64, 64, 32], "units_u": [64, 64, 32], "pool_method": "mean"}, "name": "meg_net_layer_3", "inbound_nodes": [[["dense_14", 0, 0, {}], ["dense_16", 0, 0, {}], ["dense_18", 0, 0, {}], ["input_4", 0, 0, {}], ["input_5", 0, 0, {}], ["input_6", 0, 0, {}], ["input_7", 0, 0, {}]]]}, {"class_name": "Add", "config": {"name": "add_7", "trainable": false, "dtype": "float32"}, "name": "add_7", "inbound_nodes": [[["add_4", 0, 0, {}], ["meg_net_layer_3", 0, 0, {}]]]}, {"class_name": "Add", "config": {"name": "add_8", "trainable": false, "dtype": "float32"}, "name": "add_8", "inbound_nodes": [[["add_5", 0, 0, {}], ["meg_net_layer_3", 0, 1, {}]]]}, {"class_name": "Set2Set", "config": {"name": "set2_set_1", "trainable": false, "dtype": "float32", "T": 2, "n_hidden": 16, "activation": "linear", "activation_lstm": "tanh", "recurrent_activation": "hard_sigmoid", "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "recurrent_initializer": {"class_name": "Orthogonal", "config": {"gain": 1.0, "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "use_bias": true, "unit_forget_bias": true, "kernel_regularizer": null, "recurrent_regularizer": null, "bias_regularizer": null, "kernel_constraint": null, "recurrent_constraint": null, "bias_constraint": null}, "name": "set2_set_1", "inbound_nodes": [[["add_7", 0, 0, {}], ["input_6", 0, 0, {}]]]}, {"class_name": "Set2Set", "config": {"name": "set2_set_2", "trainable": false, "dtype": "float32", "T": 2, "n_hidden": 16, "activation": "linear", "activation_lstm": "tanh", "recurrent_activation": "hard_sigmoid", "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "recurrent_initializer": {"class_name": "Orthogonal", "config": {"gain": 1.0, "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "use_bias": true, "unit_forget_bias": true, "kernel_regularizer": null, "recurrent_regularizer": null, "bias_regularizer": null, "kernel_constraint": null, "recurrent_constraint": null, "bias_constraint": null}, "name": "set2_set_2", "inbound_nodes": [[["add_8", 0, 0, {}], ["input_7", 0, 0, {}]]]}, {"class_name": "Add", "config": {"name": "add_9", "trainable": false, "dtype": "float32"}, "name": "add_9", "inbound_nodes": [[["add_6", 0, 0, {}], ["meg_net_layer_3", 0, 2, {}]]]}, {"class_name": "Concatenate", "config": {"name": "concatenate_1", "trainable": false, "dtype": "float32", "axis": -1}, "name": "concatenate_1", "inbound_nodes": [[["set2_set_1", 0, 0, {}], ["set2_set_2", 0, 0, {}], ["add_9", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_19", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_19", "inbound_nodes": [[["concatenate_1", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_20", "trainable": false, "dtype": "float32", "units": 16, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_20", "inbound_nodes": [[["dense_19", 0, 0, {}]]]}, {"class_name": "Dense", "config": {"name": "dense_21", "trainable": true, "dtype": "float32", "units": 1, "activation": "linear", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_21", "inbound_nodes": [[["dense_20", 0, 0, {}]]]}], "input_layers": [["input_1", 0, 0], ["input_2", 0, 0], ["input_3", 0, 0], ["input_4", 0, 0], ["input_5", 0, 0], ["input_6", 0, 0], ["input_7", 0, 0]], "output_layers": [["dense_21", 0, 0]]}, "shared_object_id": 99, "input_spec": [{"class_name": "InputSpec", "config": {"dtype": null, "shape": {"class_name": "__tuple__", "items": [null, null]}, "ndim": 2, "max_ndim": null, "min_ndim": null, "axes": {}}}, {"class_name": "InputSpec", "config": {"dtype": null, "shape": {"class_name": "__tuple__", "items": [null, null, 100]}, "ndim": 3, "max_ndim": null, "min_ndim": null, "axes": {}}}, {"class_name": "InputSpec", "config": {"dtype": null, "shape": {"class_name": "__tuple__", "items": [null, null, 2]}, "ndim": 3, "max_ndim": null, "min_ndim": null, "axes": {}}}, {"class_name": "InputSpec", "config": {"dtype": null, "shape": {"class_name": "__tuple__", "items": [null, null]}, "ndim": 2, "max_ndim": null, "min_ndim": null, "axes": {}}}, {"class_name": "InputSpec", "config": {"dtype": null, "shape": {"class_name": "__tuple__", "items": [null, null]}, "ndim": 2, "max_ndim": null, "min_ndim": null, "axes": {}}}, {"class_name": "InputSpec", "config": {"dtype": null, "shape": {"class_name": "__tuple__", "items": [null, null]}, "ndim": 2, "max_ndim": null, "min_ndim": null, "axes": {}}}, {"class_name": "InputSpec", "config": {"dtype": null, "shape": {"class_name": "__tuple__", "items": [null, null]}, "ndim": 2, "max_ndim": null, "min_ndim": null, "axes": {}}}], "build_input_shape": [{"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null, 100]}, {"class_name": "TensorShape", "items": [null, null, 2]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}], "is_graph_network": true, "full_save_spec": {"class_name": "__tuple__", "items": [[[{"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_1"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null, 100]}, "float32", "input_2"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null, 2]}, "float32", "input_3"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_4"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_5"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_6"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_7"]}]], {}]}, "save_spec": [{"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_1"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null, 100]}, "float32", "input_2"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null, 2]}, "float32", "input_3"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_4"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_5"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_6"]}, {"class_name": "TypeSpec", "type_spec": "tf.TensorSpec", "serialized": [{"class_name": "TensorShape", "items": [null, null]}, "int32", "input_7"]}], "keras_version": "2.6.0", "backend": "tensorflow", "model_config": {"class_name": "Functional", "config": {"name": "model_1", "layers": [{"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_1"}, "name": "input_1", "inbound_nodes": [], "shared_object_id": 0}, {"class_name": "Embedding", "config": {"name": "embedding_1", "trainable": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "float32", "input_dim": 95, "output_dim": 16, "embeddings_initializer": {"class_name": "RandomUniform", "config": {"minval": -0.05, "maxval": 0.05, "seed": null}, "shared_object_id": 1}, "embeddings_regularizer": null, "activity_regularizer": null, "embeddings_constraint": null, "mask_zero": false, "input_length": null}, "name": "embedding_1", "inbound_nodes": [[["input_1", 0, 0, {}]]], "shared_object_id": 2}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null, 100]}, "dtype": "float32", "sparse": false, "ragged": false, "name": "input_2"}, "name": "input_2", "inbound_nodes": [], "shared_object_id": 3}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null, 2]}, "dtype": "float32", "sparse": false, "ragged": false, "name": "input_3"}, "name": "input_3", "inbound_nodes": [], "shared_object_id": 4}, {"class_name": "Dense", "config": {"name": "dense_1", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 5}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 6}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_1", "inbound_nodes": [[["embedding_1", 0, 0, {}]]], "shared_object_id": 7}, {"class_name": "Dense", "config": {"name": "dense_3", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 8}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 9}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_3", "inbound_nodes": [[["input_2", 0, 0, {}]]], "shared_object_id": 10}, {"class_name": "Dense", "config": {"name": "dense_5", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 11}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 12}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_5", "inbound_nodes": [[["input_3", 0, 0, {}]]], "shared_object_id": 13}, {"class_name": "Dense", "config": {"name": "dense_2", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 14}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 15}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_2", "inbound_nodes": [[["dense_1", 0, 0, {}]]], "shared_object_id": 16}, {"class_name": "Dense", "config": {"name": "dense_4", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 17}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 18}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_4", "inbound_nodes": [[["dense_3", 0, 0, {}]]], "shared_object_id": 19}, {"class_name": "Dense", "config": {"name": "dense_6", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 20}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 21}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_6", "inbound_nodes": [[["dense_5", 0, 0, {}]]], "shared_object_id": 22}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_4"}, "name": "input_4", "inbound_nodes": [], "shared_object_id": 23}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_5"}, "name": "input_5", "inbound_nodes": [], "shared_object_id": 24}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_6"}, "name": "input_6", "inbound_nodes": [], "shared_object_id": 25}, {"class_name": "InputLayer", "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_7"}, "name": "input_7", "inbound_nodes": [], "shared_object_id": 26}, {"class_name": "MEGNetLayer", "config": {"name": "meg_net_layer_1", "trainable": false, "dtype": "float32", "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 27}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 28}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "units_e": [64, 64, 32], "units_v": [64, 64, 32], "units_u": [64, 64, 32], "pool_method": "mean"}, "name": "meg_net_layer_1", "inbound_nodes": [[["dense_2", 0, 0, {}], ["dense_4", 0, 0, {}], ["dense_6", 0, 0, {}], ["input_4", 0, 0, {}], ["input_5", 0, 0, {}], ["input_6", 0, 0, {}], ["input_7", 0, 0, {}]]], "shared_object_id": 29}, {"class_name": "Add", "config": {"name": "add_1", "trainable": false, "dtype": "float32"}, "name": "add_1", "inbound_nodes": [[["dense_2", 0, 0, {}], ["meg_net_layer_1", 0, 0, {}]]], "shared_object_id": 30}, {"class_name": "Add", "config": {"name": "add_2", "trainable": false, "dtype": "float32"}, "name": "add_2", "inbound_nodes": [[["dense_4", 0, 0, {}], ["meg_net_layer_1", 0, 1, {}]]], "shared_object_id": 31}, {"class_name": "Add", "config": {"name": "add_3", "trainable": false, "dtype": "float32"}, "name": "add_3", "inbound_nodes": [[["dense_6", 0, 0, {}], ["meg_net_layer_1", 0, 2, {}]]], "shared_object_id": 32}, {"class_name": "Dense", "config": {"name": "dense_7", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 33}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 34}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_7", "inbound_nodes": [[["add_1", 0, 0, {}]]], "shared_object_id": 35}, {"class_name": "Dense", "config": {"name": "dense_9", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 36}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 37}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_9", "inbound_nodes": [[["add_2", 0, 0, {}]]], "shared_object_id": 38}, {"class_name": "Dense", "config": {"name": "dense_11", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 39}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 40}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_11", "inbound_nodes": [[["add_3", 0, 0, {}]]], "shared_object_id": 41}, {"class_name": "Dense", "config": {"name": "dense_8", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 42}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 43}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_8", "inbound_nodes": [[["dense_7", 0, 0, {}]]], "shared_object_id": 44}, {"class_name": "Dense", "config": {"name": "dense_10", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 45}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 46}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_10", "inbound_nodes": [[["dense_9", 0, 0, {}]]], "shared_object_id": 47}, {"class_name": "Dense", "config": {"name": "dense_12", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 48}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 49}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_12", "inbound_nodes": [[["dense_11", 0, 0, {}]]], "shared_object_id": 50}, {"class_name": "MEGNetLayer", "config": {"name": "meg_net_layer_2", "trainable": false, "dtype": "float32", "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 51}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 52}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "units_e": [64, 64, 32], "units_v": [64, 64, 32], "units_u": [64, 64, 32], "pool_method": "mean"}, "name": "meg_net_layer_2", "inbound_nodes": [[["dense_8", 0, 0, {}], ["dense_10", 0, 0, {}], ["dense_12", 0, 0, {}], ["input_4", 0, 0, {}], ["input_5", 0, 0, {}], ["input_6", 0, 0, {}], ["input_7", 0, 0, {}]]], "shared_object_id": 53}, {"class_name": "Add", "config": {"name": "add_4", "trainable": false, "dtype": "float32"}, "name": "add_4", "inbound_nodes": [[["add_1", 0, 0, {}], ["meg_net_layer_2", 0, 0, {}]]], "shared_object_id": 54}, {"class_name": "Add", "config": {"name": "add_5", "trainable": false, "dtype": "float32"}, "name": "add_5", "inbound_nodes": [[["add_2", 0, 0, {}], ["meg_net_layer_2", 0, 1, {}]]], "shared_object_id": 55}, {"class_name": "Add", "config": {"name": "add_6", "trainable": false, "dtype": "float32"}, "name": "add_6", "inbound_nodes": [[["add_3", 0, 0, {}], ["meg_net_layer_2", 0, 2, {}]]], "shared_object_id": 56}, {"class_name": "Dense", "config": {"name": "dense_13", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 57}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 58}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_13", "inbound_nodes": [[["add_4", 0, 0, {}]]], "shared_object_id": 59}, {"class_name": "Dense", "config": {"name": "dense_15", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 60}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 61}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_15", "inbound_nodes": [[["add_5", 0, 0, {}]]], "shared_object_id": 62}, {"class_name": "Dense", "config": {"name": "dense_17", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 63}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 64}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_17", "inbound_nodes": [[["add_6", 0, 0, {}]]], "shared_object_id": 65}, {"class_name": "Dense", "config": {"name": "dense_14", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 66}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 67}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_14", "inbound_nodes": [[["dense_13", 0, 0, {}]]], "shared_object_id": 68}, {"class_name": "Dense", "config": {"name": "dense_16", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 69}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 70}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_16", "inbound_nodes": [[["dense_15", 0, 0, {}]]], "shared_object_id": 71}, {"class_name": "Dense", "config": {"name": "dense_18", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 72}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 73}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_18", "inbound_nodes": [[["dense_17", 0, 0, {}]]], "shared_object_id": 74}, {"class_name": "MEGNetLayer", "config": {"name": "meg_net_layer_3", "trainable": false, "dtype": "float32", "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 75}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 76}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "units_e": [64, 64, 32], "units_v": [64, 64, 32], "units_u": [64, 64, 32], "pool_method": "mean"}, "name": "meg_net_layer_3", "inbound_nodes": [[["dense_14", 0, 0, {}], ["dense_16", 0, 0, {}], ["dense_18", 0, 0, {}], ["input_4", 0, 0, {}], ["input_5", 0, 0, {}], ["input_6", 0, 0, {}], ["input_7", 0, 0, {}]]], "shared_object_id": 77}, {"class_name": "Add", "config": {"name": "add_7", "trainable": false, "dtype": "float32"}, "name": "add_7", "inbound_nodes": [[["add_4", 0, 0, {}], ["meg_net_layer_3", 0, 0, {}]]], "shared_object_id": 78}, {"class_name": "Add", "config": {"name": "add_8", "trainable": false, "dtype": "float32"}, "name": "add_8", "inbound_nodes": [[["add_5", 0, 0, {}], ["meg_net_layer_3", 0, 1, {}]]], "shared_object_id": 79}, {"class_name": "Set2Set", "config": {"name": "set2_set_1", "trainable": false, "dtype": "float32", "T": 2, "n_hidden": 16, "activation": "linear", "activation_lstm": "tanh", "recurrent_activation": "hard_sigmoid", "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 80}, "recurrent_initializer": {"class_name": "Orthogonal", "config": {"gain": 1.0, "seed": null}, "shared_object_id": 81}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 82}, "use_bias": true, "unit_forget_bias": true, "kernel_regularizer": null, "recurrent_regularizer": null, "bias_regularizer": null, "kernel_constraint": null, "recurrent_constraint": null, "bias_constraint": null}, "name": "set2_set_1", "inbound_nodes": [[["add_7", 0, 0, {}], ["input_6", 0, 0, {}]]], "shared_object_id": 83}, {"class_name": "Set2Set", "config": {"name": "set2_set_2", "trainable": false, "dtype": "float32", "T": 2, "n_hidden": 16, "activation": "linear", "activation_lstm": "tanh", "recurrent_activation": "hard_sigmoid", "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 84}, "recurrent_initializer": {"class_name": "Orthogonal", "config": {"gain": 1.0, "seed": null}, "shared_object_id": 85}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 86}, "use_bias": true, "unit_forget_bias": true, "kernel_regularizer": null, "recurrent_regularizer": null, "bias_regularizer": null, "kernel_constraint": null, "recurrent_constraint": null, "bias_constraint": null}, "name": "set2_set_2", "inbound_nodes": [[["add_8", 0, 0, {}], ["input_7", 0, 0, {}]]], "shared_object_id": 87}, {"class_name": "Add", "config": {"name": "add_9", "trainable": false, "dtype": "float32"}, "name": "add_9", "inbound_nodes": [[["add_6", 0, 0, {}], ["meg_net_layer_3", 0, 2, {}]]], "shared_object_id": 88}, {"class_name": "Concatenate", "config": {"name": "concatenate_1", "trainable": false, "dtype": "float32", "axis": -1}, "name": "concatenate_1", "inbound_nodes": [[["set2_set_1", 0, 0, {}], ["set2_set_2", 0, 0, {}], ["add_9", 0, 0, {}]]], "shared_object_id": 89}, {"class_name": "Dense", "config": {"name": "dense_19", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 90}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 91}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_19", "inbound_nodes": [[["concatenate_1", 0, 0, {}]]], "shared_object_id": 92}, {"class_name": "Dense", "config": {"name": "dense_20", "trainable": false, "dtype": "float32", "units": 16, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 93}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 94}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_20", "inbound_nodes": [[["dense_19", 0, 0, {}]]], "shared_object_id": 95}, {"class_name": "Dense", "config": {"name": "dense_21", "trainable": true, "dtype": "float32", "units": 1, "activation": "linear", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 96}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 97}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "name": "dense_21", "inbound_nodes": [[["dense_20", 0, 0, {}]]], "shared_object_id": 98}], "input_layers": [["input_1", 0, 0], ["input_2", 0, 0], ["input_3", 0, 0], ["input_4", 0, 0], ["input_5", 0, 0], ["input_6", 0, 0], ["input_7", 0, 0]], "output_layers": [["dense_21", 0, 0]]}}}2
+ root.layer-0"_tf_keras_input_layer*{"class_name": "InputLayer", "name": "input_1", "dtype": "int32", "sparse": false, "ragged": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_1"}}2
+root.layer_with_weights-0"_tf_keras_layer*{"name": "embedding_1", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "stateful": false, "must_restore_from_config": false, "class_name": "Embedding", "config": {"name": "embedding_1", "trainable": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "float32", "input_dim": 95, "output_dim": 16, "embeddings_initializer": {"class_name": "RandomUniform", "config": {"minval": -0.05, "maxval": 0.05, "seed": null}, "shared_object_id": 1}, "embeddings_regularizer": null, "activity_regularizer": null, "embeddings_constraint": null, "mask_zero": false, "input_length": null}, "inbound_nodes": [[["input_1", 0, 0, {}]]], "shared_object_id": 2, "build_input_shape": {"class_name": "TensorShape", "items": [null, null]}}2
+ root.layer-2"_tf_keras_input_layer*{"class_name": "InputLayer", "name": "input_2", "dtype": "float32", "sparse": false, "ragged": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null, 100]}, "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null, 100]}, "dtype": "float32", "sparse": false, "ragged": false, "name": "input_2"}}2
+ root.layer-3"_tf_keras_input_layer*{"class_name": "InputLayer", "name": "input_3", "dtype": "float32", "sparse": false, "ragged": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null, 2]}, "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null, 2]}, "dtype": "float32", "sparse": false, "ragged": false, "name": "input_3"}}2
+root.layer_with_weights-1"_tf_keras_layer*{"name": "dense_1", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_1", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 5}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 6}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["embedding_1", 0, 0, {}]]], "shared_object_id": 7, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 16}}, "shared_object_id": 107}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 16]}}2
+root.layer_with_weights-2"_tf_keras_layer*{"name": "dense_3", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_3", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 8}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 9}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["input_2", 0, 0, {}]]], "shared_object_id": 10, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 100}}, "shared_object_id": 108}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 100]}}2
+root.layer_with_weights-3"_tf_keras_layer*{"name": "dense_5", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_5", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 11}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 12}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["input_3", 0, 0, {}]]], "shared_object_id": 13, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 2}}, "shared_object_id": 109}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 2]}}2
+root.layer_with_weights-4"_tf_keras_layer*{"name": "dense_2", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_2", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 14}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 15}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_1", 0, 0, {}]]], "shared_object_id": 16, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 64}}, "shared_object_id": 110}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 64]}}2
+ root.layer_with_weights-5"_tf_keras_layer*{"name": "dense_4", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_4", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 17}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 18}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_3", 0, 0, {}]]], "shared_object_id": 19, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 64}}, "shared_object_id": 111}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 64]}}2
+
+root.layer_with_weights-6"_tf_keras_layer*{"name": "dense_6", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_6", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 20}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 21}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_5", 0, 0, {}]]], "shared_object_id": 22, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 64}}, "shared_object_id": 112}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 64]}}2
+
root.layer-10"_tf_keras_input_layer*{"class_name": "InputLayer", "name": "input_4", "dtype": "int32", "sparse": false, "ragged": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_4"}}2
+
root.layer-11"_tf_keras_input_layer*{"class_name": "InputLayer", "name": "input_5", "dtype": "int32", "sparse": false, "ragged": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_5"}}2
+
root.layer-12"_tf_keras_input_layer*{"class_name": "InputLayer", "name": "input_6", "dtype": "int32", "sparse": false, "ragged": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_6"}}2
+
root.layer-13"_tf_keras_input_layer*{"class_name": "InputLayer", "name": "input_7", "dtype": "int32", "sparse": false, "ragged": false, "batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "config": {"batch_input_shape": {"class_name": "__tuple__", "items": [null, null]}, "dtype": "int32", "sparse": false, "ragged": false, "name": "input_7"}}2
+ root.layer_with_weights-7"_tf_keras_layer* {"name": "meg_net_layer_1", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "MEGNetLayer", "config": {"name": "meg_net_layer_1", "trainable": false, "dtype": "float32", "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 27}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 28}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "units_e": [64, 64, 32], "units_v": [64, 64, 32], "units_u": [64, 64, 32], "pool_method": "mean"}, "inbound_nodes": [[["dense_2", 0, 0, {}], ["dense_4", 0, 0, {}], ["dense_6", 0, 0, {}], ["input_4", 0, 0, {}], ["input_5", 0, 0, {}], ["input_6", 0, 0, {}], ["input_7", 0, 0, {}]]], "shared_object_id": 29, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}]}2
+
root.layer-15"_tf_keras_layer*{"name": "add_1", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Add", "config": {"name": "add_1", "trainable": false, "dtype": "float32"}, "inbound_nodes": [[["dense_2", 0, 0, {}], ["meg_net_layer_1", 0, 0, {}]]], "shared_object_id": 30, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}]}2
+
root.layer-16"_tf_keras_layer*{"name": "add_2", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Add", "config": {"name": "add_2", "trainable": false, "dtype": "float32"}, "inbound_nodes": [[["dense_4", 0, 0, {}], ["meg_net_layer_1", 0, 1, {}]]], "shared_object_id": 31, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}]}2
+
root.layer-17"_tf_keras_layer*{"name": "add_3", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Add", "config": {"name": "add_3", "trainable": false, "dtype": "float32"}, "inbound_nodes": [[["dense_6", 0, 0, {}], ["meg_net_layer_1", 0, 2, {}]]], "shared_object_id": 32, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [1, null, 32]}]}2
+root.layer_with_weights-8"_tf_keras_layer*{"name": "dense_7", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_7", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 33}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 34}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["add_1", 0, 0, {}]]], "shared_object_id": 35, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 32}}, "shared_object_id": 113}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 32]}}2
+root.layer_with_weights-9"_tf_keras_layer*{"name": "dense_9", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_9", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 36}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 37}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["add_2", 0, 0, {}]]], "shared_object_id": 38, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 32}}, "shared_object_id": 114}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 32]}}2
+root.layer_with_weights-10"_tf_keras_layer*{"name": "dense_11", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_11", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 39}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 40}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["add_3", 0, 0, {}]]], "shared_object_id": 41, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 32}}, "shared_object_id": 115}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 32]}}2
+root.layer_with_weights-11"_tf_keras_layer*{"name": "dense_8", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_8", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 42}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 43}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_7", 0, 0, {}]]], "shared_object_id": 44, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 64}}, "shared_object_id": 116}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 64]}}2
+root.layer_with_weights-12"_tf_keras_layer*{"name": "dense_10", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_10", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 45}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 46}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_9", 0, 0, {}]]], "shared_object_id": 47, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 64}}, "shared_object_id": 117}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 64]}}2
+root.layer_with_weights-13"_tf_keras_layer*{"name": "dense_12", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_12", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 48}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 49}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_11", 0, 0, {}]]], "shared_object_id": 50, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 64}}, "shared_object_id": 118}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 64]}}2
+ root.layer_with_weights-14"_tf_keras_layer* {"name": "meg_net_layer_2", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "MEGNetLayer", "config": {"name": "meg_net_layer_2", "trainable": false, "dtype": "float32", "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 51}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 52}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "units_e": [64, 64, 32], "units_v": [64, 64, 32], "units_u": [64, 64, 32], "pool_method": "mean"}, "inbound_nodes": [[["dense_8", 0, 0, {}], ["dense_10", 0, 0, {}], ["dense_12", 0, 0, {}], ["input_4", 0, 0, {}], ["input_5", 0, 0, {}], ["input_6", 0, 0, {}], ["input_7", 0, 0, {}]]], "shared_object_id": 53, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}]}2
+
root.layer-25"_tf_keras_layer*{"name": "add_4", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Add", "config": {"name": "add_4", "trainable": false, "dtype": "float32"}, "inbound_nodes": [[["add_1", 0, 0, {}], ["meg_net_layer_2", 0, 0, {}]]], "shared_object_id": 54, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}]}2
+
root.layer-26"_tf_keras_layer*{"name": "add_5", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Add", "config": {"name": "add_5", "trainable": false, "dtype": "float32"}, "inbound_nodes": [[["add_2", 0, 0, {}], ["meg_net_layer_2", 0, 1, {}]]], "shared_object_id": 55, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}]}2
+
root.layer-27"_tf_keras_layer*{"name": "add_6", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Add", "config": {"name": "add_6", "trainable": false, "dtype": "float32"}, "inbound_nodes": [[["add_3", 0, 0, {}], ["meg_net_layer_2", 0, 2, {}]]], "shared_object_id": 56, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [1, null, 32]}]}2
+ root.layer_with_weights-15"_tf_keras_layer*{"name": "dense_13", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_13", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 57}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 58}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["add_4", 0, 0, {}]]], "shared_object_id": 59, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 32}}, "shared_object_id": 119}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 32]}}2
+ root.layer_with_weights-16"_tf_keras_layer*{"name": "dense_15", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_15", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 60}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 61}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["add_5", 0, 0, {}]]], "shared_object_id": 62, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 32}}, "shared_object_id": 120}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 32]}}2
+ root.layer_with_weights-17"_tf_keras_layer*{"name": "dense_17", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_17", "trainable": false, "dtype": "float32", "units": 64, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 63}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 64}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["add_6", 0, 0, {}]]], "shared_object_id": 65, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 32}}, "shared_object_id": 121}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 32]}}2
+ root.layer_with_weights-18"_tf_keras_layer*{"name": "dense_14", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_14", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 66}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 67}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_13", 0, 0, {}]]], "shared_object_id": 68, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 64}}, "shared_object_id": 122}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 64]}}2
+!root.layer_with_weights-19"_tf_keras_layer*{"name": "dense_16", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_16", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 69}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 70}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_15", 0, 0, {}]]], "shared_object_id": 71, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 64}}, "shared_object_id": 123}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 64]}}2
+"root.layer_with_weights-20"_tf_keras_layer*{"name": "dense_18", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_18", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 72}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 73}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_17", 0, 0, {}]]], "shared_object_id": 74, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 64}}, "shared_object_id": 124}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 64]}}2
+ #root.layer_with_weights-21"_tf_keras_layer* {"name": "meg_net_layer_3", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "MEGNetLayer", "config": {"name": "meg_net_layer_3", "trainable": false, "dtype": "float32", "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 75}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 76}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "units_e": [64, 64, 32], "units_v": [64, 64, 32], "units_u": [64, 64, 32], "pool_method": "mean"}, "inbound_nodes": [[["dense_14", 0, 0, {}], ["dense_16", 0, 0, {}], ["dense_18", 0, 0, {}], ["input_4", 0, 0, {}], ["input_5", 0, 0, {}], ["input_6", 0, 0, {}], ["input_7", 0, 0, {}]]], "shared_object_id": 77, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}, {"class_name": "TensorShape", "items": [null, null]}]}2
+$
root.layer-35"_tf_keras_layer*{"name": "add_7", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Add", "config": {"name": "add_7", "trainable": false, "dtype": "float32"}, "inbound_nodes": [[["add_4", 0, 0, {}], ["meg_net_layer_3", 0, 0, {}]]], "shared_object_id": 78, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}]}2
+%
root.layer-36"_tf_keras_layer*{"name": "add_8", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Add", "config": {"name": "add_8", "trainable": false, "dtype": "float32"}, "inbound_nodes": [[["add_5", 0, 0, {}], ["meg_net_layer_3", 0, 1, {}]]], "shared_object_id": 79, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}]}2
+ &root.layer_with_weights-22"_tf_keras_layer* {"name": "set2_set_1", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Set2Set", "config": {"name": "set2_set_1", "trainable": false, "dtype": "float32", "T": 2, "n_hidden": 16, "activation": "linear", "activation_lstm": "tanh", "recurrent_activation": "hard_sigmoid", "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 80}, "recurrent_initializer": {"class_name": "Orthogonal", "config": {"gain": 1.0, "seed": null}, "shared_object_id": 81}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 82}, "use_bias": true, "unit_forget_bias": true, "kernel_regularizer": null, "recurrent_regularizer": null, "bias_regularizer": null, "kernel_constraint": null, "recurrent_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["add_7", 0, 0, {}], ["input_6", 0, 0, {}]]], "shared_object_id": 83, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null]}]}2
+ 'root.layer_with_weights-23"_tf_keras_layer* {"name": "set2_set_2", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Set2Set", "config": {"name": "set2_set_2", "trainable": false, "dtype": "float32", "T": 2, "n_hidden": 16, "activation": "linear", "activation_lstm": "tanh", "recurrent_activation": "hard_sigmoid", "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 84}, "recurrent_initializer": {"class_name": "Orthogonal", "config": {"gain": 1.0, "seed": null}, "shared_object_id": 85}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 86}, "use_bias": true, "unit_forget_bias": true, "kernel_regularizer": null, "recurrent_regularizer": null, "bias_regularizer": null, "kernel_constraint": null, "recurrent_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["add_8", 0, 0, {}], ["input_7", 0, 0, {}]]], "shared_object_id": 87, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null]}]}2
+(
root.layer-39"_tf_keras_layer*{"name": "add_9", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Add", "config": {"name": "add_9", "trainable": false, "dtype": "float32"}, "inbound_nodes": [[["add_6", 0, 0, {}], ["meg_net_layer_3", 0, 2, {}]]], "shared_object_id": 88, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [1, null, 32]}]}2
+)
root.layer-40"_tf_keras_layer*{"name": "concatenate_1", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Concatenate", "config": {"name": "concatenate_1", "trainable": false, "dtype": "float32", "axis": -1}, "inbound_nodes": [[["set2_set_1", 0, 0, {}], ["set2_set_2", 0, 0, {}], ["add_9", 0, 0, {}]]], "shared_object_id": 89, "build_input_shape": [{"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}, {"class_name": "TensorShape", "items": [null, null, 32]}]}2
+*root.layer_with_weights-24"_tf_keras_layer*{"name": "dense_19", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_19", "trainable": false, "dtype": "float32", "units": 32, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 90}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 91}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["concatenate_1", 0, 0, {}]]], "shared_object_id": 92, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 96}}, "shared_object_id": 125}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 96]}}2
++root.layer_with_weights-25"_tf_keras_layer*{"name": "dense_20", "trainable": false, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_20", "trainable": false, "dtype": "float32", "units": 16, "activation": "softplus2", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 93}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 94}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_19", 0, 0, {}]]], "shared_object_id": 95, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 32}}, "shared_object_id": 126}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 32]}}2
+,root.layer_with_weights-26"_tf_keras_layer*{"name": "dense_21", "trainable": true, "expects_training_arg": false, "dtype": "float32", "batch_input_shape": null, "stateful": false, "must_restore_from_config": false, "class_name": "Dense", "config": {"name": "dense_21", "trainable": true, "dtype": "float32", "units": 1, "activation": "linear", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}, "shared_object_id": 96}, "bias_initializer": {"class_name": "Zeros", "config": {}, "shared_object_id": 97}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dense_20", 0, 0, {}]]], "shared_object_id": 98, "input_spec": {"class_name": "InputSpec", "config": {"dtype": null, "shape": null, "ndim": null, "max_ndim": null, "min_ndim": 2, "axes": {"-1": 16}}, "shared_object_id": 127}, "build_input_shape": {"class_name": "TensorShape", "items": [null, null, 16]}}2
+root.keras_api.metrics.0"_tf_keras_metric*{"class_name": "Mean", "name": "loss", "dtype": "float32", "config": {"name": "loss", "dtype": "float32"}, "shared_object_id": 128}2
\ No newline at end of file
diff --git a/tests/test_model/prob_e_form_unnorm/nn/saved_model.pb b/tests/test_model/prob_e_form_unnorm/nn/saved_model.pb
new file mode 100644
index 0000000..6f15bbf
Binary files /dev/null and b/tests/test_model/prob_e_form_unnorm/nn/saved_model.pb differ
diff --git a/tests/test_model/prob_e_form_unnorm/nn/variables/variables.data-00000-of-00001 b/tests/test_model/prob_e_form_unnorm/nn/variables/variables.data-00000-of-00001
new file mode 100644
index 0000000..4b13792
Binary files /dev/null and b/tests/test_model/prob_e_form_unnorm/nn/variables/variables.data-00000-of-00001 differ
diff --git a/tests/test_model/prob_e_form_unnorm/nn/variables/variables.index b/tests/test_model/prob_e_form_unnorm/nn/variables/variables.index
new file mode 100644
index 0000000..a42bc06
Binary files /dev/null and b/tests/test_model/prob_e_form_unnorm/nn/variables/variables.index differ
diff --git a/tests/test_model/prob_e_form_unnorm/weights.data-00000-of-00001 b/tests/test_model/prob_e_form_unnorm/weights.data-00000-of-00001
new file mode 100644
index 0000000..342f3a2
Binary files /dev/null and b/tests/test_model/prob_e_form_unnorm/weights.data-00000-of-00001 differ
diff --git a/tests/test_model/prob_e_form_unnorm/weights.index b/tests/test_model/prob_e_form_unnorm/weights.index
new file mode 100644
index 0000000..f640377
Binary files /dev/null and b/tests/test_model/prob_e_form_unnorm/weights.index differ
diff --git a/tests/utils.py b/tests/utils.py
new file mode 100644
index 0000000..37827b9
--- /dev/null
+++ b/tests/utils.py
@@ -0,0 +1,69 @@
+"""Testing suite shared functionality."""
+import os
+
+os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = ""
+
+from distutils import dir_util
+from math import floor
+from pathlib import Path
+from typing import List, Tuple
+
+import numpy as np
+import pandas as pd
+import pytest
+
+SplitData = Tuple[Tuple[list, list], Tuple[list, list]]
+
+@pytest.fixture
+def datadir(tmpdir, request) -> Path:
+ """Access data directory.
+
+ Fixture responsible for searching a folder with the same name of test
+ module and, if available, moving all contents to a temporary directory so
+ tests can use them freely.
+
+ Source: https://stackoverflow.com/a/29631801/
+
+ """
+ filename = request.module.__file__
+ test_dir, _ = os.path.splitext(filename)
+
+ if os.path.isdir(test_dir):
+ dir_util.copy_tree(test_dir, str(tmpdir))
+
+ return tmpdir
+
+
+def weights_equal(weights_a: List[np.ndarray], weights_b: List[np.ndarray]) -> bool:
+ """Check equality between weights."""
+ return all(
+ weight1 == pytest.approx(weight2, rel=1e-6) for weight1, weight2 in zip(weights_a, weights_b)
+ )
+
+def load_df_head(fname: Path, num_entries: int=100) -> pd.DataFrame:
+ """Load first entries of a pandas DataFrame in a backwards-compatible way.
+
+ Args:
+ fname: The pickle file to open.
+ num_entries: How many values to read.
+
+ """
+ try:
+ return pd.read_pickle(fname)[:num_entries]
+ except ValueError:
+ # Older python version
+ import pickle5 as pkl
+
+ with fname.open("rb") as f:
+ return pkl.load(f)[:num_entries]
+
+def train_test_split(
+ structures: list, targets: list, train_frac: float = 0.8
+) -> SplitData:
+ """Split structures and targets into training and testing subsets."""
+ num_train = floor(len(structures) * train_frac)
+ return (
+ (structures[:num_train], targets[:num_train]),
+ (structures[num_train:], targets[num_train:]),
+ )
diff --git a/tox.ini b/tox.ini
index 68f3a03..545a07f 100644
--- a/tox.ini
+++ b/tox.ini
@@ -13,6 +13,7 @@ deps =
tensorflow-probability>=0.10.1
pymatgen<=2021.2.8
pytest
+ pytest-mock
; pytest-xdist
pandas
pyarrow>=1.0.1
diff --git a/unlocknn/__init__.py b/unlocknn/__init__.py
index fee6176..e170ff0 100644
--- a/unlocknn/__init__.py
+++ b/unlocknn/__init__.py
@@ -1,3 +1,19 @@
"""Tools for adding uncertainty quantification to neural networks."""
+import warnings
+
+# Modules with deprecation warnings
+DEPR_MODULES = ["tensorflow", "tensorflow_probability", "megnet"]
+PENDING_DEPR_MODULES = ["ruamel", "pymatgen", "monty"]
+
+for module in DEPR_MODULES:
+ warnings.filterwarnings("ignore", "DeprecationWarning", module=module)
+
+for module in PENDING_DEPR_MODULES:
+ warnings.filterwarnings("ignore", "PendingDeprecationWarning", module=module)
+
+warnings.filterwarnings("ignore", r"UserWarning: Unable to detect statically whether the number of index_points is 1.")
+warnings.filterwarnings("ignore", r"UserWarning: Converting sparse .* to a dense Tensor of unknown shape.")
+warnings.filterwarnings("ignore", r"CustomMaskWarning: Custom mask layers require a config and must override get_config.")
+
from .kernel_layers import *
from .model import *
diff --git a/unlocknn/download.py b/unlocknn/download.py
index 325b03d..e5ebe54 100644
--- a/unlocknn/download.py
+++ b/unlocknn/download.py
@@ -7,7 +7,7 @@
try:
from typing import Literal
-except ImportError:
+except ImportError: # pragma: no cover
from typish import Literal
import pandas as pd
@@ -151,11 +151,11 @@ def _load_struct_data(fname: PathLike) -> pd.DataFrame:
return serial_df
-def _save_struct_data(df: pd.DataFrame, fname: PathLike):
- """Save data containing ``Structure``s to a file.
+def save_struct_data(df: pd.DataFrame, fname: PathLike):
+ """Save data containing ``Structure`` s to a file.
Serializes the "structure" column to a string.
- The converse of :func:`load_struct_data`.
+ The converse of :func:`_load_struct_data`.
Args:
df: The :class:`pd.DataFrame` to serialize.
diff --git a/unlocknn/kernel_layers.py b/unlocknn/kernel_layers.py
index a8c8dbc..56e3801 100644
--- a/unlocknn/kernel_layers.py
+++ b/unlocknn/kernel_layers.py
@@ -89,8 +89,10 @@ class AmpAndLengthScaleFn(KernelLayer, ABC):
"""An ABC for kernels with amplitude and length scale parameters.
Attributes:
- _amplitude (tf.Tensor): The amplitude of the kernel.
- _length_scale (tf.Tensor): The length scale of the kernel.
+ _amplitude_basis (tf.Tensor): The basis for the kernel amplitude,
+ which is passed through a softplus to calculate the actual amplitude.
+ _length_scale_basis (tf.Tensor): The basis for the length scale of the kernel.
+ which is passed through a softplus to calculate the actual amplitude.
"""
@@ -99,21 +101,33 @@ def __init__(self, **kwargs):
super().__init__(**kwargs)
dtype = kwargs.get("dtype", tf.float64)
- self._amplitude = self.add_weight(
+ self._amplitude_basis = self.add_weight(
initializer=tf.constant_initializer(0), dtype=dtype, name="amplitude"
)
- self._length_scale = self.add_weight(
+ self._length_scale_basis = self.add_weight(
initializer=tf.constant_initializer(0), dtype=dtype, name="length_scale"
)
+
+ @property
+ def amplitude(self) -> tf.Tensor:
+ """Get the current kernel amplitude."""
+ return tf.nn.softplus(0.1 * self._amplitude_basis)
+
+ @property
+ def length_scale(self) -> tf.Tensor:
+ """Get the current kernel length scale."""
+ return tf.nn.softplus(5.0 * self._length_scale_basis)
class RBFKernelFn(AmpAndLengthScaleFn):
"""A radial basis function implementation that works with keras.
Attributes:
- _amplitude (tf.Tensor): The amplitude of the kernel.
- _length_scale (tf.Tensor): The length scale of the kernel.
+ _amplitude_basis (tf.Tensor): The basis for the kernel amplitude,
+ which is passed through a softplus to calculate the actual amplitude.
+ _length_scale_basis (tf.Tensor): The basis for the length scale of the kernel.
+ which is passed through a softplus to calculate the actual amplitude.
"""
@@ -121,8 +135,8 @@ class RBFKernelFn(AmpAndLengthScaleFn):
def kernel(self) -> tfp.math.psd_kernels.PositiveSemidefiniteKernel:
"""Get a callable kernel."""
return tfp.math.psd_kernels.ExponentiatedQuadratic(
- amplitude=tf.nn.softplus(0.1 * self._amplitude),
- length_scale=tf.nn.softplus(5.0 * self._length_scale),
+ amplitude=self.amplitude,
+ length_scale=self.length_scale,
)
@property
@@ -134,8 +148,10 @@ class MaternOneHalfFn(AmpAndLengthScaleFn):
"""A Matern kernel with parameter 1/2 implementation that works with keras.
Attributes:
- _amplitude (tf.Tensor): The amplitude of the kernel.
- _length_scale (tf.Tensor): The length scale of the kernel.
+ _amplitude_basis (tf.Tensor): The basis for the kernel amplitude,
+ which is passed through a softplus to calculate the actual amplitude.
+ _length_scale_basis (tf.Tensor): The basis for the length scale of the kernel.
+ which is passed through a softplus to calculate the actual amplitude.
"""
@@ -143,8 +159,8 @@ class MaternOneHalfFn(AmpAndLengthScaleFn):
def kernel(self) -> tfp.math.psd_kernels.PositiveSemidefiniteKernel:
"""Get a callable kernel."""
return tfp.math.psd_kernels.MaternOneHalf(
- amplitude=tf.nn.softplus(0.1 * self._amplitude),
- length_scale=tf.nn.softplus(5.0 * self._length_scale),
+ amplitude=self.amplitude,
+ length_scale=self.length_scale,
)
@property
diff --git a/unlocknn/megnet_utils.py b/unlocknn/megnet_utils.py
index 990ab82..9414867 100644
--- a/unlocknn/megnet_utils.py
+++ b/unlocknn/megnet_utils.py
@@ -10,6 +10,7 @@
MEGNetGraph = Dict[str, Union[np.ndarray, List[Union[int, float]]]]
Targets = List[Union[float, np.ndarray]]
+ModelInput = Union[Structure, MEGNetGraph]
def default_megnet_config(
@@ -58,7 +59,7 @@ def scale_targets(
def create_megnet_input(
meg_model: MEGNetModel,
- structs: List[Structure],
+ inputs: List[ModelInput],
targets: Optional[Targets] = None,
batch_size: int = 128,
scrub_failed_structs: bool = False,
@@ -68,7 +69,7 @@ def create_megnet_input(
Args:
meg_model: The :class:`MEGNetModel` whose graph converter to use.
- structs: The input structures.
+ inputs: The input, either as graphs or structures.
targets: The input targets, if any.
batch_size: The batch size for the generator.
scrub_failed_structures: Whether to discard structures
@@ -84,11 +85,17 @@ def create_megnet_input(
"""
# Make some targets up for compatibility
has_targets = targets is not None
- target_buffer = targets if has_targets else [0.0] * len(structs)
+ target_buffer = targets if has_targets else [0.0] * len(inputs)
+
+ is_struct = isinstance(inputs[0], Structure)
+ if is_struct:
+ graphs, trunc_targets = meg_model.get_all_graphs_targets(
+ inputs, target_buffer, scrub_failed_structs
+ )
+ else:
+ graphs = inputs
+ trunc_targets = target_buffer
- graphs, trunc_targets = meg_model.get_all_graphs_targets(
- structs, target_buffer, scrub_failed_structs
- )
# Check dimensions of model against converted graphs
meg_model.check_dimension(graphs[0])
diff --git a/unlocknn/metrics.py b/unlocknn/metrics.py
index 30e903e..d286b95 100644
--- a/unlocknn/metrics.py
+++ b/unlocknn/metrics.py
@@ -1,5 +1,4 @@
"""Uncertainty quantification metrics and evaluation utilities."""
-from pprint import pprint
from typing import Callable, Dict, List, Optional
import numpy as np
@@ -8,7 +7,7 @@
from pymatgen.core.structure import Structure
from .megnet_utils import Targets
-from .model import ProbNN
+from .model import ProbNN, ModelInput
def neg_log_likelihood(
@@ -16,7 +15,26 @@ def neg_log_likelihood(
stddevs: Targets,
true_vals: Targets,
) -> float:
- """Calculate the negative log likelihood of true values given predictions."""
+ r"""Calculate the negative log likelihood (NLL) of true values given predictions.
+
+ NLL is given by
+
+ .. math::
+
+ \mathrm{NLL} = -\sum_i \log p_i(y_i),
+
+ where :math:`y_i` is the :math:`i^\mathrm{th}` observed (true) value and
+ :math:`p_i` is the probability density function for the
+ :math:`i^\mathrm{th}` predicted Gaussian distribution:
+
+ .. math::
+
+ p_i \sim \mathcal{N} \left( \hat{y}_i, \sigma_i^2 \right),
+
+ where :math:`\hat{y}_i` is the :math:`i^\mathrm{th}` predicted mean and
+ :math:`\sigma_i` is the :math:`i^\mathrm{th}` predicted standard deviation.
+
+ """
dist = tfp.distributions.Normal(loc=predictions, scale=stddevs)
nll = -tf.math.reduce_sum(dist.log_prob(true_vals))
return nll.numpy()
@@ -48,7 +66,7 @@ def variation(
the mean of the standard deviations and :math:`N` be the number of
predictions. The coefficient of variation is given by:
- .. math:: C_v = \sqrt{\frac{\sum_i^N{(\sigma_i - \bar{\sigma})^2}}{N - 1}}
+ .. math:: C_v = \frac{1}{\bar{\sigma}}\sqrt{\frac{\sum_i^N{(\sigma_i - \bar{\sigma})^2}}{N - 1}}
"""
stddev_mean = np.mean(stddevs)
@@ -83,7 +101,7 @@ def __call__(
true_vals: Targets,
) -> float:
"""Calculate the mean metric."""
- mean_error = np.mean(self.mean_what_func(true_vals - predictions))
+ mean_error = np.mean(self.mean_what_func(np.array(true_vals) - np.array(predictions)))
if self.take_root:
return np.sqrt(mean_error)
else:
@@ -108,7 +126,7 @@ def __call__(
def evaluate_uq_metrics(
prob_model: ProbNN,
- test_structs: List[Structure],
+ test_inputs: List[ModelInput],
test_targets: Targets,
metrics: List[str] = list(AVAILABLE_METRICS.keys()),
) -> Dict[str, float]:
@@ -116,7 +134,7 @@ def evaluate_uq_metrics(
Args:
prob_model: The probabilistic model to evaluate.
- test_structs: The input structures.
+ test_inputs: The input structures or graphs.
test_targets: The target values for the structures.
metrics: A list of metrics to compute. Defaults
to computing all of the currently implemented
@@ -149,7 +167,7 @@ def evaluate_uq_metrics(
"""
metrics_dict = {metric: AVAILABLE_METRICS[metric] for metric in metrics}
- predictions, stddevs = prob_model.predict(test_structs)
+ predictions, stddevs = prob_model.predict(test_inputs)
return {
metric: func(predictions, stddevs, test_targets)
for metric, func in metrics_dict.items()
diff --git a/unlocknn/model.py b/unlocknn/model.py
index adbbf48..0434eff 100644
--- a/unlocknn/model.py
+++ b/unlocknn/model.py
@@ -5,11 +5,11 @@
import warnings
from abc import ABC, abstractmethod
from pathlib import Path
-from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
try:
from typing import Literal, get_args
-except ImportError:
+except ImportError: # pragma: no cover
from typish import Literal, get_args
@@ -23,7 +23,7 @@
from tensorflow.python.keras.utils import losses_utils
from .kernel_layers import KernelLayer, RBFKernelFn, load_kernel
-from .megnet_utils import create_megnet_input, Targets
+from .megnet_utils import create_megnet_input, Targets, ModelInput
tfd = tfp.distributions
@@ -159,7 +159,9 @@ class ProbNN(ABC):
index_initializer: A custom initializer to use for the VGP index points.
use_normalization: Whether to use a ``BatchNormalization`` layer before
the VGP. Recommended for better training efficiency.
-
+ compile: Whether to compile the model for training. Not needed when loading
+ the model for inference only.
+
Attributes:
CONFIG_VARS: A list of attribute names, as strings, to include in metadata
when saving. These variables will be saved in a ``config.json`` file
@@ -266,9 +268,7 @@ def set_frozen(
self.model.layers[-1].trainable = not freeze
if recompile:
- self.compile(
- kl_weight=self.kl_weight, optimizer=self.optimizer, **compilation_kwargs
- )
+ self.compile(optimizer=self.optimizer, **compilation_kwargs)
@abstractmethod
def train(self, *args, **kwargs) -> None:
@@ -292,7 +292,6 @@ def _update_pred_model(self) -> None:
use_normalization=self.use_normalization,
prediction_mode=True,
)
- pred_model.compile()
pred_model.set_weights(weights)
self.pred_model = pred_model
@@ -360,7 +359,7 @@ def norm_frozen(self) -> Optional[bool]:
def compile(
self,
- kl_weight: float = 1.0,
+ new_kl_weight: Optional[float] = None,
optimizer: keras.optimizers.Optimizer = tf.optimizers.Adam(),
new_metrics: Optional[Metrics] = None,
):
@@ -369,13 +368,15 @@ def compile(
Recompilation is required whenever layers are (un)frozen.
Args:
- kl_weight: The relative weighting of the Kullback-Leibler divergence
- in the loss function.
+ new_kl_weight: The relative weighting of the Kullback-Leibler divergence
+ in the loss function. Default (``None``) is to leave unchanged.
optimizer: The model optimizer, needed for recompilation.
new_metrics: New metrics with which to compile.
"""
- loss = VariationalLoss(kl_weight)
+ loss = VariationalLoss(
+ new_kl_weight if new_kl_weight is not None else self.kl_weight
+ )
if new_metrics is not None:
self.metrics = new_metrics
self.model.compile(optimizer, loss=loss, metrics=self.metrics)
@@ -425,10 +426,12 @@ def save_kernel(self, kernel_save_path: Path):
@classmethod
def load(
- cls: "ProbNN", save_path: Path, load_ckpt: bool = True, **kwargs
+ cls: "ProbNN", save_path: PathLike, load_ckpt: bool = True, **kwargs
) -> "ProbNN":
"""Load a ProbNN from disk.
+ Loaded models must be recompiled before training.
+
Args:
save_path: The path to the model's save directory.
load_ckpt: Whether to load the best checkpoint's weights, instead
@@ -443,7 +446,7 @@ def load(
"""
# Check the save path exists
- if not save_path.exists():
+ if not Path(save_path).exists():
raise FileNotFoundError(f"`{save_path}` does not exist.")
paths = _get_save_paths(save_path)
@@ -534,25 +537,25 @@ def __init__(
def train(
self,
- structs: List[Structure],
+ inputs: List[Structure],
targets: Targets,
epochs: int,
- val_structs: Optional[List[Structure]] = None,
+ val_inputs: Optional[List[ModelInput]] = None,
val_targets: Optional[Targets] = None,
callbacks: List[tf.keras.callbacks.Callback] = [],
use_default_ckpt_handler: bool = True,
ckpt_path: PathLike = "checkpoint.h5",
- batch_size: int = 128,
+ batch_size: int = 32,
scrub_failed_structs: bool = False,
verbose: Literal[0, 1, 2] = 2,
):
"""Train the model.
Args:
- structs: A list of training crystal structures.
+ inputs: A list of training crystal structures or graphs.
targets: A list of training target values.
epochs: The number of training epochs.
- val_structs: A list of validation crystal structures.
+ val_inputs: A list of validation crystal structures or graphs.
val_targets: A list of validation target values.
callbacks: A list of additional callbacks.
use_default_ckpt_handler: Whether to use the default
@@ -568,17 +571,17 @@ def train(
"""
# Convert structures to graphs for model input
train_gen, train_graphs = create_megnet_input(
- self.meg_model, structs, targets, batch_size, scrub_failed_structs
+ self.meg_model, inputs, targets, batch_size, scrub_failed_structs
)
steps_per_train = int(np.ceil(len(train_graphs) / batch_size))
val_gen = None
val_graphs = None
steps_per_val = None
- if val_structs is not None and val_targets is not None:
+ if val_inputs is not None and val_targets is not None:
val_gen, val_graphs = create_megnet_input(
self.meg_model,
- val_structs,
+ val_inputs,
val_targets,
batch_size,
scrub_failed_structs,
@@ -633,7 +636,7 @@ def evaluate(
)
def predict(
- self, input: Union[Structure, Iterable[Structure]], batch_size: int = 128
+ self, input: Union[ModelInput, List[ModelInput]], batch_size: int = 128
) -> Tuple[np.ndarray, np.ndarray]:
"""Predict target values and standard deviations for a given input.
@@ -665,12 +668,14 @@ def predict(
"""
self.update_pred_model()
- if isinstance(input, Structure):
- # Just one to predict
+ try:
+ len(input)
+ except TypeError:
+ # Make it a sequence
input = [input]
inputs, graphs = create_megnet_input(
- self.meg_model, structs=input, batch_size=batch_size, shuffle=False
+ self.meg_model, inputs=input, batch_size=batch_size, shuffle=False
)
num_atoms = [len(graph["atom"]) for graph in graphs]