diff --git a/.gitignore b/.gitignore
index 24c71f8..cd6e613 100644
--- a/.gitignore
+++ b/.gitignore
@@ -21,4 +21,4 @@ cfg_constructor/out
 **/.venv
 **/logdat
 
-!**/demo_file_id.txt
\ No newline at end of file
+!**/demo_file_id.txt
diff --git a/GraphIsomorphismNetwork/jax_gat.py b/GraphIsomorphismNetwork/jax_gat.py
new file mode 100644
index 0000000..bd5b6d2
--- /dev/null
+++ b/GraphIsomorphismNetwork/jax_gat.py
@@ -0,0 +1,317 @@
+import jax
+import jax.numpy as jnp
+from flax import linen as nn
+from flax.training import train_state, checkpoints
+from new_loader import get_paths, load_dataset_jax_new
+import optax
+from pathlib import Path
+import pandas as pd
+import numpy as np
+import time
+import os
+import logging
+import tqdm
+
+log_file = Path(__file__).resolve().parent.parent / "logs" / "gat_training.log"
+log_file.parent.mkdir(parents=True, exist_ok=True)
+
+root_logger = logging.getLogger()
+if root_logger.handlers:
+    for handler in root_logger.handlers[:]:
+        root_logger.removeHandler(handler)
+
+logging.basicConfig(
+    filename=log_file,
+    level=logging.INFO,
+    filemode="a",
+    format="%(asctime)s - %(levelname)s - %(message)s",
+    datefmt='%Y-%m-%d %H:%M:%S'
+)
+
+logging.info("Starting GAT training...")
+
+class GraphAttentionLayer(nn.Module):
+    out_dim: int
+    dropout_rate: float = 0.5
+    
+    @nn.compact
+    def __call__(self, x, senders, receivers, training: bool):
+        W = nn.Dense(self.out_dim, use_bias=False)
+        x = W(x)
+        
+        a1 = nn.Dense(1, use_bias=False, name='attention_left')
+        a2 = nn.Dense(1, use_bias=False, name='attention_right')
+        
+        f_1 = a1(x)  # [N, 1]
+        f_2 = a2(x)  
+        
+        # Create attention logits for edges
+        edge_logits = f_1[senders] + f_2[receivers] 
+        edge_logits = nn.leaky_relu(edge_logits.squeeze(-1), negative_slope=0.2)
+        
+        # Normalize attention weights using softmax
+        max_logits = jax.ops.segment_max(
+            edge_logits, receivers, num_segments=x.shape[0], indices_are_sorted=False
+        )
+        edge_logits_normalized = edge_logits - max_logits[receivers]
+        edge_weights = jnp.exp(edge_logits_normalized)
+        
+        
+        weight_sum = jax.ops.segment_sum(
+            edge_weights, receivers, num_segments=x.shape[0], indices_are_sorted=False
+        )
+        edge_weights = edge_weights / (weight_sum[receivers] + 1e-16)
+        
+        # Apply dropout to attention weights
+        edge_weights = nn.Dropout(rate=self.dropout_rate)(
+            edge_weights, deterministic=not training
+        )
+        
+        
+        x = nn.Dropout(rate=self.dropout_rate)(x, deterministic=not training)
+        
+        
+        weighted_features = edge_weights[:, None] * x[senders] 
+        aggregated = jax.ops.segment_sum(
+            weighted_features, receivers, num_segments=x.shape[0], indices_are_sorted=False
+        )
+        
+        return nn.elu(aggregated)
+
+class MultiHeadAttention(nn.Module):
+    num_heads: int
+    hidden_dim: int
+    dropout_rate: float = 0.5
+    last_layer: bool = False
+    
+    @nn.compact
+    def __call__(self, x, senders, receivers, training: bool):
+        #multiple attention heads
+        heads = []
+        for i in range(self.num_heads):
+            head = GraphAttentionLayer(
+                out_dim=self.hidden_dim,
+                dropout_rate=self.dropout_rate,
+                name=f'attention_head_{i}'
+            )
+            head_out = head(x, senders, receivers, training)
+            heads.append(head_out)
+        
+        if self.last_layer:
+            # Average the heads for the last layer
+            return jnp.mean(jnp.stack(heads), axis=0)
+        else:
+            return jnp.concatenate(heads, axis=-1)
+
+class GAT(nn.Module):
+    in_channels: int
+    hidden_channels: int
+    out_channels: int
+    num_heads: list
+    dropout_rate: float = 0.5
+    
+    @nn.compact
+    def __call__(self, x, edge_index, batch, training: bool):
+        senders, receivers = edge_index
+        
+        x = MultiHeadAttention(
+            num_heads=self.num_heads[0],
+            hidden_dim=self.hidden_channels,
+            dropout_rate=self.dropout_rate,
+            last_layer=False,
+            name='layer_0'
+        )(x, senders, receivers, training)
+        
+  
+        x = MultiHeadAttention(
+            num_heads=self.num_heads[1],
+            hidden_dim=self.out_channels,
+            dropout_rate=self.dropout_rate,
+            last_layer=True,
+            name='layer_1'
+        )(x, senders, receivers, training)
+        
+        # Global add pooling
+        x = jax.ops.segment_sum(x, batch, num_segments=1)
+        
+        return jax.nn.log_softmax(x, axis=-1)
+
+class TrainState(train_state.TrainState):
+    batch_stats: dict
+
+def create_train_state(rng, model, learning_rate, sample_input):
+    variables = model.init(
+        rng,
+        sample_input["x"],
+        sample_input["edge_index"],
+        sample_input["batch"],
+        training=True,
+    )
+    tx = optax.adam(learning_rate)
+    return TrainState.create(
+        apply_fn=model.apply,
+        params=variables["params"],
+        tx=tx,
+        batch_stats=variables.get("batch_stats", {})
+    )
+
+@jax.jit
+def train_step(state, batch, dropout_rng):
+    def loss_fn(params):
+        variables = {"params": params, "batch_stats": state.batch_stats}
+        logits, new_model_state = state.apply_fn(
+            variables,
+            batch["x"],
+            batch["edge_index"],
+            batch["batch"],
+            training=True,
+            mutable=["batch_stats"],
+            rngs={"dropout": dropout_rng},
+        )
+        labels = batch["y"]
+        nll = -jnp.mean(jnp.take_along_axis(logits, labels[:, None], axis=-1).squeeze())
+        
+        # L2 regularization
+        l2_loss = 5e-4 * sum(jnp.sum(p**2) for p in jax.tree_leaves(params))
+        total_loss = nll + l2_loss
+        
+        return total_loss, new_model_state
+
+    grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
+    (loss, new_model_state), grads = grad_fn(state.params)
+    state = state.apply_gradients(grads=grads, batch_stats=new_model_state["batch_stats"])
+    return state, loss
+
+@jax.jit
+def test_step(state, batch):
+    variables = {"params": state.params, "batch_stats": state.batch_stats}
+    logits = state.apply_fn(
+        variables,
+        batch["x"],
+        batch["edge_index"],
+        batch["batch"],
+        training=False,
+        mutable=False
+    )
+    pred = jnp.argmax(logits, axis=-1)
+    correct = jnp.sum(pred == batch["y"])
+    return correct
+
+def save_model(state, checkpoint_dir, step):
+    os.makedirs(checkpoint_dir, exist_ok=True)
+    checkpoints.save_checkpoint(
+        ckpt_dir=checkpoint_dir,
+        target=state,
+        step=step,
+        overwrite=True
+    )
+    msg = f"Model saved at step {step} to {checkpoint_dir}"
+    print(msg)
+    logging.info(msg)
+
+def load_model_if_exists(state, checkpoint_dir):
+    latest_ckpt = checkpoints.latest_checkpoint(checkpoint_dir)
+    if latest_ckpt:
+        state = checkpoints.restore_checkpoint(
+            ckpt_dir=checkpoint_dir,
+            target=state
+        )
+        msg = f"Model restored from checkpoint: {latest_ckpt}"
+        print(msg)
+        logging.info(msg)
+    else:
+        msg = "No checkpoint found. Training from scratch."
+        print(msg)
+        logging.info(msg)
+    return state
+
+def main():
+    # print(jax.devices())
+    dev = jax.devices()[0] if jax.devices() else None
+
+    paths = get_paths(samples_2000=True)
+    data_loader, num_features, num_classes = load_dataset_jax_new(paths, max_files=200)
+
+    split_ratio = 0.8
+    n_total = len(data_loader)
+    split_idx = int(n_total * split_ratio)
+    train_loader = data_loader[:split_idx]
+    test_loader = data_loader[split_idx:]
+    print(f"Total batches: {n_total}; Training batches: {len(train_loader)}; Test batches: {len(test_loader)}")
+
+    print(f"Number of Features: {num_features}, model hidden layer dim: {max(1, int(num_features * 1e-4))}")
+    
+    # GAT config 
+    num_heads = [8, 1]  
+    hidden_dim = max(1, int(num_features * 1e-4))
+    
+    model = GAT(
+        in_channels=num_features,
+        hidden_channels=hidden_dim,
+        out_channels=num_classes,
+        num_heads=num_heads,
+        dropout_rate=0.5
+    )
+
+    rng = jax.random.PRNGKey(0)
+    dropout_rng, init_rng = jax.random.split(rng)
+    
+    print(f"Total samples: {n_total}")
+    sample_input = train_loader[0]
+    state = create_train_state(init_rng, model, learning_rate=0.005, sample_input=sample_input)
+
+    checkpoint_dir = Path(__file__).resolve().parent.parent / "weights" / "gat"
+    #state = load_model_if_exists(state, checkpoint_dir)
+    
+    start_train = time.perf_counter()
+    num_epochs = 400
+
+    state = jax.device_put(state, device=dev)
+    
+    for epoch in range(1, num_epochs + 1):
+        epoch_start = time.perf_counter()
+
+        # Train Loop
+        epoch_loss = 0.0
+        total_graphs = 0
+        for batch in tqdm.tqdm(train_loader, desc=f"training epoch {epoch}"):
+            dropout_rng, new_dropout_rng = jax.random.split(dropout_rng)
+            state, loss = train_step(state, batch, dropout_rng)
+            state = jax.device_put(state, device=dev)
+            epoch_loss += loss * batch["num_graphs"]
+            total_graphs += batch["num_graphs"]
+        avg_loss = epoch_loss / total_graphs
+
+        # Test Loop
+        total_correct = 0
+        total_test_graphs = 0
+        for batch in test_loader:
+            correct = test_step(state, batch)
+            total_correct += correct
+            total_test_graphs += batch["num_graphs"]
+        test_acc = total_correct / total_test_graphs
+
+        epoch_end = time.perf_counter()
+        msg1 = f"Time for epoch {epoch} was {(epoch_end - epoch_start):.6f} seconds"
+        msg2 = f"Epoch: {epoch:03d}, Loss: {avg_loss:.4f}, Test Acc: {test_acc:.4f}"
+        print(msg1)
+        print(msg2)
+        logging.info(msg1)
+        logging.info(msg2)
+
+        if epoch % 50 == 0:  
+            save_model(state, checkpoint_dir, epoch)
+
+    
+    save_model(state, checkpoint_dir, num_epochs)
+
+    end_train = time.perf_counter()
+    msg3 = f"Total training time : {(end_train - start_train):.6f} seconds"
+    msg4 = f"Average time per epoch: {((end_train - start_train) / num_epochs):.6f} seconds"
+    print(msg3)
+    print(msg4)
+    logging.info(msg3)
+    logging.info(msg4)
+
+if __name__ == "__main__":
+    main()
\ No newline at end of file
diff --git a/GraphIsomorphismNetwork/jax_gin.py b/GraphIsomorphismNetwork/jax_gin.py
new file mode 100644
index 0000000..ac74607
--- /dev/null
+++ b/GraphIsomorphismNetwork/jax_gin.py
@@ -0,0 +1,256 @@
+import jax
+import jax.numpy as jnp
+from flax import linen as nn
+from flax.training import train_state, checkpoints
+from new_loader import get_paths, load_dataset_jax_new
+import optax
+from pathlib import Path
+import pandas as pd
+import numpy as np
+import time
+import os
+import logging
+import tqdm
+
+log_file = Path(__file__).resolve().parent.parent / "logs" / "gin_training.log"
+log_file.parent.mkdir(parents=True, exist_ok=True)
+
+root_logger = logging.getLogger()
+if root_logger.handlers:
+    for handler in root_logger.handlers[:]:
+        root_logger.removeHandler(handler)
+
+logging.basicConfig(
+    filename=log_file,
+    level=logging.INFO,
+    filemode="a",
+    format="%(asctime)s - %(levelname)s - %(message)s",
+    datefmt='%Y-%m-%d %H:%M:%S'
+)
+
+logging.info("Starting GIN training...")
+
+# Dense -> BatchNorm -> ReLU -> Dense.
+class MLP(nn.Module):
+    hidden_dim: int
+
+    @nn.compact
+    def __call__(self, x, training: bool):
+        x = nn.Dense(2 * self.hidden_dim)(x)
+        x = nn.BatchNorm(use_running_average=not training)(x)
+        x = nn.relu(x)
+        x = nn.Dense(self.hidden_dim)(x)
+        return x
+
+# Define a GIN convolution layer: x_i' = MLP((1+eps)*x_i + sum_{j in N(i)} x_j)
+class GINConv(nn.Module):
+    hidden_dim: int
+    train_eps: bool = True
+
+    @nn.compact
+    def __call__(self, x, senders, receivers, training: bool):
+        mlp = MLP(hidden_dim=self.hidden_dim)
+        # Learnable epsilon
+        if self.train_eps:
+            eps = self.param("eps", lambda rng: jnp.zeros(()))
+        else:
+            eps = 0.0
+        # Aggregate neighbor features using segment_sum.
+        aggregated = jax.ops.segment_sum(x[senders], receivers, num_segments=x.shape[0])
+        out = mlp((1 + eps) * x + aggregated, training=training)
+        return out
+
+class GIN(nn.Module):
+    in_channels: int
+    hidden_channels: int
+    out_channels: int
+    num_layers: int
+    dropout_rate: float = 0.5
+    train_eps: bool = True
+
+    @nn.compact
+    def __call__(self, x, edge_index, batch, training: bool):
+        senders, receivers = edge_index  # edge_index is a tuple: (senders, receivers)
+        for _ in range(self.num_layers):
+            x = GINConv(hidden_dim=self.hidden_channels, train_eps=self.train_eps)(
+                x, senders, receivers, training=training
+            )
+            x = nn.BatchNorm(use_running_average=not training)(x)
+            x = nn.relu(x)
+
+        # Global add pooling: sum node features per graph.
+        x = jax.ops.segment_sum(x, batch, num_segments=1)
+
+        # Two-layer MLP for graph-level output.
+        x = nn.Dense(self.hidden_channels)(x)
+        x = nn.LayerNorm()(x)
+        x = nn.relu(x)
+        x = nn.Dropout(rate=self.dropout_rate)(x, deterministic=not training)
+        x = nn.Dense(self.out_channels)(x)
+        return jax.nn.log_softmax(x, axis=-1)
+
+class TrainState(train_state.TrainState):
+    batch_stats: dict
+
+def create_train_state(rng, model, learning_rate, sample_input):
+    variables = model.init(
+        rng,
+        sample_input["x"],
+        sample_input["edge_index"],
+        sample_input["batch"],
+        training=True,
+    )
+    tx = optax.adam(learning_rate)
+    return TrainState.create(
+        apply_fn=model.apply,
+        params=variables["params"],
+        tx=tx,
+        batch_stats=variables.get("batch_stats", {})
+    )
+
+@jax.jit
+def train_step(state, batch, dropout_rng):
+    def loss_fn(params):
+        variables = {"params": params, "batch_stats": state.batch_stats}
+        logits, new_model_state = state.apply_fn(
+            variables,
+            batch["x"],
+            batch["edge_index"],
+            batch["batch"],
+            training=True,
+            mutable=["batch_stats"],
+            rngs={"dropout": dropout_rng},
+        )
+        labels = batch["y"]
+        nll = -jnp.mean(jnp.take_along_axis(logits, labels[:, None], axis=-1).squeeze())
+        return nll, new_model_state
+
+    grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
+    (loss, new_model_state), grads = grad_fn(state.params)
+    state = state.apply_gradients(grads=grads, batch_stats=new_model_state["batch_stats"])
+    return state, loss
+
+@jax.jit
+def test_step(state, batch):
+    variables = {"params": state.params, "batch_stats": state.batch_stats}
+    logits = state.apply_fn(
+        variables,
+        batch["x"],
+        batch["edge_index"],
+        batch["batch"],
+        training=False,
+        mutable=False
+    )
+    pred = jnp.argmax(logits, axis=-1)
+    correct = jnp.sum(pred == batch["y"])
+    return correct
+
+def save_model(state, checkpoint_dir, step):
+    os.makedirs(checkpoint_dir, exist_ok=True)
+    checkpoints.save_checkpoint(
+        ckpt_dir=checkpoint_dir,
+        target=state,
+        step=step,
+        overwrite=True
+    )
+    msg = f"Model saved at step {step} to {checkpoint_dir}"
+    print(msg)
+    logging.info(msg)
+
+def load_model_if_exists(state, checkpoint_dir):
+    latest_ckpt = checkpoints.latest_checkpoint(checkpoint_dir)
+    if latest_ckpt:
+        state = checkpoints.restore_checkpoint(
+            ckpt_dir=checkpoint_dir,
+            target=state
+        )
+        msg = f"Model restored from checkpoint: {latest_ckpt}"
+        print(msg)
+        logging.info(msg)
+    else:
+        msg = "No checkpoint found. Training from scratch."
+        print(msg)
+        logging.info(msg)
+    return state
+
+def main():
+    # print(jax.devices())
+    dev = jax.devices()[0] if jax.devices() else None
+
+    paths = get_paths(samples_2000=True)
+    data_loader, num_features, num_classes = load_dataset_jax_new(paths, max_files=None)
+
+    split_ratio = 0.8
+    n_total = len(data_loader)
+    split_idx = int(n_total * split_ratio)
+    train_loader = data_loader[:split_idx]
+    test_loader = data_loader[split_idx:]
+    print(f"Total batches: {n_total}; Training batches: {len(train_loader)}; Test batches: {len(test_loader)}")
+
+    print(f"Number of Features: {num_features}, model hidden layer dim: {int(num_features * 1e-4)}")
+    model = GIN(
+        in_channels=num_features,
+        hidden_channels= max(1, int(num_features * 1e-4)),
+        out_channels=num_classes,
+        num_layers=4,
+        dropout_rate=0.5
+    )
+
+    rng = jax.random.PRNGKey(0)
+    dropout_rng, init_rng = jax.random.split(rng)
+    #print(train_loader)
+    print(n_total)
+    sample_input = train_loader[0]
+    state = create_train_state(init_rng, model, learning_rate=0.01, sample_input=sample_input)
+
+    checkpoint_dir = Path(__file__).resolve().parent.parent / "weights"
+    #state = load_model_if_exists(state, checkpoint_dir)
+    
+    start_train = time.perf_counter()
+    num_epochs = 10
+
+    state = jax.device_put(state, device=dev)
+    
+    for epoch in range(1, num_epochs + 1):
+        epoch_start = time.perf_counter()
+
+        # Train Loop
+        epoch_loss = 0.0
+        total_graphs = 0
+        for batch in tqdm.tqdm(train_loader, desc=f"training epoch {epoch}"):
+            dropout_rng, new_dropout_rng = jax.random.split(dropout_rng)
+            state, loss = train_step(state, batch, dropout_rng)
+            state = jax.device_put(state, device=dev)
+            epoch_loss += loss * batch["num_graphs"]
+            total_graphs += batch["num_graphs"]
+        avg_loss = epoch_loss / total_graphs
+
+        # Test Loop
+        total_correct = 0
+        total_test_graphs = 0
+        for batch in test_loader:
+            correct = test_step(state, batch)
+            total_correct += correct
+            total_test_graphs += batch["num_graphs"]
+        test_acc = total_correct / total_test_graphs
+
+        epoch_end = time.perf_counter()
+        msg1 = f"Time for epoch {epoch} was {(epoch_end - epoch_start):.6f} seconds"
+        msg2 = f"Epoch: {epoch:03d}, Loss: {avg_loss:.4f}, Test Acc: {test_acc:.4f}"
+        print(msg1)
+        print(msg2)
+        logging.info(msg1)
+        logging.info(msg2)
+
+        save_model(state, checkpoint_dir, epoch)
+
+    end_train = time.perf_counter()
+    msg3 = f"Total training time : {(end_train - start_train):.6f} seconds"
+    msg4 = f"Average time per epoch: {((end_train - start_train) / num_epochs):.6f} seconds"
+    print(msg3)
+    print(msg4)
+    logging.info(msg3)
+    logging.info(msg4)
+
+if __name__ == "__main__":
+    main()
\ No newline at end of file
diff --git a/GraphIsomorphismNetwork/load.py b/GraphIsomorphismNetwork/load.py
new file mode 100644
index 0000000..290302a
--- /dev/null
+++ b/GraphIsomorphismNetwork/load.py
@@ -0,0 +1,248 @@
+import time
+import torch
+import shutil
+import numpy as np
+from pathlib import Path
+import pandas as pd
+import jax.numpy as jnp
+from torch.utils import data
+from torch_geometric.data import Data, DataLoader
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def load_dataset_jax(paths: list[str]) -> tuple:
+    start_load = time.perf_counter()
+
+    """
+    Loads the dataset and returns a list of samples (each a dict with keys:
+    "x", "edge_index", "batch", "y", and "num_graphs"),
+    along with the number of features and number of classes.
+
+    This version is adapted for use with JAX/Flax.
+    """
+    num_features, num_classes = 0, 2
+    max_edgelist, max_nodes = 0, 0
+
+    # Dictionary to hold data for each file id: {file_id: [edge_index, node_features]}
+    files = {}
+    labels = None
+    for path in sorted(paths):
+        fname = Path(path).name
+        if "edgelist" in fname:
+            # Expect file names like "<id>_edgelist..."
+            curr_id = int(fname.split("_")[0])
+            print(f"Loading edgelist from: {path}")
+            edges = pd.read_csv(path, header=None).values.astype(np.int32)
+            max_edgelist = edges.shape[0] if edges.shape[0] > max_edgelist else max_edgelist
+            edge_index = edges.T  # shape: (2, num_edges)
+            if curr_id not in files:
+                files[curr_id] = [edge_index, None]
+            else:
+                _, curr_features = files[curr_id]
+                files[curr_id] = [edge_index, curr_features]
+        elif "node_names" in fname:
+            curr_id = int(fname.split("_")[0])
+            print(f"Loading node features from: {path}")
+            node_features = pd.read_csv(path, header=None).values
+            # Remove first row and second column as in your original code.
+            node_features = np.delete(np.delete(node_features, 0, axis=0), 1, axis=1).astype(np.int32)
+            if curr_id not in files:
+                files[curr_id] = [None, node_features]
+            else:
+                curr_edge, _ = files[curr_id]
+                files[curr_id] = [curr_edge, node_features]
+            # num_features = max(num_features, node_features.shape[1])
+            num_features = max(num_features, node_features.shape[0])
+            max_nodes = max(max_nodes, node_features.shape[0])
+        elif "samplekey" in fname:
+            print(f"Loading labels from: {path}")
+            labels = pd.read_csv(path, header=None)
+        elif "removed" in fname:
+            print(f"Skipping removed file: {path}")
+
+    print(f"Max number of features: {num_features}")
+    print(f"Number of classes: {num_classes}")
+    print(f"Largest edgelist (number of edges): {max_edgelist}")
+    print(f"Largest number of nodes: {max_nodes}")
+
+    data_set = []
+    for file_id, (edges, feature_vector) in files.items():
+        if feature_vector is None or edges is None:
+            print(f"Incomplete data for file id {file_id}, skipping.")
+            continue
+
+        # Pad node feature matrix to have max_nodes rows.
+        x = np.pad(feature_vector, ((0, max_nodes - feature_vector.shape[0]), (0, 0)),
+                   mode="constant").astype(np.float32)
+        if file_id in labels[0].values:
+            y_val = int(labels.loc[labels[0] == file_id].iloc[0, 12])
+        else:
+            print(f"Label for file id {file_id} not found; skipping.")
+            continue
+
+        if np.max(edges) >= x.shape[0]:
+            print("Edge list references a node outside the padded feature matrix; skipping.")
+            continue
+
+        print(f"Creating data sample from file id: {file_id}")
+        sample = {
+            "x": jnp.array(x),  # shape: [num_nodes, num_features]
+            "edge_index": (jnp.array(edges[0], dtype=jnp.int32),
+                           jnp.array(edges[1], dtype=jnp.int32)),
+            "y": jnp.array([y_val], dtype=jnp.int32),  # shape: [1]
+            "batch": jnp.zeros(x.shape[0], dtype=jnp.int32),  # all nodes in graph 0
+            "num_graphs": 1
+        }
+        data_set.append(sample)
+    print(f"Created dataset with {len(data_set)} samples.")
+
+    end_load = time.perf_counter()
+    print("\n\n")
+    print(f"Time to load data: {(end_load - start_load):.6f} seconds")
+
+    return data_set, num_features, num_classes
+
+
+def load_dataset_torch(paths: list[str]) -> tuple:
+    """
+    Loads the dataset using the new data loader and conversion methods.
+    Returns a DataLoader, number of features, and number of classes.
+    """
+    num_features, num_classes = 0, 2 # binary classification
+    max_edgelist, max_nodes = 0, 0
+    
+    # id: [edgelist, features]
+    files = {}
+
+    labels = None
+    for path in sorted(paths):
+        if "edgelist" in str(path).split("/")[-1]:
+            # add file id to file_ids lists
+            curr_id = int((str(path).split("/")[-1]).split("_")[0])
+            print(f"Loading edgelist from: {str(path)}")
+            edges = pd.read_csv(path, header=None).values.astype(np.int32)
+            max_edgelist = edges.shape[0] if edges.shape[0] > max_edgelist else max_edgelist
+            edge_index = torch.tensor(edges, dtype=torch.long).t().contiguous()
+            if curr_id not in files:
+                files[curr_id] = [edge_index, None]
+            else:
+                curr_edge, curr_features = files[curr_id]
+                files[curr_id] = [edge_index, curr_features]
+
+        elif "node_names" in str(path):
+            curr_id = int((str(path).split("/")[-1]).split("_")[0])
+            
+            print(f"Loading node features from: {str(path)}")
+            node_features = np.delete(np.delete(pd.read_csv(path, header=None).values, 0, axis=0), 1, axis=1).astype(np.int32)
+            if curr_id not in files:
+                files[curr_id] = [None, node_features]
+            else:
+                curr_edge, curr_features = files[curr_id]
+                files[curr_id] = [curr_edge, node_features]
+
+            num_features = node_features.shape[1] if node_features.shape[1] > num_features else num_features
+            max_nodes = node_features.shape[0] if node_features.shape[0] > max_nodes else max_nodes
+
+        elif "samplekey" in str(path):
+            print(f"Loading labels from: {str(path)}")
+            labels = pd.read_csv(path, header=None)
+        elif "removed" in str(path):
+            print(f"Disconnected Nodes File Found at: {str(path)}")
+            print(f"SKIPPING FILE: {str(path)}")
+            # TODO: do we want to feed this into the network?
+
+    print(f"Max number of features: {num_features}")
+    print(f"Number of classes: {num_classes}")
+    print(f"Largest edge_list: {max_edgelist}")
+    print(f"Largest number of nodes: {max_nodes}")
+    
+    data_set = []
+    for file_id, v in files.items():
+        edges, feature_vector = v
+
+        # pad node feature vectors with 0s
+        x = torch.tensor(np.pad(feature_vector, ((0, max_nodes - feature_vector.shape[0]), (0, 0))), dtype=torch.float)
+        if file_id in labels[0].values:
+            y = torch.tensor(labels.loc[labels[labels[0] == file_id].index[0], 12])
+        else:
+            print(f"{file_id} not found in labels")
+        
+        if torch.max(edges) >= x.shape[0]:
+            print("Edgelist is referencing a node that does not have a corresponding feature vector")
+            continue
+
+        print(f"Creating Data object from file with id: {file_id}")
+        data = Data(x=x, edge_index=edges, y=y).to(device)
+        data_set.append(data)
+
+    # Create a dataloader
+    print(f"Creating dataloader with data set of size: {len(data_set)}")
+    data_loader = DataLoader(data_set, batch_size=1, shuffle=True)
+    print(f"DataLoader created successfully with batch size: {data_loader.batch_size}")
+
+    return data_loader, num_features, num_classes
+
+def get_paths(edgelist=False) -> list[str]:
+    data_dir = Path(__file__).resolve().parent.parent / "data"
+    data_dir = data_dir / "newset" if edgelist else data_dir
+    print(f"Collecting file paths from directory: {data_dir}")
+    return list([str(path) for path in data_dir.iterdir()])
+
+def adjmat_to_edgelist(adjmat_df: np.ndarray) -> list:
+    """
+    Converts adjacency matrix to an edge list.
+    """
+    edges = []
+    print("Converting adjacency matrix to edge list...")
+    for i in range(adjmat_df.shape[0]):
+        for j in range(adjmat_df.shape[1]):
+            if adjmat_df[i, j] != 0:
+                edges.append([i, j])
+    print("Edge list created with", len(edges), "edges.")
+    return edges
+
+def save_edgelist(edge_list: list[list[int]], file_path: str) -> None:
+    """
+    Saves the edge list to a CSV file.
+    """
+    print("Saving edge list to CSV file...")
+    df = pd.DataFrame(edge_list, columns=["Source", "Target"])
+    df.to_csv(file_path, index=False)
+    print(f"Edge list saved to {file_path}")
+
+def convert_data_to_edgelists(paths: list[str]):
+    edgelist_dir = Path(__file__).resolve().parent.parent / "data" / "edgelists"
+    for path in paths:
+        adjmat = pd.read_csv(path, header=None).values
+        print(f"Converting adjacency matrix for path: {path}")
+        edge_list = adjmat_to_edgelist(adjmat)
+        file_name = Path(path).name
+        save_edgelist(edge_list, edgelist_dir / file_name)
+
+def convert_data_to_edgelists():
+    paths = get_paths()
+    edgelist_dir = Path(__file__).resolve().parent.parent / "data" / "edgelists"
+    edgelist_dir.mkdir(parents=True, exist_ok=True) 
+    for path in paths:
+        if "adjacency_matrix" in str(path):
+            adjmat = pd.read_csv(path, header=None).values
+            print(f"Converting adjacency matrix for path: {path} to edgelist")
+            edge_list = adjmat_to_edgelist(adjmat)
+            file_name = "".join([i for i in Path(path).stem if i in "1234567890"]) + "_edgelist" + ".csv"
+            save_edgelist(edge_list, edgelist_dir / file_name)
+        else:
+            print(f"Copying non-adjacency matrix file: {path}")
+            shutil.copy(path, edgelist_dir / Path(path).name)
+
+if __name__ == "__main__":
+    """
+    NOTE: the function call below will attempt to read adjmat files and convert them to edgelists
+    It should be used to generate a few edgelist for testing the load_dataset_new function
+    """
+    # convert_data_to_edgelists()
+    """
+    The function call below reads the files in the edgelist directory and constructs a dataloader
+    """
+    # load_dataset_new(get_paths(edgelist=True))
+    pass
diff --git a/GraphIsomorphismNetwork/new_loader.py b/GraphIsomorphismNetwork/new_loader.py
new file mode 100644
index 0000000..a6b15a8
--- /dev/null
+++ b/GraphIsomorphismNetwork/new_loader.py
@@ -0,0 +1,249 @@
+import time
+# import torch
+import shutil
+import numpy as np
+from pathlib import Path
+import pandas as pd
+import jax.numpy as jnp
+# from torch.utils import data
+# from torch_geometric.data import Data, DataLoader
+import tqdm
+
+# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+def load_dataset_jax(paths: list[str]) -> tuple:
+    start_load = time.perf_counter()
+
+    """
+    Loads the dataset and returns a list of samples (each a dict with keys:
+    "x", "edge_index", "batch", "y", and "num_graphs"),
+    along with the number of features and number of classes.
+
+    This version is adapted for use with JAX/Flax.
+    """
+    num_features, num_classes = 0, 2
+    max_edgelist, max_nodes = 0, 0
+
+    # Dictionary to hold data for each file id: {file_id: [edge_index, node_features]}
+    files = {}
+    labels = None
+    for path in sorted(paths):
+        fname = Path(path).name
+        if "edgelist" in fname:
+            # Expect file names like "<id>_edgelist..."
+            curr_id = int(fname.split("_")[0])
+            print(f"Loading edgelist from: {path}")
+            edges = pd.read_csv(path, header=None).values.astype(np.int32)
+            max_edgelist = edges.shape[0] if edges.shape[0] > max_edgelist else max_edgelist
+            edge_index = edges.T  # shape: (2, num_edges)
+            if curr_id not in files:
+                files[curr_id] = [edge_index, None]
+            else:
+                _, curr_features = files[curr_id]
+                files[curr_id] = [edge_index, curr_features]
+        elif "node_names" in fname:
+            curr_id = int(fname.split("_")[0])
+            print(f"Loading node features from: {path}")
+            node_features = pd.read_csv(path, header=None).values
+            # Remove first row and second column as in your original code.
+            node_features = np.delete(np.delete(node_features, 0, axis=0), 1, axis=1).astype(np.int32)
+            if curr_id not in files:
+                files[curr_id] = [None, node_features]
+            else:
+                curr_edge, _ = files[curr_id]
+                files[curr_id] = [curr_edge, node_features]
+            # num_features = max(num_features, node_features.shape[1])
+            num_features = max(num_features, node_features.shape[0])
+            max_nodes = max(max_nodes, node_features.shape[0])
+        elif "samplekey" in fname:
+            print(f"Loading labels from: {path}")
+            labels = pd.read_csv(path, header=None)
+        elif "removed" in fname:
+            print(f"Skipping removed file: {path}")
+
+    print(f"Max number of features: {num_features}")
+    print(f"Number of classes: {num_classes}")
+    print(f"Largest edgelist (number of edges): {max_edgelist}")
+    print(f"Largest number of nodes: {max_nodes}")
+
+    data_set = []
+    for file_id, (edges, feature_vector) in files.items():
+        if feature_vector is None or edges is None:
+            print(f"Incomplete data for file id {file_id}, skipping.")
+            continue
+
+        # Pad node feature matrix to have max_nodes rows.
+        x = np.pad(feature_vector, ((0, max_nodes - feature_vector.shape[0]), (0, 0)),
+                   mode="constant").astype(np.float32)
+        if file_id in labels[1].values:
+            y_val = int(labels.loc[labels[1] == file_id].iloc[0, 12])
+        else:
+            print(f"Label for file id {file_id} not found; skipping.")
+            continue
+
+        if np.max(edges) >= x.shape[0]:
+            print("Edge list references a node outside the padded feature matrix; skipping.")
+            continue
+
+        print(f"Creating data sample from file id: {file_id}")
+        sample = {
+            "x": jnp.array(x),  # shape: [num_nodes, num_features]
+            "edge_index": (jnp.array(edges[0], dtype=jnp.int32),
+                           jnp.array(edges[1], dtype=jnp.int32)),
+            "y": jnp.array([y_val], dtype=jnp.int32),  # shape: [1]
+            "batch": jnp.zeros(x.shape[0], dtype=jnp.int32),  # all nodes in graph 0
+            "num_graphs": 1
+        }
+        data_set.append(sample)
+    print(f"Created dataset with {len(data_set)} samples.")
+
+    end_load = time.perf_counter()
+    print("\n\n")
+    print(f"Time to load data: {(end_load - start_load):.6f} seconds")
+
+    return data_set, num_features, num_classes
+
+def load_dataset_jax_new(paths: list[str], max_files=20) -> tuple:
+    start_load = time.perf_counter()
+    
+    paths = [Path(path) for path in paths] # NOTE: this could just be done in the get_paths function
+    paths_set = set([path.name for path in paths]) # for quick verification that all 3 files for any given id exist
+    
+    labels_path = None
+    if max_files is None:
+        max_files = float('inf')
+
+    ids = {} # {id: [node_names_path, edgelist_path, block_lens_path]}
+    num_files = 0
+    for path in paths:
+        file_name = path.name
+    
+        if "sample" in file_name:
+            labels_path = path
+        
+        if "edgelist.csv" in file_name and num_files < max_files:
+            #print(file_name)
+            file_id = int(file_name.split("/")[-1].split("_")[0])
+            if (file_id not in ids) and (str(file_id) + "_edgelist.csv") in paths_set and (str(file_id) + "_block_lens.csv") in paths_set:
+                base_path = path.parents[0]
+                ids[file_id] = [base_path / file_name, base_path / (str(file_id) + "_edgelist.csv"), base_path / (str(file_id) + "_block_lens.csv")]
+                num_files += 1
+            # if num_files >= max_files: # NOTE: if you uncomment this line you will have to add a line of code to set labels_path manually
+                # break
+    
+    samples = {} # {id: [node_features, edges]}
+    max_nodes, max_edges, max_blocks = 0, 0, 0
+    num_classes = 2
+    for file_id, (node_names_path, edgelist_path, block_lens_path) in tqdm.tqdm(ids.items(), desc="Building Dataset"):
+        # print("#"*30)
+        bp = str(block_lens_path)
+        bp = bp.replace("out_edgelists", "out_file_feats")
+        block_lens_path = Path(bp)
+        
+        # print(node_names_path, edgelist_path, block_lens_path)
+
+        # node_features = pd.read_csv(node_names_path, header=None, low_memory=False).values
+        # print("read feats")
+        # node_features = np.delete(np.delete(node_features, 0, axis=0), 1, axis=1).astype(np.int32) # shape: num_nodes x 1 data: (node_id)
+        # print(node_features.shape)
+        # max_nodes = max(max_nodes, node_features.shape[0])
+
+        I32_MAX = (2**(31)) - 1
+        
+        block_lens = pd.read_csv(block_lens_path, header=None, low_memory=False).values
+        # print("read blocklens")
+        block_lens = np.delete(block_lens, 0, axis=0)
+        block_lens = block_lens.astype(np.uint64) # shape: num_blocks x 2 data: (node_id, num_instructions)
+
+        block_lens = block_lens % I32_MAX
+        block_lens = block_lens.astype(np.int32)
+        
+        # print(np.max(block_lens, axis=0)[1], np.min(block_lens, axis=0)[1])
+
+        # print(block_lens.shape)
+        
+        edges = None
+        try:
+            edges = pd.read_csv(edgelist_path, header=None) # shape: num_edges x 2 data: (node_id, node_id)
+        except:
+            # most likely "No columns to parse"; single big block in file i think (or series of big blocks)
+            # introduce series of self-edges
+            edges = pd.DataFrame([[node,node] for node,_ in block_lens])
+
+        edges = edges % I32_MAX
+        edges = edges.astype(np.int32).values
+        
+        # node_features = np.delete(np.delete(node_features, 0, axis=0), 1, axis=1).astype(np.int32) # shape: num_nodes x 1 data: (node_id)
+        # print(node_features.shape)
+        # max_nodes = max(max_nodes, node_features.shape[0])
+            
+        # print(edges.shape)
+        # print("read edges")
+        # max_edges = max(max_edges, edges.shape[0])
+        # samples[file_id] = [node_features, edges]
+
+        # if node_features.shape[0] != block_lens.shape[0]:
+        #     print(f"node_features shape {node_features.shape} is not equal to block_lens shape {block_lens.shape}, file_id {file_id}")
+
+        samples[file_id] = [block_lens, edges]
+            
+    print(labels_path)
+    print(f"{len(samples.items())} files in sample dict")
+    labels = pd.read_csv(labels_path)
+    
+    dataset = []
+    
+    max_nodes, max_edges, max_blocks = 0, 0, 0
+
+    for file_id, (features, edges) in tqdm.tqdm(samples.items(), desc="Dataset Sample Adding"):
+        matching_row = labels[labels['id'] == file_id]
+        if not matching_row.empty:
+            # print(matching_row)
+            # file_label = matching_row.iloc[1, 7]
+            file_label = matching_row.iloc[0]['list']
+            # print(f"File id: {file_id} has label {file_label}")
+            y = 0 if file_label == "Whitelist" else 1
+            max_edges = max(max_edges, edges.shape[0])
+    
+            max_blocks = max(max_blocks, features.shape[0])
+            max_nodes = max_blocks # each block should be distinct node
+    
+            # TODO: do we even need to pad?
+            # x = np.pad(features, ((0, max_nodes - features.shape[0]), (0, 0)), mode="constant").astype(np.float32)
+            x = features
+            #print(file_id)
+            # print(edges)
+            #second_idx = 1 if len(edges) > 1 else 0
+            # print(type(edges))
+            # print(edges.shape)
+            curr_sample = {
+                "x": jnp.array(x),
+                "edge_index": (jnp.array(edges[:,0], dtype=jnp.int32), jnp.array(edges[:,1], dtype=jnp.int32)),
+                "y": jnp.array([y], dtype=jnp.int32),
+                "batch": jnp.zeros(x.shape[0], dtype=jnp.int32),
+                "num_graphs": 1
+            }
+    
+            dataset.append(curr_sample)
+
+        else:
+            # print(f"File id: {file_id} does not have a label. Skipping.")
+            continue
+
+        
+    print(f"Created dataset with {len(dataset)} samples.")
+
+    end_load = time.perf_counter()
+    print("\n\n")
+    print(f"Time to load data: {(end_load - start_load):.6f} seconds")
+    print(max_nodes, max_edges, max_blocks)
+
+
+    return dataset, dataset[0]["x"].shape[1], num_classes
+
+
+def get_paths(samples_2000=False) -> list[str]:
+    data_dir = Path(__file__).resolve().parent.parent / "out"
+    #data_dir = data_dir / "2kds" / "found_files" if samples_2000 else data_dir
+    print(f"Collecting file paths from directory: {data_dir}")
+    return list([str(path) for path in data_dir.rglob("*")])
\ No newline at end of file
diff --git a/GraphIsomorphismNetwork/old_gin.py b/GraphIsomorphismNetwork/old_gin.py
new file mode 100644
index 0000000..76099e3
--- /dev/null
+++ b/GraphIsomorphismNetwork/old_gin.py
@@ -0,0 +1,254 @@
+import jax
+import jax.numpy as jnp
+from flax import linen as nn
+from flax.training import train_state, checkpoints
+# from load import get_paths, load_dataset_jax
+from new_loader import get_paths, load_dataset_jax_new
+import optax
+from pathlib import Path
+import pandas as pd
+import numpy as np
+import time
+import os
+import logging
+
+log_file = Path(__file__).resolve().parent.parent / "logs" / "gin_training.log"
+log_file.parent.mkdir(parents=True, exist_ok=True)
+
+root_logger = logging.getLogger()
+if root_logger.handlers:
+    for handler in root_logger.handlers[:]:
+        root_logger.removeHandler(handler)
+
+logging.basicConfig(
+    filename=log_file,
+    level=logging.INFO,
+    filemode="a",
+    format="%(asctime)s - %(levelname)s - %(message)s",
+    datefmt='%Y-%m-%d %H:%M:%S'
+)
+
+logging.info("Starting GIN training...")
+
+# Dense -> BatchNorm -> ReLU -> Dense.
+class MLP(nn.Module):
+    hidden_dim: int
+
+    @nn.compact
+    def __call__(self, x, training: bool):
+        x = nn.Dense(2 * self.hidden_dim)(x)
+        x = nn.BatchNorm(use_running_average=not training)(x)
+        x = nn.relu(x)
+        x = nn.Dense(self.hidden_dim)(x)
+        return x
+
+# Define a GIN convolution layer: x_i' = MLP((1+eps)*x_i + sum_{j in N(i)} x_j)
+class GINConv(nn.Module):
+    hidden_dim: int
+    train_eps: bool = True
+
+    @nn.compact
+    def __call__(self, x, senders, receivers, training: bool):
+        mlp = MLP(hidden_dim=self.hidden_dim)
+        # Learnable epsilon
+        if self.train_eps:
+            eps = self.param("eps", lambda rng: jnp.zeros(()))
+        else:
+            eps = 0.0
+        # Aggregate neighbor features using segment_sum.
+        aggregated = jax.ops.segment_sum(x[senders], receivers, num_segments=x.shape[0])
+        out = mlp((1 + eps) * x + aggregated, training=training)
+        return out
+
+class GIN(nn.Module):
+    in_channels: int
+    hidden_channels: int
+    out_channels: int
+    num_layers: int
+    dropout_rate: float = 0.5
+    train_eps: bool = True
+
+    @nn.compact
+    def __call__(self, x, edge_index, batch, training: bool):
+        senders, receivers = edge_index  # edge_index is a tuple: (senders, receivers)
+        for _ in range(self.num_layers):
+            x = GINConv(hidden_dim=self.hidden_channels, train_eps=self.train_eps)(
+                x, senders, receivers, training=training
+            )
+            x = nn.BatchNorm(use_running_average=not training)(x)
+            x = nn.relu(x)
+
+        # Global add pooling: sum node features per graph.
+        x = jax.ops.segment_sum(x, batch, num_segments=1)
+
+        # Two-layer MLP for graph-level output.
+        x = nn.Dense(self.hidden_channels)(x)
+        x = nn.LayerNorm()(x)
+        x = nn.relu(x)
+        x = nn.Dropout(rate=self.dropout_rate)(x, deterministic=not training)
+        x = nn.Dense(self.out_channels)(x)
+        return jax.nn.log_softmax(x, axis=-1)
+
+class TrainState(train_state.TrainState):
+    batch_stats: dict
+
+def create_train_state(rng, model, learning_rate, sample_input):
+    variables = model.init(
+        rng,
+        sample_input["x"],
+        sample_input["edge_index"],
+        sample_input["batch"],
+        training=True,
+    )
+    tx = optax.adam(learning_rate)
+    return TrainState.create(
+        apply_fn=model.apply,
+        params=variables["params"],
+        tx=tx,
+        batch_stats=variables.get("batch_stats", {})
+    )
+
+@jax.jit
+def train_step(state, batch, dropout_rng):
+    def loss_fn(params):
+        variables = {"params": params, "batch_stats": state.batch_stats}
+        logits, new_model_state = state.apply_fn(
+            variables,
+            batch["x"],
+            batch["edge_index"],
+            batch["batch"],
+            training=True,
+            mutable=["batch_stats"],
+            rngs={"dropout": dropout_rng},
+        )
+        labels = batch["y"]
+        nll = -jnp.mean(jnp.take_along_axis(logits, labels[:, None], axis=-1).squeeze())
+        return nll, new_model_state
+
+    grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
+    (loss, new_model_state), grads = grad_fn(state.params)
+    state = state.apply_gradients(grads=grads, batch_stats=new_model_state["batch_stats"])
+    return state, loss
+
+@jax.jit
+def test_step(state, batch):
+    variables = {"params": state.params, "batch_stats": state.batch_stats}
+    logits = state.apply_fn(
+        variables,
+        batch["x"],
+        batch["edge_index"],
+        batch["batch"],
+        training=False,
+        mutable=False
+    )
+    pred = jnp.argmax(logits, axis=-1)
+    correct = jnp.sum(pred == batch["y"])
+    return correct
+
+def save_model(state, checkpoint_dir, step):
+    os.makedirs(checkpoint_dir, exist_ok=True)
+    checkpoints.save_checkpoint(
+        ckpt_dir=checkpoint_dir,
+        target=state,
+        step=step,
+        overwrite=True
+    )
+    msg = f"Model saved at step {step} to {checkpoint_dir}"
+    print(msg)
+    logging.info(msg)
+
+def load_model_if_exists(state, checkpoint_dir):
+    latest_ckpt = checkpoints.latest_checkpoint(checkpoint_dir)
+    if latest_ckpt:
+        state = checkpoints.restore_checkpoint(
+            ckpt_dir=checkpoint_dir,
+            target=state
+        )
+        msg = f"Model restored from checkpoint: {latest_ckpt}"
+        print(msg)
+        logging.info(msg)
+    else:
+        msg = "No checkpoint found. Training from scratch."
+        print(msg)
+        logging.info(msg)
+    return state
+
+def main():
+    # paths = get_paths(edgelist=True)
+    # data_loader, num_features, num_classes = load_dataset_jax(paths)
+
+    paths = get_paths(samples_2000=True)
+    data_loader, num_features, num_classes = load_dataset_jax_new(paths, max_files=20)
+    
+    """
+    model = GIN(
+        in_channels=num_features,
+        hidden_channels=int(num_features * 1.5),
+        out_channels=num_classes,
+        num_layers=4,
+        dropout_rate=0.5
+    )
+    """ 
+
+    print(f"Number of Features: {num_features}, model hidden layer: {int(num_features * 1e-4)}")
+    model = GIN(
+        in_channels=num_features,
+        hidden_channels=int(num_features * 2e-4),
+        out_channels=num_classes,
+        num_layers=4,
+        dropout_rate=0.5
+    )
+
+    rng = jax.random.PRNGKey(0)
+    dropout_rng, init_rng = jax.random.split(rng)
+
+    sample_input = data_loader[0]
+    state = create_train_state(init_rng, model, learning_rate=0.01, sample_input=sample_input)
+
+    checkpoint_dir = Path(__file__).resolve().parent.parent / "weights"
+
+    state = load_model_if_exists(state, checkpoint_dir)
+
+    start_train = time.perf_counter()
+
+    num_epochs = 10
+    for epoch in range(1, num_epochs + 1):
+        epoch_start = time.perf_counter()
+
+        epoch_loss = 0.0
+        total_graphs = 0
+        for batch in data_loader:
+            dropout_rng, new_dropout_rng = jax.random.split(dropout_rng)
+            state, loss = train_step(state, batch, dropout_rng)
+            epoch_loss += loss * batch["num_graphs"]
+            total_graphs += batch["num_graphs"]
+        avg_loss = epoch_loss / total_graphs
+
+        total_correct = 0
+        total_samples = 0
+        for batch in data_loader:
+            correct = test_step(state, batch)
+            total_correct += correct
+            total_samples += batch["num_graphs"]
+        train_acc = total_correct / total_samples
+
+        epoch_end = time.perf_counter()
+        msg1 = f"Time for epoch {epoch} was {(epoch_end - epoch_start):.6f} seconds"
+        msg2 = f"Epoch: {epoch:03d}, Loss: {avg_loss:.4f}, Train Acc: {train_acc:.4f}"
+        print(msg1)
+        print(msg2)
+        logging.info(msg1)
+        logging.info(msg2)
+
+        save_model(state, checkpoint_dir, epoch)
+
+    end_train = time.perf_counter()
+    msg3 = f"Total training time : {(end_train - start_train):.6f} seconds"
+    msg4 = f"Average time per epoch: {((end_train - start_train) / num_epochs):.6f} seconds"
+    print(msg3)
+    print(msg4)
+    logging.info(msg3)
+    logging.info(msg4)
+
+if __name__ == "__main__":
+    main()
diff --git a/poetry.lock b/poetry.lock
index 7304519..6f6efc8 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -1,4 +1,16 @@
-# This file is automatically @generated by Poetry 1.8.3 and should not be changed by hand.
+# This file is automatically @generated by Poetry 2.1.3 and should not be changed by hand.
+
+[[package]]
+name = "absl-py"
+version = "2.2.0"
+description = "Abseil Python Common Libraries, see https://github.com/abseil/abseil-py."
+optional = false
+python-versions = ">=3.8"
+groups = ["main"]
+files = [
+    {file = "absl_py-2.2.0-py3-none-any.whl", hash = "sha256:5c432cdf7b045f89c4ddc3bba196cabb389c0c321322f8dec68eecdfa732fdad"},
+    {file = "absl_py-2.2.0.tar.gz", hash = "sha256:2aabeae1403380e338fba88d4f8c9bf9925c20ad04c1c96d4a26930d034c507b"},
+]
 
 [[package]]
 name = "asttokens"
@@ -6,6 +18,7 @@ version = "2.4.1"
 description = "Annotate AST trees with source code positions"
 optional = false
 python-versions = "*"
+groups = ["main"]
 files = [
     {file = "asttokens-2.4.1-py2.py3-none-any.whl", hash = "sha256:051ed49c3dcae8913ea7cd08e46a606dba30b79993209636c4875bc1d637bc24"},
     {file = "asttokens-2.4.1.tar.gz", hash = "sha256:b03869718ba9a6eb027e134bfdf69f38a236d681c83c160d510768af11254ba0"},
@@ -15,8 +28,29 @@ files = [
 six = ">=1.12.0"
 
 [package.extras]
-astroid = ["astroid (>=1,<2)", "astroid (>=2,<4)"]
-test = ["astroid (>=1,<2)", "astroid (>=2,<4)", "pytest"]
+astroid = ["astroid (>=1,<2) ; python_version < \"3\"", "astroid (>=2,<4) ; python_version >= \"3\""]
+test = ["astroid (>=1,<2) ; python_version < \"3\"", "astroid (>=2,<4) ; python_version >= \"3\"", "pytest"]
+
+[[package]]
+name = "chex"
+version = "0.1.89"
+description = "Chex: Testing made fun, in JAX!"
+optional = false
+python-versions = ">=3.9"
+groups = ["main"]
+files = [
+    {file = "chex-0.1.89-py3-none-any.whl", hash = "sha256:145241c27d8944adb634fb7d472a460e1c1b643f561507d4031ad5156ef82dfa"},
+    {file = "chex-0.1.89.tar.gz", hash = "sha256:78f856e6a0a8459edfcbb402c2c044d2b8102eac4b633838cbdfdcdb09c6c8e0"},
+]
+
+[package.dependencies]
+absl-py = ">=0.9.0"
+jax = ">=0.4.27"
+jaxlib = ">=0.4.27"
+numpy = ">=1.24.1"
+setuptools = {version = "*", markers = "python_version >= \"3.12\""}
+toolz = ">=0.9.0"
+typing_extensions = ">=4.2.0"
 
 [[package]]
 name = "click"
@@ -24,6 +58,7 @@ version = "8.1.7"
 description = "Composable command line interface toolkit"
 optional = false
 python-versions = ">=3.7"
+groups = ["main"]
 files = [
     {file = "click-8.1.7-py3-none-any.whl", hash = "sha256:ae74fb96c20a0277a1d615f1e4d73c8414f5a98db8b799a7931d1582f3390c28"},
     {file = "click-8.1.7.tar.gz", hash = "sha256:ca9853ad459e787e2192211578cc907e7594e294c7ccc834310722b41b9ca6de"},
@@ -38,6 +73,8 @@ version = "0.4.6"
 description = "Cross-platform colored terminal text."
 optional = false
 python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*,>=2.7"
+groups = ["main"]
+markers = "platform_system == \"Windows\" or sys_platform == \"win32\""
 files = [
     {file = "colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6"},
     {file = "colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44"},
@@ -49,6 +86,7 @@ version = "1.3.0"
 description = "Python library for calculating contours of 2D quadrilateral grids"
 optional = false
 python-versions = ">=3.9"
+groups = ["main"]
 files = [
     {file = "contourpy-1.3.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:880ea32e5c774634f9fcd46504bf9f080a41ad855f4fef54f5380f5133d343c7"},
     {file = "contourpy-1.3.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:76c905ef940a4474a6289c71d53122a4f77766eef23c03cd57016ce19d0f7b42"},
@@ -133,6 +171,7 @@ version = "0.12.1"
 description = "Composable style cycles"
 optional = false
 python-versions = ">=3.8"
+groups = ["main"]
 files = [
     {file = "cycler-0.12.1-py3-none-any.whl", hash = "sha256:85cef7cff222d8644161529808465972e51340599459b8ac3ccbac5a854e0d30"},
     {file = "cycler-0.12.1.tar.gz", hash = "sha256:88bb128f02ba341da8ef447245a9e138fae777f6a23943da4540077d3601eb1c"},
@@ -148,24 +187,111 @@ version = "5.1.1"
 description = "Decorators for Humans"
 optional = false
 python-versions = ">=3.5"
+groups = ["main"]
 files = [
     {file = "decorator-5.1.1-py3-none-any.whl", hash = "sha256:b8c3f85900b9dc423225913c5aace94729fe1fa9763b38939a95226f02d37186"},
     {file = "decorator-5.1.1.tar.gz", hash = "sha256:637996211036b6385ef91435e4fae22989472f9d571faba8927ba8253acbc330"},
 ]
 
+[[package]]
+name = "etils"
+version = "1.12.2"
+description = "Collection of common python utils"
+optional = false
+python-versions = ">=3.10"
+groups = ["main"]
+files = [
+    {file = "etils-1.12.2-py3-none-any.whl", hash = "sha256:4600bec9de6cf5cb043a171e1856e38b5f273719cf3ecef90199f7091a6b3912"},
+    {file = "etils-1.12.2.tar.gz", hash = "sha256:c6b9e1f0ce66d1bbf54f99201b08a60ba396d3446d9eb18d4bc39b26a2e1a5ee"},
+]
+
+[package.dependencies]
+fsspec = {version = "*", optional = true, markers = "extra == \"epath\""}
+importlib_resources = {version = "*", optional = true, markers = "extra == \"epath\""}
+typing_extensions = {version = "*", optional = true, markers = "extra == \"epath\" or extra == \"epy\""}
+zipp = {version = "*", optional = true, markers = "extra == \"epath\""}
+
+[package.extras]
+all = ["etils[array-types]", "etils[eapp]", "etils[ecolab]", "etils[edc]", "etils[enp]", "etils[epath-gcs]", "etils[epath-s3]", "etils[epath]", "etils[epy]", "etils[etqdm]", "etils[etree-dm]", "etils[etree-jax]", "etils[etree-tf]", "etils[etree]"]
+array-types = ["etils[enp]"]
+dev = ["chex", "fiddle", "optree", "pydantic", "pyink", "pylint (>=2.6.0)", "pytest", "pytest-subtests", "pytest-xdist", "tensorflow_datasets", "torch"]
+docs = ["etils[all,dev]", "sphinx-apitree[ext]"]
+eapp = ["absl-py", "etils[epy]", "simple_parsing"]
+ecolab = ["etils[enp]", "etils[epy]", "etils[etree]", "jupyter", "mediapy", "numpy", "packaging", "protobuf"]
+edc = ["etils[epy]"]
+enp = ["einops", "etils[epy]", "numpy"]
+epath = ["etils[epy]", "fsspec", "importlib_resources", "typing_extensions", "zipp"]
+epath-gcs = ["etils[epath]", "gcsfs"]
+epath-s3 = ["etils[epath]", "s3fs"]
+epy = ["typing_extensions"]
+etqdm = ["absl-py", "etils[epy]", "tqdm"]
+etree = ["etils[array-types]", "etils[enp]", "etils[epy]", "etils[etqdm]"]
+etree-dm = ["dm-tree", "etils[etree]"]
+etree-jax = ["etils[etree]", "jax[cpu]"]
+etree-tf = ["etils[etree]", "tensorflow"]
+lazy-imports = ["etils[ecolab]"]
+
 [[package]]
 name = "executing"
 version = "2.1.0"
 description = "Get the currently executing AST node of a frame, and other information"
 optional = false
 python-versions = ">=3.8"
+groups = ["main"]
 files = [
     {file = "executing-2.1.0-py2.py3-none-any.whl", hash = "sha256:8d63781349375b5ebccc3142f4b30350c0cd9c79f921cde38be2be4637e98eaf"},
     {file = "executing-2.1.0.tar.gz", hash = "sha256:8ea27ddd260da8150fa5a708269c4a10e76161e2496ec3e587da9e3c0fe4b9ab"},
 ]
 
 [package.extras]
-tests = ["asttokens (>=2.1.0)", "coverage", "coverage-enable-subprocess", "ipython", "littleutils", "pytest", "rich"]
+tests = ["asttokens (>=2.1.0)", "coverage", "coverage-enable-subprocess", "ipython", "littleutils", "pytest", "rich ; python_version >= \"3.11\""]
+
+[[package]]
+name = "filelock"
+version = "3.18.0"
+description = "A platform independent file lock."
+optional = false
+python-versions = ">=3.9"
+groups = ["main"]
+files = [
+    {file = "filelock-3.18.0-py3-none-any.whl", hash = "sha256:c401f4f8377c4464e6db25fff06205fd89bdd83b65eb0488ed1b160f780e21de"},
+    {file = "filelock-3.18.0.tar.gz", hash = "sha256:adbc88eabb99d2fec8c9c1b229b171f18afa655400173ddc653d5d01501fb9f2"},
+]
+
+[package.extras]
+docs = ["furo (>=2024.8.6)", "sphinx (>=8.1.3)", "sphinx-autodoc-typehints (>=3)"]
+testing = ["covdefaults (>=2.3)", "coverage (>=7.6.10)", "diff-cover (>=9.2.1)", "pytest (>=8.3.4)", "pytest-asyncio (>=0.25.2)", "pytest-cov (>=6)", "pytest-mock (>=3.14)", "pytest-timeout (>=2.3.1)", "virtualenv (>=20.28.1)"]
+typing = ["typing-extensions (>=4.12.2) ; python_version < \"3.11\""]
+
+[[package]]
+name = "flax"
+version = "0.10.4"
+description = "Flax: A neural network library for JAX designed for flexibility"
+optional = false
+python-versions = ">=3.10"
+groups = ["main"]
+files = [
+    {file = "flax-0.10.4-py3-none-any.whl", hash = "sha256:8cc83d91654ff943909730e02e858b4cd4577531373f83abe6597c58c581032d"},
+    {file = "flax-0.10.4.tar.gz", hash = "sha256:57ae44d3f111fc85cff9049adb9684ce8ebd44e87bd8ca776ed52422c2d85021"},
+]
+
+[package.dependencies]
+jax = ">=0.4.27"
+msgpack = "*"
+numpy = {version = ">=1.26.0", markers = "python_version >= \"3.12\""}
+optax = "*"
+orbax-checkpoint = "*"
+PyYAML = ">=5.4.1"
+rich = ">=11.1"
+tensorstore = "*"
+treescope = ">=0.1.7"
+typing_extensions = ">=4.2"
+
+[package.extras]
+all = ["matplotlib"]
+dev = ["pre-commit (>=3.8.0)"]
+docs = ["Pygments (>=2.6.1)", "dm-haiku", "docutils (==0.16)", "einops", "ipykernel", "ipython_genutils", "ipywidgets (>=8.1.5)", "jupytext (==1.13.8)", "kagglehub (>=0.3.3)", "matplotlib", "ml_collections", "myst_nb", "nbstripout", "recommonmark", "scikit-learn", "sphinx (>=3.3.1)", "sphinx-book-theme", "sphinx-design"]
+testing = ["cloudpickle (>=3.0.0)", "clu", "clu (<=0.0.9) ; python_version < \"3.10\"", "einops", "gymnasium[accept-rom-license,atari]", "jaxlib", "jaxtyping", "jraph (>=0.0.6dev0)", "ml-collections", "mypy", "opencv-python", "pytest", "pytest-cov", "pytest-custom_exit_code", "pytest-xdist", "pytype", "sentencepiece", "tensorflow (>=2.12.0)", "tensorflow_datasets", "tensorflow_text (>=2.11.0) ; platform_system != \"Darwin\"", "torch", "treescope (>=0.1.1) ; python_version >= \"3.10\""]
 
 [[package]]
 name = "fonttools"
@@ -173,6 +299,7 @@ version = "4.53.1"
 description = "Tools to manipulate font files"
 optional = false
 python-versions = ">=3.8"
+groups = ["main"]
 files = [
     {file = "fonttools-4.53.1-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:0679a30b59d74b6242909945429dbddb08496935b82f91ea9bf6ad240ec23397"},
     {file = "fonttools-4.53.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:e8bf06b94694251861ba7fdeea15c8ec0967f84c3d4143ae9daf42bbc7717fe3"},
@@ -219,18 +346,73 @@ files = [
 ]
 
 [package.extras]
-all = ["brotli (>=1.0.1)", "brotlicffi (>=0.8.0)", "fs (>=2.2.0,<3)", "lxml (>=4.0)", "lz4 (>=1.7.4.2)", "matplotlib", "munkres", "pycairo", "scipy", "skia-pathops (>=0.5.0)", "sympy", "uharfbuzz (>=0.23.0)", "unicodedata2 (>=15.1.0)", "xattr", "zopfli (>=0.1.4)"]
+all = ["brotli (>=1.0.1) ; platform_python_implementation == \"CPython\"", "brotlicffi (>=0.8.0) ; platform_python_implementation != \"CPython\"", "fs (>=2.2.0,<3)", "lxml (>=4.0)", "lz4 (>=1.7.4.2)", "matplotlib", "munkres ; platform_python_implementation == \"PyPy\"", "pycairo", "scipy ; platform_python_implementation != \"PyPy\"", "skia-pathops (>=0.5.0)", "sympy", "uharfbuzz (>=0.23.0)", "unicodedata2 (>=15.1.0) ; python_version <= \"3.12\"", "xattr ; sys_platform == \"darwin\"", "zopfli (>=0.1.4)"]
 graphite = ["lz4 (>=1.7.4.2)"]
-interpolatable = ["munkres", "pycairo", "scipy"]
+interpolatable = ["munkres ; platform_python_implementation == \"PyPy\"", "pycairo", "scipy ; platform_python_implementation != \"PyPy\""]
 lxml = ["lxml (>=4.0)"]
 pathops = ["skia-pathops (>=0.5.0)"]
 plot = ["matplotlib"]
 repacker = ["uharfbuzz (>=0.23.0)"]
 symfont = ["sympy"]
-type1 = ["xattr"]
+type1 = ["xattr ; sys_platform == \"darwin\""]
 ufo = ["fs (>=2.2.0,<3)"]
-unicode = ["unicodedata2 (>=15.1.0)"]
-woff = ["brotli (>=1.0.1)", "brotlicffi (>=0.8.0)", "zopfli (>=0.1.4)"]
+unicode = ["unicodedata2 (>=15.1.0) ; python_version <= \"3.12\""]
+woff = ["brotli (>=1.0.1) ; platform_python_implementation == \"CPython\"", "brotlicffi (>=0.8.0) ; platform_python_implementation != \"CPython\"", "zopfli (>=0.1.4)"]
+
+[[package]]
+name = "fsspec"
+version = "2025.3.0"
+description = "File-system specification"
+optional = false
+python-versions = ">=3.8"
+groups = ["main"]
+files = [
+    {file = "fsspec-2025.3.0-py3-none-any.whl", hash = "sha256:efb87af3efa9103f94ca91a7f8cb7a4df91af9f74fc106c9c7ea0efd7277c1b3"},
+    {file = "fsspec-2025.3.0.tar.gz", hash = "sha256:a935fd1ea872591f2b5148907d103488fc523295e6c64b835cfad8c3eca44972"},
+]
+
+[package.extras]
+abfs = ["adlfs"]
+adl = ["adlfs"]
+arrow = ["pyarrow (>=1)"]
+dask = ["dask", "distributed"]
+dev = ["pre-commit", "ruff"]
+doc = ["numpydoc", "sphinx", "sphinx-design", "sphinx-rtd-theme", "yarl"]
+dropbox = ["dropbox", "dropboxdrivefs", "requests"]
+full = ["adlfs", "aiohttp (!=4.0.0a0,!=4.0.0a1)", "dask", "distributed", "dropbox", "dropboxdrivefs", "fusepy", "gcsfs", "libarchive-c", "ocifs", "panel", "paramiko", "pyarrow (>=1)", "pygit2", "requests", "s3fs", "smbprotocol", "tqdm"]
+fuse = ["fusepy"]
+gcs = ["gcsfs"]
+git = ["pygit2"]
+github = ["requests"]
+gs = ["gcsfs"]
+gui = ["panel"]
+hdfs = ["pyarrow (>=1)"]
+http = ["aiohttp (!=4.0.0a0,!=4.0.0a1)"]
+libarchive = ["libarchive-c"]
+oci = ["ocifs"]
+s3 = ["s3fs"]
+sftp = ["paramiko"]
+smb = ["smbprotocol"]
+ssh = ["paramiko"]
+test = ["aiohttp (!=4.0.0a0,!=4.0.0a1)", "numpy", "pytest", "pytest-asyncio (!=0.22.0)", "pytest-benchmark", "pytest-cov", "pytest-mock", "pytest-recording", "pytest-rerunfailures", "requests"]
+test-downstream = ["aiobotocore (>=2.5.4,<3.0.0)", "dask[dataframe,test]", "moto[server] (>4,<5)", "pytest-timeout", "xarray"]
+test-full = ["adlfs", "aiohttp (!=4.0.0a0,!=4.0.0a1)", "cloudpickle", "dask", "distributed", "dropbox", "dropboxdrivefs", "fastparquet", "fusepy", "gcsfs", "jinja2", "kerchunk", "libarchive-c", "lz4", "notebook", "numpy", "ocifs", "pandas", "panel", "paramiko", "pyarrow", "pyarrow (>=1)", "pyftpdlib", "pygit2", "pytest", "pytest-asyncio (!=0.22.0)", "pytest-benchmark", "pytest-cov", "pytest-mock", "pytest-recording", "pytest-rerunfailures", "python-snappy", "requests", "smbprotocol", "tqdm", "urllib3", "zarr", "zstandard"]
+tqdm = ["tqdm"]
+
+[[package]]
+name = "humanize"
+version = "4.12.2"
+description = "Python humanize utilities"
+optional = false
+python-versions = ">=3.9"
+groups = ["main"]
+files = [
+    {file = "humanize-4.12.2-py3-none-any.whl", hash = "sha256:e4e44dced598b7e03487f3b1c6fd5b1146c30ea55a110e71d5d4bca3e094259e"},
+    {file = "humanize-4.12.2.tar.gz", hash = "sha256:ce0715740e9caacc982bb89098182cf8ded3552693a433311c6a4ce6f4e12a2c"},
+]
+
+[package.extras]
+tests = ["freezegun", "pytest", "pytest-cov"]
 
 [[package]]
 name = "iced-x86"
@@ -238,6 +420,7 @@ version = "1.21.0"
 description = "iced-x86 is a blazing fast and correct x86/x64 disassembler, assembler and instruction decoder"
 optional = false
 python-versions = "~=3.8"
+groups = ["main"]
 files = [
     {file = "iced-x86-1.21.0.tar.gz", hash = "sha256:1248b00647dd83b911214895eafbeb2642caa3007de2796ffad1c69aa6d8cd87"},
     {file = "iced_x86-1.21.0-cp38-abi3-macosx_10_12_x86_64.whl", hash = "sha256:bf4ff01b891dac585f40ee86ba482becdf7208d33c788f590dbff19409a02a41"},
@@ -250,12 +433,33 @@ files = [
     {file = "iced_x86-1.21.0-pp38-pypy38_pp73-win_amd64.whl", hash = "sha256:3090e252de96fef2a628eb9cb43629e462fceb1eaeb9e6a73f5ec4925dc5f142"},
 ]
 
+[[package]]
+name = "importlib-resources"
+version = "6.5.2"
+description = "Read resources from Python packages"
+optional = false
+python-versions = ">=3.9"
+groups = ["main"]
+files = [
+    {file = "importlib_resources-6.5.2-py3-none-any.whl", hash = "sha256:789cfdc3ed28c78b67a06acb8126751ced69a3d5f79c095a98298cd8a760ccec"},
+    {file = "importlib_resources-6.5.2.tar.gz", hash = "sha256:185f87adef5bcc288449d98fb4fba07cea78bc036455dd44c5fc4a2fe78fed2c"},
+]
+
+[package.extras]
+check = ["pytest-checkdocs (>=2.4)", "pytest-ruff (>=0.2.1) ; sys_platform != \"cygwin\""]
+cover = ["pytest-cov"]
+doc = ["furo", "jaraco.packaging (>=9.3)", "jaraco.tidelift (>=1.4)", "rst.linker (>=1.9)", "sphinx (>=3.5)", "sphinx-lint"]
+enabler = ["pytest-enabler (>=2.2)"]
+test = ["jaraco.test (>=5.4)", "pytest (>=6,!=8.1.*)", "zipp (>=3.17)"]
+type = ["pytest-mypy"]
+
 [[package]]
 name = "ipython"
 version = "8.27.0"
 description = "IPython: Productive Interactive Computing"
 optional = false
 python-versions = ">=3.10"
+groups = ["main"]
 files = [
     {file = "ipython-8.27.0-py3-none-any.whl", hash = "sha256:f68b3cb8bde357a5d7adc9598d57e22a45dfbea19eb6b98286fa3b288c9cd55c"},
     {file = "ipython-8.27.0.tar.gz", hash = "sha256:0b99a2dc9f15fd68692e898e5568725c6d49c527d36a9fb5960ffbdeaa82ff7e"},
@@ -275,7 +479,7 @@ traitlets = ">=5.13.0"
 [package.extras]
 all = ["ipython[black,doc,kernel,matplotlib,nbconvert,nbformat,notebook,parallel,qtconsole]", "ipython[test,test-extra]"]
 black = ["black"]
-doc = ["docrepr", "exceptiongroup", "intersphinx-registry", "ipykernel", "ipython[test]", "matplotlib", "setuptools (>=18.5)", "sphinx (>=1.3)", "sphinx-rtd-theme", "sphinxcontrib-jquery", "tomli", "typing-extensions"]
+doc = ["docrepr", "exceptiongroup", "intersphinx-registry", "ipykernel", "ipython[test]", "matplotlib", "setuptools (>=18.5)", "sphinx (>=1.3)", "sphinx-rtd-theme", "sphinxcontrib-jquery", "tomli ; python_version < \"3.11\"", "typing-extensions"]
 kernel = ["ipykernel"]
 matplotlib = ["matplotlib"]
 nbconvert = ["nbconvert"]
@@ -286,12 +490,75 @@ qtconsole = ["qtconsole"]
 test = ["packaging", "pickleshare", "pytest", "pytest-asyncio (<0.22)", "testpath"]
 test-extra = ["curio", "ipython[test]", "matplotlib (!=3.2.0)", "nbformat", "numpy (>=1.23)", "pandas", "trio"]
 
+[[package]]
+name = "jax"
+version = "0.5.3"
+description = "Differentiate, compile, and transform Numpy code."
+optional = false
+python-versions = ">=3.10"
+groups = ["main"]
+files = [
+    {file = "jax-0.5.3-py3-none-any.whl", hash = "sha256:1483dc237b4f47e41755d69429e8c3c138736716147cd43bb2b99b259d4e3c41"},
+    {file = "jax-0.5.3.tar.gz", hash = "sha256:f17fcb0fd61dc289394af6ce4de2dada2312f2689bb0d73642c6f026a95fbb2c"},
+]
+
+[package.dependencies]
+jaxlib = "0.5.3"
+ml_dtypes = ">=0.4.0"
+numpy = {version = ">=1.26.0", markers = "python_version >= \"3.12\""}
+opt_einsum = "*"
+scipy = ">=1.11.1"
+
+[package.extras]
+ci = ["jaxlib (==0.5.1)"]
+cuda = ["jax-cuda12-plugin[with-cuda] (==0.5.3)", "jaxlib (==0.5.3)"]
+cuda12 = ["jax-cuda12-plugin[with-cuda] (==0.5.3)", "jaxlib (==0.5.3)"]
+cuda12-local = ["jax-cuda12-plugin (==0.5.3)", "jaxlib (==0.5.3)"]
+cuda12-pip = ["jax-cuda12-plugin[with-cuda] (==0.5.3)", "jaxlib (==0.5.3)"]
+k8s = ["kubernetes"]
+minimum-jaxlib = ["jaxlib (==0.5.3)"]
+rocm = ["jax-rocm60-plugin (==0.5.3)", "jaxlib (==0.5.3)"]
+tpu = ["jaxlib (==0.5.3)", "libtpu (==0.0.11.*)", "requests"]
+
+[[package]]
+name = "jaxlib"
+version = "0.5.3"
+description = "XLA library for JAX"
+optional = false
+python-versions = ">=3.10"
+groups = ["main"]
+files = [
+    {file = "jaxlib-0.5.3-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:48ff5c89fb8a0fe04d475e9ddc074b4879a91d7ab68a51cec5cd1e87f81e6c47"},
+    {file = "jaxlib-0.5.3-cp310-cp310-manylinux2014_aarch64.whl", hash = "sha256:972400db4af6e85270d81db5e6e620d31395f0472e510c50dfcd4cb3f72b7220"},
+    {file = "jaxlib-0.5.3-cp310-cp310-manylinux2014_x86_64.whl", hash = "sha256:52be6c9775aff738a61170d8c047505c75bb799a45518e66a7a0908127b11785"},
+    {file = "jaxlib-0.5.3-cp310-cp310-win_amd64.whl", hash = "sha256:b41a6fcaeb374fabc4ee7e74cfed60843bdab607cd54f60a68b7f7655cde2b66"},
+    {file = "jaxlib-0.5.3-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:b62bd8b29e5a4f9bfaa57c8daf6e04820b2c994f448f3dec602d64255545e9f2"},
+    {file = "jaxlib-0.5.3-cp311-cp311-manylinux2014_aarch64.whl", hash = "sha256:a4666f81d72c060ed3e581ded116a9caa9b0a70a148a54cb12a1d3afca3624b5"},
+    {file = "jaxlib-0.5.3-cp311-cp311-manylinux2014_x86_64.whl", hash = "sha256:29e1530fc81833216f1e28b578d0c59697654f72ee31c7a44ed7753baf5ac466"},
+    {file = "jaxlib-0.5.3-cp311-cp311-win_amd64.whl", hash = "sha256:8eb54e38d789557579f900ea3d70f104a440f8555a9681ed45f4a122dcbfd92e"},
+    {file = "jaxlib-0.5.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:d394dbde4a1c6bd67501cfb29d3819a10b900cb534cc0fc603319f7092f24cfa"},
+    {file = "jaxlib-0.5.3-cp312-cp312-manylinux2014_aarch64.whl", hash = "sha256:bddf6360377aa1c792e47fd87f307c342e331e5ff3582f940b1bca00f6b4bc73"},
+    {file = "jaxlib-0.5.3-cp312-cp312-manylinux2014_x86_64.whl", hash = "sha256:5a5e88ab1cd6fdf78d69abe3544e8f09cce200dd339bb85fbe3c2ea67f2a5e68"},
+    {file = "jaxlib-0.5.3-cp312-cp312-win_amd64.whl", hash = "sha256:520665929649f29f7d948d4070dbaf3e032a4c1f7c11f2863eac73320fcee784"},
+    {file = "jaxlib-0.5.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:31321c25282a06a6dfc940507bc14d0a0ac838d8ced6c07aa00a7fae34ce7b3f"},
+    {file = "jaxlib-0.5.3-cp313-cp313-manylinux2014_aarch64.whl", hash = "sha256:e904b92dedfbc7e545725a8d7676987030ae9c069001d94701bc109c6dab4100"},
+    {file = "jaxlib-0.5.3-cp313-cp313-manylinux2014_x86_64.whl", hash = "sha256:bb7593cb7fffcb13963f22fa5229ed960b8fb4ae5ec3b0820048cbd67f1e8e31"},
+    {file = "jaxlib-0.5.3-cp313-cp313-win_amd64.whl", hash = "sha256:8019f73a10b1290f988dd3768c684f3a8a147239091c3b790ce7e47e3bbc00bd"},
+    {file = "jaxlib-0.5.3-cp313-cp313t-manylinux2014_x86_64.whl", hash = "sha256:4c9a9d4cda091a3ef068ace8379fff9e98eea2fc51dbdd7c3386144a1bdf715d"},
+]
+
+[package.dependencies]
+ml_dtypes = ">=0.2.0"
+numpy = ">=1.25"
+scipy = ">=1.11.1"
+
 [[package]]
 name = "jedi"
 version = "0.19.1"
 description = "An autocompletion tool for Python that can be used for text editors."
 optional = false
 python-versions = ">=3.6"
+groups = ["main"]
 files = [
     {file = "jedi-0.19.1-py2.py3-none-any.whl", hash = "sha256:e983c654fe5c02867aef4cdfce5a2fbb4a50adc0af145f70504238f18ef5e7e0"},
     {file = "jedi-0.19.1.tar.gz", hash = "sha256:cf0496f3651bc65d7174ac1b7d043eff454892c708a87d1b683e57b569927ffd"},
@@ -311,6 +578,7 @@ version = "3.1.4"
 description = "A very fast and expressive template engine."
 optional = false
 python-versions = ">=3.7"
+groups = ["main"]
 files = [
     {file = "jinja2-3.1.4-py3-none-any.whl", hash = "sha256:bc5dd2abb727a5319567b7a813e6a2e7318c39f4f487cfe6c89c6f9c7d25197d"},
     {file = "jinja2-3.1.4.tar.gz", hash = "sha256:4a3aee7acbbe7303aede8e9648d13b8bf88a429282aa6122a993f0ac800cb369"},
@@ -328,6 +596,7 @@ version = "3.3.0"
 description = "Python library for serializing arbitrary object graphs into JSON"
 optional = false
 python-versions = ">=3.7"
+groups = ["main"]
 files = [
     {file = "jsonpickle-3.3.0-py3-none-any.whl", hash = "sha256:287c12143f35571ab00e224fa323aa4b090d5a7f086f5f494d7ee9c7eb1a380a"},
     {file = "jsonpickle-3.3.0.tar.gz", hash = "sha256:ab467e601e5b1a1cd76f1819d014795165da071744ef30bf3786e9bc549de25a"},
@@ -336,7 +605,7 @@ files = [
 [package.extras]
 docs = ["furo", "rst.linker (>=1.9)", "sphinx"]
 packaging = ["build", "twine"]
-testing = ["bson", "ecdsa", "feedparser", "gmpy2", "numpy", "pandas", "pymongo", "pytest (>=3.5,!=3.7.3)", "pytest-benchmark", "pytest-benchmark[histogram]", "pytest-checkdocs (>=1.2.3)", "pytest-cov", "pytest-enabler (>=1.0.1)", "pytest-ruff (>=0.2.1)", "scikit-learn", "scipy", "scipy (>=1.9.3)", "simplejson", "sqlalchemy", "ujson"]
+testing = ["bson", "ecdsa", "feedparser", "gmpy2 ; python_version < \"3.12\"", "numpy", "pandas", "pymongo", "pytest (>=3.5,!=3.7.3)", "pytest-benchmark", "pytest-benchmark[histogram]", "pytest-checkdocs (>=1.2.3)", "pytest-cov", "pytest-enabler (>=1.0.1)", "pytest-ruff (>=0.2.1)", "scikit-learn", "scipy (>=1.9.3) ; python_version > \"3.10\"", "scipy ; python_version <= \"3.10\"", "simplejson", "sqlalchemy", "ujson"]
 
 [[package]]
 name = "kiwisolver"
@@ -344,6 +613,7 @@ version = "1.4.7"
 description = "A fast implementation of the Cassowary constraint solver"
 optional = false
 python-versions = ">=3.8"
+groups = ["main"]
 files = [
     {file = "kiwisolver-1.4.7-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:8a9c83f75223d5e48b0bc9cb1bf2776cf01563e00ade8775ffe13b0b6e1af3a6"},
     {file = "kiwisolver-1.4.7-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:58370b1ffbd35407444d57057b57da5d6549d2d854fa30249771775c63b5fe17"},
@@ -461,12 +731,38 @@ files = [
     {file = "kiwisolver-1.4.7.tar.gz", hash = "sha256:9893ff81bd7107f7b685d3017cc6583daadb4fc26e4a888350df530e41980a60"},
 ]
 
+[[package]]
+name = "markdown-it-py"
+version = "3.0.0"
+description = "Python port of markdown-it. Markdown parsing, done right!"
+optional = false
+python-versions = ">=3.8"
+groups = ["main"]
+files = [
+    {file = "markdown-it-py-3.0.0.tar.gz", hash = "sha256:e3f60a94fa066dc52ec76661e37c851cb232d92f9886b15cb560aaada2df8feb"},
+    {file = "markdown_it_py-3.0.0-py3-none-any.whl", hash = "sha256:355216845c60bd96232cd8d8c40e8f9765cc86f46880e43a8fd22dc1a1a8cab1"},
+]
+
+[package.dependencies]
+mdurl = ">=0.1,<1.0"
+
+[package.extras]
+benchmarking = ["psutil", "pytest", "pytest-benchmark"]
+code-style = ["pre-commit (>=3.0,<4.0)"]
+compare = ["commonmark (>=0.9,<1.0)", "markdown (>=3.4,<4.0)", "mistletoe (>=1.0,<2.0)", "mistune (>=2.0,<3.0)", "panflute (>=2.3,<3.0)"]
+linkify = ["linkify-it-py (>=1,<3)"]
+plugins = ["mdit-py-plugins"]
+profiling = ["gprof2dot"]
+rtd = ["jupyter_sphinx", "mdit-py-plugins", "myst-parser", "pyyaml", "sphinx", "sphinx-copybutton", "sphinx-design", "sphinx_book_theme"]
+testing = ["coverage", "pytest", "pytest-cov", "pytest-regressions"]
+
 [[package]]
 name = "markupsafe"
 version = "2.1.5"
 description = "Safely add untrusted strings to HTML/XML markup."
 optional = false
 python-versions = ">=3.7"
+groups = ["main"]
 files = [
     {file = "MarkupSafe-2.1.5-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:a17a92de5231666cfbe003f0e4b9b3a7ae3afb1ec2845aadc2bacc93ff85febc"},
     {file = "MarkupSafe-2.1.5-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:72b6be590cc35924b02c78ef34b467da4ba07e4e0f0454a2c5907f473fc50ce5"},
@@ -536,6 +832,7 @@ version = "3.9.2"
 description = "Python plotting package"
 optional = false
 python-versions = ">=3.9"
+groups = ["main"]
 files = [
     {file = "matplotlib-3.9.2-cp310-cp310-macosx_10_12_x86_64.whl", hash = "sha256:9d78bbc0cbc891ad55b4f39a48c22182e9bdaea7fc0e5dbd364f49f729ca1bbb"},
     {file = "matplotlib-3.9.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:c375cc72229614632c87355366bdf2570c2dac01ac66b8ad048d2dabadf2d0d4"},
@@ -599,6 +896,7 @@ version = "0.1.7"
 description = "Inline Matplotlib backend for Jupyter"
 optional = false
 python-versions = ">=3.8"
+groups = ["main"]
 files = [
     {file = "matplotlib_inline-0.1.7-py3-none-any.whl", hash = "sha256:df192d39a4ff8f21b1895d72e6a13f5fcc5099f00fa84384e0ea28c2cc0653ca"},
     {file = "matplotlib_inline-0.1.7.tar.gz", hash = "sha256:8423b23ec666be3d16e16b60bdd8ac4e86e840ebd1dd11a30b9f117f2fa0ab90"},
@@ -607,12 +905,172 @@ files = [
 [package.dependencies]
 traitlets = "*"
 
+[[package]]
+name = "mdurl"
+version = "0.1.2"
+description = "Markdown URL utilities"
+optional = false
+python-versions = ">=3.7"
+groups = ["main"]
+files = [
+    {file = "mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8"},
+    {file = "mdurl-0.1.2.tar.gz", hash = "sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba"},
+]
+
+[[package]]
+name = "ml-dtypes"
+version = "0.5.1"
+description = ""
+optional = false
+python-versions = ">=3.9"
+groups = ["main"]
+files = [
+    {file = "ml_dtypes-0.5.1-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:bd73f51957949069573ff783563486339a9285d72e2f36c18e0c1aa9ca7eb190"},
+    {file = "ml_dtypes-0.5.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:810512e2eccdfc3b41eefa3a27402371a3411453a1efc7e9c000318196140fed"},
+    {file = "ml_dtypes-0.5.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:141b2ea2f20bb10802ddca55d91fe21231ef49715cfc971998e8f2a9838f3dbe"},
+    {file = "ml_dtypes-0.5.1-cp310-cp310-win_amd64.whl", hash = "sha256:26ebcc69d7b779c8f129393e99732961b5cc33fcff84090451f448c89b0e01b4"},
+    {file = "ml_dtypes-0.5.1-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:023ce2f502efd4d6c1e0472cc58ce3640d051d40e71e27386bed33901e201327"},
+    {file = "ml_dtypes-0.5.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7000b6e4d8ef07542c05044ec5d8bbae1df083b3f56822c3da63993a113e716f"},
+    {file = "ml_dtypes-0.5.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c09526488c3a9e8b7a23a388d4974b670a9a3dd40c5c8a61db5593ce9b725bab"},
+    {file = "ml_dtypes-0.5.1-cp311-cp311-win_amd64.whl", hash = "sha256:15ad0f3b0323ce96c24637a88a6f44f6713c64032f27277b069f285c3cf66478"},
+    {file = "ml_dtypes-0.5.1-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:6f462f5eca22fb66d7ff9c4744a3db4463af06c49816c4b6ac89b16bfcdc592e"},
+    {file = "ml_dtypes-0.5.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6f76232163b5b9c34291b54621ee60417601e2e4802a188a0ea7157cd9b323f4"},
+    {file = "ml_dtypes-0.5.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ad4953c5eb9c25a56d11a913c2011d7e580a435ef5145f804d98efa14477d390"},
+    {file = "ml_dtypes-0.5.1-cp312-cp312-win_amd64.whl", hash = "sha256:9626d0bca1fb387d5791ca36bacbba298c5ef554747b7ebeafefb4564fc83566"},
+    {file = "ml_dtypes-0.5.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:12651420130ee7cc13059fc56dac6ad300c3af3848b802d475148c9defd27c23"},
+    {file = "ml_dtypes-0.5.1-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c9945669d3dadf8acb40ec2e57d38c985d8c285ea73af57fc5b09872c516106d"},
+    {file = "ml_dtypes-0.5.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bf9975bda82a99dc935f2ae4c83846d86df8fd6ba179614acac8e686910851da"},
+    {file = "ml_dtypes-0.5.1-cp313-cp313-win_amd64.whl", hash = "sha256:fd918d4e6a4e0c110e2e05be7a7814d10dc1b95872accbf6512b80a109b71ae1"},
+    {file = "ml_dtypes-0.5.1-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:05f23447a1c20ddf4dc7c2c661aa9ed93fcb2658f1017c204d1e758714dc28a8"},
+    {file = "ml_dtypes-0.5.1-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1b7fbe5571fdf28fd3aaab3ef4aafc847de9ebf263be959958c1ca58ec8eadf5"},
+    {file = "ml_dtypes-0.5.1-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d13755f8e8445b3870114e5b6240facaa7cb0c3361e54beba3e07fa912a6e12b"},
+    {file = "ml_dtypes-0.5.1-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:b8a9d46b4df5ae2135a8e8e72b465448ebbc1559997f4f9304a9ecc3413efb5b"},
+    {file = "ml_dtypes-0.5.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:afb2009ac98da274e893e03162f6269398b2b00d947e7057ee2469a921d58135"},
+    {file = "ml_dtypes-0.5.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:aefedc579ece2f8fb38f876aa7698204ee4c372d0e54f1c1ffa8ca580b54cc60"},
+    {file = "ml_dtypes-0.5.1-cp39-cp39-win_amd64.whl", hash = "sha256:8f2c028954f16ede77902b223a8da2d9cbb3892375b85809a5c3cfb1587960c4"},
+    {file = "ml_dtypes-0.5.1.tar.gz", hash = "sha256:ac5b58559bb84a95848ed6984eb8013249f90b6bab62aa5acbad876e256002c9"},
+]
+
+[package.dependencies]
+numpy = [
+    {version = ">=2.1.0", markers = "python_version >= \"3.13\""},
+    {version = ">=1.26.0", markers = "python_version == \"3.12\""},
+]
+
+[package.extras]
+dev = ["absl-py", "pyink", "pylint (>=2.6.0)", "pytest", "pytest-xdist"]
+
+[[package]]
+name = "mpmath"
+version = "1.3.0"
+description = "Python library for arbitrary-precision floating-point arithmetic"
+optional = false
+python-versions = "*"
+groups = ["main"]
+files = [
+    {file = "mpmath-1.3.0-py3-none-any.whl", hash = "sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c"},
+    {file = "mpmath-1.3.0.tar.gz", hash = "sha256:7a28eb2a9774d00c7bc92411c19a89209d5da7c4c9a9e227be8330a23a25b91f"},
+]
+
+[package.extras]
+develop = ["codecov", "pycodestyle", "pytest (>=4.6)", "pytest-cov", "wheel"]
+docs = ["sphinx"]
+gmpy = ["gmpy2 (>=2.1.0a4) ; platform_python_implementation != \"PyPy\""]
+tests = ["pytest (>=4.6)"]
+
+[[package]]
+name = "msgpack"
+version = "1.1.0"
+description = "MessagePack serializer"
+optional = false
+python-versions = ">=3.8"
+groups = ["main"]
+files = [
+    {file = "msgpack-1.1.0-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:7ad442d527a7e358a469faf43fda45aaf4ac3249c8310a82f0ccff9164e5dccd"},
+    {file = "msgpack-1.1.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:74bed8f63f8f14d75eec75cf3d04ad581da6b914001b474a5d3cd3372c8cc27d"},
+    {file = "msgpack-1.1.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:914571a2a5b4e7606997e169f64ce53a8b1e06f2cf2c3a7273aa106236d43dd5"},
+    {file = "msgpack-1.1.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c921af52214dcbb75e6bdf6a661b23c3e6417f00c603dd2070bccb5c3ef499f5"},
+    {file = "msgpack-1.1.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d8ce0b22b890be5d252de90d0e0d119f363012027cf256185fc3d474c44b1b9e"},
+    {file = "msgpack-1.1.0-cp310-cp310-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:73322a6cc57fcee3c0c57c4463d828e9428275fb85a27aa2aa1a92fdc42afd7b"},
+    {file = "msgpack-1.1.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:e1f3c3d21f7cf67bcf2da8e494d30a75e4cf60041d98b3f79875afb5b96f3a3f"},
+    {file = "msgpack-1.1.0-cp310-cp310-musllinux_1_2_i686.whl", hash = "sha256:64fc9068d701233effd61b19efb1485587560b66fe57b3e50d29c5d78e7fef68"},
+    {file = "msgpack-1.1.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:42f754515e0f683f9c79210a5d1cad631ec3d06cea5172214d2176a42e67e19b"},
+    {file = "msgpack-1.1.0-cp310-cp310-win32.whl", hash = "sha256:3df7e6b05571b3814361e8464f9304c42d2196808e0119f55d0d3e62cd5ea044"},
+    {file = "msgpack-1.1.0-cp310-cp310-win_amd64.whl", hash = "sha256:685ec345eefc757a7c8af44a3032734a739f8c45d1b0ac45efc5d8977aa4720f"},
+    {file = "msgpack-1.1.0-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:3d364a55082fb2a7416f6c63ae383fbd903adb5a6cf78c5b96cc6316dc1cedc7"},
+    {file = "msgpack-1.1.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:79ec007767b9b56860e0372085f8504db5d06bd6a327a335449508bbee9648fa"},
+    {file = "msgpack-1.1.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:6ad622bf7756d5a497d5b6836e7fc3752e2dd6f4c648e24b1803f6048596f701"},
+    {file = "msgpack-1.1.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8e59bca908d9ca0de3dc8684f21ebf9a690fe47b6be93236eb40b99af28b6ea6"},
+    {file = "msgpack-1.1.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:5e1da8f11a3dd397f0a32c76165cf0c4eb95b31013a94f6ecc0b280c05c91b59"},
+    {file = "msgpack-1.1.0-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:452aff037287acb1d70a804ffd022b21fa2bb7c46bee884dbc864cc9024128a0"},
+    {file = "msgpack-1.1.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:8da4bf6d54ceed70e8861f833f83ce0814a2b72102e890cbdfe4b34764cdd66e"},
+    {file = "msgpack-1.1.0-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:41c991beebf175faf352fb940bf2af9ad1fb77fd25f38d9142053914947cdbf6"},
+    {file = "msgpack-1.1.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:a52a1f3a5af7ba1c9ace055b659189f6c669cf3657095b50f9602af3a3ba0fe5"},
+    {file = "msgpack-1.1.0-cp311-cp311-win32.whl", hash = "sha256:58638690ebd0a06427c5fe1a227bb6b8b9fdc2bd07701bec13c2335c82131a88"},
+    {file = "msgpack-1.1.0-cp311-cp311-win_amd64.whl", hash = "sha256:fd2906780f25c8ed5d7b323379f6138524ba793428db5d0e9d226d3fa6aa1788"},
+    {file = "msgpack-1.1.0-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:d46cf9e3705ea9485687aa4001a76e44748b609d260af21c4ceea7f2212a501d"},
+    {file = "msgpack-1.1.0-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:5dbad74103df937e1325cc4bfeaf57713be0b4f15e1c2da43ccdd836393e2ea2"},
+    {file = "msgpack-1.1.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:58dfc47f8b102da61e8949708b3eafc3504509a5728f8b4ddef84bd9e16ad420"},
+    {file = "msgpack-1.1.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4676e5be1b472909b2ee6356ff425ebedf5142427842aa06b4dfd5117d1ca8a2"},
+    {file = "msgpack-1.1.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:17fb65dd0bec285907f68b15734a993ad3fc94332b5bb21b0435846228de1f39"},
+    {file = "msgpack-1.1.0-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:a51abd48c6d8ac89e0cfd4fe177c61481aca2d5e7ba42044fd218cfd8ea9899f"},
+    {file = "msgpack-1.1.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:2137773500afa5494a61b1208619e3871f75f27b03bcfca7b3a7023284140247"},
+    {file = "msgpack-1.1.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:398b713459fea610861c8a7b62a6fec1882759f308ae0795b5413ff6a160cf3c"},
+    {file = "msgpack-1.1.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:06f5fd2f6bb2a7914922d935d3b8bb4a7fff3a9a91cfce6d06c13bc42bec975b"},
+    {file = "msgpack-1.1.0-cp312-cp312-win32.whl", hash = "sha256:ad33e8400e4ec17ba782f7b9cf868977d867ed784a1f5f2ab46e7ba53b6e1e1b"},
+    {file = "msgpack-1.1.0-cp312-cp312-win_amd64.whl", hash = "sha256:115a7af8ee9e8cddc10f87636767857e7e3717b7a2e97379dc2054712693e90f"},
+    {file = "msgpack-1.1.0-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:071603e2f0771c45ad9bc65719291c568d4edf120b44eb36324dcb02a13bfddf"},
+    {file = "msgpack-1.1.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:0f92a83b84e7c0749e3f12821949d79485971f087604178026085f60ce109330"},
+    {file = "msgpack-1.1.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:4a1964df7b81285d00a84da4e70cb1383f2e665e0f1f2a7027e683956d04b734"},
+    {file = "msgpack-1.1.0-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:59caf6a4ed0d164055ccff8fe31eddc0ebc07cf7326a2aaa0dbf7a4001cd823e"},
+    {file = "msgpack-1.1.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0907e1a7119b337971a689153665764adc34e89175f9a34793307d9def08e6ca"},
+    {file = "msgpack-1.1.0-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:65553c9b6da8166e819a6aa90ad15288599b340f91d18f60b2061f402b9a4915"},
+    {file = "msgpack-1.1.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:7a946a8992941fea80ed4beae6bff74ffd7ee129a90b4dd5cf9c476a30e9708d"},
+    {file = "msgpack-1.1.0-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:4b51405e36e075193bc051315dbf29168d6141ae2500ba8cd80a522964e31434"},
+    {file = "msgpack-1.1.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:b4c01941fd2ff87c2a934ee6055bda4ed353a7846b8d4f341c428109e9fcde8c"},
+    {file = "msgpack-1.1.0-cp313-cp313-win32.whl", hash = "sha256:7c9a35ce2c2573bada929e0b7b3576de647b0defbd25f5139dcdaba0ae35a4cc"},
+    {file = "msgpack-1.1.0-cp313-cp313-win_amd64.whl", hash = "sha256:bce7d9e614a04d0883af0b3d4d501171fbfca038f12c77fa838d9f198147a23f"},
+    {file = "msgpack-1.1.0-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c40ffa9a15d74e05ba1fe2681ea33b9caffd886675412612d93ab17b58ea2fec"},
+    {file = "msgpack-1.1.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f1ba6136e650898082d9d5a5217d5906d1e138024f836ff48691784bbe1adf96"},
+    {file = "msgpack-1.1.0-cp38-cp38-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:e0856a2b7e8dcb874be44fea031d22e5b3a19121be92a1e098f46068a11b0870"},
+    {file = "msgpack-1.1.0-cp38-cp38-musllinux_1_2_aarch64.whl", hash = "sha256:471e27a5787a2e3f974ba023f9e265a8c7cfd373632247deb225617e3100a3c7"},
+    {file = "msgpack-1.1.0-cp38-cp38-musllinux_1_2_i686.whl", hash = "sha256:646afc8102935a388ffc3914b336d22d1c2d6209c773f3eb5dd4d6d3b6f8c1cb"},
+    {file = "msgpack-1.1.0-cp38-cp38-musllinux_1_2_x86_64.whl", hash = "sha256:13599f8829cfbe0158f6456374e9eea9f44eee08076291771d8ae93eda56607f"},
+    {file = "msgpack-1.1.0-cp38-cp38-win32.whl", hash = "sha256:8a84efb768fb968381e525eeeb3d92857e4985aacc39f3c47ffd00eb4509315b"},
+    {file = "msgpack-1.1.0-cp38-cp38-win_amd64.whl", hash = "sha256:879a7b7b0ad82481c52d3c7eb99bf6f0645dbdec5134a4bddbd16f3506947feb"},
+    {file = "msgpack-1.1.0-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:53258eeb7a80fc46f62fd59c876957a2d0e15e6449a9e71842b6d24419d88ca1"},
+    {file = "msgpack-1.1.0-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:7e7b853bbc44fb03fbdba34feb4bd414322180135e2cb5164f20ce1c9795ee48"},
+    {file = "msgpack-1.1.0-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:f3e9b4936df53b970513eac1758f3882c88658a220b58dcc1e39606dccaaf01c"},
+    {file = "msgpack-1.1.0-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:46c34e99110762a76e3911fc923222472c9d681f1094096ac4102c18319e6468"},
+    {file = "msgpack-1.1.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8a706d1e74dd3dea05cb54580d9bd8b2880e9264856ce5068027eed09680aa74"},
+    {file = "msgpack-1.1.0-cp39-cp39-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:534480ee5690ab3cbed89d4c8971a5c631b69a8c0883ecfea96c19118510c846"},
+    {file = "msgpack-1.1.0-cp39-cp39-musllinux_1_2_aarch64.whl", hash = "sha256:8cf9e8c3a2153934a23ac160cc4cba0ec035f6867c8013cc6077a79823370346"},
+    {file = "msgpack-1.1.0-cp39-cp39-musllinux_1_2_i686.whl", hash = "sha256:3180065ec2abbe13a4ad37688b61b99d7f9e012a535b930e0e683ad6bc30155b"},
+    {file = "msgpack-1.1.0-cp39-cp39-musllinux_1_2_x86_64.whl", hash = "sha256:c5a91481a3cc573ac8c0d9aace09345d989dc4a0202b7fcb312c88c26d4e71a8"},
+    {file = "msgpack-1.1.0-cp39-cp39-win32.whl", hash = "sha256:f80bc7d47f76089633763f952e67f8214cb7b3ee6bfa489b3cb6a84cfac114cd"},
+    {file = "msgpack-1.1.0-cp39-cp39-win_amd64.whl", hash = "sha256:4d1b7ff2d6146e16e8bd665ac726a89c74163ef8cd39fa8c1087d4e52d3a2325"},
+    {file = "msgpack-1.1.0.tar.gz", hash = "sha256:dd432ccc2c72b914e4cb77afce64aab761c1137cc698be3984eee260bcb2896e"},
+]
+
+[[package]]
+name = "nest-asyncio"
+version = "1.6.0"
+description = "Patch asyncio to allow nested event loops"
+optional = false
+python-versions = ">=3.5"
+groups = ["main"]
+files = [
+    {file = "nest_asyncio-1.6.0-py3-none-any.whl", hash = "sha256:87af6efd6b5e897c81050477ef65c62e2b2f35d51703cae01aff2905b1852e1c"},
+    {file = "nest_asyncio-1.6.0.tar.gz", hash = "sha256:6f172d5449aca15afd6c646851f4e31e02c598d553a667e38cafa997cfec55fe"},
+]
+
 [[package]]
 name = "networkx"
 version = "3.3"
 description = "Python package for creating and manipulating graphs and networks"
 optional = false
 python-versions = ">=3.10"
+groups = ["main"]
 files = [
     {file = "networkx-3.3-py3-none-any.whl", hash = "sha256:28575580c6ebdaf4505b22c6256a2b9de86b316dc63ba9e93abde3d78dfdbcf2"},
     {file = "networkx-3.3.tar.gz", hash = "sha256:0c127d8b2f4865f59ae9cb8aafcd60b5c70f3241ebd66f7defad7c4ab90126c9"},
@@ -631,6 +1089,7 @@ version = "2.1.1"
 description = "Fundamental package for array computing in Python"
 optional = false
 python-versions = ">=3.10"
+groups = ["main"]
 files = [
     {file = "numpy-2.1.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:c8a0e34993b510fc19b9a2ce7f31cb8e94ecf6e924a40c0c9dd4f62d0aac47d9"},
     {file = "numpy-2.1.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:7dd86dfaf7c900c0bbdcb8b16e2f6ddf1eb1fe39c6c8cca6e94844ed3152a8fd"},
@@ -687,12 +1146,303 @@ files = [
     {file = "numpy-2.1.1.tar.gz", hash = "sha256:d0cf7d55b1051387807405b3898efafa862997b4cba8aa5dbe657be794afeafd"},
 ]
 
+[[package]]
+name = "nvidia-cublas-cu12"
+version = "12.6.4.1"
+description = "CUBLAS native runtime libraries"
+optional = false
+python-versions = ">=3"
+groups = ["main"]
+markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""
+files = [
+    {file = "nvidia_cublas_cu12-12.6.4.1-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:08ed2686e9875d01b58e3cb379c6896df8e76c75e0d4a7f7dace3d7b6d9ef8eb"},
+    {file = "nvidia_cublas_cu12-12.6.4.1-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:235f728d6e2a409eddf1df58d5b0921cf80cfa9e72b9f2775ccb7b4a87984668"},
+    {file = "nvidia_cublas_cu12-12.6.4.1-py3-none-win_amd64.whl", hash = "sha256:9e4fa264f4d8a4eb0cdbd34beadc029f453b3bafae02401e999cf3d5a5af75f8"},
+]
+
+[[package]]
+name = "nvidia-cuda-cupti-cu12"
+version = "12.6.80"
+description = "CUDA profiling tools runtime libs."
+optional = false
+python-versions = ">=3"
+groups = ["main"]
+markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""
+files = [
+    {file = "nvidia_cuda_cupti_cu12-12.6.80-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:166ee35a3ff1587f2490364f90eeeb8da06cd867bd5b701bf7f9a02b78bc63fc"},
+    {file = "nvidia_cuda_cupti_cu12-12.6.80-py3-none-manylinux2014_aarch64.whl", hash = "sha256:358b4a1d35370353d52e12f0a7d1769fc01ff74a191689d3870b2123156184c4"},
+    {file = "nvidia_cuda_cupti_cu12-12.6.80-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6768bad6cab4f19e8292125e5f1ac8aa7d1718704012a0e3272a6f61c4bce132"},
+    {file = "nvidia_cuda_cupti_cu12-12.6.80-py3-none-manylinux2014_x86_64.whl", hash = "sha256:a3eff6cdfcc6a4c35db968a06fcadb061cbc7d6dde548609a941ff8701b98b73"},
+    {file = "nvidia_cuda_cupti_cu12-12.6.80-py3-none-win_amd64.whl", hash = "sha256:bbe6ae76e83ce5251b56e8c8e61a964f757175682bbad058b170b136266ab00a"},
+]
+
+[[package]]
+name = "nvidia-cuda-nvrtc-cu12"
+version = "12.6.77"
+description = "NVRTC native runtime libraries"
+optional = false
+python-versions = ">=3"
+groups = ["main"]
+markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""
+files = [
+    {file = "nvidia_cuda_nvrtc_cu12-12.6.77-py3-none-manylinux2014_aarch64.whl", hash = "sha256:5847f1d6e5b757f1d2b3991a01082a44aad6f10ab3c5c0213fa3e25bddc25a13"},
+    {file = "nvidia_cuda_nvrtc_cu12-12.6.77-py3-none-manylinux2014_x86_64.whl", hash = "sha256:35b0cc6ee3a9636d5409133e79273ce1f3fd087abb0532d2d2e8fff1fe9efc53"},
+    {file = "nvidia_cuda_nvrtc_cu12-12.6.77-py3-none-win_amd64.whl", hash = "sha256:f7007dbd914c56bd80ea31bc43e8e149da38f68158f423ba845fc3292684e45a"},
+]
+
+[[package]]
+name = "nvidia-cuda-runtime-cu12"
+version = "12.6.77"
+description = "CUDA Runtime native Libraries"
+optional = false
+python-versions = ">=3"
+groups = ["main"]
+markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""
+files = [
+    {file = "nvidia_cuda_runtime_cu12-12.6.77-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:6116fad3e049e04791c0256a9778c16237837c08b27ed8c8401e2e45de8d60cd"},
+    {file = "nvidia_cuda_runtime_cu12-12.6.77-py3-none-manylinux2014_aarch64.whl", hash = "sha256:d461264ecb429c84c8879a7153499ddc7b19b5f8d84c204307491989a365588e"},
+    {file = "nvidia_cuda_runtime_cu12-12.6.77-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ba3b56a4f896141e25e19ab287cd71e52a6a0f4b29d0d31609f60e3b4d5219b7"},
+    {file = "nvidia_cuda_runtime_cu12-12.6.77-py3-none-manylinux2014_x86_64.whl", hash = "sha256:a84d15d5e1da416dd4774cb42edf5e954a3e60cc945698dc1d5be02321c44dc8"},
+    {file = "nvidia_cuda_runtime_cu12-12.6.77-py3-none-win_amd64.whl", hash = "sha256:86c58044c824bf3c173c49a2dbc7a6c8b53cb4e4dca50068be0bf64e9dab3f7f"},
+]
+
+[[package]]
+name = "nvidia-cudnn-cu12"
+version = "9.5.1.17"
+description = "cuDNN runtime libraries"
+optional = false
+python-versions = ">=3"
+groups = ["main"]
+markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""
+files = [
+    {file = "nvidia_cudnn_cu12-9.5.1.17-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:9fd4584468533c61873e5fda8ca41bac3a38bcb2d12350830c69b0a96a7e4def"},
+    {file = "nvidia_cudnn_cu12-9.5.1.17-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:30ac3869f6db17d170e0e556dd6cc5eee02647abc31ca856634d5a40f82c15b2"},
+    {file = "nvidia_cudnn_cu12-9.5.1.17-py3-none-win_amd64.whl", hash = "sha256:d7af0f8a4f3b4b9dbb3122f2ef553b45694ed9c384d5a75bab197b8eefb79ab8"},
+]
+
+[package.dependencies]
+nvidia-cublas-cu12 = "*"
+
+[[package]]
+name = "nvidia-cufft-cu12"
+version = "11.3.0.4"
+description = "CUFFT native runtime libraries"
+optional = false
+python-versions = ">=3"
+groups = ["main"]
+markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""
+files = [
+    {file = "nvidia_cufft_cu12-11.3.0.4-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d16079550df460376455cba121db6564089176d9bac9e4f360493ca4741b22a6"},
+    {file = "nvidia_cufft_cu12-11.3.0.4-py3-none-manylinux2014_aarch64.whl", hash = "sha256:8510990de9f96c803a051822618d42bf6cb8f069ff3f48d93a8486efdacb48fb"},
+    {file = "nvidia_cufft_cu12-11.3.0.4-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ccba62eb9cef5559abd5e0d54ceed2d9934030f51163df018532142a8ec533e5"},
+    {file = "nvidia_cufft_cu12-11.3.0.4-py3-none-manylinux2014_x86_64.whl", hash = "sha256:768160ac89f6f7b459bee747e8d175dbf53619cfe74b2a5636264163138013ca"},
+    {file = "nvidia_cufft_cu12-11.3.0.4-py3-none-win_amd64.whl", hash = "sha256:6048ebddfb90d09d2707efb1fd78d4e3a77cb3ae4dc60e19aab6be0ece2ae464"},
+]
+
+[package.dependencies]
+nvidia-nvjitlink-cu12 = "*"
+
+[[package]]
+name = "nvidia-cufile-cu12"
+version = "1.11.1.6"
+description = "cuFile GPUDirect libraries"
+optional = false
+python-versions = ">=3"
+groups = ["main"]
+markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""
+files = [
+    {file = "nvidia_cufile_cu12-1.11.1.6-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:cc23469d1c7e52ce6c1d55253273d32c565dd22068647f3aa59b3c6b005bf159"},
+    {file = "nvidia_cufile_cu12-1.11.1.6-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:8f57a0051dcf2543f6dc2b98a98cb2719c37d3cee1baba8965d57f3bbc90d4db"},
+]
+
+[[package]]
+name = "nvidia-curand-cu12"
+version = "10.3.7.77"
+description = "CURAND native runtime libraries"
+optional = false
+python-versions = ">=3"
+groups = ["main"]
+markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""
+files = [
+    {file = "nvidia_curand_cu12-10.3.7.77-py3-none-manylinux2014_aarch64.whl", hash = "sha256:6e82df077060ea28e37f48a3ec442a8f47690c7499bff392a5938614b56c98d8"},
+    {file = "nvidia_curand_cu12-10.3.7.77-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:a42cd1344297f70b9e39a1e4f467a4e1c10f1da54ff7a85c12197f6c652c8bdf"},
+    {file = "nvidia_curand_cu12-10.3.7.77-py3-none-manylinux2014_x86_64.whl", hash = "sha256:99f1a32f1ac2bd134897fc7a203f779303261268a65762a623bf30cc9fe79117"},
+    {file = "nvidia_curand_cu12-10.3.7.77-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:7b2ed8e95595c3591d984ea3603dd66fe6ce6812b886d59049988a712ed06b6e"},
+    {file = "nvidia_curand_cu12-10.3.7.77-py3-none-win_amd64.whl", hash = "sha256:6d6d935ffba0f3d439b7cd968192ff068fafd9018dbf1b85b37261b13cfc9905"},
+]
+
+[[package]]
+name = "nvidia-cusolver-cu12"
+version = "11.7.1.2"
+description = "CUDA solver native runtime libraries"
+optional = false
+python-versions = ">=3"
+groups = ["main"]
+markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""
+files = [
+    {file = "nvidia_cusolver_cu12-11.7.1.2-py3-none-manylinux2014_aarch64.whl", hash = "sha256:0ce237ef60acde1efc457335a2ddadfd7610b892d94efee7b776c64bb1cac9e0"},
+    {file = "nvidia_cusolver_cu12-11.7.1.2-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:e9e49843a7707e42022babb9bcfa33c29857a93b88020c4e4434656a655b698c"},
+    {file = "nvidia_cusolver_cu12-11.7.1.2-py3-none-manylinux2014_x86_64.whl", hash = "sha256:6cf28f17f64107a0c4d7802be5ff5537b2130bfc112f25d5a30df227058ca0e6"},
+    {file = "nvidia_cusolver_cu12-11.7.1.2-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:dbbe4fc38ec1289c7e5230e16248365e375c3673c9c8bac5796e2e20db07f56e"},
+    {file = "nvidia_cusolver_cu12-11.7.1.2-py3-none-win_amd64.whl", hash = "sha256:6813f9d8073f555444a8705f3ab0296d3e1cb37a16d694c5fc8b862a0d8706d7"},
+]
+
+[package.dependencies]
+nvidia-cublas-cu12 = "*"
+nvidia-cusparse-cu12 = "*"
+nvidia-nvjitlink-cu12 = "*"
+
+[[package]]
+name = "nvidia-cusparse-cu12"
+version = "12.5.4.2"
+description = "CUSPARSE native runtime libraries"
+optional = false
+python-versions = ">=3"
+groups = ["main"]
+markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""
+files = [
+    {file = "nvidia_cusparse_cu12-12.5.4.2-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d25b62fb18751758fe3c93a4a08eff08effedfe4edf1c6bb5afd0890fe88f887"},
+    {file = "nvidia_cusparse_cu12-12.5.4.2-py3-none-manylinux2014_aarch64.whl", hash = "sha256:7aa32fa5470cf754f72d1116c7cbc300b4e638d3ae5304cfa4a638a5b87161b1"},
+    {file = "nvidia_cusparse_cu12-12.5.4.2-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7556d9eca156e18184b94947ade0fba5bb47d69cec46bf8660fd2c71a4b48b73"},
+    {file = "nvidia_cusparse_cu12-12.5.4.2-py3-none-manylinux2014_x86_64.whl", hash = "sha256:23749a6571191a215cb74d1cdbff4a86e7b19f1200c071b3fcf844a5bea23a2f"},
+    {file = "nvidia_cusparse_cu12-12.5.4.2-py3-none-win_amd64.whl", hash = "sha256:4acb8c08855a26d737398cba8fb6f8f5045d93f82612b4cfd84645a2332ccf20"},
+]
+
+[package.dependencies]
+nvidia-nvjitlink-cu12 = "*"
+
+[[package]]
+name = "nvidia-cusparselt-cu12"
+version = "0.6.3"
+description = "NVIDIA cuSPARSELt"
+optional = false
+python-versions = "*"
+groups = ["main"]
+markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""
+files = [
+    {file = "nvidia_cusparselt_cu12-0.6.3-py3-none-manylinux2014_aarch64.whl", hash = "sha256:8371549623ba601a06322af2133c4a44350575f5a3108fb75f3ef20b822ad5f1"},
+    {file = "nvidia_cusparselt_cu12-0.6.3-py3-none-manylinux2014_x86_64.whl", hash = "sha256:e5c8a26c36445dd2e6812f1177978a24e2d37cacce7e090f297a688d1ec44f46"},
+    {file = "nvidia_cusparselt_cu12-0.6.3-py3-none-win_amd64.whl", hash = "sha256:3b325bcbd9b754ba43df5a311488fca11a6b5dc3d11df4d190c000cf1a0765c7"},
+]
+
+[[package]]
+name = "nvidia-nccl-cu12"
+version = "2.26.2"
+description = "NVIDIA Collective Communication Library (NCCL) Runtime"
+optional = false
+python-versions = ">=3"
+groups = ["main"]
+markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""
+files = [
+    {file = "nvidia_nccl_cu12-2.26.2-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:5c196e95e832ad30fbbb50381eb3cbd1fadd5675e587a548563993609af19522"},
+    {file = "nvidia_nccl_cu12-2.26.2-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:694cf3879a206553cc9d7dbda76b13efaf610fdb70a50cba303de1b0d1530ac6"},
+]
+
+[[package]]
+name = "nvidia-nvjitlink-cu12"
+version = "12.6.85"
+description = "Nvidia JIT LTO Library"
+optional = false
+python-versions = ">=3"
+groups = ["main"]
+markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""
+files = [
+    {file = "nvidia_nvjitlink_cu12-12.6.85-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:eedc36df9e88b682efe4309aa16b5b4e78c2407eac59e8c10a6a47535164369a"},
+    {file = "nvidia_nvjitlink_cu12-12.6.85-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:cf4eaa7d4b6b543ffd69d6abfb11efdeb2db48270d94dfd3a452c24150829e41"},
+    {file = "nvidia_nvjitlink_cu12-12.6.85-py3-none-win_amd64.whl", hash = "sha256:e61120e52ed675747825cdd16febc6a0730537451d867ee58bee3853b1b13d1c"},
+]
+
+[[package]]
+name = "nvidia-nvtx-cu12"
+version = "12.6.77"
+description = "NVIDIA Tools Extension"
+optional = false
+python-versions = ">=3"
+groups = ["main"]
+markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""
+files = [
+    {file = "nvidia_nvtx_cu12-12.6.77-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:f44f8d86bb7d5629988d61c8d3ae61dddb2015dee142740536bc7481b022fe4b"},
+    {file = "nvidia_nvtx_cu12-12.6.77-py3-none-manylinux2014_aarch64.whl", hash = "sha256:adcaabb9d436c9761fca2b13959a2d237c5f9fd406c8e4b723c695409ff88059"},
+    {file = "nvidia_nvtx_cu12-12.6.77-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:b90bed3df379fa79afbd21be8e04a0314336b8ae16768b58f2d34cb1d04cd7d2"},
+    {file = "nvidia_nvtx_cu12-12.6.77-py3-none-manylinux2014_x86_64.whl", hash = "sha256:6574241a3ec5fdc9334353ab8c479fe75841dbe8f4532a8fc97ce63503330ba1"},
+    {file = "nvidia_nvtx_cu12-12.6.77-py3-none-win_amd64.whl", hash = "sha256:2fb11a4af04a5e6c84073e6404d26588a34afd35379f0855a99797897efa75c0"},
+]
+
+[[package]]
+name = "opt-einsum"
+version = "3.4.0"
+description = "Path optimization of einsum functions."
+optional = false
+python-versions = ">=3.8"
+groups = ["main"]
+files = [
+    {file = "opt_einsum-3.4.0-py3-none-any.whl", hash = "sha256:69bb92469f86a1565195ece4ac0323943e83477171b91d24c35afe028a90d7cd"},
+    {file = "opt_einsum-3.4.0.tar.gz", hash = "sha256:96ca72f1b886d148241348783498194c577fa30a8faac108586b14f1ba4473ac"},
+]
+
+[[package]]
+name = "optax"
+version = "0.2.4"
+description = "A gradient processing and optimization library in JAX."
+optional = false
+python-versions = ">=3.9"
+groups = ["main"]
+files = [
+    {file = "optax-0.2.4-py3-none-any.whl", hash = "sha256:db35c04e50b52596662efb002334de08c2a0a74971e4da33f467e84fac08886a"},
+    {file = "optax-0.2.4.tar.gz", hash = "sha256:4e05d3d5307e6dde4c319187ae36e6cd3a0c035d4ed25e9e992449a304f47336"},
+]
+
+[package.dependencies]
+absl-py = ">=0.7.1"
+chex = ">=0.1.87"
+etils = {version = "*", extras = ["epy"]}
+jax = ">=0.4.27"
+jaxlib = ">=0.4.27"
+numpy = ">=1.18.0"
+
+[package.extras]
+docs = ["flax", "ipython (>=8.8.0)", "matplotlib (>=3.5.0)", "myst-nb (>=1.0.0)", "sphinx (>=6.0.0)", "sphinx-autodoc-typehints", "sphinx-book-theme (>=1.0.1)", "sphinx-collections (>=0.0.1)", "sphinx-gallery (>=0.14.0)", "sphinx_contributors", "sphinxcontrib-katex", "tensorflow (>=2.4.0)", "tensorflow-datasets (>=4.2.0)"]
+dp-accounting = ["absl-py (>=1.0.0)", "attrs (>=21.4.0)", "mpmath (>=1.2.1)", "numpy (>=1.21.4)", "scipy (>=1.7.1)"]
+examples = ["dp_accounting (>=0.4)", "flax", "ipywidgets", "tensorflow (>=2.4.0)", "tensorflow-datasets (>=4.2.0)"]
+test = ["dm-tree (>=0.1.7)", "flax (>=0.5.3)", "scikit-learn", "scipy (>=1.7.1)"]
+
+[[package]]
+name = "orbax-checkpoint"
+version = "0.11.10"
+description = "Orbax Checkpoint"
+optional = false
+python-versions = ">=3.10"
+groups = ["main"]
+files = [
+    {file = "orbax_checkpoint-0.11.10-py3-none-any.whl", hash = "sha256:11e20aa97a3b0ddef79a24cf192fd997298604ab5541dc4f3ec7512ecbe94bdf"},
+    {file = "orbax_checkpoint-0.11.10.tar.gz", hash = "sha256:9e415b0d041b4c256ff2e126df9c6b056f0155f322de0a69befd73d1657fb9e5"},
+]
+
+[package.dependencies]
+absl-py = "*"
+etils = {version = "*", extras = ["epath", "epy"]}
+humanize = "*"
+jax = ">=0.5.0"
+msgpack = "*"
+nest_asyncio = "*"
+numpy = "*"
+protobuf = "*"
+pyyaml = "*"
+simplejson = ">=3.16.0"
+tensorstore = ">=0.1.71"
+typing_extensions = "*"
+
+[package.extras]
+docs = ["flax", "google-cloud-logging"]
+testing = ["chex", "flax", "google-cloud-logging", "mock", "pytest", "pytest-xdist"]
+
 [[package]]
 name = "packaging"
 version = "24.1"
 description = "Core utilities for Python packages"
 optional = false
 python-versions = ">=3.8"
+groups = ["main"]
 files = [
     {file = "packaging-24.1-py3-none-any.whl", hash = "sha256:5b8f2217dbdbd2f7f384c41c628544e6d52f2d0f53c6d0c3ea61aa5d1d7ff124"},
     {file = "packaging-24.1.tar.gz", hash = "sha256:026ed72c8ed3fcce5bf8950572258698927fd1dbda10a5e981cdf0ac37f4f002"},
@@ -704,6 +1454,7 @@ version = "2.2.2"
 description = "Powerful data structures for data analysis, time series, and statistics"
 optional = false
 python-versions = ">=3.9"
+groups = ["main"]
 files = [
     {file = "pandas-2.2.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:90c6fca2acf139569e74e8781709dccb6fe25940488755716d1d354d6bc58bce"},
     {file = "pandas-2.2.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:c7adfc142dac335d8c1e0dcbd37eb8617eac386596eb9e1a1b77791cf2498238"},
@@ -773,6 +1524,7 @@ version = "0.8.4"
 description = "A Python Parser"
 optional = false
 python-versions = ">=3.6"
+groups = ["main"]
 files = [
     {file = "parso-0.8.4-py2.py3-none-any.whl", hash = "sha256:a418670a20291dacd2dddc80c377c5c3791378ee1e8d12bffc35420643d43f18"},
     {file = "parso-0.8.4.tar.gz", hash = "sha256:eb3a7b58240fb99099a345571deecc0f9540ea5f4dd2fe14c2a99d6b281ab92d"},
@@ -788,6 +1540,8 @@ version = "4.9.0"
 description = "Pexpect allows easy control of interactive console applications."
 optional = false
 python-versions = "*"
+groups = ["main"]
+markers = "sys_platform != \"win32\" and sys_platform != \"emscripten\""
 files = [
     {file = "pexpect-4.9.0-py2.py3-none-any.whl", hash = "sha256:7236d1e080e4936be2dc3e326cec0af72acf9212a7e1d060210e70a47e253523"},
     {file = "pexpect-4.9.0.tar.gz", hash = "sha256:ee7d41123f3c9911050ea2c2dac107568dc43b2d3b0c7557a33212c398ead30f"},
@@ -802,6 +1556,7 @@ version = "10.4.0"
 description = "Python Imaging Library (Fork)"
 optional = false
 python-versions = ">=3.8"
+groups = ["main"]
 files = [
     {file = "pillow-10.4.0-cp310-cp310-macosx_10_10_x86_64.whl", hash = "sha256:4d9667937cfa347525b319ae34375c37b9ee6b525440f3ef48542fcf66f2731e"},
     {file = "pillow-10.4.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:543f3dc61c18dafb755773efc89aae60d06b6596a63914107f75459cf984164d"},
@@ -890,7 +1645,7 @@ docs = ["furo", "olefile", "sphinx (>=7.3)", "sphinx-copybutton", "sphinx-inline
 fpx = ["olefile"]
 mic = ["olefile"]
 tests = ["check-manifest", "coverage", "defusedxml", "markdown2", "olefile", "packaging", "pyroma", "pytest", "pytest-cov", "pytest-timeout"]
-typing = ["typing-extensions"]
+typing = ["typing-extensions ; python_version < \"3.10\""]
 xmp = ["defusedxml"]
 
 [[package]]
@@ -899,6 +1654,7 @@ version = "3.0.47"
 description = "Library for building powerful interactive command lines in Python"
 optional = false
 python-versions = ">=3.7.0"
+groups = ["main"]
 files = [
     {file = "prompt_toolkit-3.0.47-py3-none-any.whl", hash = "sha256:0d7bfa67001d5e39d02c224b663abc33687405033a8c422d0d675a5a13361d10"},
     {file = "prompt_toolkit-3.0.47.tar.gz", hash = "sha256:1e1b29cb58080b1e69f207c893a1a7bf16d127a5c30c9d17a25a5d77792e5360"},
@@ -907,12 +1663,33 @@ files = [
 [package.dependencies]
 wcwidth = "*"
 
+[[package]]
+name = "protobuf"
+version = "6.30.1"
+description = ""
+optional = false
+python-versions = ">=3.9"
+groups = ["main"]
+files = [
+    {file = "protobuf-6.30.1-cp310-abi3-win32.whl", hash = "sha256:ba0706f948d0195f5cac504da156d88174e03218d9364ab40d903788c1903d7e"},
+    {file = "protobuf-6.30.1-cp310-abi3-win_amd64.whl", hash = "sha256:ed484f9ddd47f0f1bf0648806cccdb4fe2fb6b19820f9b79a5adf5dcfd1b8c5f"},
+    {file = "protobuf-6.30.1-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:aa4f7dfaed0d840b03d08d14bfdb41348feaee06a828a8c455698234135b4075"},
+    {file = "protobuf-6.30.1-cp39-abi3-manylinux2014_aarch64.whl", hash = "sha256:47cd320b7db63e8c9ac35f5596ea1c1e61491d8a8eb6d8b45edc44760b53a4f6"},
+    {file = "protobuf-6.30.1-cp39-abi3-manylinux2014_x86_64.whl", hash = "sha256:e3083660225fa94748ac2e407f09a899e6a28bf9c0e70c75def8d15706bf85fc"},
+    {file = "protobuf-6.30.1-cp39-cp39-win32.whl", hash = "sha256:554d7e61cce2aa4c63ca27328f757a9f3867bce8ec213bf09096a8d16bcdcb6a"},
+    {file = "protobuf-6.30.1-cp39-cp39-win_amd64.whl", hash = "sha256:b510f55ce60f84dc7febc619b47215b900466e3555ab8cb1ba42deb4496d6cc0"},
+    {file = "protobuf-6.30.1-py3-none-any.whl", hash = "sha256:3c25e51e1359f1f5fa3b298faa6016e650d148f214db2e47671131b9063c53be"},
+    {file = "protobuf-6.30.1.tar.gz", hash = "sha256:535fb4e44d0236893d5cf1263a0f706f1160b689a7ab962e9da8a9ce4050b780"},
+]
+
 [[package]]
 name = "ptyprocess"
 version = "0.7.0"
 description = "Run a subprocess in a pseudo terminal"
 optional = false
 python-versions = "*"
+groups = ["main"]
+markers = "sys_platform != \"win32\" and sys_platform != \"emscripten\""
 files = [
     {file = "ptyprocess-0.7.0-py2.py3-none-any.whl", hash = "sha256:4b41f3967fce3af57cc7e94b888626c18bf37a083e3651ca8feeb66d492fef35"},
     {file = "ptyprocess-0.7.0.tar.gz", hash = "sha256:5c5d0a3b48ceee0b48485e0c26037c0acd7d29765ca3fbb5cb3831d347423220"},
@@ -924,6 +1701,7 @@ version = "0.2.3"
 description = "Safely evaluate AST nodes without side effects"
 optional = false
 python-versions = "*"
+groups = ["main"]
 files = [
     {file = "pure_eval-0.2.3-py3-none-any.whl", hash = "sha256:1db8e35b67b3d218d818ae653e27f06c3aa420901fa7b081ca98cbedc874e0d0"},
     {file = "pure_eval-0.2.3.tar.gz", hash = "sha256:5f4e983f40564c576c7c8635ae88db5956bb2229d7e9237d03b3c0b0190eaf42"},
@@ -938,6 +1716,7 @@ version = "2.18.0"
 description = "Pygments is a syntax highlighting package written in Python."
 optional = false
 python-versions = ">=3.8"
+groups = ["main"]
 files = [
     {file = "pygments-2.18.0-py3-none-any.whl", hash = "sha256:b8e6aca0523f3ab76fee51799c488e38782ac06eafcf95e7ba832985c8e7b13a"},
     {file = "pygments-2.18.0.tar.gz", hash = "sha256:786ff802f32e91311bff3889f6e9a86e81505fe99f2735bb6d60ae0c5004f199"},
@@ -952,6 +1731,7 @@ version = "3.1.4"
 description = "pyparsing module - Classes and methods to define and execute parsing grammars"
 optional = false
 python-versions = ">=3.6.8"
+groups = ["main"]
 files = [
     {file = "pyparsing-3.1.4-py3-none-any.whl", hash = "sha256:a6a7ee4235a3f944aa1fa2249307708f893fe5717dc603503c6c7969c070fb7c"},
     {file = "pyparsing-3.1.4.tar.gz", hash = "sha256:f86ec8d1a83f11977c9a6ea7598e8c27fc5cddfa5b07ea2241edbbde1d7bc032"},
@@ -966,6 +1746,7 @@ version = "2.9.0.post0"
 description = "Extensions to the standard Python datetime module"
 optional = false
 python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,>=2.7"
+groups = ["main"]
 files = [
     {file = "python-dateutil-2.9.0.post0.tar.gz", hash = "sha256:37dd54208da7e1cd875388217d5e00ebd4179249f90fb72437e91a35459a0ad3"},
     {file = "python_dateutil-2.9.0.post0-py2.py3-none-any.whl", hash = "sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427"},
@@ -980,6 +1761,7 @@ version = "2024.2"
 description = "World timezone definitions, modern and historical"
 optional = false
 python-versions = "*"
+groups = ["main"]
 files = [
     {file = "pytz-2024.2-py2.py3-none-any.whl", hash = "sha256:31c7c1817eb7fae7ca4b8c7ee50c72f93aa2dd863de768e1ef4245d426aa0725"},
     {file = "pytz-2024.2.tar.gz", hash = "sha256:2aa355083c50a0f93fa581709deac0c9ad65cca8a9e9beac660adcbd493c798a"},
@@ -991,6 +1773,7 @@ version = "0.3.2"
 description = "A Python network graph visualization library"
 optional = false
 python-versions = ">3.6"
+groups = ["main"]
 files = [
     {file = "pyvis-0.3.2-py3-none-any.whl", hash = "sha256:5720c4ca8161dc5d9ab352015723abb7a8bb8fb443edeb07f7a322db34a97555"},
 ]
@@ -1001,12 +1784,95 @@ jinja2 = ">=2.9.6"
 jsonpickle = ">=1.4.1"
 networkx = ">=1.11"
 
+[[package]]
+name = "pyyaml"
+version = "6.0.2"
+description = "YAML parser and emitter for Python"
+optional = false
+python-versions = ">=3.8"
+groups = ["main"]
+files = [
+    {file = "PyYAML-6.0.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:0a9a2848a5b7feac301353437eb7d5957887edbf81d56e903999a75a3d743086"},
+    {file = "PyYAML-6.0.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:29717114e51c84ddfba879543fb232a6ed60086602313ca38cce623c1d62cfbf"},
+    {file = "PyYAML-6.0.2-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8824b5a04a04a047e72eea5cec3bc266db09e35de6bdfe34c9436ac5ee27d237"},
+    {file = "PyYAML-6.0.2-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:7c36280e6fb8385e520936c3cb3b8042851904eba0e58d277dca80a5cfed590b"},
+    {file = "PyYAML-6.0.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ec031d5d2feb36d1d1a24380e4db6d43695f3748343d99434e6f5f9156aaa2ed"},
+    {file = "PyYAML-6.0.2-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:936d68689298c36b53b29f23c6dbb74de12b4ac12ca6cfe0e047bedceea56180"},
+    {file = "PyYAML-6.0.2-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:23502f431948090f597378482b4812b0caae32c22213aecf3b55325e049a6c68"},
+    {file = "PyYAML-6.0.2-cp310-cp310-win32.whl", hash = "sha256:2e99c6826ffa974fe6e27cdb5ed0021786b03fc98e5ee3c5bfe1fd5015f42b99"},
+    {file = "PyYAML-6.0.2-cp310-cp310-win_amd64.whl", hash = "sha256:a4d3091415f010369ae4ed1fc6b79def9416358877534caf6a0fdd2146c87a3e"},
+    {file = "PyYAML-6.0.2-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:cc1c1159b3d456576af7a3e4d1ba7e6924cb39de8f67111c735f6fc832082774"},
+    {file = "PyYAML-6.0.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:1e2120ef853f59c7419231f3bf4e7021f1b936f6ebd222406c3b60212205d2ee"},
+    {file = "PyYAML-6.0.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:5d225db5a45f21e78dd9358e58a98702a0302f2659a3c6cd320564b75b86f47c"},
+    {file = "PyYAML-6.0.2-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:5ac9328ec4831237bec75defaf839f7d4564be1e6b25ac710bd1a96321cc8317"},
+    {file = "PyYAML-6.0.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3ad2a3decf9aaba3d29c8f537ac4b243e36bef957511b4766cb0057d32b0be85"},
+    {file = "PyYAML-6.0.2-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:ff3824dc5261f50c9b0dfb3be22b4567a6f938ccce4587b38952d85fd9e9afe4"},
+    {file = "PyYAML-6.0.2-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:797b4f722ffa07cc8d62053e4cff1486fa6dc094105d13fea7b1de7d8bf71c9e"},
+    {file = "PyYAML-6.0.2-cp311-cp311-win32.whl", hash = "sha256:11d8f3dd2b9c1207dcaf2ee0bbbfd5991f571186ec9cc78427ba5bd32afae4b5"},
+    {file = "PyYAML-6.0.2-cp311-cp311-win_amd64.whl", hash = "sha256:e10ce637b18caea04431ce14fabcf5c64a1c61ec9c56b071a4b7ca131ca52d44"},
+    {file = "PyYAML-6.0.2-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:c70c95198c015b85feafc136515252a261a84561b7b1d51e3384e0655ddf25ab"},
+    {file = "PyYAML-6.0.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:ce826d6ef20b1bc864f0a68340c8b3287705cae2f8b4b1d932177dcc76721725"},
+    {file = "PyYAML-6.0.2-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1f71ea527786de97d1a0cc0eacd1defc0985dcf6b3f17bb77dcfc8c34bec4dc5"},
+    {file = "PyYAML-6.0.2-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:9b22676e8097e9e22e36d6b7bda33190d0d400f345f23d4065d48f4ca7ae0425"},
+    {file = "PyYAML-6.0.2-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:80bab7bfc629882493af4aa31a4cfa43a4c57c83813253626916b8c7ada83476"},
+    {file = "PyYAML-6.0.2-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:0833f8694549e586547b576dcfaba4a6b55b9e96098b36cdc7ebefe667dfed48"},
+    {file = "PyYAML-6.0.2-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8b9c7197f7cb2738065c481a0461e50ad02f18c78cd75775628afb4d7137fb3b"},
+    {file = "PyYAML-6.0.2-cp312-cp312-win32.whl", hash = "sha256:ef6107725bd54b262d6dedcc2af448a266975032bc85ef0172c5f059da6325b4"},
+    {file = "PyYAML-6.0.2-cp312-cp312-win_amd64.whl", hash = "sha256:7e7401d0de89a9a855c839bc697c079a4af81cf878373abd7dc625847d25cbd8"},
+    {file = "PyYAML-6.0.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:efdca5630322a10774e8e98e1af481aad470dd62c3170801852d752aa7a783ba"},
+    {file = "PyYAML-6.0.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:50187695423ffe49e2deacb8cd10510bc361faac997de9efef88badc3bb9e2d1"},
+    {file = "PyYAML-6.0.2-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0ffe8360bab4910ef1b9e87fb812d8bc0a308b0d0eef8c8f44e0254ab3b07133"},
+    {file = "PyYAML-6.0.2-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:17e311b6c678207928d649faa7cb0d7b4c26a0ba73d41e99c4fff6b6c3276484"},
+    {file = "PyYAML-6.0.2-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:70b189594dbe54f75ab3a1acec5f1e3faa7e8cf2f1e08d9b561cb41b845f69d5"},
+    {file = "PyYAML-6.0.2-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:41e4e3953a79407c794916fa277a82531dd93aad34e29c2a514c2c0c5fe971cc"},
+    {file = "PyYAML-6.0.2-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:68ccc6023a3400877818152ad9a1033e3db8625d899c72eacb5a668902e4d652"},
+    {file = "PyYAML-6.0.2-cp313-cp313-win32.whl", hash = "sha256:bc2fa7c6b47d6bc618dd7fb02ef6fdedb1090ec036abab80d4681424b84c1183"},
+    {file = "PyYAML-6.0.2-cp313-cp313-win_amd64.whl", hash = "sha256:8388ee1976c416731879ac16da0aff3f63b286ffdd57cdeb95f3f2e085687563"},
+    {file = "PyYAML-6.0.2-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:24471b829b3bf607e04e88d79542a9d48bb037c2267d7927a874e6c205ca7e9a"},
+    {file = "PyYAML-6.0.2-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d7fded462629cfa4b685c5416b949ebad6cec74af5e2d42905d41e257e0869f5"},
+    {file = "PyYAML-6.0.2-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:d84a1718ee396f54f3a086ea0a66d8e552b2ab2017ef8b420e92edbc841c352d"},
+    {file = "PyYAML-6.0.2-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9056c1ecd25795207ad294bcf39f2db3d845767be0ea6e6a34d856f006006083"},
+    {file = "PyYAML-6.0.2-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:82d09873e40955485746739bcb8b4586983670466c23382c19cffecbf1fd8706"},
+    {file = "PyYAML-6.0.2-cp38-cp38-win32.whl", hash = "sha256:43fa96a3ca0d6b1812e01ced1044a003533c47f6ee8aca31724f78e93ccc089a"},
+    {file = "PyYAML-6.0.2-cp38-cp38-win_amd64.whl", hash = "sha256:01179a4a8559ab5de078078f37e5c1a30d76bb88519906844fd7bdea1b7729ff"},
+    {file = "PyYAML-6.0.2-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:688ba32a1cffef67fd2e9398a2efebaea461578b0923624778664cc1c914db5d"},
+    {file = "PyYAML-6.0.2-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:a8786accb172bd8afb8be14490a16625cbc387036876ab6ba70912730faf8e1f"},
+    {file = "PyYAML-6.0.2-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d8e03406cac8513435335dbab54c0d385e4a49e4945d2909a581c83647ca0290"},
+    {file = "PyYAML-6.0.2-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:f753120cb8181e736c57ef7636e83f31b9c0d1722c516f7e86cf15b7aa57ff12"},
+    {file = "PyYAML-6.0.2-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3b1fdb9dc17f5a7677423d508ab4f243a726dea51fa5e70992e59a7411c89d19"},
+    {file = "PyYAML-6.0.2-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:0b69e4ce7a131fe56b7e4d770c67429700908fc0752af059838b1cfb41960e4e"},
+    {file = "PyYAML-6.0.2-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:a9f8c2e67970f13b16084e04f134610fd1d374bf477b17ec1599185cf611d725"},
+    {file = "PyYAML-6.0.2-cp39-cp39-win32.whl", hash = "sha256:6395c297d42274772abc367baaa79683958044e5d3835486c16da75d2a694631"},
+    {file = "PyYAML-6.0.2-cp39-cp39-win_amd64.whl", hash = "sha256:39693e1f8320ae4f43943590b49779ffb98acb81f788220ea932a6b6c51004d8"},
+    {file = "pyyaml-6.0.2.tar.gz", hash = "sha256:d584d9ec91ad65861cc08d42e834324ef890a082e591037abe114850ff7bbc3e"},
+]
+
+[[package]]
+name = "rich"
+version = "13.9.4"
+description = "Render rich text, tables, progress bars, syntax highlighting, markdown and more to the terminal"
+optional = false
+python-versions = ">=3.8.0"
+groups = ["main"]
+files = [
+    {file = "rich-13.9.4-py3-none-any.whl", hash = "sha256:6049d5e6ec054bf2779ab3358186963bac2ea89175919d699e378b99738c2a90"},
+    {file = "rich-13.9.4.tar.gz", hash = "sha256:439594978a49a09530cff7ebc4b5c7103ef57baf48d5ea3184f21d9a2befa098"},
+]
+
+[package.dependencies]
+markdown-it-py = ">=2.2.0"
+pygments = ">=2.13.0,<3.0.0"
+
+[package.extras]
+jupyter = ["ipywidgets (>=7.5.1,<9)"]
+
 [[package]]
 name = "scipy"
 version = "1.14.1"
 description = "Fundamental algorithms for scientific computing in Python"
 optional = false
 python-versions = ">=3.10"
+groups = ["main"]
 files = [
     {file = "scipy-1.14.1-cp310-cp310-macosx_10_13_x86_64.whl", hash = "sha256:b28d2ca4add7ac16ae8bb6632a3c86e4b9e4d52d3e34267f6e1b0c1f8d87e389"},
     {file = "scipy-1.14.1-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:d0d2821003174de06b69e58cef2316a6622b60ee613121199cb2852a873f8cf3"},
@@ -1049,7 +1915,7 @@ numpy = ">=1.23.5,<2.3"
 [package.extras]
 dev = ["cython-lint (>=0.12.2)", "doit (>=0.36.0)", "mypy (==1.10.0)", "pycodestyle", "pydevtool", "rich-click", "ruff (>=0.0.292)", "types-psutil", "typing_extensions"]
 doc = ["jupyterlite-pyodide-kernel", "jupyterlite-sphinx (>=0.13.1)", "jupytext", "matplotlib (>=3.5)", "myst-nb", "numpydoc", "pooch", "pydata-sphinx-theme (>=0.15.2)", "sphinx (>=5.0.0,<=7.3.7)", "sphinx-design (>=0.4.0)"]
-test = ["Cython", "array-api-strict (>=2.0)", "asv", "gmpy2", "hypothesis (>=6.30)", "meson", "mpmath", "ninja", "pooch", "pytest", "pytest-cov", "pytest-timeout", "pytest-xdist", "scikit-umfpack", "threadpoolctl"]
+test = ["Cython", "array-api-strict (>=2.0)", "asv", "gmpy2", "hypothesis (>=6.30)", "meson", "mpmath", "ninja ; sys_platform != \"emscripten\"", "pooch", "pytest", "pytest-cov", "pytest-timeout", "pytest-xdist", "scikit-umfpack", "threadpoolctl"]
 
 [[package]]
 name = "semantic-version"
@@ -1057,13 +1923,14 @@ version = "2.10.0"
 description = "A library implementing the 'SemVer' scheme."
 optional = false
 python-versions = ">=2.7"
+groups = ["main"]
 files = [
     {file = "semantic_version-2.10.0-py2.py3-none-any.whl", hash = "sha256:de78a3b8e0feda74cabc54aab2da702113e33ac9d9eb9d2389bcf1f58b7d9177"},
     {file = "semantic_version-2.10.0.tar.gz", hash = "sha256:bdabb6d336998cbb378d4b9db3a4b56a1e3235701dc05ea2690d9a997ed5041c"},
 ]
 
 [package.extras]
-dev = ["Django (>=1.11)", "check-manifest", "colorama (<=0.4.1)", "coverage", "flake8", "nose2", "readme-renderer (<25.0)", "tox", "wheel", "zest.releaser[recommended]"]
+dev = ["Django (>=1.11)", "check-manifest", "colorama (<=0.4.1) ; python_version == \"3.4\"", "coverage", "flake8", "nose2", "readme-renderer (<25.0) ; python_version == \"3.4\"", "tox", "wheel", "zest.releaser[recommended]"]
 doc = ["Sphinx", "sphinx-rtd-theme"]
 
 [[package]]
@@ -1072,15 +1939,16 @@ version = "72.2.0"
 description = "Easily download, build, install, upgrade, and uninstall Python packages"
 optional = false
 python-versions = ">=3.8"
+groups = ["main"]
 files = [
     {file = "setuptools-72.2.0-py3-none-any.whl", hash = "sha256:f11dd94b7bae3a156a95ec151f24e4637fb4fa19c878e4d191bfb8b2d82728c4"},
     {file = "setuptools-72.2.0.tar.gz", hash = "sha256:80aacbf633704e9c8bfa1d99fa5dd4dc59573efcf9e4042c13d3bcef91ac2ef9"},
 ]
 
 [package.extras]
-core = ["importlib-metadata (>=6)", "importlib-resources (>=5.10.2)", "jaraco.text (>=3.7)", "more-itertools (>=8.8)", "ordered-set (>=3.1.1)", "packaging (>=24)", "platformdirs (>=2.6.2)", "tomli (>=2.0.1)", "wheel (>=0.43.0)"]
+core = ["importlib-metadata (>=6) ; python_version < \"3.10\"", "importlib-resources (>=5.10.2) ; python_version < \"3.9\"", "jaraco.text (>=3.7)", "more-itertools (>=8.8)", "ordered-set (>=3.1.1)", "packaging (>=24)", "platformdirs (>=2.6.2)", "tomli (>=2.0.1) ; python_version < \"3.11\"", "wheel (>=0.43.0)"]
 doc = ["furo", "jaraco.packaging (>=9.3)", "jaraco.tidelift (>=1.4)", "pygments-github-lexers (==0.0.5)", "pyproject-hooks (!=1.1)", "rst.linker (>=1.9)", "sphinx (>=3.5)", "sphinx-favicon", "sphinx-inline-tabs", "sphinx-lint", "sphinx-notfound-page (>=1,<2)", "sphinx-reredirects", "sphinxcontrib-towncrier", "towncrier (<24.7)"]
-test = ["build[virtualenv] (>=1.0.3)", "filelock (>=3.4.0)", "importlib-metadata", "ini2toml[lite] (>=0.14)", "jaraco.develop (>=7.21)", "jaraco.envs (>=2.2)", "jaraco.path (>=3.2.0)", "jaraco.test", "mypy (==1.11.*)", "packaging (>=23.2)", "pip (>=19.1)", "pyproject-hooks (!=1.1)", "pytest (>=6,!=8.1.*)", "pytest-checkdocs (>=2.4)", "pytest-cov", "pytest-enabler (>=2.2)", "pytest-home (>=0.5)", "pytest-mypy", "pytest-perf", "pytest-ruff (<0.4)", "pytest-ruff (>=0.2.1)", "pytest-ruff (>=0.3.2)", "pytest-subprocess", "pytest-timeout", "pytest-xdist (>=3)", "tomli", "tomli-w (>=1.0.0)", "virtualenv (>=13.0.0)", "wheel"]
+test = ["build[virtualenv] (>=1.0.3)", "filelock (>=3.4.0)", "importlib-metadata", "ini2toml[lite] (>=0.14)", "jaraco.develop (>=7.21) ; python_version >= \"3.9\" and sys_platform != \"cygwin\"", "jaraco.envs (>=2.2)", "jaraco.path (>=3.2.0)", "jaraco.test", "mypy (==1.11.*)", "packaging (>=23.2)", "pip (>=19.1)", "pyproject-hooks (!=1.1)", "pytest (>=6,!=8.1.*)", "pytest-checkdocs (>=2.4)", "pytest-cov", "pytest-enabler (>=2.2)", "pytest-home (>=0.5)", "pytest-mypy", "pytest-perf ; sys_platform != \"cygwin\"", "pytest-ruff (<0.4) ; platform_system == \"Windows\"", "pytest-ruff (>=0.2.1) ; sys_platform != \"cygwin\"", "pytest-ruff (>=0.3.2) ; sys_platform != \"cygwin\"", "pytest-subprocess", "pytest-timeout", "pytest-xdist (>=3)", "tomli", "tomli-w (>=1.0.0)", "virtualenv (>=13.0.0)", "wheel"]
 
 [[package]]
 name = "setuptools-rust"
@@ -1088,6 +1956,7 @@ version = "1.10.1"
 description = "Setuptools Rust extension plugin"
 optional = false
 python-versions = ">=3.8"
+groups = ["main"]
 files = [
     {file = "setuptools_rust-1.10.1-py3-none-any.whl", hash = "sha256:3837616cc0a7705b2c44058f626c97f774eeb980f28427c16ece562661bc20c5"},
     {file = "setuptools_rust-1.10.1.tar.gz", hash = "sha256:d79035fc54cdf9342e9edf4b009491ecab06c3a652b37c3c137c7ba85547d3e6"},
@@ -1097,12 +1966,133 @@ files = [
 semantic-version = ">=2.8.2,<3"
 setuptools = ">=62.4"
 
+[[package]]
+name = "simplejson"
+version = "3.20.1"
+description = "Simple, fast, extensible JSON encoder/decoder for Python"
+optional = false
+python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,>=2.5"
+groups = ["main"]
+files = [
+    {file = "simplejson-3.20.1-cp27-cp27m-manylinux1_i686.whl", hash = "sha256:f5272b5866b259fe6c33c4a8c5073bf8b359c3c97b70c298a2f09a69b52c7c41"},
+    {file = "simplejson-3.20.1-cp27-cp27m-manylinux1_x86_64.whl", hash = "sha256:5c0de368f3052a59a1acf21f8b2dd28686a9e4eba2da7efae7ed9554cb31e7bc"},
+    {file = "simplejson-3.20.1-cp27-cp27m-manylinux2010_i686.whl", hash = "sha256:0821871404a537fd0e22eba240c74c0467c28af6cc435903eca394cfc74a0497"},
+    {file = "simplejson-3.20.1-cp27-cp27m-manylinux2010_x86_64.whl", hash = "sha256:c939a1e576bded47d7d03aa2afc2ae90b928b2cf1d9dc2070ceec51fd463f430"},
+    {file = "simplejson-3.20.1-cp27-cp27mu-manylinux1_i686.whl", hash = "sha256:3c4f0a61cdc05550782ca4a2cdb311ea196c2e6be6b24a09bf71360ca8c3ca9b"},
+    {file = "simplejson-3.20.1-cp27-cp27mu-manylinux1_x86_64.whl", hash = "sha256:6c21f5c026ca633cfffcb6bc1fac2e99f65cb2b24657d3bef21aed9916cc3bbf"},
+    {file = "simplejson-3.20.1-cp27-cp27mu-manylinux2010_i686.whl", hash = "sha256:8d23b7f8d6b72319d6d55a0261089ff621ce87e54731c2d3de6a9bf7be5c028c"},
+    {file = "simplejson-3.20.1-cp27-cp27mu-manylinux2010_x86_64.whl", hash = "sha256:cda5c32a98f392909088111ecec23f2b0d39346ceae1a0fea23ab2d1f84ec21d"},
+    {file = "simplejson-3.20.1-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:e580aa65d5f6c3bf41b9b4afe74be5d5ddba9576701c107c772d936ea2b5043a"},
+    {file = "simplejson-3.20.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:4a586ce4f78cec11f22fe55c5bee0f067e803aab9bad3441afe2181693b5ebb5"},
+    {file = "simplejson-3.20.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:74a1608f9e6e8c27a4008d70a54270868306d80ed48c9df7872f9f4b8ac87808"},
+    {file = "simplejson-3.20.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:03db8cb64154189a92a7786209f24e391644f3a3fa335658be2df2af1960b8d8"},
+    {file = "simplejson-3.20.1-cp310-cp310-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:eea7e2b7d858f6fdfbf0fe3cb846d6bd8a45446865bc09960e51f3d473c2271b"},
+    {file = "simplejson-3.20.1-cp310-cp310-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:e66712b17d8425bb7ff8968d4c7c7fd5a2dd7bd63728b28356223c000dd2f91f"},
+    {file = "simplejson-3.20.1-cp310-cp310-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a2cc4f6486f9f515b62f5831ff1888886619b84fc837de68f26d919ba7bbdcbc"},
+    {file = "simplejson-3.20.1-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:a3c2df555ee4016148fa192e2b9cd9e60bc1d40769366134882685e90aee2a1e"},
+    {file = "simplejson-3.20.1-cp310-cp310-musllinux_1_2_i686.whl", hash = "sha256:78520f04b7548a5e476b5396c0847e066f1e0a4c0c5e920da1ad65e95f410b11"},
+    {file = "simplejson-3.20.1-cp310-cp310-musllinux_1_2_ppc64le.whl", hash = "sha256:f4bd49ecde87b0fe9f55cc971449a32832bca9910821f7072bbfae1155eaa007"},
+    {file = "simplejson-3.20.1-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:7eaae2b88eb5da53caaffdfa50e2e12022553949b88c0df4f9a9663609373f72"},
+    {file = "simplejson-3.20.1-cp310-cp310-win32.whl", hash = "sha256:e836fb88902799eac8debc2b642300748f4860a197fa3d9ea502112b6bb8e142"},
+    {file = "simplejson-3.20.1-cp310-cp310-win_amd64.whl", hash = "sha256:b122a19b552b212fc3b5b96fc5ce92333d4a9ac0a800803e1f17ebb16dac4be5"},
+    {file = "simplejson-3.20.1-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:325b8c107253d3217e89d7b50c71015b5b31e2433e6c5bf38967b2f80630a8ca"},
+    {file = "simplejson-3.20.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:88a7baa8211089b9e58d78fbc1b0b322103f3f3d459ff16f03a36cece0d0fcf0"},
+    {file = "simplejson-3.20.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:299b1007b8101d50d95bc0db1bf5c38dc372e85b504cf77f596462083ee77e3f"},
+    {file = "simplejson-3.20.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:03ec618ed65caab48e81e3ed29586236a8e57daef792f1f3bb59504a7e98cd10"},
+    {file = "simplejson-3.20.1-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:cd2cdead1d3197f0ff43373cf4730213420523ba48697743e135e26f3d179f38"},
+    {file = "simplejson-3.20.1-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:3466d2839fdc83e1af42e07b90bc8ff361c4e8796cd66722a40ba14e458faddd"},
+    {file = "simplejson-3.20.1-cp311-cp311-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d492ed8e92f3a9f9be829205f44b1d0a89af6582f0cf43e0d129fa477b93fe0c"},
+    {file = "simplejson-3.20.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:f924b485537b640dc69434565463fd6fc0c68c65a8c6e01a823dd26c9983cf79"},
+    {file = "simplejson-3.20.1-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:9e8eacf6a3491bf76ea91a8d46726368a6be0eb94993f60b8583550baae9439e"},
+    {file = "simplejson-3.20.1-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:d34d04bf90b4cea7c22d8b19091633908f14a096caa301b24c2f3d85b5068fb8"},
+    {file = "simplejson-3.20.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:69dd28d4ce38390ea4aaf212902712c0fd1093dc4c1ff67e09687c3c3e15a749"},
+    {file = "simplejson-3.20.1-cp311-cp311-win32.whl", hash = "sha256:dfe7a9da5fd2a3499436cd350f31539e0a6ded5da6b5b3d422df016444d65e43"},
+    {file = "simplejson-3.20.1-cp311-cp311-win_amd64.whl", hash = "sha256:896a6c04d7861d507d800da7642479c3547060bf97419d9ef73d98ced8258766"},
+    {file = "simplejson-3.20.1-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:f31c4a3a7ab18467ee73a27f3e59158255d1520f3aad74315edde7a940f1be23"},
+    {file = "simplejson-3.20.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:884e6183d16b725e113b83a6fc0230152ab6627d4d36cb05c89c2c5bccfa7bc6"},
+    {file = "simplejson-3.20.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:03d7a426e416fe0d3337115f04164cd9427eb4256e843a6b8751cacf70abc832"},
+    {file = "simplejson-3.20.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:000602141d0bddfcff60ea6a6e97d5e10c9db6b17fd2d6c66199fa481b6214bb"},
+    {file = "simplejson-3.20.1-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:af8377a8af78226e82e3a4349efdde59ffa421ae88be67e18cef915e4023a595"},
+    {file = "simplejson-3.20.1-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:15c7de4c88ab2fbcb8781a3b982ef883696736134e20b1210bca43fb42ff1acf"},
+    {file = "simplejson-3.20.1-cp312-cp312-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:455a882ff3f97d810709f7b620007d4e0aca8da71d06fc5c18ba11daf1c4df49"},
+    {file = "simplejson-3.20.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:fc0f523ce923e7f38eb67804bc80e0a028c76d7868500aa3f59225574b5d0453"},
+    {file = "simplejson-3.20.1-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:76461ec929282dde4a08061071a47281ad939d0202dc4e63cdd135844e162fbc"},
+    {file = "simplejson-3.20.1-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:ab19c2da8c043607bde4d4ef3a6b633e668a7d2e3d56f40a476a74c5ea71949f"},
+    {file = "simplejson-3.20.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:b2578bedaedf6294415197b267d4ef678fea336dd78ee2a6d2f4b028e9d07be3"},
+    {file = "simplejson-3.20.1-cp312-cp312-win32.whl", hash = "sha256:339f407373325a36b7fd744b688ba5bae0666b5d340ec6d98aebc3014bf3d8ea"},
+    {file = "simplejson-3.20.1-cp312-cp312-win_amd64.whl", hash = "sha256:627d4486a1ea7edf1f66bb044ace1ce6b4c1698acd1b05353c97ba4864ea2e17"},
+    {file = "simplejson-3.20.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:71e849e7ceb2178344998cbe5ade101f1b329460243c79c27fbfc51c0447a7c3"},
+    {file = "simplejson-3.20.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:b63fdbab29dc3868d6f009a59797cefaba315fd43cd32ddd998ee1da28e50e29"},
+    {file = "simplejson-3.20.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:1190f9a3ce644fd50ec277ac4a98c0517f532cfebdcc4bd975c0979a9f05e1fb"},
+    {file = "simplejson-3.20.1-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c1336ba7bcb722ad487cd265701ff0583c0bb6de638364ca947bb84ecc0015d1"},
+    {file = "simplejson-3.20.1-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:e975aac6a5acd8b510eba58d5591e10a03e3d16c1cf8a8624ca177491f7230f0"},
+    {file = "simplejson-3.20.1-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:6a6dd11ee282937ad749da6f3b8d87952ad585b26e5edfa10da3ae2536c73078"},
+    {file = "simplejson-3.20.1-cp313-cp313-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ab980fcc446ab87ea0879edad41a5c28f2d86020014eb035cf5161e8de4474c6"},
+    {file = "simplejson-3.20.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:f5aee2a4cb6b146bd17333ac623610f069f34e8f31d2f4f0c1a2186e50c594f0"},
+    {file = "simplejson-3.20.1-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:652d8eecbb9a3b6461b21ec7cf11fd0acbab144e45e600c817ecf18e4580b99e"},
+    {file = "simplejson-3.20.1-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:8c09948f1a486a89251ee3a67c9f8c969b379f6ffff1a6064b41fea3bce0a112"},
+    {file = "simplejson-3.20.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:cbbd7b215ad4fc6f058b5dd4c26ee5c59f72e031dfda3ac183d7968a99e4ca3a"},
+    {file = "simplejson-3.20.1-cp313-cp313-win32.whl", hash = "sha256:ae81e482476eaa088ef9d0120ae5345de924f23962c0c1e20abbdff597631f87"},
+    {file = "simplejson-3.20.1-cp313-cp313-win_amd64.whl", hash = "sha256:1b9fd15853b90aec3b1739f4471efbf1ac05066a2c7041bf8db821bb73cd2ddc"},
+    {file = "simplejson-3.20.1-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:c7edf279c1376f28bf41e916c015a2a08896597869d57d621f55b6a30c7e1e6d"},
+    {file = "simplejson-3.20.1-cp36-cp36m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d9202b9de38f12e99a40addd1a8d508a13c77f46d87ab1f9095f154667f4fe81"},
+    {file = "simplejson-3.20.1-cp36-cp36m-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:391345b4157cc4e120027e013bd35c45e2c191e2bf48b8913af488cdc3b9243c"},
+    {file = "simplejson-3.20.1-cp36-cp36m-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:c6fdcc9debb711ddd2ad6d69f9386a3d9e8e253234bbb30513e0a7caa9510c51"},
+    {file = "simplejson-3.20.1-cp36-cp36m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9daf8cdc7ee8a9e9f7a3b313ba0a003391857e90d0e82fbcd4d614aa05cb7c3b"},
+    {file = "simplejson-3.20.1-cp36-cp36m-musllinux_1_2_aarch64.whl", hash = "sha256:c02f4868a3a46ffe284a51a88d134dc96feff6079a7115164885331a1ba8ed9f"},
+    {file = "simplejson-3.20.1-cp36-cp36m-musllinux_1_2_i686.whl", hash = "sha256:3d7310172d5340febd258cb147f46aae30ad57c445f4d7e1ae8461c10aaf43b0"},
+    {file = "simplejson-3.20.1-cp36-cp36m-musllinux_1_2_ppc64le.whl", hash = "sha256:4762e05577955312a4c6802f58dd02e040cc79ae59cda510aa1564d84449c102"},
+    {file = "simplejson-3.20.1-cp36-cp36m-musllinux_1_2_x86_64.whl", hash = "sha256:8bb98fdf318c05aefd08a92583bd6ee148e93c6756fb1befb7b2d5f27824be78"},
+    {file = "simplejson-3.20.1-cp36-cp36m-win32.whl", hash = "sha256:9a74e70818818981294b8e6956ce3496c5e1bd4726ac864fae473197671f7b85"},
+    {file = "simplejson-3.20.1-cp36-cp36m-win_amd64.whl", hash = "sha256:e041add470e8f8535cc05509485eb7205729a84441f03b25cde80ad48823792e"},
+    {file = "simplejson-3.20.1-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:7e9d73f46119240e4f4f07868241749d67d09873f40cb968d639aa9ccc488b86"},
+    {file = "simplejson-3.20.1-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ae6e637dc24f8fee332ed23dd070e81394138e42cd4fd9d0923e5045ba122e27"},
+    {file = "simplejson-3.20.1-cp37-cp37m-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:efd3bc6c6b17e3d4620eb6be5196f0d1c08b6ce7c3101fa8e292b79e0908944b"},
+    {file = "simplejson-3.20.1-cp37-cp37m-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:87fc623d457173a0213bc9ca4e346b83c9d443f63ed5cca847fb0cacea3cfc95"},
+    {file = "simplejson-3.20.1-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ec6a1e0a7aff76f0e008bebfa950188b9c50b58c1885d898145f48fc8e189a56"},
+    {file = "simplejson-3.20.1-cp37-cp37m-musllinux_1_2_aarch64.whl", hash = "sha256:9c079606f461a6e950099167e21e13985147c8a24be8eea66c9ad68f73fad744"},
+    {file = "simplejson-3.20.1-cp37-cp37m-musllinux_1_2_i686.whl", hash = "sha256:9faceb68fba27ef17eda306e4cd97a7b4b14fdadca5fbb15790ba8b26ebeec0c"},
+    {file = "simplejson-3.20.1-cp37-cp37m-musllinux_1_2_ppc64le.whl", hash = "sha256:7ceed598e4bacbf5133fe7a418f7991bb2df0683f3ac11fbf9e36a2bc7aa4b85"},
+    {file = "simplejson-3.20.1-cp37-cp37m-musllinux_1_2_x86_64.whl", hash = "sha256:ede69c765e9901861ad7c6139023b7b7d5807c48a2539d817b4ab40018002d5f"},
+    {file = "simplejson-3.20.1-cp37-cp37m-win32.whl", hash = "sha256:d8853c269a4c5146ddca4aa7c70e631795e9d11239d5fedb1c6bbc91ffdebcac"},
+    {file = "simplejson-3.20.1-cp37-cp37m-win_amd64.whl", hash = "sha256:ed6a17fd397f0e2b3ad668fc9e19253ed2e3875ad9086bd7f795c29a3223f4a1"},
+    {file = "simplejson-3.20.1-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:7551682b60bba3a9e2780742e101cf0a64250e76de7d09b1c4b0c8a7c7cc6834"},
+    {file = "simplejson-3.20.1-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:bd9577ec1c8c3a43040e3787711e4c257c70035b7551a21854b5dec88dad09e1"},
+    {file = "simplejson-3.20.1-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:4a8e197e4cf6d42c2c57e7c52cd7c1e7b3e37c5911df1314fb393320131e2101"},
+    {file = "simplejson-3.20.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6bd09c8c75666e7f62a33d2f1fb57f81da1fcbb19a9fe7d7910b5756e1dd6048"},
+    {file = "simplejson-3.20.1-cp38-cp38-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:1bd6bfe5678d73fbd5328eea6a35216503796428fc47f1237432522febaf3a0c"},
+    {file = "simplejson-3.20.1-cp38-cp38-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:71b75d448fd0ceb2e7c90e72bb82c41f8462550d48529980bc0bab1d2495bfbb"},
+    {file = "simplejson-3.20.1-cp38-cp38-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a7e15b716d09f318c8cda3e20f82fae81684ce3d3acd1d7770fa3007df1769de"},
+    {file = "simplejson-3.20.1-cp38-cp38-musllinux_1_2_aarch64.whl", hash = "sha256:3e7963197d958fcf9e98b212b80977d56c022384621ff463d98afc3b6b1ce7e8"},
+    {file = "simplejson-3.20.1-cp38-cp38-musllinux_1_2_i686.whl", hash = "sha256:2e671dd62051129185d3a9a92c60101f56cbc174854a1a3dfb69114ebd9e1699"},
+    {file = "simplejson-3.20.1-cp38-cp38-musllinux_1_2_ppc64le.whl", hash = "sha256:e25b2a0c396f3b84fb89573d07b0e1846ed563eb364f2ea8230ca92b8a8cb786"},
+    {file = "simplejson-3.20.1-cp38-cp38-musllinux_1_2_x86_64.whl", hash = "sha256:489c3a43116082bad56795215786313832ba3991cca1f55838e52a553f451ab6"},
+    {file = "simplejson-3.20.1-cp38-cp38-win32.whl", hash = "sha256:4a92e948bad8df7fa900ba2ba0667a98303f3db206cbaac574935c332838208e"},
+    {file = "simplejson-3.20.1-cp38-cp38-win_amd64.whl", hash = "sha256:49d059b8363327eee3c94799dd96782314b2dbd7bcc293b4ad48db69d6f4d362"},
+    {file = "simplejson-3.20.1-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:a8011f1dd1d676befcd4d675ebdbfdbbefd3bf350052b956ba8c699fca7d8cef"},
+    {file = "simplejson-3.20.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:e91703a4c5fec53e36875ae426ad785f4120bd1d93b65bed4752eeccd1789e0c"},
+    {file = "simplejson-3.20.1-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:e39eaa57c7757daa25bcd21f976c46be443b73dd6c3da47fe5ce7b7048ccefe2"},
+    {file = "simplejson-3.20.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ceab2ce2acdc7fbaa433a93006758db6ba9a659e80c4faa13b80b9d2318e9b17"},
+    {file = "simplejson-3.20.1-cp39-cp39-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:6d4f320c33277a5b715db5bf5b10dae10c19076bd6d66c2843e04bd12d1f1ea5"},
+    {file = "simplejson-3.20.1-cp39-cp39-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:2b6436c48e64378fa844d8c9e58a5ed0352bbcfd4028369a9b46679b7ab79d2d"},
+    {file = "simplejson-3.20.1-cp39-cp39-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6e18345c8dda5d699be8166b61f9d80aaee4545b709f1363f60813dc032dac53"},
+    {file = "simplejson-3.20.1-cp39-cp39-musllinux_1_2_aarch64.whl", hash = "sha256:90b573693d1526bed576f6817e2a492eaaef68f088b57d7a9e83d122bbb49e51"},
+    {file = "simplejson-3.20.1-cp39-cp39-musllinux_1_2_i686.whl", hash = "sha256:272cc767826e924a6bd369ea3dbf18e166ded29059c7a4d64d21a9a22424b5b5"},
+    {file = "simplejson-3.20.1-cp39-cp39-musllinux_1_2_ppc64le.whl", hash = "sha256:51b41f284d603c4380732d7d619f8b34bd04bc4aa0ed0ed5f4ffd0539b14da44"},
+    {file = "simplejson-3.20.1-cp39-cp39-musllinux_1_2_x86_64.whl", hash = "sha256:6e6697a3067d281f01de0fe96fc7cba4ea870d96d7deb7bfcf85186d74456503"},
+    {file = "simplejson-3.20.1-cp39-cp39-win32.whl", hash = "sha256:6dd3a1d5aca87bf947f3339b0f8e8e329f1badf548bdbff37fac63c17936da8e"},
+    {file = "simplejson-3.20.1-cp39-cp39-win_amd64.whl", hash = "sha256:463f1fca8fbf23d088e5850fdd0dd4d5faea8900a9f9680270bd98fd649814ca"},
+    {file = "simplejson-3.20.1-py3-none-any.whl", hash = "sha256:8a6c1bbac39fa4a79f83cbf1df6ccd8ff7069582a9fd8db1e52cea073bc2c697"},
+    {file = "simplejson-3.20.1.tar.gz", hash = "sha256:e64139b4ec4f1f24c142ff7dcafe55a22b811a74d86d66560c8815687143037d"},
+]
+
 [[package]]
 name = "six"
 version = "1.16.0"
 description = "Python 2 and 3 compatibility utilities"
 optional = false
 python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*"
+groups = ["main"]
 files = [
     {file = "six-1.16.0-py2.py3-none-any.whl", hash = "sha256:8abb2f1d86890a2dfb989f9a77cfcfd3e47c2a354b01111771326f8aa26e0254"},
     {file = "six-1.16.0.tar.gz", hash = "sha256:1e61c37477a1626458e36f7b1d82aa5c9b094fa4802892072e49de9c60c4c926"},
@@ -1114,6 +2104,7 @@ version = "0.6.3"
 description = "Extract data from python stack frames and tracebacks for informative displays"
 optional = false
 python-versions = "*"
+groups = ["main"]
 files = [
     {file = "stack_data-0.6.3-py3-none-any.whl", hash = "sha256:d5558e0c25a4cb0853cddad3d77da9891a08cb85dd9f9f91b9f8cd66e511e695"},
     {file = "stack_data-0.6.3.tar.gz", hash = "sha256:836a778de4fec4dcd1dcd89ed8abff8a221f58308462e1c4aa2a3cf30148f0b9"},
@@ -1127,12 +2118,140 @@ pure-eval = "*"
 [package.extras]
 tests = ["cython", "littleutils", "pygments", "pytest", "typeguard"]
 
+[[package]]
+name = "sympy"
+version = "1.14.0"
+description = "Computer algebra system (CAS) in Python"
+optional = false
+python-versions = ">=3.9"
+groups = ["main"]
+files = [
+    {file = "sympy-1.14.0-py3-none-any.whl", hash = "sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5"},
+    {file = "sympy-1.14.0.tar.gz", hash = "sha256:d3d3fe8df1e5a0b42f0e7bdf50541697dbe7d23746e894990c030e2b05e72517"},
+]
+
+[package.dependencies]
+mpmath = ">=1.1.0,<1.4"
+
+[package.extras]
+dev = ["hypothesis (>=6.70.0)", "pytest (>=7.1.0)"]
+
+[[package]]
+name = "tensorstore"
+version = "0.1.72"
+description = "Read and write large, multi-dimensional arrays"
+optional = false
+python-versions = ">=3.10"
+groups = ["main"]
+files = [
+    {file = "tensorstore-0.1.72-cp310-cp310-macosx_10_14_x86_64.whl", hash = "sha256:a41b4fe0603943d23472619a8ada70b8d2c9458747fad88b0ce7b29f1ccf4e74"},
+    {file = "tensorstore-0.1.72-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:170172b698fefb4b5507c6cb339ca0b75d56d12ba6a43d9569c61800c1eeb121"},
+    {file = "tensorstore-0.1.72-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b71134b85f540e17a1ae65da1fb906781b7470ef0ed71d98d29459325897f574"},
+    {file = "tensorstore-0.1.72-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:08c5318535aac5e20e247c6e9b43f5887b2293f548de7279650bc73804ccf3ed"},
+    {file = "tensorstore-0.1.72-cp310-cp310-win_amd64.whl", hash = "sha256:9113d3fcf78c1366688aa90ee7efdc86b57962ea72276944cc57e916a6180749"},
+    {file = "tensorstore-0.1.72-cp311-cp311-macosx_10_14_x86_64.whl", hash = "sha256:599cc7b26b0c96373e89ff5bcf9b76e832802169229680bef985b10011f9bae7"},
+    {file = "tensorstore-0.1.72-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:a7e7b02da26ca5c95b3c613efd0fe10c082dfa4dc3e9818fefc69e30fe70ea1e"},
+    {file = "tensorstore-0.1.72-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4a6825cdb6751663ca0bd9abd528ea354ad2199f549bf1f36feac79a6c06efe2"},
+    {file = "tensorstore-0.1.72-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ed916b9aeca242a3f367679f65ba376149251ebb28b873becd76c73b688399b6"},
+    {file = "tensorstore-0.1.72-cp311-cp311-win_amd64.whl", hash = "sha256:5d410c879dc4b34036ec38e20ff05c7e3b0ad5d1eb595412b27a9dbb5e435035"},
+    {file = "tensorstore-0.1.72-cp312-cp312-macosx_10_14_x86_64.whl", hash = "sha256:92fac5e2cbc90e5ca8fc72c5bf112816d981e266a3cf9fb1681ba8b3f59537ef"},
+    {file = "tensorstore-0.1.72-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:7c9413f8318a4fa259ec5325f569c0759bccee936df44bd2f7bb35c8afdcdfc8"},
+    {file = "tensorstore-0.1.72-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c0f722218f494b1631dbec451b9863f579054e27da2f39aab418db4493694abe"},
+    {file = "tensorstore-0.1.72-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d5dced3f367308e9fa8e7b72e9e57a4c491fa47c066e035ac33421e2b2408e3f"},
+    {file = "tensorstore-0.1.72-cp312-cp312-win_amd64.whl", hash = "sha256:721d599db0113d75ab6ba1365989bbaf2ab752d7a6268f975c8bfd3a8eb6084b"},
+    {file = "tensorstore-0.1.72-cp313-cp313-macosx_10_14_x86_64.whl", hash = "sha256:9c3a36f681ffcc104ba931d471447e8901e64e8cc6913b61792870ff59529961"},
+    {file = "tensorstore-0.1.72-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:0cd951e593a17babbbde1410cfadb4a04e1cddfa5ace0de5ccb41029223f96b9"},
+    {file = "tensorstore-0.1.72-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:2fdfa0118be0721c110bcbe7e464758f78d3e14ee8c30a911eb8f4465e6c2e81"},
+    {file = "tensorstore-0.1.72-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:5ed6fe937b0433b573c3d6805d0759d33ccc24aa2aba720e4b8ba689c2f9775f"},
+    {file = "tensorstore-0.1.72-cp313-cp313-win_amd64.whl", hash = "sha256:66c0658689243af0825fff222fb56fdf05a8553bcb3b471dbf18830161302986"},
+    {file = "tensorstore-0.1.72.tar.gz", hash = "sha256:763d7f6898711783f199c8226a9c0b259546f5c6d9b4dc0ad3c9e39627060022"},
+]
+
+[package.dependencies]
+ml_dtypes = ">=0.3.1"
+numpy = ">=1.22.0"
+
+[[package]]
+name = "toolz"
+version = "1.0.0"
+description = "List processing tools and functional utilities"
+optional = false
+python-versions = ">=3.8"
+groups = ["main"]
+files = [
+    {file = "toolz-1.0.0-py3-none-any.whl", hash = "sha256:292c8f1c4e7516bf9086f8850935c799a874039c8bcf959d47b600e4c44a6236"},
+    {file = "toolz-1.0.0.tar.gz", hash = "sha256:2c86e3d9a04798ac556793bced838816296a2f085017664e4995cb40a1047a02"},
+]
+
+[[package]]
+name = "torch"
+version = "2.7.1"
+description = "Tensors and Dynamic neural networks in Python with strong GPU acceleration"
+optional = false
+python-versions = ">=3.9.0"
+groups = ["main"]
+files = [
+    {file = "torch-2.7.1-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:a103b5d782af5bd119b81dbcc7ffc6fa09904c423ff8db397a1e6ea8fd71508f"},
+    {file = "torch-2.7.1-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:fe955951bdf32d182ee8ead6c3186ad54781492bf03d547d31771a01b3d6fb7d"},
+    {file = "torch-2.7.1-cp310-cp310-win_amd64.whl", hash = "sha256:885453d6fba67d9991132143bf7fa06b79b24352f4506fd4d10b309f53454162"},
+    {file = "torch-2.7.1-cp310-none-macosx_11_0_arm64.whl", hash = "sha256:d72acfdb86cee2a32c0ce0101606f3758f0d8bb5f8f31e7920dc2809e963aa7c"},
+    {file = "torch-2.7.1-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:236f501f2e383f1cb861337bdf057712182f910f10aeaf509065d54d339e49b2"},
+    {file = "torch-2.7.1-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:06eea61f859436622e78dd0cdd51dbc8f8c6d76917a9cf0555a333f9eac31ec1"},
+    {file = "torch-2.7.1-cp311-cp311-win_amd64.whl", hash = "sha256:8273145a2e0a3c6f9fd2ac36762d6ee89c26d430e612b95a99885df083b04e52"},
+    {file = "torch-2.7.1-cp311-none-macosx_11_0_arm64.whl", hash = "sha256:aea4fc1bf433d12843eb2c6b2204861f43d8364597697074c8d38ae2507f8730"},
+    {file = "torch-2.7.1-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:27ea1e518df4c9de73af7e8a720770f3628e7f667280bce2be7a16292697e3fa"},
+    {file = "torch-2.7.1-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:c33360cfc2edd976c2633b3b66c769bdcbbf0e0b6550606d188431c81e7dd1fc"},
+    {file = "torch-2.7.1-cp312-cp312-win_amd64.whl", hash = "sha256:d8bf6e1856ddd1807e79dc57e54d3335f2b62e6f316ed13ed3ecfe1fc1df3d8b"},
+    {file = "torch-2.7.1-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:787687087412c4bd68d315e39bc1223f08aae1d16a9e9771d95eabbb04ae98fb"},
+    {file = "torch-2.7.1-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:03563603d931e70722dce0e11999d53aa80a375a3d78e6b39b9f6805ea0a8d28"},
+    {file = "torch-2.7.1-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:d632f5417b6980f61404a125b999ca6ebd0b8b4bbdbb5fbbba44374ab619a412"},
+    {file = "torch-2.7.1-cp313-cp313-win_amd64.whl", hash = "sha256:23660443e13995ee93e3d844786701ea4ca69f337027b05182f5ba053ce43b38"},
+    {file = "torch-2.7.1-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:0da4f4dba9f65d0d203794e619fe7ca3247a55ffdcbd17ae8fb83c8b2dc9b585"},
+    {file = "torch-2.7.1-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:e08d7e6f21a617fe38eeb46dd2213ded43f27c072e9165dc27300c9ef9570934"},
+    {file = "torch-2.7.1-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:30207f672328a42df4f2174b8f426f354b2baa0b7cca3a0adb3d6ab5daf00dc8"},
+    {file = "torch-2.7.1-cp313-cp313t-win_amd64.whl", hash = "sha256:79042feca1c634aaf6603fe6feea8c6b30dfa140a6bbc0b973e2260c7e79a22e"},
+    {file = "torch-2.7.1-cp313-none-macosx_11_0_arm64.whl", hash = "sha256:988b0cbc4333618a1056d2ebad9eb10089637b659eb645434d0809d8d937b946"},
+    {file = "torch-2.7.1-cp39-cp39-manylinux_2_28_aarch64.whl", hash = "sha256:e0d81e9a12764b6f3879a866607c8ae93113cbcad57ce01ebde63eb48a576369"},
+    {file = "torch-2.7.1-cp39-cp39-manylinux_2_28_x86_64.whl", hash = "sha256:8394833c44484547ed4a47162318337b88c97acdb3273d85ea06e03ffff44998"},
+    {file = "torch-2.7.1-cp39-cp39-win_amd64.whl", hash = "sha256:df41989d9300e6e3c19ec9f56f856187a6ef060c3662fe54f4b6baf1fc90bd19"},
+    {file = "torch-2.7.1-cp39-none-macosx_11_0_arm64.whl", hash = "sha256:a737b5edd1c44a5c1ece2e9f3d00df9d1b3fb9541138bee56d83d38293fb6c9d"},
+]
+
+[package.dependencies]
+filelock = "*"
+fsspec = "*"
+jinja2 = "*"
+networkx = "*"
+nvidia-cublas-cu12 = {version = "12.6.4.1", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""}
+nvidia-cuda-cupti-cu12 = {version = "12.6.80", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""}
+nvidia-cuda-nvrtc-cu12 = {version = "12.6.77", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""}
+nvidia-cuda-runtime-cu12 = {version = "12.6.77", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""}
+nvidia-cudnn-cu12 = {version = "9.5.1.17", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""}
+nvidia-cufft-cu12 = {version = "11.3.0.4", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""}
+nvidia-cufile-cu12 = {version = "1.11.1.6", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""}
+nvidia-curand-cu12 = {version = "10.3.7.77", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""}
+nvidia-cusolver-cu12 = {version = "11.7.1.2", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""}
+nvidia-cusparse-cu12 = {version = "12.5.4.2", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""}
+nvidia-cusparselt-cu12 = {version = "0.6.3", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""}
+nvidia-nccl-cu12 = {version = "2.26.2", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""}
+nvidia-nvjitlink-cu12 = {version = "12.6.85", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""}
+nvidia-nvtx-cu12 = {version = "12.6.77", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""}
+setuptools = {version = "*", markers = "python_version >= \"3.12\""}
+sympy = ">=1.13.3"
+triton = {version = "3.3.1", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""}
+typing-extensions = ">=4.10.0"
+
+[package.extras]
+opt-einsum = ["opt-einsum (>=3.3)"]
+optree = ["optree (>=0.13.0)"]
+
 [[package]]
 name = "tqdm"
 version = "4.66.5"
 description = "Fast, Extensible Progress Meter"
 optional = false
 python-versions = ">=3.7"
+groups = ["main"]
 files = [
     {file = "tqdm-4.66.5-py3-none-any.whl", hash = "sha256:90279a3770753eafc9194a0364852159802111925aa30eb3f9d85b0e805ac7cd"},
     {file = "tqdm-4.66.5.tar.gz", hash = "sha256:e1020aef2e5096702d8a025ac7d16b1577279c9d63f8375b63083e9a5f0fcbad"},
@@ -1153,6 +2272,7 @@ version = "5.14.3"
 description = "Traitlets Python configuration system"
 optional = false
 python-versions = ">=3.8"
+groups = ["main"]
 files = [
     {file = "traitlets-5.14.3-py3-none-any.whl", hash = "sha256:b74e89e397b1ed28cc831db7aea759ba6640cb3de13090ca145426688ff1ac4f"},
     {file = "traitlets-5.14.3.tar.gz", hash = "sha256:9ed0579d3502c94b4b3732ac120375cda96f923114522847de4b3bb98b96b6b7"},
@@ -1162,12 +2282,71 @@ files = [
 docs = ["myst-parser", "pydata-sphinx-theme", "sphinx"]
 test = ["argcomplete (>=3.0.3)", "mypy (>=1.7.0)", "pre-commit", "pytest (>=7.0,<8.2)", "pytest-mock", "pytest-mypy-testing"]
 
+[[package]]
+name = "treescope"
+version = "0.1.9"
+description = "Treescope: An interactive HTML pretty-printer for ML research in IPython notebooks."
+optional = false
+python-versions = ">=3.10"
+groups = ["main"]
+files = [
+    {file = "treescope-0.1.9-py3-none-any.whl", hash = "sha256:68677013a9f0228212fccf835f3fb037be07ae8b4c5f6f58eefab11198f83cf7"},
+    {file = "treescope-0.1.9.tar.gz", hash = "sha256:ba6cdbdc9c5b52691d5f3bb4c5d5c7daa5627119acac8640b46d37e6aabe63a6"},
+]
+
+[package.dependencies]
+numpy = ">=1.25.2"
+
+[package.extras]
+dev = ["ipython", "jupyter", "pyink (>=24.3.0)", "pylint (>=2.6.0)", "pytest (>=8.2.2)", "pytype"]
+docs = ["ipython", "ipython (>=8.8.0)", "jax[cpu] (>=0.4.23)", "matplotlib (>=3.5.0)", "myst-nb (>=1.0.0)", "myst-parser (>=3.0.1)", "packaging (==24.1)", "palettable (==3.3.3)", "pandas (==2.2.2)", "penzai (>=0.2.4,<0.3.0)", "plotly (==5.22.0)", "sphinx (>=6.0.0,<7.3.0)", "sphinx-book-theme (>=1.0.1)", "sphinx-hoverxref", "sphinx_contributors", "sphinxcontrib-katex", "torch (==2.3.1)"]
+notebook = ["ipython", "jax (>=0.4.23)", "palettable"]
+test = ["absl-py (>=1.4.0)", "jax (>=0.4.23)", "omegaconf (>=2.0.0)", "pydantic (>=2.0.0)", "pytest (>=8.2.2)", "torch (>=2.0.0)"]
+
+[[package]]
+name = "triton"
+version = "3.3.1"
+description = "A language and compiler for custom Deep Learning operations"
+optional = false
+python-versions = "*"
+groups = ["main"]
+markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""
+files = [
+    {file = "triton-3.3.1-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b74db445b1c562844d3cfad6e9679c72e93fdfb1a90a24052b03bb5c49d1242e"},
+    {file = "triton-3.3.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b31e3aa26f8cb3cc5bf4e187bf737cbacf17311e1112b781d4a059353dfd731b"},
+    {file = "triton-3.3.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9999e83aba21e1a78c1f36f21bce621b77bcaa530277a50484a7cb4a822f6e43"},
+    {file = "triton-3.3.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b89d846b5a4198317fec27a5d3a609ea96b6d557ff44b56c23176546023c4240"},
+    {file = "triton-3.3.1-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a3198adb9d78b77818a5388bff89fa72ff36f9da0bc689db2f0a651a67ce6a42"},
+    {file = "triton-3.3.1-cp39-cp39-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f6139aeb04a146b0b8e0fbbd89ad1e65861c57cfed881f21d62d3cb94a36bab7"},
+]
+
+[package.dependencies]
+setuptools = ">=40.8.0"
+
+[package.extras]
+build = ["cmake (>=3.20)", "lit"]
+tests = ["autopep8", "isort", "llnl-hatchet", "numpy", "pytest", "pytest-forked", "pytest-xdist", "scipy (>=1.7.1)"]
+tutorials = ["matplotlib", "pandas", "tabulate"]
+
+[[package]]
+name = "typing-extensions"
+version = "4.12.2"
+description = "Backported and Experimental Type Hints for Python 3.8+"
+optional = false
+python-versions = ">=3.8"
+groups = ["main"]
+files = [
+    {file = "typing_extensions-4.12.2-py3-none-any.whl", hash = "sha256:04e5ca0351e0f3f85c6853954072df659d0d13fac324d0072316b67d7794700d"},
+    {file = "typing_extensions-4.12.2.tar.gz", hash = "sha256:1a7ead55c7e559dd4dee8856e3a88b41225abfe1ce8df57b7c13915fe121ffb8"},
+]
+
 [[package]]
 name = "tzdata"
 version = "2024.1"
 description = "Provider of IANA time zone data"
 optional = false
 python-versions = ">=2"
+groups = ["main"]
 files = [
     {file = "tzdata-2024.1-py2.py3-none-any.whl", hash = "sha256:9068bc196136463f5245e51efda838afa15aaeca9903f49050dfa2679db4d252"},
     {file = "tzdata-2024.1.tar.gz", hash = "sha256:2674120f8d891909751c38abcdfd386ac0a5a1127954fbc332af6b5ceae07efd"},
@@ -1179,6 +2358,7 @@ version = "0.2.13"
 description = "Measures the displayed width of unicode strings in a terminal"
 optional = false
 python-versions = "*"
+groups = ["main"]
 files = [
     {file = "wcwidth-0.2.13-py2.py3-none-any.whl", hash = "sha256:3da69048e4540d84af32131829ff948f1e022c1c6bdb8d6102117aac784f6859"},
     {file = "wcwidth-0.2.13.tar.gz", hash = "sha256:72ea0c06399eb286d978fdedb6923a9eb47e1c486ce63e9b4e64fc18303972b5"},
@@ -1190,6 +2370,7 @@ version = "0.44.0"
 description = "A built-package format for Python"
 optional = false
 python-versions = ">=3.8"
+groups = ["main"]
 files = [
     {file = "wheel-0.44.0-py3-none-any.whl", hash = "sha256:2376a90c98cc337d18623527a97c31797bd02bad0033d41547043a1cbfbe448f"},
     {file = "wheel-0.44.0.tar.gz", hash = "sha256:a29c3f2817e95ab89aa4660681ad547c0e9547f20e75b0562fe7723c9a2a9d49"},
@@ -1198,7 +2379,27 @@ files = [
 [package.extras]
 test = ["pytest (>=6.0.0)", "setuptools (>=65)"]
 
+[[package]]
+name = "zipp"
+version = "3.21.0"
+description = "Backport of pathlib-compatible object wrapper for zip files"
+optional = false
+python-versions = ">=3.9"
+groups = ["main"]
+files = [
+    {file = "zipp-3.21.0-py3-none-any.whl", hash = "sha256:ac1bbe05fd2991f160ebce24ffbac5f6d11d83dc90891255885223d42b3cd931"},
+    {file = "zipp-3.21.0.tar.gz", hash = "sha256:2c9958f6430a2040341a52eb608ed6dd93ef4392e02ffe219417c1b28b5dd1f4"},
+]
+
+[package.extras]
+check = ["pytest-checkdocs (>=2.4)", "pytest-ruff (>=0.2.1) ; sys_platform != \"cygwin\""]
+cover = ["pytest-cov"]
+doc = ["furo", "jaraco.packaging (>=9.3)", "jaraco.tidelift (>=1.4)", "rst.linker (>=1.9)", "sphinx (>=3.5)", "sphinx-lint"]
+enabler = ["pytest-enabler (>=2.2)"]
+test = ["big-O", "importlib-resources ; python_version < \"3.9\"", "jaraco.functools", "jaraco.itertools", "jaraco.test", "more-itertools", "pytest (>=6,!=8.1.*)", "pytest-ignore-flaky"]
+type = ["pytest-mypy"]
+
 [metadata]
-lock-version = "2.0"
+lock-version = "2.1"
 python-versions = "^3.12"
-content-hash = "f802131d57996774dddc2b3f1467aa5df28e90f21b6236274f68d192967663ff"
+content-hash = "89ae70124a07e45ef527d1b86b94a1d2587de7688c187bd3f12b8fac7f2342e5"
diff --git a/pyproject.toml b/pyproject.toml
index 815a4d7..6cee724 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -18,6 +18,9 @@ wheel = "^0.44.0"
 setuptools-rust = "^1.10.1"
 iced-x86 = "^1.21.0"
 scipy = "^1.14.0"
+jax = "^0.5.3"
+flax = "^0.10.4"
+torch = "^2.7.1"
 
 
 [build-system]
diff --git a/weights/checkpoint_10/_CHECKPOINT_METADATA b/weights/checkpoint_10/_CHECKPOINT_METADATA
new file mode 100644
index 0000000..3b9ec13
--- /dev/null
+++ b/weights/checkpoint_10/_CHECKPOINT_METADATA
@@ -0,0 +1 @@
+{"item_handlers": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler", "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1742949866281792936, "commit_timestamp_nsecs": 1742949866840282865, "custom_metadata": {}}
\ No newline at end of file
diff --git a/weights/checkpoint_10/_METADATA b/weights/checkpoint_10/_METADATA
new file mode 100644
index 0000000..56ac739
--- /dev/null
+++ b/weights/checkpoint_10/_METADATA
@@ -0,0 +1 @@
+{"tree_metadata": {"('step',)": {"key_metadata": [{"key": "step", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('params', 'BatchNorm_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('params', 'BatchNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('params', 'BatchNorm_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "BatchNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('params', 'BatchNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "BatchNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('params', 'BatchNorm_2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "BatchNorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('params', 'BatchNorm_2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "BatchNorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('params', 'BatchNorm_3', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "BatchNorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('params', 'BatchNorm_3', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "BatchNorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('params', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('params', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96, 96]}}, "('params', 'Dense_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('params', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96, 2]}}, "('params', 'GINConv_0', 'MLP_0', 'BatchNorm_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "GINConv_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('params', 'GINConv_0', 'MLP_0', 'BatchNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "GINConv_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('params', 'GINConv_0', 'MLP_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "GINConv_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('params', 'GINConv_0', 'MLP_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "GINConv_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 192]}}, "('params', 'GINConv_0', 'MLP_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "GINConv_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('params', 'GINConv_0', 'MLP_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "GINConv_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 96]}}, "('params', 'GINConv_0', 'eps')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "GINConv_0", "key_type": 2}, {"key": "eps", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('params', 'GINConv_1', 'MLP_0', 'BatchNorm_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "GINConv_1", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('params', 'GINConv_1', 'MLP_0', 'BatchNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "GINConv_1", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('params', 'GINConv_1', 'MLP_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "GINConv_1", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('params', 'GINConv_1', 'MLP_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "GINConv_1", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96, 192]}}, "('params', 'GINConv_1', 'MLP_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "GINConv_1", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('params', 'GINConv_1', 'MLP_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "GINConv_1", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 96]}}, "('params', 'GINConv_1', 'eps')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "GINConv_1", "key_type": 2}, {"key": "eps", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('params', 'GINConv_2', 'MLP_0', 'BatchNorm_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "GINConv_2", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('params', 'GINConv_2', 'MLP_0', 'BatchNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "GINConv_2", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('params', 'GINConv_2', 'MLP_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "GINConv_2", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('params', 'GINConv_2', 'MLP_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "GINConv_2", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96, 192]}}, "('params', 'GINConv_2', 'MLP_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "GINConv_2", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('params', 'GINConv_2', 'MLP_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "GINConv_2", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 96]}}, "('params', 'GINConv_2', 'eps')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "GINConv_2", "key_type": 2}, {"key": "eps", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('params', 'GINConv_3', 'MLP_0', 'BatchNorm_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "GINConv_3", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('params', 'GINConv_3', 'MLP_0', 'BatchNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "GINConv_3", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('params', 'GINConv_3', 'MLP_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "GINConv_3", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('params', 'GINConv_3', 'MLP_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "GINConv_3", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96, 192]}}, "('params', 'GINConv_3', 'MLP_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "GINConv_3", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('params', 'GINConv_3', 'MLP_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "GINConv_3", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 96]}}, "('params', 'GINConv_3', 'eps')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "GINConv_3", "key_type": 2}, {"key": "eps", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('params', 'LayerNorm_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('params', 'LayerNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('opt_state', '0', 'count')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "count", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('opt_state', '0', 'mu', 'BatchNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('opt_state', '0', 'mu', 'BatchNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('opt_state', '0', 'mu', 'BatchNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "BatchNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('opt_state', '0', 'mu', 'BatchNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "BatchNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('opt_state', '0', 'mu', 'BatchNorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "BatchNorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('opt_state', '0', 'mu', 'BatchNorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "BatchNorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('opt_state', '0', 'mu', 'BatchNorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "BatchNorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('opt_state', '0', 'mu', 'BatchNorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "BatchNorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('opt_state', '0', 'mu', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('opt_state', '0', 'mu', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96, 96]}}, "('opt_state', '0', 'mu', 'Dense_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('opt_state', '0', 'mu', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96, 2]}}, "('opt_state', '0', 'mu', 'GINConv_0', 'MLP_0', 'BatchNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "GINConv_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('opt_state', '0', 'mu', 'GINConv_0', 'MLP_0', 'BatchNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "GINConv_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('opt_state', '0', 'mu', 'GINConv_0', 'MLP_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "GINConv_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('opt_state', '0', 'mu', 'GINConv_0', 'MLP_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "GINConv_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 192]}}, "('opt_state', '0', 'mu', 'GINConv_0', 'MLP_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "GINConv_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('opt_state', '0', 'mu', 'GINConv_0', 'MLP_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "GINConv_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 96]}}, "('opt_state', '0', 'mu', 'GINConv_0', 'eps')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "GINConv_0", "key_type": 2}, {"key": "eps", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('opt_state', '0', 'mu', 'GINConv_1', 'MLP_0', 'BatchNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "GINConv_1", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('opt_state', '0', 'mu', 'GINConv_1', 'MLP_0', 'BatchNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "GINConv_1", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('opt_state', '0', 'mu', 'GINConv_1', 'MLP_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "GINConv_1", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('opt_state', '0', 'mu', 'GINConv_1', 'MLP_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "GINConv_1", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96, 192]}}, "('opt_state', '0', 'mu', 'GINConv_1', 'MLP_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "GINConv_1", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('opt_state', '0', 'mu', 'GINConv_1', 'MLP_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "GINConv_1", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 96]}}, "('opt_state', '0', 'mu', 'GINConv_1', 'eps')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "GINConv_1", "key_type": 2}, {"key": "eps", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('opt_state', '0', 'mu', 'GINConv_2', 'MLP_0', 'BatchNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "GINConv_2", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('opt_state', '0', 'mu', 'GINConv_2', 'MLP_0', 'BatchNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "GINConv_2", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('opt_state', '0', 'mu', 'GINConv_2', 'MLP_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "GINConv_2", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('opt_state', '0', 'mu', 'GINConv_2', 'MLP_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "GINConv_2", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96, 192]}}, "('opt_state', '0', 'mu', 'GINConv_2', 'MLP_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "GINConv_2", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('opt_state', '0', 'mu', 'GINConv_2', 'MLP_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "GINConv_2", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 96]}}, "('opt_state', '0', 'mu', 'GINConv_2', 'eps')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "GINConv_2", "key_type": 2}, {"key": "eps", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('opt_state', '0', 'mu', 'GINConv_3', 'MLP_0', 'BatchNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "GINConv_3", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('opt_state', '0', 'mu', 'GINConv_3', 'MLP_0', 'BatchNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "GINConv_3", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('opt_state', '0', 'mu', 'GINConv_3', 'MLP_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "GINConv_3", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('opt_state', '0', 'mu', 'GINConv_3', 'MLP_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "GINConv_3", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96, 192]}}, "('opt_state', '0', 'mu', 'GINConv_3', 'MLP_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "GINConv_3", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('opt_state', '0', 'mu', 'GINConv_3', 'MLP_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "GINConv_3", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 96]}}, "('opt_state', '0', 'mu', 'GINConv_3', 'eps')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "GINConv_3", "key_type": 2}, {"key": "eps", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('opt_state', '0', 'mu', 'LayerNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('opt_state', '0', 'mu', 'LayerNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('opt_state', '0', 'nu', 'BatchNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('opt_state', '0', 'nu', 'BatchNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('opt_state', '0', 'nu', 'BatchNorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "BatchNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('opt_state', '0', 'nu', 'BatchNorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "BatchNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('opt_state', '0', 'nu', 'BatchNorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "BatchNorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('opt_state', '0', 'nu', 'BatchNorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "BatchNorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('opt_state', '0', 'nu', 'BatchNorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "BatchNorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('opt_state', '0', 'nu', 'BatchNorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "BatchNorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('opt_state', '0', 'nu', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('opt_state', '0', 'nu', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96, 96]}}, "('opt_state', '0', 'nu', 'Dense_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('opt_state', '0', 'nu', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96, 2]}}, "('opt_state', '0', 'nu', 'GINConv_0', 'MLP_0', 'BatchNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "GINConv_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('opt_state', '0', 'nu', 'GINConv_0', 'MLP_0', 'BatchNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "GINConv_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('opt_state', '0', 'nu', 'GINConv_0', 'MLP_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "GINConv_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('opt_state', '0', 'nu', 'GINConv_0', 'MLP_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "GINConv_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 192]}}, "('opt_state', '0', 'nu', 'GINConv_0', 'MLP_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "GINConv_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('opt_state', '0', 'nu', 'GINConv_0', 'MLP_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "GINConv_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 96]}}, "('opt_state', '0', 'nu', 'GINConv_0', 'eps')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "GINConv_0", "key_type": 2}, {"key": "eps", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('opt_state', '0', 'nu', 'GINConv_1', 'MLP_0', 'BatchNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "GINConv_1", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('opt_state', '0', 'nu', 'GINConv_1', 'MLP_0', 'BatchNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "GINConv_1", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('opt_state', '0', 'nu', 'GINConv_1', 'MLP_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "GINConv_1", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('opt_state', '0', 'nu', 'GINConv_1', 'MLP_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "GINConv_1", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96, 192]}}, "('opt_state', '0', 'nu', 'GINConv_1', 'MLP_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "GINConv_1", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('opt_state', '0', 'nu', 'GINConv_1', 'MLP_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "GINConv_1", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 96]}}, "('opt_state', '0', 'nu', 'GINConv_1', 'eps')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "GINConv_1", "key_type": 2}, {"key": "eps", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('opt_state', '0', 'nu', 'GINConv_2', 'MLP_0', 'BatchNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "GINConv_2", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('opt_state', '0', 'nu', 'GINConv_2', 'MLP_0', 'BatchNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "GINConv_2", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('opt_state', '0', 'nu', 'GINConv_2', 'MLP_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "GINConv_2", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('opt_state', '0', 'nu', 'GINConv_2', 'MLP_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "GINConv_2", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96, 192]}}, "('opt_state', '0', 'nu', 'GINConv_2', 'MLP_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "GINConv_2", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('opt_state', '0', 'nu', 'GINConv_2', 'MLP_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "GINConv_2", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 96]}}, "('opt_state', '0', 'nu', 'GINConv_2', 'eps')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "GINConv_2", "key_type": 2}, {"key": "eps", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('opt_state', '0', 'nu', 'GINConv_3', 'MLP_0', 'BatchNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "GINConv_3", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('opt_state', '0', 'nu', 'GINConv_3', 'MLP_0', 'BatchNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "GINConv_3", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('opt_state', '0', 'nu', 'GINConv_3', 'MLP_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "GINConv_3", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('opt_state', '0', 'nu', 'GINConv_3', 'MLP_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "GINConv_3", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96, 192]}}, "('opt_state', '0', 'nu', 'GINConv_3', 'MLP_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "GINConv_3", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('opt_state', '0', 'nu', 'GINConv_3', 'MLP_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "GINConv_3", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 96]}}, "('opt_state', '0', 'nu', 'GINConv_3', 'eps')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "GINConv_3", "key_type": 2}, {"key": "eps", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('opt_state', '0', 'nu', 'LayerNorm_0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('opt_state', '0', 'nu', 'LayerNorm_0', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('opt_state', '1')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('batch_stats', 'BatchNorm_0', 'mean')": {"key_metadata": [{"key": "batch_stats", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('batch_stats', 'BatchNorm_0', 'var')": {"key_metadata": [{"key": "batch_stats", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('batch_stats', 'BatchNorm_1', 'mean')": {"key_metadata": [{"key": "batch_stats", "key_type": 2}, {"key": "BatchNorm_1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('batch_stats', 'BatchNorm_1', 'var')": {"key_metadata": [{"key": "batch_stats", "key_type": 2}, {"key": "BatchNorm_1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('batch_stats', 'BatchNorm_2', 'mean')": {"key_metadata": [{"key": "batch_stats", "key_type": 2}, {"key": "BatchNorm_2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('batch_stats', 'BatchNorm_2', 'var')": {"key_metadata": [{"key": "batch_stats", "key_type": 2}, {"key": "BatchNorm_2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('batch_stats', 'BatchNorm_3', 'mean')": {"key_metadata": [{"key": "batch_stats", "key_type": 2}, {"key": "BatchNorm_3", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('batch_stats', 'BatchNorm_3', 'var')": {"key_metadata": [{"key": "batch_stats", "key_type": 2}, {"key": "BatchNorm_3", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('batch_stats', 'GINConv_0', 'MLP_0', 'BatchNorm_0', 'mean')": {"key_metadata": [{"key": "batch_stats", "key_type": 2}, {"key": "GINConv_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('batch_stats', 'GINConv_0', 'MLP_0', 'BatchNorm_0', 'var')": {"key_metadata": [{"key": "batch_stats", "key_type": 2}, {"key": "GINConv_0", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('batch_stats', 'GINConv_1', 'MLP_0', 'BatchNorm_0', 'mean')": {"key_metadata": [{"key": "batch_stats", "key_type": 2}, {"key": "GINConv_1", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('batch_stats', 'GINConv_1', 'MLP_0', 'BatchNorm_0', 'var')": {"key_metadata": [{"key": "batch_stats", "key_type": 2}, {"key": "GINConv_1", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('batch_stats', 'GINConv_2', 'MLP_0', 'BatchNorm_0', 'mean')": {"key_metadata": [{"key": "batch_stats", "key_type": 2}, {"key": "GINConv_2", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('batch_stats', 'GINConv_2', 'MLP_0', 'BatchNorm_0', 'var')": {"key_metadata": [{"key": "batch_stats", "key_type": 2}, {"key": "GINConv_2", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('batch_stats', 'GINConv_3', 'MLP_0', 'BatchNorm_0', 'mean')": {"key_metadata": [{"key": "batch_stats", "key_type": 2}, {"key": "GINConv_3", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('batch_stats', 'GINConv_3', 'MLP_0', 'BatchNorm_0', 'var')": {"key_metadata": [{"key": "batch_stats", "key_type": 2}, {"key": "GINConv_3", "key_type": 2}, {"key": "MLP_0", "key_type": 2}, {"key": "BatchNorm_0", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}}, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}
\ No newline at end of file
diff --git a/weights/checkpoint_10/_sharding b/weights/checkpoint_10/_sharding
new file mode 100644
index 0000000..bb414b2
--- /dev/null
+++ b/weights/checkpoint_10/_sharding
@@ -0,0 +1 @@
+{"YmF0Y2hfc3RhdHMuQmF0Y2hOb3JtXzAubWVhbg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","YmF0Y2hfc3RhdHMuQmF0Y2hOb3JtXzAudmFy":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","YmF0Y2hfc3RhdHMuQmF0Y2hOb3JtXzEubWVhbg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","YmF0Y2hfc3RhdHMuQmF0Y2hOb3JtXzEudmFy":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","YmF0Y2hfc3RhdHMuQmF0Y2hOb3JtXzIubWVhbg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","YmF0Y2hfc3RhdHMuQmF0Y2hOb3JtXzIudmFy":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","YmF0Y2hfc3RhdHMuQmF0Y2hOb3JtXzMubWVhbg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","YmF0Y2hfc3RhdHMuQmF0Y2hOb3JtXzMudmFy":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","YmF0Y2hfc3RhdHMuR0lOQ29udl8wLk1MUF8wLkJhdGNoTm9ybV8wLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","YmF0Y2hfc3RhdHMuR0lOQ29udl8wLk1MUF8wLkJhdGNoTm9ybV8wLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","YmF0Y2hfc3RhdHMuR0lOQ29udl8xLk1MUF8wLkJhdGNoTm9ybV8wLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","YmF0Y2hfc3RhdHMuR0lOQ29udl8xLk1MUF8wLkJhdGNoTm9ybV8wLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","YmF0Y2hfc3RhdHMuR0lOQ29udl8yLk1MUF8wLkJhdGNoTm9ybV8wLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","YmF0Y2hfc3RhdHMuR0lOQ29udl8yLk1MUF8wLkJhdGNoTm9ybV8wLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","YmF0Y2hfc3RhdHMuR0lOQ29udl8zLk1MUF8wLkJhdGNoTm9ybV8wLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","YmF0Y2hfc3RhdHMuR0lOQ29udl8zLk1MUF8wLkJhdGNoTm9ybV8wLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAuY291bnQ=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuQmF0Y2hOb3JtXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuQmF0Y2hOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuQmF0Y2hOb3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuQmF0Y2hOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuQmF0Y2hOb3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuQmF0Y2hOb3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuQmF0Y2hOb3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuQmF0Y2hOb3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuR0lOQ29udl8wLk1MUF8wLkJhdGNoTm9ybV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuR0lOQ29udl8wLk1MUF8wLkJhdGNoTm9ybV8wLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuR0lOQ29udl8wLk1MUF8wLkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuR0lOQ29udl8wLk1MUF8wLkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuR0lOQ29udl8wLk1MUF8wLkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuR0lOQ29udl8wLk1MUF8wLkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuR0lOQ29udl8wLmVwcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuR0lOQ29udl8xLk1MUF8wLkJhdGNoTm9ybV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuR0lOQ29udl8xLk1MUF8wLkJhdGNoTm9ybV8wLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuR0lOQ29udl8xLk1MUF8wLkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuR0lOQ29udl8xLk1MUF8wLkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuR0lOQ29udl8xLk1MUF8wLkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuR0lOQ29udl8xLk1MUF8wLkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuR0lOQ29udl8xLmVwcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuR0lOQ29udl8yLk1MUF8wLkJhdGNoTm9ybV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuR0lOQ29udl8yLk1MUF8wLkJhdGNoTm9ybV8wLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuR0lOQ29udl8yLk1MUF8wLkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuR0lOQ29udl8yLk1MUF8wLkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuR0lOQ29udl8yLk1MUF8wLkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuR0lOQ29udl8yLk1MUF8wLkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuR0lOQ29udl8yLmVwcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuR0lOQ29udl8zLk1MUF8wLkJhdGNoTm9ybV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuR0lOQ29udl8zLk1MUF8wLkJhdGNoTm9ybV8wLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuR0lOQ29udl8zLk1MUF8wLkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuR0lOQ29udl8zLk1MUF8wLkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuR0lOQ29udl8zLk1MUF8wLkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuR0lOQ29udl8zLk1MUF8wLkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuR0lOQ29udl8zLmVwcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuTGF5ZXJOb3JtXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubXUuTGF5ZXJOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuQmF0Y2hOb3JtXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuQmF0Y2hOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuQmF0Y2hOb3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuQmF0Y2hOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuQmF0Y2hOb3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuQmF0Y2hOb3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuQmF0Y2hOb3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuQmF0Y2hOb3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuR0lOQ29udl8wLk1MUF8wLkJhdGNoTm9ybV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuR0lOQ29udl8wLk1MUF8wLkJhdGNoTm9ybV8wLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuR0lOQ29udl8wLk1MUF8wLkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuR0lOQ29udl8wLk1MUF8wLkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuR0lOQ29udl8wLk1MUF8wLkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuR0lOQ29udl8wLk1MUF8wLkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuR0lOQ29udl8wLmVwcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuR0lOQ29udl8xLk1MUF8wLkJhdGNoTm9ybV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuR0lOQ29udl8xLk1MUF8wLkJhdGNoTm9ybV8wLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuR0lOQ29udl8xLk1MUF8wLkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuR0lOQ29udl8xLk1MUF8wLkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuR0lOQ29udl8xLk1MUF8wLkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuR0lOQ29udl8xLk1MUF8wLkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuR0lOQ29udl8xLmVwcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuR0lOQ29udl8yLk1MUF8wLkJhdGNoTm9ybV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuR0lOQ29udl8yLk1MUF8wLkJhdGNoTm9ybV8wLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuR0lOQ29udl8yLk1MUF8wLkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuR0lOQ29udl8yLk1MUF8wLkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuR0lOQ29udl8yLk1MUF8wLkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuR0lOQ29udl8yLk1MUF8wLkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuR0lOQ29udl8yLmVwcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuR0lOQ29udl8zLk1MUF8wLkJhdGNoTm9ybV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuR0lOQ29udl8zLk1MUF8wLkJhdGNoTm9ybV8wLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuR0lOQ29udl8zLk1MUF8wLkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuR0lOQ29udl8zLk1MUF8wLkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuR0lOQ29udl8zLk1MUF8wLkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuR0lOQ29udl8zLk1MUF8wLkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuR0lOQ29udl8zLmVwcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuTGF5ZXJOb3JtXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","b3B0X3N0YXRlLjAubnUuTGF5ZXJOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","c3RlcA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkJhdGNoTm9ybV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkJhdGNoTm9ybV8wLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkJhdGNoTm9ybV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkJhdGNoTm9ybV8xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkJhdGNoTm9ybV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkJhdGNoTm9ybV8yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkJhdGNoTm9ybV8zLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkJhdGNoTm9ybV8zLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkdJTkNvbnZfMC5NTFBfMC5CYXRjaE5vcm1fMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkdJTkNvbnZfMC5NTFBfMC5CYXRjaE5vcm1fMC5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkdJTkNvbnZfMC5NTFBfMC5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkdJTkNvbnZfMC5NTFBfMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkdJTkNvbnZfMC5NTFBfMC5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkdJTkNvbnZfMC5NTFBfMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkdJTkNvbnZfMC5lcHM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkdJTkNvbnZfMS5NTFBfMC5CYXRjaE5vcm1fMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkdJTkNvbnZfMS5NTFBfMC5CYXRjaE5vcm1fMC5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkdJTkNvbnZfMS5NTFBfMC5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkdJTkNvbnZfMS5NTFBfMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkdJTkNvbnZfMS5NTFBfMC5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkdJTkNvbnZfMS5NTFBfMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkdJTkNvbnZfMS5lcHM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkdJTkNvbnZfMi5NTFBfMC5CYXRjaE5vcm1fMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkdJTkNvbnZfMi5NTFBfMC5CYXRjaE5vcm1fMC5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkdJTkNvbnZfMi5NTFBfMC5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkdJTkNvbnZfMi5NTFBfMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkdJTkNvbnZfMi5NTFBfMC5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkdJTkNvbnZfMi5NTFBfMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkdJTkNvbnZfMi5lcHM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkdJTkNvbnZfMy5NTFBfMC5CYXRjaE5vcm1fMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkdJTkNvbnZfMy5NTFBfMC5CYXRjaE5vcm1fMC5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkdJTkNvbnZfMy5NTFBfMC5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkdJTkNvbnZfMy5NTFBfMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkdJTkNvbnZfMy5NTFBfMC5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkdJTkNvbnZfMy5NTFBfMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkdJTkNvbnZfMy5lcHM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkxheWVyTm9ybV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}","cGFyYW1zLkxheWVyTm9ybV8wLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TFRT_CPU_0\"}"}
\ No newline at end of file
diff --git a/weights/checkpoint_10/array_metadatas/process_0 b/weights/checkpoint_10/array_metadatas/process_0
new file mode 100644
index 0000000..823cce8
--- /dev/null
+++ b/weights/checkpoint_10/array_metadatas/process_0
@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "step", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.BatchNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.BatchNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.BatchNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.BatchNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.BatchNorm_2.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.BatchNorm_2.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.BatchNorm_3.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.BatchNorm_3.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.Dense_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.Dense_0.kernel", "write_shape": [96, 96], "chunk_shape": [96, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.Dense_1.bias", "write_shape": [2], "chunk_shape": [2], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.Dense_1.kernel", "write_shape": [96, 2], "chunk_shape": [96, 2], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.GINConv_0.MLP_0.BatchNorm_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.GINConv_0.MLP_0.BatchNorm_0.scale", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.GINConv_0.MLP_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.GINConv_0.MLP_0.Dense_0.kernel", "write_shape": [1, 192], "chunk_shape": [1, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.GINConv_0.MLP_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.GINConv_0.MLP_0.Dense_1.kernel", "write_shape": [192, 96], "chunk_shape": [192, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.GINConv_0.eps", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.GINConv_1.MLP_0.BatchNorm_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.GINConv_1.MLP_0.BatchNorm_0.scale", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.GINConv_1.MLP_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.GINConv_1.MLP_0.Dense_0.kernel", "write_shape": [96, 192], "chunk_shape": [96, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.GINConv_1.MLP_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.GINConv_1.MLP_0.Dense_1.kernel", "write_shape": [192, 96], "chunk_shape": [192, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.GINConv_1.eps", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.GINConv_2.MLP_0.BatchNorm_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.GINConv_2.MLP_0.BatchNorm_0.scale", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.GINConv_2.MLP_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.GINConv_2.MLP_0.Dense_0.kernel", "write_shape": [96, 192], "chunk_shape": [96, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.GINConv_2.MLP_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.GINConv_2.MLP_0.Dense_1.kernel", "write_shape": [192, 96], "chunk_shape": [192, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.GINConv_2.eps", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.GINConv_3.MLP_0.BatchNorm_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.GINConv_3.MLP_0.BatchNorm_0.scale", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.GINConv_3.MLP_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.GINConv_3.MLP_0.Dense_0.kernel", "write_shape": [96, 192], "chunk_shape": [96, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.GINConv_3.MLP_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.GINConv_3.MLP_0.Dense_1.kernel", "write_shape": [192, 96], "chunk_shape": [192, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.GINConv_3.eps", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.count", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.BatchNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.BatchNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.BatchNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.BatchNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.BatchNorm_2.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.BatchNorm_2.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.BatchNorm_3.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.BatchNorm_3.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.Dense_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.Dense_0.kernel", "write_shape": [96, 96], "chunk_shape": [96, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.Dense_1.bias", "write_shape": [2], "chunk_shape": [2], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.Dense_1.kernel", "write_shape": [96, 2], "chunk_shape": [96, 2], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.GINConv_0.MLP_0.BatchNorm_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.GINConv_0.MLP_0.BatchNorm_0.scale", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.GINConv_0.MLP_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.GINConv_0.MLP_0.Dense_0.kernel", "write_shape": [1, 192], "chunk_shape": [1, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.GINConv_0.MLP_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.GINConv_0.MLP_0.Dense_1.kernel", "write_shape": [192, 96], "chunk_shape": [192, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.GINConv_0.eps", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.GINConv_1.MLP_0.BatchNorm_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.GINConv_1.MLP_0.BatchNorm_0.scale", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.GINConv_1.MLP_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.GINConv_1.MLP_0.Dense_0.kernel", "write_shape": [96, 192], "chunk_shape": [96, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.GINConv_1.MLP_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.GINConv_1.MLP_0.Dense_1.kernel", "write_shape": [192, 96], "chunk_shape": [192, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.GINConv_1.eps", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.GINConv_2.MLP_0.BatchNorm_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.GINConv_2.MLP_0.BatchNorm_0.scale", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.GINConv_2.MLP_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.GINConv_2.MLP_0.Dense_0.kernel", "write_shape": [96, 192], "chunk_shape": [96, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.GINConv_2.MLP_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.GINConv_2.MLP_0.Dense_1.kernel", "write_shape": [192, 96], "chunk_shape": [192, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.GINConv_2.eps", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.GINConv_3.MLP_0.BatchNorm_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.GINConv_3.MLP_0.BatchNorm_0.scale", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.GINConv_3.MLP_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.GINConv_3.MLP_0.Dense_0.kernel", "write_shape": [96, 192], "chunk_shape": [96, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.GINConv_3.MLP_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.GINConv_3.MLP_0.Dense_1.kernel", "write_shape": [192, 96], "chunk_shape": [192, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.GINConv_3.eps", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.mu.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.BatchNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.BatchNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.BatchNorm_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.BatchNorm_1.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.BatchNorm_2.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.BatchNorm_2.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.BatchNorm_3.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.BatchNorm_3.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.Dense_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.Dense_0.kernel", "write_shape": [96, 96], "chunk_shape": [96, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.Dense_1.bias", "write_shape": [2], "chunk_shape": [2], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.Dense_1.kernel", "write_shape": [96, 2], "chunk_shape": [96, 2], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.GINConv_0.MLP_0.BatchNorm_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.GINConv_0.MLP_0.BatchNorm_0.scale", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.GINConv_0.MLP_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.GINConv_0.MLP_0.Dense_0.kernel", "write_shape": [1, 192], "chunk_shape": [1, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.GINConv_0.MLP_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.GINConv_0.MLP_0.Dense_1.kernel", "write_shape": [192, 96], "chunk_shape": [192, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.GINConv_0.eps", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.GINConv_1.MLP_0.BatchNorm_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.GINConv_1.MLP_0.BatchNorm_0.scale", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.GINConv_1.MLP_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.GINConv_1.MLP_0.Dense_0.kernel", "write_shape": [96, 192], "chunk_shape": [96, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.GINConv_1.MLP_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.GINConv_1.MLP_0.Dense_1.kernel", "write_shape": [192, 96], "chunk_shape": [192, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.GINConv_1.eps", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.GINConv_2.MLP_0.BatchNorm_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.GINConv_2.MLP_0.BatchNorm_0.scale", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.GINConv_2.MLP_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.GINConv_2.MLP_0.Dense_0.kernel", "write_shape": [96, 192], "chunk_shape": [96, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.GINConv_2.MLP_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.GINConv_2.MLP_0.Dense_1.kernel", "write_shape": [192, 96], "chunk_shape": [192, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.GINConv_2.eps", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.GINConv_3.MLP_0.BatchNorm_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.GINConv_3.MLP_0.BatchNorm_0.scale", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.GINConv_3.MLP_0.Dense_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.GINConv_3.MLP_0.Dense_0.kernel", "write_shape": [96, 192], "chunk_shape": [96, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.GINConv_3.MLP_0.Dense_1.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.GINConv_3.MLP_0.Dense_1.kernel", "write_shape": [192, 96], "chunk_shape": [192, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.GINConv_3.eps", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.LayerNorm_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.0.nu.LayerNorm_0.scale", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "batch_stats.BatchNorm_0.mean", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "batch_stats.BatchNorm_0.var", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "batch_stats.BatchNorm_1.mean", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "batch_stats.BatchNorm_1.var", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "batch_stats.BatchNorm_2.mean", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "batch_stats.BatchNorm_2.var", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "batch_stats.BatchNorm_3.mean", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "batch_stats.BatchNorm_3.var", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "batch_stats.GINConv_0.MLP_0.BatchNorm_0.mean", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "batch_stats.GINConv_0.MLP_0.BatchNorm_0.var", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "batch_stats.GINConv_1.MLP_0.BatchNorm_0.mean", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "batch_stats.GINConv_1.MLP_0.BatchNorm_0.var", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "batch_stats.GINConv_2.MLP_0.BatchNorm_0.mean", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "batch_stats.GINConv_2.MLP_0.BatchNorm_0.var", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "batch_stats.GINConv_3.MLP_0.BatchNorm_0.mean", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "batch_stats.GINConv_3.MLP_0.BatchNorm_0.var", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}]}
\ No newline at end of file
diff --git a/weights/checkpoint_10/d/85bb797189b91de645cc09202167b979 b/weights/checkpoint_10/d/85bb797189b91de645cc09202167b979
new file mode 100644
index 0000000..b0ad5b4
Binary files /dev/null and b/weights/checkpoint_10/d/85bb797189b91de645cc09202167b979 differ
diff --git a/weights/checkpoint_10/manifest.ocdbt b/weights/checkpoint_10/manifest.ocdbt
new file mode 100644
index 0000000..590ee58
Binary files /dev/null and b/weights/checkpoint_10/manifest.ocdbt differ
diff --git a/weights/checkpoint_10/ocdbt.process_0/d/1d2ac04c6579f3d8c6d751ec125a08ca b/weights/checkpoint_10/ocdbt.process_0/d/1d2ac04c6579f3d8c6d751ec125a08ca
new file mode 100644
index 0000000..32aeb3b
Binary files /dev/null and b/weights/checkpoint_10/ocdbt.process_0/d/1d2ac04c6579f3d8c6d751ec125a08ca differ
diff --git a/weights/checkpoint_10/ocdbt.process_0/d/592716ea24d3ab6a96012664b908c6df b/weights/checkpoint_10/ocdbt.process_0/d/592716ea24d3ab6a96012664b908c6df
new file mode 100644
index 0000000..e142bff
Binary files /dev/null and b/weights/checkpoint_10/ocdbt.process_0/d/592716ea24d3ab6a96012664b908c6df differ
diff --git a/weights/checkpoint_10/ocdbt.process_0/d/b2da5b3705e76f0908e019f11e6c4fbc b/weights/checkpoint_10/ocdbt.process_0/d/b2da5b3705e76f0908e019f11e6c4fbc
new file mode 100644
index 0000000..3a529f3
Binary files /dev/null and b/weights/checkpoint_10/ocdbt.process_0/d/b2da5b3705e76f0908e019f11e6c4fbc differ
diff --git a/weights/checkpoint_10/ocdbt.process_0/d/dbc05a4fec106f56e6164ec6ca36e830 b/weights/checkpoint_10/ocdbt.process_0/d/dbc05a4fec106f56e6164ec6ca36e830
new file mode 100644
index 0000000..5a3cd8b
Binary files /dev/null and b/weights/checkpoint_10/ocdbt.process_0/d/dbc05a4fec106f56e6164ec6ca36e830 differ
diff --git a/weights/checkpoint_10/ocdbt.process_0/manifest.ocdbt b/weights/checkpoint_10/ocdbt.process_0/manifest.ocdbt
new file mode 100644
index 0000000..b0ecd66
Binary files /dev/null and b/weights/checkpoint_10/ocdbt.process_0/manifest.ocdbt differ