From 89826b6b81ac6b5b2a94abbd0eafddc23831b03a Mon Sep 17 00:00:00 2001
From: mariaschuld <mariaschuld@gmail.com>
Date: Wed, 28 Feb 2024 19:03:32 +0200
Subject: [PATCH 001/100] implemented branching logic for different use cases

---
 src/qml_benchmarks/model_utils.py            | 153 +++++++++++++++++-
 src/qml_benchmarks/models/iqp_kernel.py      |  46 ++++--
 src/qml_benchmarks/models/iqp_variational.py | 161 +++++++++++++------
 3 files changed, 298 insertions(+), 62 deletions(-)

diff --git a/src/qml_benchmarks/model_utils.py b/src/qml_benchmarks/model_utils.py
index 312118e2..14e03875 100644
--- a/src/qml_benchmarks/model_utils.py
+++ b/src/qml_benchmarks/model_utils.py
@@ -22,12 +22,13 @@
 import optax
 import jax
 import jax.numpy as jnp
+from pennylane import numpy as pnp
 from sklearn.exceptions import ConvergenceWarning
 from sklearn.utils import gen_batches
 
 
 def train(
-    model, loss_fn, optimizer, X, y, random_key_generator, convergence_interval=200
+        model, loss_fn, optimizer, X, y, random_key_generator, convergence_interval=200
 ):
     """
     Trains a model using an optimizer and a loss function via gradient descent. We assume that the loss function
@@ -99,7 +100,7 @@ def update(params, opt_state, x, y):
             # get means of last two intervals and standard deviation of last interval
             average1 = np.mean(loss_history[-convergence_interval:])
             average2 = np.mean(
-                loss_history[-2 * convergence_interval : -convergence_interval]
+                loss_history[-2 * convergence_interval: -convergence_interval]
             )
             std1 = np.std(loss_history[-convergence_interval:])
             # if the difference in averages is small compared to the statistical fluctuations, stop training.
@@ -124,6 +125,84 @@ def update(params, opt_state, x, y):
     return params
 
 
+def train_without_jax(
+        model,
+        loss_fn,
+        optimizer,
+        X,
+        y,
+        random_key_generator,
+        convergence_interval=200
+):
+    """Trains a model using an optimizer and a loss function.
+
+    We assume that the `model.forward` method works with the `model.params_` to
+    give predictions that the loss function uses for training as:
+
+        >> preds = model.forward(model.params_, X)
+        >> loss = loss_fn(preds, y)
+
+    Args:
+        model (_type_): _description_
+        loss_fn (Callable): _description_
+        optimizer (_type_): _description_
+        X (_type_): _description_
+        y (_type_): _description_
+        key_generator (_type_): _description_
+        interval (int, optional): _description_ Size of interval to average over loss values. The average of the last
+            two intervals is compared to establish convergence.
+        tol (float, optional): _description_ Tolerance criterion for convergence.
+
+    Returns:
+        params (dict): _description_
+    """
+
+    params = list(model.params_.values())
+    opt = optimizer(stepsize=model.learning_rate)
+
+    loss_history = []
+    converged = False
+    start = time.time()
+    for step in range(model.max_steps):
+        key = random_key_generator()
+        X_batch, y_batch = get_batch_without_jax(X, y, key, batch_size=model.batch_size)
+
+        X_batch = pnp.array(X_batch, requires_grad=False)
+        y_batch = pnp.array(y_batch, requires_grad=False)
+        loss_val = loss_fn(*params, X_batch, y_batch)
+        params = opt.step(loss_fn, *params, X_batch, y_batch)[:len(params)]
+        loss_history.append(loss_val)
+
+        logging.debug(f"{step} - loss: {loss_val}")
+
+        if np.isnan(loss_val):
+            logging.info(f"nan encountered. Training aborted.")
+            break
+
+        if step > 2 * convergence_interval:
+            average1 = np.mean(loss_history[-convergence_interval:])
+            average2 = np.mean(loss_history[-2 * convergence_interval:-convergence_interval])
+            std1 = np.std(loss_history[-convergence_interval:])
+            if np.abs(average2 - average1) <= std1 / np.sqrt(convergence_interval) / 2:
+                logging.info(f"Model {model.__class__.__name__} converged after {step} steps.")
+                converged = True
+                break
+
+    end = time.time()
+    loss_history = np.array(loss_history)
+    model.loss_history_ = loss_history / np.max(np.abs(loss_history))
+    model.training_time_ = end - start
+
+    if not converged:
+        raise ConvergenceWarning(
+            f"Model {model.__class__.__name__} has not converged after the maximum number of {model.max_steps} steps.")
+
+    for i, key in enumerate(model.params_.keys()):
+        model.params_[key] = params[i]
+
+    return model.params_
+
+
 def get_batch(X, y, rnd_key, batch_size=32):
     """
     A generator to get random batches of the data (X, y)
@@ -145,6 +224,25 @@ def get_batch(X, y, rnd_key, batch_size=32):
     return X[rnd_indices], y[rnd_indices]
 
 
+def get_batch_without_jax(X, y, rnd_key, batch_size=32):
+    """
+    A generator to get random batches of the data (X, y)
+
+    Args:
+        X (array[float]): Input data with shape (n_samples, n_features).
+        y (array[float]): Target labels with shape (n_samples,)
+        rnd_key: A jax random key object
+        batch_size (int): Number of elements in batch
+
+    Returns:
+        array[float]: A batch of input data shape (batch_size, n_features)
+        array[float]: A batch of target labels shaped (batch_size,)
+    """
+    all_indices = list(range(len(X)))
+    rnd_indices = np.random.choice(all_indices, size=(batch_size,), replace=True)
+    return X[rnd_indices], y[rnd_indices]
+
+
 def get_from_dict(dict, key_list):
     """
     Access a value from a nested dictionary.
@@ -292,3 +390,54 @@ def chunked_loss(params, X, y):
         return jnp.mean(res)
 
     return chunked_loss
+
+
+####### LOSS UTILS WITHOUT JAX
+
+def l2_loss(pred, y):
+    """
+    The square loss function. 0.5 is there to match optax.l2_loss.
+    """
+    return 0.5 * (pred - y) ** 2
+
+
+def softmax(x, axis=-1):
+    """
+    copied from JAX: https://jax.readthedocs.io/en/latest/_modules/jax/_src/nn/functions.html#softmax
+    """
+    x_max = pnp.max(x, axis, keepdims=True)
+    unnormalized = pnp.exp(x - x_max)
+    result = unnormalized / pnp.sum(unnormalized, axis, keepdims=True)
+    return result
+
+
+def one_hot(a, num_classes=2):
+    """
+    convert an array to a one hot encoded array.
+    Taken from https://stackoverflow.com/questions/29831489/convert-array-of-indices-to-one-hot-encoded-array-in-numpy
+    """
+    b = pnp.zeros((a.size, num_classes))
+    b[pnp.arange(a.size), a] = 1
+    return b
+
+
+def log_softmax(x, axis=-1):
+    """
+    taken from jax.nn.log_softmax:
+    https://jax.readthedocs.io/en/latest/_modules/jax/_src/nn/functions.html#log_softmax
+    """
+    x_arr = pnp.asarray(x)
+    x_max = pnp.max(x_arr, axis, keepdims=True)
+    x_max = pnp.array(x_max, requires_grad=False)
+    shifted = x_arr - x_max
+    shifted_logsumexp = pnp.log(
+        pnp.sum(pnp.exp(shifted), axis, keepdims=True))
+    result = shifted - shifted_logsumexp
+    return result
+
+
+def softmax_cross_entropy(logits, labels):
+    """taken from optax source:
+    https: // github.com / google - deepmind / optax / blob / master / optax / losses / _classification.py  # L78%23L103
+    """
+    return -pnp.sum(labels * log_softmax(logits, axis=-1), axis=-1)
diff --git a/src/qml_benchmarks/models/iqp_kernel.py b/src/qml_benchmarks/models/iqp_kernel.py
index e8f5853d..65cfc2a8 100644
--- a/src/qml_benchmarks/models/iqp_kernel.py
+++ b/src/qml_benchmarks/models/iqp_kernel.py
@@ -16,7 +16,6 @@
 import pennylane as qml
 import numpy as np
 import jax
-import jax.numpy as jnp
 from sklearn.base import BaseEstimator, ClassifierMixin
 from sklearn.svm import SVC
 from sklearn.preprocessing import MinMaxScaler
@@ -31,12 +30,14 @@ def __init__(
         svm=SVC(kernel="precomputed", probability=True),
         repeats=2,
         C=1.0,
+        use_jax=False,
+        vmap=False,
         jit=False,
         random_state=42,
         scaling=1.0,
         max_vmap=250,
-        dev_type="default.qubit.jax",
-        qnode_kwargs={"interface": "jax-jit", "diff_method": None},
+        dev_type="default.qubit",
+        qnode_kwargs={},
     ):
         r"""
         Kernel version of the classifier from https://arxiv.org/pdf/1804.11326v2.pdf.
@@ -58,10 +59,12 @@ def __init__(
             svm (sklearn.svm.SVC): scikit-learn SVM class object used to fit the model from the kernel matrix
             repeats (int): number of times the IQP structure is repeated in the embedding circuit.
             C (float): regularization parameter for SVC. Lower values imply stronger regularization.
+            use_jax (bool): Whether to use jax. If False, no jitting and vmapping is performed either
             jit (bool): Whether to use just in time compilation.
-            dev_type (str): string specifying the pennylane device type; e.g. 'default.qubit'.
+            vmap (bool): Whether to use jax.vmap.
             max_vmap (int or None): The maximum size of a chunk to vectorise over. Lower values use less memory.
                 must divide batch_size.
+            dev_type (str): string specifying the pennylane device type; e.g. 'default.qubit'.
             qnode_kwargs (str): the key word arguments passed to the circuit qnode.
             scaling (float): Factor by which to scale the input data.
             random_state (int): seed used for reproducibility.
@@ -69,6 +72,8 @@ def __init__(
         # attributes that do not depend on data
         self.repeats = repeats
         self.C = C
+        self.use_jax = use_jax
+        self.vmap = vmap
         self.jit = jit
         self.max_vmap = max_vmap
         self.svm = svm
@@ -86,7 +91,9 @@ def __init__(
         self.circuit = None
 
     def generate_key(self):
-        return jax.random.PRNGKey(self.rng.integers(1000000))
+        if self.use_jax:
+            return jax.random.PRNGKey(self.rng.integers(1000000))
+        return self.rng.integers(1000000)
 
     def construct_circuit(self):
         dev = qml.device(self.dev_type, wires=self.n_qubits_)
@@ -115,7 +122,7 @@ def circuit(x):
 
         self.circuit = circuit
 
-        if self.jit:
+        if self.use_jax and self.jit:
             circuit = jax.jit(circuit)
         return circuit
 
@@ -132,15 +139,19 @@ def precompute_kernel(self, X1, X2):
         dim2 = len(X2)
 
         # concatenate all pairs of vectors
-        Z = jnp.array(
+        Z = np.array(
             [np.concatenate((X1[i], X2[j])) for i in range(dim1) for j in range(dim2)]
         )
 
         circuit = self.construct_circuit()
-        self.batched_circuit = chunk_vmapped_fn(
-            jax.vmap(circuit, 0), start=0, max_vmap=self.max_vmap
-        )
-        kernel_values = self.batched_circuit(Z)[:, 0]
+
+        if self.use_jax and self.vmap:
+            self.batched_circuit = chunk_vmapped_fn(
+                jax.vmap(circuit, 0), start=0, max_vmap=self.max_vmap
+            )
+            kernel_values = self.batched_circuit(Z)[:, 0]
+        else:
+            kernel_values = np.array([circuit(z)[0] for z in Z])
 
         # reshape the values into the kernel matrix
         kernel_matrix = np.reshape(kernel_values, (dim1, dim2))
@@ -174,11 +185,14 @@ def fit(self, X, y):
             y (np.ndarray): Labels of shape (n_samples,)
         """
 
-        self.svm.random_state = int(
-            jax.random.randint(
-                self.generate_key(), shape=(1,), minval=0, maxval=1000000
+        if self.use_jax:
+            self.svm.random_state = int(
+                jax.random.randint(
+                    self.generate_key(), shape=(1,), minval=0, maxval=1000000
+                )
             )
-        )
+        else:
+            self.svm.random_state = self.generate_key()
 
         self.initialize(X.shape[1], np.unique(y))
 
@@ -244,3 +258,5 @@ def transform(self, X, preprocess=True):
             X = self.scaler.transform(X)
 
         return X * self.scaling
+
+
diff --git a/src/qml_benchmarks/models/iqp_variational.py b/src/qml_benchmarks/models/iqp_variational.py
index 09883456..8ccc4dc1 100644
--- a/src/qml_benchmarks/models/iqp_variational.py
+++ b/src/qml_benchmarks/models/iqp_variational.py
@@ -26,14 +26,16 @@ def __init__(
         n_layers=10,
         learning_rate=0.001,
         batch_size=32,
-        max_vmap=None,
+        use_jax=False,
         jit=True,
+        vmap=False,
+        max_vmap=None,
         max_steps=10000,
         convergence_interval=200,
         random_state=42,
         scaling=1.0,
-        dev_type="default.qubit.jax",
-        qnode_kwargs={"interface": "jax"},
+        dev_type="default.qubit",
+        qnode_kwargs={},
     ):
         r"""
         Variational verison of the classifier from https://arxiv.org/pdf/1804.11326v2.pdf.
@@ -54,10 +56,12 @@ def __init__(
             n_layers (int): Number of layers in the variational part of the circuit.
             learning_rate (float): Learning rate for gradient descent.
             batch_size (int): Size of batches used for computing paraemeter updates.
+            use_jax (bool): Whether to use jax. If False, no jitting and vmapping is performed either.
+            jit (bool): Whether to use just in time compilation.
+            vmap (bool): Whether to use jax.vmap.
             max_vmap (int or None): The maximum size of a chunk to vectorise over. Lower values use less memory.
                 must divide batch_size.
-            jit (bool): Whether to use just in time compilation.
-            convergence_threshold (float): If loss changes less than this threshold for 10 consecutive steps we stop training.
+            convergence_interval (int): If loss does not change significantly in this interval, stop training.
             max_steps (int): Maximum number of training steps. A warning will be raised if training did not converge.
             dev_type (str): string specifying the pennylane device type; e.g. 'default.qubit'.
             qnode_kwargs (str): the keyword arguments passed to the circuit qnode.
@@ -74,6 +78,8 @@ def __init__(
         self.batch_size = batch_size
         self.dev_type = dev_type
         self.qnode_kwargs = qnode_kwargs
+        self.use_jax = use_jax
+        self.vmap = vmap
         self.jit = jit
         self.scaling = scaling
         self.random_state = random_state
@@ -92,29 +98,60 @@ def __init__(
         self.circuit = None
 
     def generate_key(self):
-        return jax.random.PRNGKey(self.rng.integers(1000000))
+        if self.use_jax:
+            return jax.random.PRNGKey(self.rng.integers(1000000))
+        return self.rng.integers(1000000)
 
     def construct_model(self):
         dev = qml.device(self.dev_type, wires=self.n_qubits_)
 
-        @qml.qnode(dev, **self.qnode_kwargs)
-        def circuit(params, x):
-            """
-            The variational circuit from the plots. Uses an IQP data embedding.
-            We use the same observable as in the plots.
-            """
-            qml.IQPEmbedding(x, wires=range(self.n_qubits_), n_repeats=self.repeats)
-            qml.StronglyEntanglingLayers(
-                params["weights"], wires=range(self.n_qubits_), imprimitive=qml.CZ
-            )
-            return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
+        if self.use_jax:
+            @qml.qnode(dev, **self.qnode_kwargs)
+            def circuit(params, x):
+                """
+                The variational circuit from the plots. Uses an IQP data embedding.
+                We use the same observable as in the plots.
+                """
+                qml.IQPEmbedding(x, wires=range(self.n_qubits_), n_repeats=self.repeats)
+                qml.StronglyEntanglingLayers(
+                    params["weights"], wires=range(self.n_qubits_), imprimitive=qml.CZ
+                )
+                return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
+        else:
+            @qml.qnode(dev, **self.qnode_kwargs)
+            def circuit(weights, x):
+                """
+                The variational circuit from the plots. Uses an IQP data embedding.
+                We use the same observable as in the plots.
+                """
+                qml.IQPEmbedding(x, wires=range(self.n_qubits_), n_repeats=self.repeats)
+                qml.StronglyEntanglingLayers(
+                    weights, wires=range(self.n_qubits_), imprimitive=qml.CZ
+                )
+                return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
 
         self.circuit = circuit
 
-        if self.jit:
+        if self.use_jax and self.jit:
             circuit = jax.jit(circuit)
-        self.forward = jax.vmap(circuit, in_axes=(None, 0))
-        self.chunked_forward = chunk_vmapped_fn(self.forward, 1, self.max_vmap)
+
+        if self.use_jax:
+            if self.vmap:
+                # use jax and batch feed the circuit
+                self.forward = jax.vmap(circuit, in_axes=(None, 0))
+                self.chunked_forward = chunk_vmapped_fn(self.forward, 1, self.max_vmap)
+            else:
+                # use jax but do not batch feed the circuit
+                def forward(params, X):
+                    return jnp.stack([circuit(params, x) for x in X])
+
+                self.forward = forward
+        else:
+            # use autograd and do not batch feed the circuit
+            def forward(params, X):
+                return pnp.array([circuit(params, x) for x in X])
+
+            self.forward = forward
 
         return self.forward
 
@@ -139,13 +176,24 @@ def initialize(self, n_features, classes=None):
 
     def initialize_params(self):
         # initialise the trainable parameters
-        weights = (
-            2
-            * jnp.pi
-            * jax.random.uniform(
-                shape=(self.n_layers, self.n_qubits_, 3), key=self.generate_key()
+
+        if self.use_jax:
+            weights = (
+                2
+                * jnp.pi
+                * jax.random.uniform(
+                    shape=(self.n_layers, self.n_qubits_, 3), key=self.generate_key()
+                )
+            )
+        else:
+            weights = (
+                2
+                * np.pi
+                * np.random.uniform(
+                    size=(self.n_layers, self.n_qubits_, 3)
+                )
             )
-        )
+            weights = pnp.array(weights, requires_grad=True)
         self.params_ = {"weights": weights}
 
     def fit(self, X, y):
@@ -162,24 +210,41 @@ def fit(self, X, y):
         self.scaler.fit(X)
         X = self.transform(X)
 
-        optimizer = optax.adam
-
-        def loss_fn(params, X, y):
-            expvals = self.forward(params, X)
-            probs = (1 - expvals * y) / 2  # the probs of incorrect classification
-            return jnp.mean(probs)
-
-        if self.jit:
-            loss_fn = jax.jit(loss_fn)
-        self.params_ = train(
-            self,
-            loss_fn,
-            optimizer,
-            X,
-            y,
-            self.generate_key,
-            convergence_interval=self.convergence_interval,
-        )
+        if self.use_jax:
+
+            optimizer = optax.adam
+
+            def loss_fn(params, X, y):
+                expvals = self.forward(params, X)
+                probs = (1 - expvals * y) / 2  # the probs of incorrect classification
+                if self.use_jax:
+                    return jnp.mean(probs)
+                np.mean(probs)
+
+            if self.use_jax and self.jit:
+                loss_fn = jax.jit(loss_fn)
+
+            self.params_ = train(
+                self,
+                loss_fn,
+                optimizer,
+                X,
+                y,
+                self.generate_key,
+                convergence_interval=self.convergence_interval,
+            )
+
+        else:
+            X = pnp.array(X, requires_grad=False)
+            y = pnp.array(y, requires_grad=False)
+            optimizer = qml.AdamOptimizer
+
+            def loss_fn(weights, X, y):
+                expvals = self.forward(weights, X)
+                probs = (1 - expvals * y) / 2  # the probs of incorrect classification
+                return pnp.mean(probs)
+
+            self.params_ = train_without_jax(self, loss_fn, optimizer, X, y, self.generate_key)
 
         return self
 
@@ -207,7 +272,13 @@ def predict_proba(self, X):
             (n_samples, n_classes)
         """
         X = self.transform(X)
-        predictions = self.chunked_forward(self.params_, X)
+        if self.use_jax:
+            if self.vmap:
+                predictions = self.chunked_forward(self.params_, X)
+            else:
+                predictions = self.forward(self.params_, X)
+        else:
+            predictions = self.forward(self.params_["weights"], X)
         predictions_2d = np.c_[(1 - predictions) / 2, (1 + predictions) / 2]
         return predictions_2d
 

From ac10674ac6fa01f3a1d487a1ac8ea53bc9ea409c Mon Sep 17 00:00:00 2001
From: mariaschuld <mariaschuld@gmail.com>
Date: Wed, 28 Feb 2024 19:34:40 +0200
Subject: [PATCH 002/100] some minor corrections

---
 src/qml_benchmarks/model_utils.py            | 27 +++-----------------
 src/qml_benchmarks/models/iqp_variational.py |  7 +++--
 2 files changed, 6 insertions(+), 28 deletions(-)

diff --git a/src/qml_benchmarks/model_utils.py b/src/qml_benchmarks/model_utils.py
index 14e03875..2f893f21 100644
--- a/src/qml_benchmarks/model_utils.py
+++ b/src/qml_benchmarks/model_utils.py
@@ -28,7 +28,7 @@
 
 
 def train(
-        model, loss_fn, optimizer, X, y, random_key_generator, convergence_interval=200
+    model, loss_fn, optimizer, X, y, random_key_generator, convergence_interval=200
 ):
     """
     Trains a model using an optimizer and a loss function via gradient descent. We assume that the loss function
@@ -100,7 +100,7 @@ def update(params, opt_state, x, y):
             # get means of last two intervals and standard deviation of last interval
             average1 = np.mean(loss_history[-convergence_interval:])
             average2 = np.mean(
-                loss_history[-2 * convergence_interval: -convergence_interval]
+                loss_history[-2 * convergence_interval : -convergence_interval]
             )
             std1 = np.std(loss_history[-convergence_interval:])
             # if the difference in averages is small compared to the statistical fluctuations, stop training.
@@ -134,27 +134,7 @@ def train_without_jax(
         random_key_generator,
         convergence_interval=200
 ):
-    """Trains a model using an optimizer and a loss function.
-
-    We assume that the `model.forward` method works with the `model.params_` to
-    give predictions that the loss function uses for training as:
-
-        >> preds = model.forward(model.params_, X)
-        >> loss = loss_fn(preds, y)
-
-    Args:
-        model (_type_): _description_
-        loss_fn (Callable): _description_
-        optimizer (_type_): _description_
-        X (_type_): _description_
-        y (_type_): _description_
-        key_generator (_type_): _description_
-        interval (int, optional): _description_ Size of interval to average over loss values. The average of the last
-            two intervals is compared to establish convergence.
-        tol (float, optional): _description_ Tolerance criterion for convergence.
-
-    Returns:
-        params (dict): _description_
+    """Trains a model using an optimizer and a loss function, using PennyLane's autograd interface.
     """
 
     params = list(model.params_.values())
@@ -166,7 +146,6 @@ def train_without_jax(
     for step in range(model.max_steps):
         key = random_key_generator()
         X_batch, y_batch = get_batch_without_jax(X, y, key, batch_size=model.batch_size)
-
         X_batch = pnp.array(X_batch, requires_grad=False)
         y_batch = pnp.array(y_batch, requires_grad=False)
         loss_val = loss_fn(*params, X_batch, y_batch)
diff --git a/src/qml_benchmarks/models/iqp_variational.py b/src/qml_benchmarks/models/iqp_variational.py
index 8ccc4dc1..554bf4e5 100644
--- a/src/qml_benchmarks/models/iqp_variational.py
+++ b/src/qml_benchmarks/models/iqp_variational.py
@@ -194,6 +194,7 @@ def initialize_params(self):
                 )
             )
             weights = pnp.array(weights, requires_grad=True)
+
         self.params_ = {"weights": weights}
 
     def fit(self, X, y):
@@ -217,11 +218,9 @@ def fit(self, X, y):
             def loss_fn(params, X, y):
                 expvals = self.forward(params, X)
                 probs = (1 - expvals * y) / 2  # the probs of incorrect classification
-                if self.use_jax:
-                    return jnp.mean(probs)
-                np.mean(probs)
+                return jnp.mean(probs)
 
-            if self.use_jax and self.jit:
+            if self.jit:
                 loss_fn = jax.jit(loss_fn)
 
             self.params_ = train(

From 8ff542dbb64180ad0a4e708ccd706cd63feaa3cb Mon Sep 17 00:00:00 2001
From: mariaschuld <mariaschuld@gmail.com>
Date: Wed, 28 Feb 2024 19:42:15 +0200
Subject: [PATCH 003/100] simplify if logic a little

---
 src/qml_benchmarks/models/iqp_kernel.py      |  3 +-
 src/qml_benchmarks/models/iqp_variational.py | 35 +++++++++++---------
 2 files changed, 20 insertions(+), 18 deletions(-)

diff --git a/src/qml_benchmarks/models/iqp_kernel.py b/src/qml_benchmarks/models/iqp_kernel.py
index 65cfc2a8..d7559996 100644
--- a/src/qml_benchmarks/models/iqp_kernel.py
+++ b/src/qml_benchmarks/models/iqp_kernel.py
@@ -146,6 +146,7 @@ def precompute_kernel(self, X1, X2):
         circuit = self.construct_circuit()
 
         if self.use_jax and self.vmap:
+            # if batched circuit is used
             self.batched_circuit = chunk_vmapped_fn(
                 jax.vmap(circuit, 0), start=0, max_vmap=self.max_vmap
             )
@@ -258,5 +259,3 @@ def transform(self, X, preprocess=True):
             X = self.scaler.transform(X)
 
         return X * self.scaling
-
-
diff --git a/src/qml_benchmarks/models/iqp_variational.py b/src/qml_benchmarks/models/iqp_variational.py
index 554bf4e5..2fa45fd6 100644
--- a/src/qml_benchmarks/models/iqp_variational.py
+++ b/src/qml_benchmarks/models/iqp_variational.py
@@ -117,36 +117,39 @@ def circuit(params, x):
                     params["weights"], wires=range(self.n_qubits_), imprimitive=qml.CZ
                 )
                 return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
-        else:
-            @qml.qnode(dev, **self.qnode_kwargs)
-            def circuit(weights, x):
-                """
-                The variational circuit from the plots. Uses an IQP data embedding.
-                We use the same observable as in the plots.
-                """
-                qml.IQPEmbedding(x, wires=range(self.n_qubits_), n_repeats=self.repeats)
-                qml.StronglyEntanglingLayers(
-                    weights, wires=range(self.n_qubits_), imprimitive=qml.CZ
-                )
-                return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
 
-        self.circuit = circuit
+            self.circuit = circuit
 
-        if self.use_jax and self.jit:
-            circuit = jax.jit(circuit)
+            if self.jit:
+                circuit = jax.jit(circuit)
 
-        if self.use_jax:
             if self.vmap:
                 # use jax and batch feed the circuit
                 self.forward = jax.vmap(circuit, in_axes=(None, 0))
                 self.chunked_forward = chunk_vmapped_fn(self.forward, 1, self.max_vmap)
+
             else:
                 # use jax but do not batch feed the circuit
                 def forward(params, X):
                     return jnp.stack([circuit(params, x) for x in X])
 
                 self.forward = forward
+
         else:
+            @qml.qnode(dev, **self.qnode_kwargs)
+            def circuit(weights, x):
+                """
+                The variational circuit from the plots. Uses an IQP data embedding.
+                We use the same observable as in the plots.
+                """
+                qml.IQPEmbedding(x, wires=range(self.n_qubits_), n_repeats=self.repeats)
+                qml.StronglyEntanglingLayers(
+                    weights, wires=range(self.n_qubits_), imprimitive=qml.CZ
+                )
+                return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
+
+            self.circuit = circuit
+
             # use autograd and do not batch feed the circuit
             def forward(params, X):
                 return pnp.array([circuit(params, x) for x in X])

From c57a18499194274feae8069a9b9bd84611e3b353 Mon Sep 17 00:00:00 2001
From: Maria Schuld <mariaschuld@gmail.com>
Date: Thu, 29 Feb 2024 08:53:28 +0200
Subject: [PATCH 004/100] Update src/qml_benchmarks/models/iqp_kernel.py

Co-authored-by: Vincent Michaud-Rioux <vincent.michaud-rioux@xanadu.ai>
---
 src/qml_benchmarks/models/iqp_kernel.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/src/qml_benchmarks/models/iqp_kernel.py b/src/qml_benchmarks/models/iqp_kernel.py
index d7559996..3d899b58 100644
--- a/src/qml_benchmarks/models/iqp_kernel.py
+++ b/src/qml_benchmarks/models/iqp_kernel.py
@@ -36,7 +36,11 @@ def __init__(
         random_state=42,
         scaling=1.0,
         max_vmap=250,
-        dev_type="default.qubit",
+        # device-related attributes
+        self.dev_type = dev_type
+        self.use_jax = use_jax if use_jax is not None else self.dev_type == "default.qubit.jax"
+        self.vmap = vmap if vmap is not None else self.use_jax
+        self.jit = jit if jit is not None else self.use_jax
         qnode_kwargs={},
     ):
         r"""

From bf67629be49e63e631ab38b3b96725deb83ab022 Mon Sep 17 00:00:00 2001
From: Vincent Michaud-Rioux <vincent.michaud-rioux@xanadu.ai>
Date: Wed, 28 Feb 2024 21:27:02 +0000
Subject: [PATCH 005/100] Add makefile. Run make format. Add catalyst qjit as
 an option for lightning devices.

---
 Makefile                                      | 29 ++++++
 requirements.txt                              |  2 +
 src/qml_benchmarks/data/__init__.py           |  1 -
 src/qml_benchmarks/data/mnist.py              | 16 +---
 src/qml_benchmarks/model_utils.py             | 50 +++-------
 src/qml_benchmarks/models/__init__.py         | 68 +++++++-------
 src/qml_benchmarks/models/circuit_centric.py  | 16 +---
 .../models/convolutional_neural_network.py    |  8 +-
 src/qml_benchmarks/models/data_reuploading.py | 27 ++----
 .../models/dressed_quantum_circuit.py         | 22 ++---
 src/qml_benchmarks/models/iqp_kernel.py       | 40 ++++----
 src/qml_benchmarks/models/iqp_variational.py  | 10 +-
 .../models/projected_quantum_kernel.py        | 21 +----
 .../models/quantum_boltzmann_machine.py       | 91 +++++++++++--------
 .../models/quantum_kitchen_sinks.py           | 10 +-
 .../models/quantum_metric_learning.py         | 20 +---
 .../models/quanvolutional_neural_network.py   | 12 +--
 src/qml_benchmarks/models/separable.py        | 12 +--
 src/qml_benchmarks/models/tree_tensor.py      | 17 +---
 src/qml_benchmarks/models/weinet.py           | 19 +---
 20 files changed, 205 insertions(+), 286 deletions(-)
 create mode 100644 Makefile

diff --git a/Makefile b/Makefile
new file mode 100644
index 00000000..50577e2a
--- /dev/null
+++ b/Makefile
@@ -0,0 +1,29 @@
+PYTHON3 := $(shell which python3 2>/dev/null)
+
+PYTHON := python3
+
+.PHONY: help
+help:
+	@echo "Please use \`make <target>' where <target> is one of"
+	@echo "  clean              to delete all temporary, cache, and build files"
+	@echo "  format [check=1]   to apply black formatter; use with 'check=1' to check instead of modify (requires black)"
+	@echo "  lint               to run pylint on source files"
+
+.PHONY : clean
+clean:
+	rm -rf src/qml_benchmarks.egg-info/
+	rm -rf src/qml_benchmarks/__pycache__/
+	rm -rf src/qml_benchmarks/models/__pycache__/
+
+.PHONY:format
+format:
+ifdef check
+	black -l 100 ./src/qml_benchmarks --check
+else
+	black -l 100 ./src/qml_benchmarks
+endif
+
+.PHONY: lint
+lint:
+	pylint src/qml_benchmarks --rcfile .pylintrc
+
diff --git a/requirements.txt b/requirements.txt
index 9cc6c785..7eaa0af3 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -10,3 +10,5 @@ pyyaml~=6.0
 pennyLane~=0.34
 scipy~=1.11
 pandas~=2.2
+black~=24.0
+pylint
\ No newline at end of file
diff --git a/src/qml_benchmarks/data/__init__.py b/src/qml_benchmarks/data/__init__.py
index 654ac570..feb22334 100644
--- a/src/qml_benchmarks/data/__init__.py
+++ b/src/qml_benchmarks/data/__init__.py
@@ -19,4 +19,3 @@
 from qml_benchmarks.data.hyperplanes import generate_hyperplanes_parity
 from qml_benchmarks.data.linearly_separable import generate_linearly_separable
 from qml_benchmarks.data.two_curves import generate_two_curves
-           
\ No newline at end of file
diff --git a/src/qml_benchmarks/data/mnist.py b/src/qml_benchmarks/data/mnist.py
index df4bae01..15b31593 100644
--- a/src/qml_benchmarks/data/mnist.py
+++ b/src/qml_benchmarks/data/mnist.py
@@ -29,29 +29,21 @@
 from sklearn.preprocessing import StandardScaler
 
 
-def generate_mnist(
-    digitA, digitB, preprocessing, n_features=None, n_samples=None, height=None
-):
+def generate_mnist(digitA, digitB, preprocessing, n_features=None, n_samples=None, height=None):
     if preprocessing == "cg":
 
-        mnist_train = torchvision.datasets.MNIST(
-            "mnist_original/", download=True, train=True
-        )
+        mnist_train = torchvision.datasets.MNIST("mnist_original/", download=True, train=True)
         X_train = mnist_train.data
         y_train = mnist_train.targets
 
-        mnist_test = torchvision.datasets.MNIST(
-            "mnist_original/", download=True, train=False
-        )
+        mnist_test = torchvision.datasets.MNIST("mnist_original/", download=True, train=False)
         X_test = mnist_test.data
         y_test = mnist_test.targets
 
         idxs_train = np.concatenate(
             (np.where(y_train == digitA)[0], np.where(y_train == digitB)[0])
         )
-        idxs_test = np.concatenate(
-            (np.where(y_test == digitA)[0], np.where(y_test == digitB)[0])
-        )
+        idxs_test = np.concatenate((np.where(y_test == digitA)[0], np.where(y_test == digitB)[0]))
 
         X_train = X_train[idxs_train]
         y_train = y_train[idxs_train]
diff --git a/src/qml_benchmarks/model_utils.py b/src/qml_benchmarks/model_utils.py
index 2f893f21..d35ab9df 100644
--- a/src/qml_benchmarks/model_utils.py
+++ b/src/qml_benchmarks/model_utils.py
@@ -27,9 +27,7 @@
 from sklearn.utils import gen_batches
 
 
-def train(
-    model, loss_fn, optimizer, X, y, random_key_generator, convergence_interval=200
-):
+def train(model, loss_fn, optimizer, X, y, random_key_generator, convergence_interval=200):
     """
     Trains a model using an optimizer and a loss function via gradient descent. We assume that the loss function
     is of the form `loss(params, X, y)` and that the trainable parameters are stored in model.params_ as a dictionary
@@ -99,15 +97,11 @@ def update(params, opt_state, x, y):
         if step > 2 * convergence_interval:
             # get means of last two intervals and standard deviation of last interval
             average1 = np.mean(loss_history[-convergence_interval:])
-            average2 = np.mean(
-                loss_history[-2 * convergence_interval : -convergence_interval]
-            )
+            average2 = np.mean(loss_history[-2 * convergence_interval : -convergence_interval])
             std1 = np.std(loss_history[-convergence_interval:])
             # if the difference in averages is small compared to the statistical fluctuations, stop training.
             if np.abs(average2 - average1) <= std1 / np.sqrt(convergence_interval) / 2:
-                logging.info(
-                    f"Model {model.__class__.__name__} converged after {step} steps."
-                )
+                logging.info(f"Model {model.__class__.__name__} converged after {step} steps.")
                 converged = True
                 break
 
@@ -126,16 +120,9 @@ def update(params, opt_state, x, y):
 
 
 def train_without_jax(
-        model,
-        loss_fn,
-        optimizer,
-        X,
-        y,
-        random_key_generator,
-        convergence_interval=200
+    model, loss_fn, optimizer, X, y, random_key_generator, convergence_interval=200
 ):
-    """Trains a model using an optimizer and a loss function, using PennyLane's autograd interface.
-    """
+    """Trains a model using an optimizer and a loss function, using PennyLane's autograd interface."""
 
     params = list(model.params_.values())
     opt = optimizer(stepsize=model.learning_rate)
@@ -149,7 +136,7 @@ def train_without_jax(
         X_batch = pnp.array(X_batch, requires_grad=False)
         y_batch = pnp.array(y_batch, requires_grad=False)
         loss_val = loss_fn(*params, X_batch, y_batch)
-        params = opt.step(loss_fn, *params, X_batch, y_batch)[:len(params)]
+        params = opt.step(loss_fn, *params, X_batch, y_batch)[: len(params)]
         loss_history.append(loss_val)
 
         logging.debug(f"{step} - loss: {loss_val}")
@@ -160,7 +147,7 @@ def train_without_jax(
 
         if step > 2 * convergence_interval:
             average1 = np.mean(loss_history[-convergence_interval:])
-            average2 = np.mean(loss_history[-2 * convergence_interval:-convergence_interval])
+            average2 = np.mean(loss_history[-2 * convergence_interval : -convergence_interval])
             std1 = np.std(loss_history[-convergence_interval:])
             if np.abs(average2 - average1) <= std1 / np.sqrt(convergence_interval) / 2:
                 logging.info(f"Model {model.__class__.__name__} converged after {step} steps.")
@@ -174,7 +161,8 @@ def train_without_jax(
 
     if not converged:
         raise ConvergenceWarning(
-            f"Model {model.__class__.__name__} has not converged after the maximum number of {model.max_steps} steps.")
+            f"Model {model.__class__.__name__} has not converged after the maximum number of {model.max_steps} steps."
+        )
 
     for i, key in enumerate(model.params_.keys()):
         model.params_[key] = params[i]
@@ -197,9 +185,7 @@ def get_batch(X, y, rnd_key, batch_size=32):
         array[float]: A batch of target labels shaped (batch_size,)
     """
     all_indices = jnp.array(range(len(X)))
-    rnd_indices = jax.random.choice(
-        key=rnd_key, a=all_indices, shape=(batch_size,), replace=True
-    )
+    rnd_indices = jax.random.choice(key=rnd_key, a=all_indices, shape=(batch_size,), replace=True)
     return X[rnd_indices], y[rnd_indices]
 
 
@@ -301,9 +287,7 @@ def chunked_fn(*args):
         # jnp.concatenate needs to act on arrays with the same shape, so pad the last array if necessary
         if batch_len / max_vmap % 1 != 0.0:
             diff = max_vmap - len(res[-1])
-            res[-1] = jnp.pad(
-                res[-1], [(0, diff), *[(0, 0)] * (len(res[-1].shape) - 1)]
-            )
+            res[-1] = jnp.pad(res[-1], [(0, diff), *[(0, 0)] * (len(res[-1].shape) - 1)])
             return jnp.concatenate(res)[:-diff]
         else:
             return jnp.concatenate(res)
@@ -338,9 +322,7 @@ def chunked_grad(params, X, y):
             set_in_dict(
                 grad_dict,
                 key_list,
-                jnp.mean(
-                    jnp.array([get_from_dict(grad, key_list) for grad in grads]), axis=0
-                ),
+                jnp.mean(jnp.array([get_from_dict(grad, key_list) for grad in grads]), axis=0),
             )
         return grad_dict
 
@@ -363,9 +345,7 @@ def chunk_loss(loss_fn, max_vmap):
 
     def chunked_loss(params, X, y):
         batch_slices = list(gen_batches(len(X), max_vmap))
-        res = jnp.array(
-            [loss_fn(params, *[X[slice], y[slice]]) for slice in batch_slices]
-        )
+        res = jnp.array([loss_fn(params, *[X[slice], y[slice]]) for slice in batch_slices])
         return jnp.mean(res)
 
     return chunked_loss
@@ -373,6 +353,7 @@ def chunked_loss(params, X, y):
 
 ####### LOSS UTILS WITHOUT JAX
 
+
 def l2_loss(pred, y):
     """
     The square loss function. 0.5 is there to match optax.l2_loss.
@@ -409,8 +390,7 @@ def log_softmax(x, axis=-1):
     x_max = pnp.max(x_arr, axis, keepdims=True)
     x_max = pnp.array(x_max, requires_grad=False)
     shifted = x_arr - x_max
-    shifted_logsumexp = pnp.log(
-        pnp.sum(pnp.exp(shifted), axis, keepdims=True))
+    shifted_logsumexp = pnp.log(pnp.sum(pnp.exp(shifted), axis, keepdims=True))
     result = shifted - shifted_logsumexp
     return result
 
diff --git a/src/qml_benchmarks/models/__init__.py b/src/qml_benchmarks/models/__init__.py
index a84e4fae..a35b4de2 100644
--- a/src/qml_benchmarks/models/__init__.py
+++ b/src/qml_benchmarks/models/__init__.py
@@ -80,30 +80,30 @@
 
 class MLPClassifier(MLP):
     def __init__(
-            self,
-            hidden_layer_sizes=(100, 100),
-            activation="relu",
-            solver="adam",
-            alpha=0.0001,
-            batch_size="auto",
-            learning_rate="constant",
-            learning_rate_init=0.001,
-            power_t=0.5,
-            max_iter=3000,
-            shuffle=True,
-            random_state=None,
-            tol=1e-4,
-            verbose=False,
-            warm_start=False,
-            momentum=0.9,
-            nesterovs_momentum=True,
-            early_stopping=False,
-            validation_fraction=0.1,
-            beta_1=0.9,
-            beta_2=0.999,
-            epsilon=1e-8,
-            n_iter_no_change=10,
-            max_fun=15000,
+        self,
+        hidden_layer_sizes=(100, 100),
+        activation="relu",
+        solver="adam",
+        alpha=0.0001,
+        batch_size="auto",
+        learning_rate="constant",
+        learning_rate_init=0.001,
+        power_t=0.5,
+        max_iter=3000,
+        shuffle=True,
+        random_state=None,
+        tol=1e-4,
+        verbose=False,
+        warm_start=False,
+        momentum=0.9,
+        nesterovs_momentum=True,
+        early_stopping=False,
+        validation_fraction=0.1,
+        beta_1=0.9,
+        beta_2=0.999,
+        epsilon=1e-8,
+        n_iter_no_change=10,
+        max_fun=15000,
     ):
         super().__init__(
             hidden_layer_sizes=hidden_layer_sizes,
@@ -134,16 +134,16 @@ def __init__(
 
 class SVC(SVC_base):
     def __init__(
-            self,
-            C=1.0,
-            degree=3,
-            gamma="scale",
-            coef0=0.0,
-            shrinking=True,
-            probability=False,
-            tol=0.001,
-            max_iter=-1,
-            random_state=None,
+        self,
+        C=1.0,
+        degree=3,
+        gamma="scale",
+        coef0=0.0,
+        shrinking=True,
+        probability=False,
+        tol=0.001,
+        max_iter=-1,
+        random_state=None,
     ):
         super().__init__(
             C=C,
diff --git a/src/qml_benchmarks/models/circuit_centric.py b/src/qml_benchmarks/models/circuit_centric.py
index 497a2bd6..1c51e883 100644
--- a/src/qml_benchmarks/models/circuit_centric.py
+++ b/src/qml_benchmarks/models/circuit_centric.py
@@ -170,12 +170,8 @@ def initialize(self, n_features, classes=None):
 
     def initialize_params(self):
         # initialise the trainable parameters
-        shape = qml.StronglyEntanglingLayers.shape(
-            n_layers=self.n_layers, n_wires=self.n_qubits_
-        )
-        weights = jax.random.uniform(
-            self.generate_key(), minval=0, maxval=2 * np.pi, shape=shape
-        )
+        shape = qml.StronglyEntanglingLayers.shape(n_layers=self.n_layers, n_wires=self.n_qubits_)
+        weights = jax.random.uniform(self.generate_key(), minval=0, maxval=2 * np.pi, shape=shape)
         b = jnp.array(0.01)
         self.params_ = {"weights": weights, "b": b}
 
@@ -194,9 +190,7 @@ def fit(self, X, y):
         optimizer = optax.adam
 
         def loss_fn(params, X, y):
-            pred = self.forward(
-                params, X
-            )  # jnp.stack([self.forward(params, x) for x in X])
+            pred = self.forward(params, X)  # jnp.stack([self.forward(params, x) for x in X])
             return jnp.mean(optax.l2_loss(pred, y))
 
         if self.jit:
@@ -254,7 +248,5 @@ def transform(self, X, preprocess=False):
         padding = np.ones(shape=(len(X), n_padding)) / max_n_features
 
         X_padded = np.c_[X, padding]
-        X_normalised = np.divide(
-            X_padded, np.expand_dims(np.linalg.norm(X_padded, axis=1), axis=1)
-        )
+        X_normalised = np.divide(X_padded, np.expand_dims(np.linalg.norm(X_padded, axis=1), axis=1))
         return X_normalised
diff --git a/src/qml_benchmarks/models/convolutional_neural_network.py b/src/qml_benchmarks/models/convolutional_neural_network.py
index f4670d20..08f28c70 100644
--- a/src/qml_benchmarks/models/convolutional_neural_network.py
+++ b/src/qml_benchmarks/models/convolutional_neural_network.py
@@ -37,13 +37,9 @@ class CNN(nn.Module):
 
         @nn.compact
         def __call__(self, x):
-            x = nn.Conv(
-                features=output_channels[0], kernel_size=(kernel_shape, kernel_shape)
-            )(x)
+            x = nn.Conv(features=output_channels[0], kernel_size=(kernel_shape, kernel_shape))(x)
             x = nn.max_pool(x, window_shape=(2, 2), strides=(2, 2))
-            x = nn.Conv(
-                features=output_channels[1], kernel_size=(kernel_shape, kernel_shape)
-            )(x)
+            x = nn.Conv(features=output_channels[1], kernel_size=(kernel_shape, kernel_shape))(x)
             x = nn.max_pool(x, window_shape=(2, 2), strides=(2, 2))
             x = x.reshape((x.shape[0], -1))  # flatten
             x = nn.Dense(features=output_channels[1] * 2)(x)
diff --git a/src/qml_benchmarks/models/data_reuploading.py b/src/qml_benchmarks/models/data_reuploading.py
index 9d80f6ad..8466a12a 100644
--- a/src/qml_benchmarks/models/data_reuploading.py
+++ b/src/qml_benchmarks/models/data_reuploading.py
@@ -138,10 +138,7 @@ def circuit(params, x):
                 x_idx = 0  # to keep track of the data index
                 for i in range(self.n_qubits_):
                     # scaled inputs
-                    angles = (
-                        x[x_idx : x_idx + 3]
-                        * params["omegas"][layer, x_idx : x_idx + 3]
-                    )
+                    angles = x[x_idx : x_idx + 3] * params["omegas"][layer, x_idx : x_idx + 3]
                     qml.Rot(*angles, wires=i)
 
                     # variational
@@ -158,8 +155,7 @@ def circuit(params, x):
             x_idx = 0
             for i in range(self.n_qubits_):
                 angles = (
-                    x[x_idx : x_idx + 3]
-                    * params["omegas"][self.n_layers, x_idx : x_idx + 3]
+                    x[x_idx : x_idx + 3] * params["omegas"][self.n_layers, x_idx : x_idx + 3]
                     + params["thetas"][i, self.n_layers, :]
                 )
                 qml.Rot(*angles, wires=i)
@@ -254,8 +250,7 @@ def loss_fn(params, x, y):
                 / 2
                 * (
                     jnp.sum(
-                        (alpha_mat0 * probs0 - y_mat0) ** 2
-                        + (alpha_mat1 * probs1 - y_mat1) ** 2
+                        (alpha_mat0 * probs0 - y_mat0) ** 2 + (alpha_mat1 * probs1 - y_mat1) ** 2
                     )
                 )
             )  # eqn 23 in plots
@@ -409,20 +404,14 @@ def circuit(params, x):
             for layer in range(self.n_layers):
                 x_idx = 0  # to keep track of the data index
                 for i in range(self.n_qubits_):
-                    angles = (
-                        x[x_idx : x_idx + 3]
-                        * params["omegas"][layer, x_idx : x_idx + 3]
-                    )
+                    angles = x[x_idx : x_idx + 3] * params["omegas"][layer, x_idx : x_idx + 3]
                     qml.Rot(*angles, wires=i)
                     x_idx += 3
 
             # final reupload without CZs
             x_idx = 0
             for i in range(self.n_qubits_):
-                angles = (
-                    x[x_idx : x_idx + 3]
-                    * params["omegas"][self.n_layers, x_idx : x_idx + 3]
-                )
+                angles = x[x_idx : x_idx + 3] * params["omegas"][self.n_layers, x_idx : x_idx + 3]
                 qml.Rot(*angles, wires=i)
                 x_idx += 3
 
@@ -521,8 +510,7 @@ def circuit(params, x):
                 x_idx = 0  # to keep track of the data index
                 for i in range(self.n_qubits_):
                     angles = (
-                        x[x_idx : x_idx + 3]
-                        * params["omegas"][layer, x_idx : x_idx + 3]
+                        x[x_idx : x_idx + 3] * params["omegas"][layer, x_idx : x_idx + 3]
                         + params["thetas"][i, layer, :]
                     )
                     qml.Rot(*angles, wires=i)
@@ -532,8 +520,7 @@ def circuit(params, x):
             x_idx = 0
             for i in range(self.n_qubits_):
                 angles = (
-                    x[x_idx : x_idx + 3]
-                    * params["omegas"][self.n_layers, x_idx : x_idx + 3]
+                    x[x_idx : x_idx + 3] * params["omegas"][self.n_layers, x_idx : x_idx + 3]
                     + params["thetas"][i, self.n_layers, :]
                 )
                 qml.Rot(*angles, wires=i)
diff --git a/src/qml_benchmarks/models/dressed_quantum_circuit.py b/src/qml_benchmarks/models/dressed_quantum_circuit.py
index 6ccca60d..ed0c7fc8 100644
--- a/src/qml_benchmarks/models/dressed_quantum_circuit.py
+++ b/src/qml_benchmarks/models/dressed_quantum_circuit.py
@@ -161,19 +161,14 @@ def initialize_params(self):
         circuit_weights = (
             2
             * jnp.pi
-            * jax.random.uniform(
-                shape=(self.n_layers, self.n_qubits_), key=self.generate_key()
-            )
+            * jax.random.uniform(shape=(self.n_layers, self.n_qubits_), key=self.generate_key())
         )
         input_weights = (
-            jax.random.normal(
-                shape=(self.n_qubits_, self.n_qubits_), key=self.generate_key()
-            )
+            jax.random.normal(shape=(self.n_qubits_, self.n_qubits_), key=self.generate_key())
             / self.n_features_
         )
         output_weights = (
-            jax.random.normal(shape=(2, self.n_qubits_), key=self.generate_key())
-            / self.n_features_
+            jax.random.normal(shape=(2, self.n_qubits_), key=self.generate_key()) / self.n_features_
         )
         self.params_ = {
             "circuit_weights": circuit_weights,
@@ -277,20 +272,15 @@ def dressed_circuit(params, x):
     def initialize_params(self):
         # initialise the trainable parameters
         mid_weights = (
-            jax.random.normal(
-                shape=(self.n_qubits_, self.n_qubits_), key=self.generate_key()
-            )
+            jax.random.normal(shape=(self.n_qubits_, self.n_qubits_), key=self.generate_key())
             / self.n_features_
         )
         input_weights = (
-            jax.random.normal(
-                shape=(self.n_qubits_, self.n_qubits_), key=self.generate_key()
-            )
+            jax.random.normal(shape=(self.n_qubits_, self.n_qubits_), key=self.generate_key())
             / self.n_features_
         )
         output_weights = (
-            jax.random.normal(shape=(2, self.n_qubits_), key=self.generate_key())
-            / self.n_features_
+            jax.random.normal(shape=(2, self.n_qubits_), key=self.generate_key()) / self.n_features_
         )
         self.params_ = {
             "mid_weights": mid_weights,
diff --git a/src/qml_benchmarks/models/iqp_kernel.py b/src/qml_benchmarks/models/iqp_kernel.py
index 3d899b58..6682abd2 100644
--- a/src/qml_benchmarks/models/iqp_kernel.py
+++ b/src/qml_benchmarks/models/iqp_kernel.py
@@ -30,9 +30,9 @@ def __init__(
         svm=SVC(kernel="precomputed", probability=True),
         repeats=2,
         C=1.0,
-        use_jax=False,
-        vmap=False,
-        jit=False,
+        use_jax=None,
+        vmap=None,
+        jit=None,
         random_state=42,
         scaling=1.0,
         max_vmap=250,
@@ -76,16 +76,17 @@ def __init__(
         # attributes that do not depend on data
         self.repeats = repeats
         self.C = C
-        self.use_jax = use_jax
-        self.vmap = vmap
-        self.jit = jit
         self.max_vmap = max_vmap
         self.svm = svm
-        self.dev_type = dev_type
         self.qnode_kwargs = qnode_kwargs
         self.scaling = scaling
         self.random_state = random_state
         self.rng = np.random.default_rng(random_state)
+        # device-related attributes
+        self.dev_type = dev_type
+        self.use_jax = use_jax if use_jax is not None else self.dev_type == "default.qubit.jax"
+        self.vmap = vmap if vmap is not None else self.use_jax
+        self.jit = jit if jit is not None else self.use_jax
 
         # data-dependant attributes
         # which will be initialised by calling "fit"
@@ -128,6 +129,10 @@ def circuit(x):
 
         if self.use_jax and self.jit:
             circuit = jax.jit(circuit)
+        elif "lightning" in self.dev_type and self.jit:
+            from catalyst import qjit
+
+            circuit = qjit(circuit)
         return circuit
 
     def precompute_kernel(self, X1, X2):
@@ -142,24 +147,23 @@ def precompute_kernel(self, X1, X2):
         dim1 = len(X1)
         dim2 = len(X2)
 
-        # concatenate all pairs of vectors
-        Z = np.array(
-            [np.concatenate((X1[i], X2[j])) for i in range(dim1) for j in range(dim2)]
-        )
-
         circuit = self.construct_circuit()
 
         if self.use_jax and self.vmap:
+            # concatenate all pairs of vectors
+            Z = np.array([np.concatenate((X1[i], X2[j])) for i in range(dim1) for j in range(dim2)])
             # if batched circuit is used
             self.batched_circuit = chunk_vmapped_fn(
                 jax.vmap(circuit, 0), start=0, max_vmap=self.max_vmap
             )
             kernel_values = self.batched_circuit(Z)[:, 0]
+            # reshape the values into the kernel matrix
+            kernel_matrix = np.reshape(kernel_values, (dim1, dim2))
         else:
-            kernel_values = np.array([circuit(z)[0] for z in Z])
-
-        # reshape the values into the kernel matrix
-        kernel_matrix = np.reshape(kernel_values, (dim1, dim2))
+            kernel_matrix = np.empty((dim1, dim2))
+            for i, x in enumerate(X1):
+                for j, y in enumerate(X2):
+                    kernel_matrix[i, j] = circuit(np.concatenate((x, y)))[0]
 
         return kernel_matrix
 
@@ -192,9 +196,7 @@ def fit(self, X, y):
 
         if self.use_jax:
             self.svm.random_state = int(
-                jax.random.randint(
-                    self.generate_key(), shape=(1,), minval=0, maxval=1000000
-                )
+                jax.random.randint(self.generate_key(), shape=(1,), minval=0, maxval=1000000)
             )
         else:
             self.svm.random_state = self.generate_key()
diff --git a/src/qml_benchmarks/models/iqp_variational.py b/src/qml_benchmarks/models/iqp_variational.py
index 2fa45fd6..866c35ea 100644
--- a/src/qml_benchmarks/models/iqp_variational.py
+++ b/src/qml_benchmarks/models/iqp_variational.py
@@ -106,6 +106,7 @@ def construct_model(self):
         dev = qml.device(self.dev_type, wires=self.n_qubits_)
 
         if self.use_jax:
+
             @qml.qnode(dev, **self.qnode_kwargs)
             def circuit(params, x):
                 """
@@ -136,6 +137,7 @@ def forward(params, X):
                 self.forward = forward
 
         else:
+
             @qml.qnode(dev, **self.qnode_kwargs)
             def circuit(weights, x):
                 """
@@ -189,13 +191,7 @@ def initialize_params(self):
                 )
             )
         else:
-            weights = (
-                2
-                * np.pi
-                * np.random.uniform(
-                    size=(self.n_layers, self.n_qubits_, 3)
-                )
-            )
+            weights = 2 * np.pi * np.random.uniform(size=(self.n_layers, self.n_qubits_, 3))
             weights = pnp.array(weights, requires_grad=True)
 
         self.params_ = {"weights": weights}
diff --git a/src/qml_benchmarks/models/projected_quantum_kernel.py b/src/qml_benchmarks/models/projected_quantum_kernel.py
index 83ae06ed..43b30798 100644
--- a/src/qml_benchmarks/models/projected_quantum_kernel.py
+++ b/src/qml_benchmarks/models/projected_quantum_kernel.py
@@ -202,24 +202,9 @@ def precompute_kernel(self, X1, X2):
 
         for i in range(dim1):
             for j in range(dim2):
-                sumX = sum(
-                    [
-                        (valsX_X1[i, q] - valsX_X2[j, q]) ** 2
-                        for q in range(self.n_qubits_)
-                    ]
-                )
-                sumY = sum(
-                    [
-                        (valsY_X1[i, q] - valsY_X2[j, q]) ** 2
-                        for q in range(self.n_qubits_)
-                    ]
-                )
-                sumZ = sum(
-                    [
-                        (valsZ_X1[i, q] - valsZ_X2[j, q]) ** 2
-                        for q in range(self.n_qubits_)
-                    ]
-                )
+                sumX = sum([(valsX_X1[i, q] - valsX_X2[j, q]) ** 2 for q in range(self.n_qubits_)])
+                sumY = sum([(valsY_X1[i, q] - valsY_X2[j, q]) ** 2 for q in range(self.n_qubits_)])
+                sumZ = sum([(valsZ_X1[i, q] - valsZ_X2[j, q]) ** 2 for q in range(self.n_qubits_)])
 
                 kernel_matrix[i, j] = np.exp(
                     -default_gamma * self.gamma_factor * (sumX + sumY + sumZ)
diff --git a/src/qml_benchmarks/models/quantum_boltzmann_machine.py b/src/qml_benchmarks/models/quantum_boltzmann_machine.py
index 843c3b97..ba394e9a 100644
--- a/src/qml_benchmarks/models/quantum_boltzmann_machine.py
+++ b/src/qml_benchmarks/models/quantum_boltzmann_machine.py
@@ -7,17 +7,19 @@
 from sklearn.preprocessing import StandardScaler
 import itertools
 from qml_benchmarks.model_utils import chunk_vmapped_fn
+
 jax.config.update("jax_enable_x64", True)
 
 sigmaZ = jnp.array([[1, 0], [0, -1]])
 sigmaX = jnp.array([[0, 1], [1, 0]])
 sigmaY = jnp.array([[0, -1j], [1j, 0]])
 
+
 def tensor_ops(ops, idxs, n_qubits):
     """
     Returns a tensor product of two operators acting at indexes idxs in an n_qubit system
     """
-    tensor_op = 1.
+    tensor_op = 1.0
     for i in range(n_qubits):
         if i in idxs:
             j = idxs.index(i)
@@ -26,21 +28,22 @@ def tensor_ops(ops, idxs, n_qubits):
             tensor_op = jnp.kron(tensor_op, jnp.eye(2))
     return tensor_op
 
+
 class QuantumBoltzmannMachine(BaseEstimator, ClassifierMixin):
 
     def __init__(
-            self,
-            visible_qubits='single',
-            observable_type='sum',
-            temperature=1,
-            learning_rate=0.001,
-            batch_size=32,
-            max_vmap=None,
-            jit=True,
-            max_steps=10000,
-            convergence_threshold=1e-6,
-            random_state=42,
-            scaling=1.0
+        self,
+        visible_qubits="single",
+        observable_type="sum",
+        temperature=1,
+        learning_rate=0.001,
+        batch_size=32,
+        max_vmap=None,
+        jit=True,
+        max_steps=10000,
+        convergence_threshold=1e-6,
+        random_state=42,
+        scaling=1.0,
     ):
         """
         Variational Quantum Boltzmann Machine from https://arxiv.org/abs/2006.06004
@@ -109,7 +112,7 @@ def __init__(
         # which will be initialised by calling "fit"
         self.params_ = None  # Dictionary containing the trainable parameters
         self.n_qubits = None
-        self.n_visible = None #number of visible qubits
+        self.n_visible = None  # number of visible qubits
         self.scaler = None  # data scaler will be fitted on training data
         self.circuit = None
 
@@ -121,16 +124,19 @@ def construct_model(self):
         doubles = list(itertools.combinations(np.arange(self.n_qubits), 2))
         self.n_params_ = 2 * len(singles) + len(doubles)
 
-        if self.observable_type == 'sum':
-            obs = sum([tensor_ops([sigmaZ], (i,), self.n_qubits) for i in range(self.n_visible)])/self.n_visible
-        elif self.observable_type == 'product':
-            obs = 1.
+        if self.observable_type == "sum":
+            obs = (
+                sum([tensor_ops([sigmaZ], (i,), self.n_qubits) for i in range(self.n_visible)])
+                / self.n_visible
+            )
+        elif self.observable_type == "product":
+            obs = 1.0
             for i in range(self.n_visible):
                 obs = jnp.kron(obs, sigmaZ)
 
         def gibbs_state(thetas, x):
 
-            H = jnp.zeros([2 ** self.n_qubits, 2 ** self.n_qubits])
+            H = jnp.zeros([2**self.n_qubits, 2**self.n_qubits])
             count = 0
             for idxs in singles:
                 H = H + tensor_ops([sigmaZ], idxs, self.n_qubits) * jnp.dot(thetas[count], x)
@@ -139,7 +145,9 @@ def gibbs_state(thetas, x):
                 count = count + 1
 
             for idxs in doubles:
-                H = H + tensor_ops([sigmaZ, sigmaZ], idxs, self.n_qubits) * jnp.dot(thetas[count], x)
+                H = H + tensor_ops([sigmaZ, sigmaZ], idxs, self.n_qubits) * jnp.dot(
+                    thetas[count], x
+                )
 
             state = jax.scipy.linalg.expm(-H / self.temperature, max_squarings=32)
             return state / jnp.trace(state)
@@ -148,8 +156,9 @@ def model(thetas, x):
             state = gibbs_state(thetas, x)
             return jnp.trace(jnp.matmul(state, obs))
 
-        if self.jit: model = jax.jit(model)
-        self.forward = jax.vmap(model,in_axes=(None,0))
+        if self.jit:
+            model = jax.jit(model)
+        self.forward = jax.vmap(model, in_axes=(None, 0))
         self.chunked_forward = chunk_vmapped_fn(self.forward, 1, self.max_vmap)
 
         return self.forward
@@ -171,11 +180,11 @@ def initialize(self, n_features, classes=None):
 
         self.n_qubits = n_features
 
-        if self.visible_qubits == 'single':
+        if self.visible_qubits == "single":
             self.n_visible = 1
-        elif self.visible_qubits == 'half':
-            self.n_visible = self.n_qubits//2
-        elif self.visible_qubits == 'all':
+        elif self.visible_qubits == "half":
+            self.n_visible = self.n_qubits // 2
+        elif self.visible_qubits == "all":
             self.n_visible = self.n_qubits
 
         self.construct_model()
@@ -203,12 +212,13 @@ def fit(self, X, y):
 
         def loss_fn(params, X, y):
             # binary cross entropy loss
-            vals = self.forward(params['thetas'], X)
+            vals = self.forward(params["thetas"], X)
             probs = (1 + vals) / 2
             y = jax.nn.relu(y)  # convert to 0,1
-            return jnp.mean(-y*jnp.log(probs) - (1-y)*jnp.log(1-probs))
+            return jnp.mean(-y * jnp.log(probs) - (1 - y) * jnp.log(1 - probs))
 
-        if self.jit: loss_fn = jax.jit(loss_fn)
+        if self.jit:
+            loss_fn = jax.jit(loss_fn)
         self.params_ = train(self, loss_fn, optimizer, X, y, self.generate_key)
 
         return self
@@ -237,7 +247,7 @@ def predict_proba(self, X):
             (n_samples, n_classes)
         """
         X = self.transform(X)
-        predictions = self.forward(self.params_['thetas'], X)
+        predictions = self.forward(self.params_["thetas"], X)
         predictions_2d = np.c_[(1 - predictions) / 2, (1 + predictions) / 2]
         return predictions_2d
 
@@ -258,25 +268,30 @@ class QuantumBoltzmannMachineSeparable(QuantumBoltzmannMachine):
 
     def construct_model(self):
         def qubit_gibbs_state(thetas, x):
-            H = sigmaZ*jnp.dot(thetas[0],x)+sigmaX*jnp.dot(thetas[1],x)
+            H = sigmaZ * jnp.dot(thetas[0], x) + sigmaX * jnp.dot(thetas[1], x)
             state = jax.scipy.linalg.expm(-H / self.temperature, max_squarings=32)
             return state / jnp.trace(state)
 
         def model(thetas, x):
-            gibbs_states = [qubit_gibbs_state(thetas[2*i:2*i+2,:],x) for i in range(self.n_visible)]
-            expvals = jnp.array([jnp.trace(jnp.matmul(state, sigmaZ)) for state in gibbs_states ])
-            if self.observable_type=='sum':
+            gibbs_states = [
+                qubit_gibbs_state(thetas[2 * i : 2 * i + 2, :], x) for i in range(self.n_visible)
+            ]
+            expvals = jnp.array([jnp.trace(jnp.matmul(state, sigmaZ)) for state in gibbs_states])
+            if self.observable_type == "sum":
                 return jnp.mean(expvals)
-            elif self.observable_type=='product':
+            elif self.observable_type == "product":
                 return jnp.prod(expvals)
 
-        if self.jit: model = jax.jit(model)
-        self.forward = jax.vmap(model,in_axes=(None,0))
+        if self.jit:
+            model = jax.jit(model)
+        self.forward = jax.vmap(model, in_axes=(None, 0))
         self.chunked_forward = chunk_vmapped_fn(self.forward, 1, self.max_vmap)
 
         return self.forward
 
     def initialize_params(self):
         # initialise the trainable parameters
-        params = jax.random.normal(shape=(2*self.n_qubits, self.n_qubits), key=self.generate_key())
+        params = jax.random.normal(
+            shape=(2 * self.n_qubits, self.n_qubits), key=self.generate_key()
+        )
         self.params_ = {"thetas": params}
diff --git a/src/qml_benchmarks/models/quantum_kitchen_sinks.py b/src/qml_benchmarks/models/quantum_kitchen_sinks.py
index 6dd6d34d..68c389b5 100644
--- a/src/qml_benchmarks/models/quantum_kitchen_sinks.py
+++ b/src/qml_benchmarks/models/quantum_kitchen_sinks.py
@@ -180,9 +180,7 @@ def initialize_params(self):
         betas = (
             2
             * np.pi
-            * jax.random.uniform(
-                key=self.generate_key(), shape=(self.n_episodes, self.n_qubits_)
-            )
+            * jax.random.uniform(key=self.generate_key(), shape=(self.n_episodes, self.n_qubits_))
         )
         self.params_ = {"omegas": np.array(omegas), "betas": np.array(betas)}
 
@@ -241,11 +239,7 @@ def transform(self, X, preprocess=True):
         Args:
             X (np.ndarray): Data of shape (n_samples, n_features)
         """
-        if (
-            self.params_["betas"] is None
-            or self.params_["omegas"] is None
-            or self.circuit is None
-        ):
+        if self.params_["betas"] is None or self.params_["omegas"] is None or self.circuit is None:
             raise ValueError("Model cannot predict without fitting to data first.")
 
         if preprocess:
diff --git a/src/qml_benchmarks/models/quantum_metric_learning.py b/src/qml_benchmarks/models/quantum_metric_learning.py
index 2dcfda5b..dcca5191 100644
--- a/src/qml_benchmarks/models/quantum_metric_learning.py
+++ b/src/qml_benchmarks/models/quantum_metric_learning.py
@@ -129,12 +129,8 @@ def construct_model(self):
         @qml.qnode(dev, **self.qnode_kwargs)
         def circuit(params, x1, x2):
             qml.QAOAEmbedding(features=x1, weights=params["weights"], wires=wires)
-            qml.adjoint(qml.QAOAEmbedding)(
-                features=x2, weights=params["weights"], wires=wires
-            )
-            return qml.expval(
-                qml.Projector(np.array([0] * self.n_qubits_), wires=wires)
-            )
+            qml.adjoint(qml.QAOAEmbedding)(features=x2, weights=params["weights"], wires=wires)
+            return qml.expval(qml.Projector(np.array([0] * self.n_qubits_), wires=wires))
 
         self.circuit = circuit
 
@@ -202,9 +198,7 @@ def fit(self, X, y):
         B = jnp.array(X[y == 1])
 
         if self.batch_size > min(len(A), len(B)):
-            warnings.warn(
-                "batch size too large, setting to " + str(min(len(A), len(B)))
-            )
+            warnings.warn("batch size too large, setting to " + str(min(len(A), len(B))))
             self.batch_size = min(len(A), len(B))
 
         def loss_fn(params, A=None, B=None):
@@ -260,13 +254,9 @@ def predict_proba(self, X):
 
         X = self.transform(X)
 
-        max_examples = min(
-            len(self.params_["examples_-1"]), len(self.params_["examples_+1"])
-        )
+        max_examples = min(len(self.params_["examples_-1"]), len(self.params_["examples_+1"]))
         if self.n_examples_predict > max_examples:
-            warnings.warn(
-                "n_examples_predict too large, setting to " + str(max_examples)
-            )
+            warnings.warn("n_examples_predict too large, setting to " + str(max_examples))
             self.n_examples_predict = max_examples
 
         A_examples, B_examples = get_batch(
diff --git a/src/qml_benchmarks/models/quanvolutional_neural_network.py b/src/qml_benchmarks/models/quanvolutional_neural_network.py
index d08d4f07..790417ac 100644
--- a/src/qml_benchmarks/models/quanvolutional_neural_network.py
+++ b/src/qml_benchmarks/models/quanvolutional_neural_network.py
@@ -174,9 +174,7 @@ def construct_random_circuit(self):
             jnp.pi
             * 2
             * jnp.array(
-                jax.random.uniform(
-                    self.generate_key(), shape=(self.rand_depth, self.rand_rot)
-                )
+                jax.random.uniform(self.generate_key(), shape=(self.rand_depth, self.rand_rot))
             )
         )
 
@@ -203,14 +201,10 @@ def construct_quanvolutional_layer(self):
         """
         construct the quantum feature map.
         """
-        random_circuits = [
-            self.construct_random_circuit() for __ in range(self.n_qchannels)
-        ]
+        random_circuits = [self.construct_random_circuit() for __ in range(self.n_qchannels)]
 
         # construct an array that specifies the indices of the 'windows' of the image used for the convolution.
-        idx_mat = jnp.array(
-            [[(i, j) for j in range(self.width)] for i in range(self.height)]
-        )
+        idx_mat = jnp.array([[(i, j) for j in range(self.width)] for i in range(self.height)])
         idxs = jnp.array(
             [
                 idx_mat[j : j + self.qkernel_shape, k : k + self.qkernel_shape]
diff --git a/src/qml_benchmarks/models/separable.py b/src/qml_benchmarks/models/separable.py
index 56e852f0..f2a773fe 100644
--- a/src/qml_benchmarks/models/separable.py
+++ b/src/qml_benchmarks/models/separable.py
@@ -101,9 +101,7 @@ def single_qubit_circuit(weights, x):
             below to get the full circuit
             """
             for layer in range(self.encoding_layers):
-                qml.Rot(
-                    weights[layer, 0], weights[layer, 1], weights[layer, 2], wires=0
-                )
+                qml.Rot(weights[layer, 0], weights[layer, 1], weights[layer, 2], wires=0)
                 qml.RY(x, wires=0)
             qml.Rot(
                 weights[self.encoding_layers, 0],
@@ -341,9 +339,7 @@ def precompute_kernel(self, X1, X2):
         dim2 = len(X2)
 
         # concatenate all pairs of vectors
-        Z = np.array(
-            [np.concatenate((X1[i], X2[j])) for i in range(dim1) for j in range(dim2)]
-        )
+        Z = np.array([np.concatenate((X1[i], X2[j])) for i in range(dim1) for j in range(dim2)])
         self.construct_circuit()
         kernel_values = [self.forward(z) for z in Z]
         # reshape the values into the kernel matrix
@@ -379,9 +375,7 @@ def fit(self, X, y):
         """
 
         self.svm.random_state = int(
-            jax.random.randint(
-                self.generate_key(), shape=(1,), minval=0, maxval=1000000
-            )
+            jax.random.randint(self.generate_key(), shape=(1,), minval=0, maxval=1000000)
         )
 
         self.initialize(X.shape[1], np.unique(y))
diff --git a/src/qml_benchmarks/models/tree_tensor.py b/src/qml_benchmarks/models/tree_tensor.py
index b4199ac6..9c386d49 100644
--- a/src/qml_benchmarks/models/tree_tensor.py
+++ b/src/qml_benchmarks/models/tree_tensor.py
@@ -110,8 +110,7 @@ def circuit(params, x):
                     qml.CNOT,
                     wires=range(self.n_qubits),
                     pattern=[
-                        ((i + 2**layer), i)
-                        for i in range(0, self.n_qubits, 2 ** (layer + 1))
+                        ((i + 2**layer), i) for i in range(0, self.n_qubits, 2 ** (layer + 1))
                     ],
                 )
             qml.RY(params["weights"][count], wires=0)
@@ -162,11 +161,7 @@ def initialize(self, n_features, classes=None):
     def initialize_params(self):
         # initialise the trainable parameters
         weights = (
-            2
-            * jnp.pi
-            * jax.random.uniform(
-                shape=(2 * self.n_qubits - 1,), key=self.generate_key()
-            )
+            2 * jnp.pi * jax.random.uniform(shape=(2 * self.n_qubits - 1,), key=self.generate_key())
         )
 
         bias = 0.1 * jax.random.normal(shape=(1,), key=self.generate_key())
@@ -241,9 +236,7 @@ def transform(self, X, preprocess=True):
         n_features = X.shape[1]
         X = X * self.scaling
 
-        n_qubits_ae = int(
-            np.ceil(np.log2(n_features))
-        )  # the num qubits needed to amplitude encode
+        n_qubits_ae = int(np.ceil(np.log2(n_features)))  # the num qubits needed to amplitude encode
         n_qubits = 2 ** int(
             np.ceil(np.log2(n_qubits_ae))
         )  # the model needs 2**m qubits, for some m
@@ -252,7 +245,5 @@ def transform(self, X, preprocess=True):
         padding = np.ones(shape=(len(X), n_padding)) / max_n_features
 
         X_padded = np.c_[X, padding]
-        X_normalised = np.divide(
-            X_padded, np.expand_dims(np.linalg.norm(X_padded, axis=1), axis=1)
-        )
+        X_normalised = np.divide(X_padded, np.expand_dims(np.linalg.norm(X_padded, axis=1), axis=1))
         return X_normalised
diff --git a/src/qml_benchmarks/models/weinet.py b/src/qml_benchmarks/models/weinet.py
index 57a8193b..7669df46 100644
--- a/src/qml_benchmarks/models/weinet.py
+++ b/src/qml_benchmarks/models/weinet.py
@@ -152,16 +152,12 @@ def circuit(x):
                         jnp.reshape(x, -1), wires=wires, normalize=True, pad_with=0.0
                     )
                     qml.QubitUnitary(
-                        jnp.kron(
-                            self.unitaries[nu][nu], jnp.array(self.unitaries[nu][mu])
-                        ),
+                        jnp.kron(self.unitaries[nu][nu], jnp.array(self.unitaries[nu][mu])),
                         wires=wires,
                     )
                     return [qml.expval(op) for op in operators]
 
-                self.circuit = (
-                    circuit  # we use the last one of the circuits here as an example
-                )
+                self.circuit = circuit  # we use the last one of the circuits here as an example
 
                 if self.jit:
                     circuit = jax.jit(circuit)
@@ -176,9 +172,7 @@ def forward_fn(self, params, x):
         is equivalent to classically sampling one of the unitaries Q_i, parameterised by params['s'].
         """
         probs = jax.nn.softmax(params["s"])
-        expvals = jnp.array(
-            [probs[i] * jnp.array(self.circuits[i](x)).T for i in range(9)]
-        )
+        expvals = jnp.array([probs[i] * jnp.array(self.circuits[i](x)).T for i in range(9)])
         expvals = jnp.sum(expvals, axis=0)
         out = jnp.sum(params["weights"] * expvals)
         # out = jax.nn.sigmoid(out)  # convert to a probability
@@ -215,15 +209,12 @@ def initialize_params(self):
         """
         # no of expvals that are combined with weights
         n_expvals = int(
-            self.n_qubits_
-            - 1
-            + factorial(self.n_qubits_ - 2) / 2 / factorial(self.n_qubits_ - 4)
+            self.n_qubits_ - 1 + factorial(self.n_qubits_ - 2) / 2 / factorial(self.n_qubits_ - 4)
         )
 
         self.params_ = {
             "s": jax.random.normal(self.generate_key(), shape=(9,)),
-            "weights": jax.random.normal(self.generate_key(), shape=(n_expvals,))
-            / n_expvals,
+            "weights": jax.random.normal(self.generate_key(), shape=(n_expvals,)) / n_expvals,
         }
 
     def fit(self, X, y):

From c39a4f8c6699902b4976aa8fc32d5c46a9968ec6 Mon Sep 17 00:00:00 2001
From: Vincent Michaud-Rioux <vincent.michaud-rioux@xanadu.ai>
Date: Thu, 29 Feb 2024 16:39:12 +0000
Subject: [PATCH 006/100] Add qjit in iqp_variational.py module.

---
 src/qml_benchmarks/models/iqp_variational.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/src/qml_benchmarks/models/iqp_variational.py b/src/qml_benchmarks/models/iqp_variational.py
index 866c35ea..ab4e8b98 100644
--- a/src/qml_benchmarks/models/iqp_variational.py
+++ b/src/qml_benchmarks/models/iqp_variational.py
@@ -150,7 +150,11 @@ def circuit(weights, x):
                 )
                 return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
 
-            self.circuit = circuit
+            if self.jit:
+                from catalyst import qjit
+                circuit = qjit(circuit)
+
+            self.circuit = circuit 
 
             # use autograd and do not batch feed the circuit
             def forward(params, X):

From 98dcf678068e3e2e5f19984167169a37845c9c56 Mon Sep 17 00:00:00 2001
From: Vincent Michaud-Rioux <vincent.michaud-rioux@xanadu.ai>
Date: Thu, 29 Feb 2024 18:34:14 +0000
Subject: [PATCH 007/100] Fix iqp_var solver.

---
 src/qml_benchmarks/models/iqp_kernel.py      |  5 -----
 src/qml_benchmarks/models/iqp_variational.py | 15 ++++++++++-----
 2 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/src/qml_benchmarks/models/iqp_kernel.py b/src/qml_benchmarks/models/iqp_kernel.py
index 6682abd2..bc731f01 100644
--- a/src/qml_benchmarks/models/iqp_kernel.py
+++ b/src/qml_benchmarks/models/iqp_kernel.py
@@ -36,11 +36,6 @@ def __init__(
         random_state=42,
         scaling=1.0,
         max_vmap=250,
-        # device-related attributes
-        self.dev_type = dev_type
-        self.use_jax = use_jax if use_jax is not None else self.dev_type == "default.qubit.jax"
-        self.vmap = vmap if vmap is not None else self.use_jax
-        self.jit = jit if jit is not None else self.use_jax
         qnode_kwargs={},
     ):
         r"""
diff --git a/src/qml_benchmarks/models/iqp_variational.py b/src/qml_benchmarks/models/iqp_variational.py
index ab4e8b98..0b9af398 100644
--- a/src/qml_benchmarks/models/iqp_variational.py
+++ b/src/qml_benchmarks/models/iqp_variational.py
@@ -153,12 +153,13 @@ def circuit(weights, x):
             if self.jit:
                 from catalyst import qjit
                 circuit = qjit(circuit)
+                # circuit(np.random.rand(self.n_layers, self.n_qubits_, 3), np.random.rand(self.n_qubits_))
 
             self.circuit = circuit 
 
             # use autograd and do not batch feed the circuit
             def forward(params, X):
-                return pnp.array([circuit(params, x) for x in X])
+                return jnp.array([circuit(params, x) for x in X])
 
             self.forward = forward
 
@@ -196,7 +197,7 @@ def initialize_params(self):
             )
         else:
             weights = 2 * np.pi * np.random.uniform(size=(self.n_layers, self.n_qubits_, 3))
-            weights = pnp.array(weights, requires_grad=True)
+            weights = jnp.array(weights)
 
         self.params_ = {"weights": weights}
 
@@ -237,14 +238,18 @@ def loss_fn(params, X, y):
             )
 
         else:
-            X = pnp.array(X, requires_grad=False)
-            y = pnp.array(y, requires_grad=False)
+            X = np.array(X)
+            y = np.array(y)
             optimizer = qml.AdamOptimizer
 
             def loss_fn(weights, X, y):
                 expvals = self.forward(weights, X)
                 probs = (1 - expvals * y) / 2  # the probs of incorrect classification
-                return pnp.mean(probs)
+                return jnp.mean(probs)
+
+            if self.jit:
+                from catalyst import qjit
+                loss_fn = qjit(loss_fn)
 
             self.params_ = train_without_jax(self, loss_fn, optimizer, X, y, self.generate_key)
 

From c7af8b6a43f4806568514149e2befdbd0a269abe Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Tue, 26 Mar 2024 16:40:38 +0100
Subject: [PATCH 008/100] faster qjitting

---
 src/qml_benchmarks/models/iqp_kernel.py | 72 +++++++++++++++----------
 1 file changed, 45 insertions(+), 27 deletions(-)

diff --git a/src/qml_benchmarks/models/iqp_kernel.py b/src/qml_benchmarks/models/iqp_kernel.py
index bc731f01..3c83d744 100644
--- a/src/qml_benchmarks/models/iqp_kernel.py
+++ b/src/qml_benchmarks/models/iqp_kernel.py
@@ -13,13 +13,17 @@
 # limitations under the License.
 
 import time
+
+import catalyst
 import pennylane as qml
 import numpy as np
 import jax
+import jax.numpy as jnp
 from sklearn.base import BaseEstimator, ClassifierMixin
 from sklearn.svm import SVC
 from sklearn.preprocessing import MinMaxScaler
 from qml_benchmarks.model_utils import chunk_vmapped_fn
+from catalyst import qjit
 
 jax.config.update("jax_enable_x64", True)
 
@@ -30,6 +34,7 @@ def __init__(
         svm=SVC(kernel="precomputed", probability=True),
         repeats=2,
         C=1.0,
+        dev_type = 'default.qubit.jax',
         use_jax=None,
         vmap=None,
         jit=None,
@@ -98,36 +103,40 @@ def generate_key(self):
     def construct_circuit(self):
         dev = qml.device(self.dev_type, wires=self.n_qubits_)
 
-        @qml.qnode(dev, **self.qnode_kwargs)
-        def circuit(x):
-            """
-            circuit used to precomute the kernel matrix K(x_1,x_2).
-            Args:
-                x (np.array): vector of length 2*num_feature that is the concatenation of x_1 and x_2
-
-            Returns:
-                (float) the value of the kernel fucntion K(x_1,x_2)
-            """
-            qml.IQPEmbedding(
-                x[: self.n_qubits_], wires=range(self.n_qubits_), n_repeats=self.repeats
-            )
-            qml.adjoint(
+        def wrapped_circuit(x):
+            @qml.qnode(dev, **self.qnode_kwargs)
+            def circuit(x):
+                """
+                circuit used to precomute the kernel matrix K(x_1,x_2).
+                Args:
+                    x (np.array): vector of length 2*num_feature that is the concatenation of x_1 and x_2
+
+                Returns:
+                    (float) the value of the kernel fucntion K(x_1,x_2)
+                """
                 qml.IQPEmbedding(
-                    x[self.n_qubits_ :],
-                    wires=range(self.n_qubits_),
-                    n_repeats=self.repeats,
+                    x[: self.n_qubits_], wires=range(self.n_qubits_), n_repeats=self.repeats
+                )
+                qml.adjoint(
+                    qml.IQPEmbedding(
+                        x[self.n_qubits_ :],
+                        wires=range(self.n_qubits_),
+                        n_repeats=self.repeats,
+                    )
                 )
-            )
-            return qml.probs()
 
-        self.circuit = circuit
+                return qml.probs()
+            return circuit(x)[0]
+
+        circuit = wrapped_circuit
 
         if self.use_jax and self.jit:
             circuit = jax.jit(circuit)
         elif "lightning" in self.dev_type and self.jit:
-            from catalyst import qjit
-
             circuit = qjit(circuit)
+
+        self.circuit = circuit
+
         return circuit
 
     def precompute_kernel(self, X1, X2):
@@ -151,14 +160,23 @@ def precompute_kernel(self, X1, X2):
             self.batched_circuit = chunk_vmapped_fn(
                 jax.vmap(circuit, 0), start=0, max_vmap=self.max_vmap
             )
-            kernel_values = self.batched_circuit(Z)[:, 0]
+            kernel_values = self.batched_circuit(Z)
             # reshape the values into the kernel matrix
             kernel_matrix = np.reshape(kernel_values, (dim1, dim2))
         else:
-            kernel_matrix = np.empty((dim1, dim2))
-            for i, x in enumerate(X1):
-                for j, y in enumerate(X2):
-                    kernel_matrix[i, j] = circuit(np.concatenate((x, y)))[0]
+            # could also use catalyst.vmap like as above, although I think it does basically the same thing as this.
+            X1 = jnp.array(X1)
+            X2 = jnp.array(X2)
+            @qjit(autograph=True)
+            def construct_kernel(X1,X2):
+                dim1 = len(X1)
+                dim2 = len(X2)
+                kernel_matrix = jnp.zeros([dim1, dim2])
+                for i, x in enumerate(X1):
+                    for j, y in enumerate(X2):
+                       kernel_matrix = kernel_matrix.at[i,j].set(circuit(jnp.concatenate((x, y))))
+                return  kernel_matrix
+            kernel_matrix = construct_kernel(X1,X2)
 
         return kernel_matrix
 

From d357dffc6a534b7e1d60bb8f9f50b8c661dfe9d0 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Wed, 27 Mar 2024 17:18:59 +0100
Subject: [PATCH 009/100] catalyst training

---
 src/qml_benchmarks/model_utils.py            | 75 ++++++++++++++++++++
 src/qml_benchmarks/models/iqp_variational.py | 20 +++---
 2 files changed, 84 insertions(+), 11 deletions(-)

diff --git a/src/qml_benchmarks/model_utils.py b/src/qml_benchmarks/model_utils.py
index d35ab9df..2645a128 100644
--- a/src/qml_benchmarks/model_utils.py
+++ b/src/qml_benchmarks/model_utils.py
@@ -18,6 +18,10 @@
 from functools import reduce
 import logging
 import time
+
+import pennylane as qml
+import catalyst
+from catalyst import qjit
 import numpy as np
 import optax
 import jax
@@ -89,6 +93,8 @@ def update(params, opt_state, x, y):
         loss_history.append(loss_val)
         logging.debug(f"{step} - loss: {loss_val}")
 
+        print(loss_val)
+
         if np.isnan(loss_val):
             logging.info(f"nan encountered. Training aborted.")
             break
@@ -170,6 +176,75 @@ def train_without_jax(
     return model.params_
 
 
+def train_with_catalyst(model, loss_fn, optimizer, X, y, random_key_generator, convergence_interval=200):
+
+    params = model.params_
+    opt = optimizer(learning_rate=model.learning_rate)
+
+    @qjit
+    def update(i, args):
+        params, opt_state, X, y, loss_history, key = args
+        X_batch, y_batch = get_batch(X, y, key, batch_size=model.batch_size)
+        loss = loss_fn(params, X_batch, y_batch)
+        loss_history =  loss_history.at[i].set(loss)
+        #backprop is not supported in catalyst yet, falling back on finite diff
+        grads = catalyst.grad(loss_fn, method="fd")(params,X_batch, y_batch)
+        updates, opt_state = opt.update(grads, opt_state)
+        params = optax.apply_updates(params, updates)
+        key, subkey = jax.random.split(key)
+        return (params, opt_state, X, y, loss_history, subkey)
+
+    @qjit
+    def optimize(params, X, y, steps, loss_history, opt_state, key):
+        args = (params, opt_state, X, y, loss_history, key)
+        (params, opt_state, _, _, loss_history, key) = catalyst.for_loop(0,steps,1)(update)(args)
+        return params, loss_history, opt_state
+
+    def train_until_convergence(params, X, y):
+        converged = False
+        current_steps = 0
+        opt_state = opt.init(params)
+        loss_histories = []
+        while current_steps < model.max_steps:
+
+            key = random_key_generator()
+            loss_history = jnp.zeros(convergence_interval)
+            params, loss_history, opt_state = optimize(params, X, y, convergence_interval, loss_history, opt_state, key)
+            loss_histories.append(loss_history)
+            current_steps += convergence_interval
+
+            if True in jnp.isnan(loss_history):
+                logging.info(f"nan encountered. Training aborted.")
+                break
+
+            if len(loss_histories)>=2:
+                average1 = jnp.mean(loss_histories[-1])
+                average2 = jnp.mean(loss_histories[-2])
+                std1 = jnp.std(loss_history[-1])
+                if jnp.abs(average2 - average1) <= std1 / jnp.sqrt(convergence_interval) / 2:
+                    logging.info(f"Model {model.__class__.__name__} converged after {step} steps.")
+                    converged = True
+                    break
+
+        if not converged:
+            print("Loss did not converge:", loss_history)
+            raise ConvergenceWarning(
+                f"Model {model.__class__.__name__} has not converged after the maximum number of {model.max_steps} steps."
+            )
+
+        return params, jnp.concatenate(loss_histories)
+
+    start = time.time()
+    params, loss_history = train_until_convergence(params, X, y)
+    end = time.time()
+    loss_history = np.array(loss_history)
+    model.loss_history_ = loss_history / np.max(np.abs(loss_history))
+    model.training_time_ = end - start
+
+    return params
+
+
+
 def get_batch(X, y, rnd_key, batch_size=32):
     """
     A generator to get random batches of the data (X, y)
diff --git a/src/qml_benchmarks/models/iqp_variational.py b/src/qml_benchmarks/models/iqp_variational.py
index 0b9af398..015c3022 100644
--- a/src/qml_benchmarks/models/iqp_variational.py
+++ b/src/qml_benchmarks/models/iqp_variational.py
@@ -16,6 +16,7 @@
 import numpy as np
 from sklearn.base import BaseEstimator, ClassifierMixin
 from sklearn.preprocessing import MinMaxScaler
+from catalyst import qjit
 from qml_benchmarks.model_utils import *
 
 
@@ -98,9 +99,7 @@ def __init__(
         self.circuit = None
 
     def generate_key(self):
-        if self.use_jax:
-            return jax.random.PRNGKey(self.rng.integers(1000000))
-        return self.rng.integers(1000000)
+        return jax.random.PRNGKey(self.rng.integers(1000000))
 
     def construct_model(self):
         dev = qml.device(self.dev_type, wires=self.n_qubits_)
@@ -139,19 +138,18 @@ def forward(params, X):
         else:
 
             @qml.qnode(dev, **self.qnode_kwargs)
-            def circuit(weights, x):
+            def circuit(params, x):
                 """
                 The variational circuit from the plots. Uses an IQP data embedding.
                 We use the same observable as in the plots.
                 """
                 qml.IQPEmbedding(x, wires=range(self.n_qubits_), n_repeats=self.repeats)
                 qml.StronglyEntanglingLayers(
-                    weights, wires=range(self.n_qubits_), imprimitive=qml.CZ
+                    params["weights"], wires=range(self.n_qubits_), imprimitive=qml.CZ
                 )
                 return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
 
             if self.jit:
-                from catalyst import qjit
                 circuit = qjit(circuit)
                 # circuit(np.random.rand(self.n_layers, self.n_qubits_, 3), np.random.rand(self.n_qubits_))
 
@@ -240,10 +238,10 @@ def loss_fn(params, X, y):
         else:
             X = np.array(X)
             y = np.array(y)
-            optimizer = qml.AdamOptimizer
+            optimizer = optax.adam
 
-            def loss_fn(weights, X, y):
-                expvals = self.forward(weights, X)
+            def loss_fn(params, X, y):
+                expvals = self.forward(params, X)
                 probs = (1 - expvals * y) / 2  # the probs of incorrect classification
                 return jnp.mean(probs)
 
@@ -251,7 +249,7 @@ def loss_fn(weights, X, y):
                 from catalyst import qjit
                 loss_fn = qjit(loss_fn)
 
-            self.params_ = train_without_jax(self, loss_fn, optimizer, X, y, self.generate_key)
+            self.params_ = train_with_catalyst(self, loss_fn, optimizer, X, y, self.generate_key)
 
         return self
 
@@ -285,7 +283,7 @@ def predict_proba(self, X):
             else:
                 predictions = self.forward(self.params_, X)
         else:
-            predictions = self.forward(self.params_["weights"], X)
+            predictions = self.forward(self.params_, X)
         predictions_2d = np.c_[(1 - predictions) / 2, (1 + predictions) / 2]
         return predictions_2d
 

From 27b4d9dc42feae4175c95faea66cce9b25769be0 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Thu, 28 Mar 2024 12:53:05 +0100
Subject: [PATCH 010/100] catalyst port attempt

---
 .../models/projected_quantum_kernel.py        | 123 ++++++++++++++----
 1 file changed, 96 insertions(+), 27 deletions(-)

diff --git a/src/qml_benchmarks/models/projected_quantum_kernel.py b/src/qml_benchmarks/models/projected_quantum_kernel.py
index 43b30798..48658d9b 100644
--- a/src/qml_benchmarks/models/projected_quantum_kernel.py
+++ b/src/qml_benchmarks/models/projected_quantum_kernel.py
@@ -13,6 +13,9 @@
 # limitations under the License.
 
 import time
+
+import catalyst
+from catalyst import qjit
 import pennylane as qml
 import numpy as np
 import jax
@@ -34,11 +37,13 @@ def __init__(
         embedding="Hamiltonian",
         t=1.0 / 3,
         trotter_steps=5,
+        use_jax=False,
         jit=True,
+        vmap = False,
         max_vmap=None,
         scaling=1.0,
         dev_type="default.qubit.jax",
-        qnode_kwargs={"interface": "jax-jit", "diff_method": None},
+        qnode_kwargs={},
         random_state=42,
     ):
         r"""
@@ -91,6 +96,8 @@ def __init__(
         self.embedding = embedding
         self.t = t
         self.trotter_steps = trotter_steps
+        self.use_jax = use_jax
+        self.vmap = vmap
         self.jit = jit
         self.dev_type = dev_type
         self.qnode_kwargs = qnode_kwargs
@@ -146,24 +153,58 @@ def embedding(x):
 
         dev = qml.device(self.dev_type, wires=self.n_qubits_)
 
-        @qml.qnode(dev, **self.qnode_kwargs)
-        def circuit(x):
-            embedding(x)
-            return (
-                [qml.expval(qml.PauliX(wires=i)) for i in range(self.n_qubits_)]
-                + [qml.expval(qml.PauliY(wires=i)) for i in range(self.n_qubits_)]
-                + [qml.expval(qml.PauliZ(wires=i)) for i in range(self.n_qubits_)]
-            )
+        if "lightning" in self.dev_type and self.jit:
+            # currently only support for returning expvals of qubit-wise commuting observables,
+            # so we split into three circuits
+            @qjit(autograph=True)
+            @qml.qnode(dev, **self.qnode_kwargs)
+            def circuitX(x):
+                embedding(x)
+                return [qml.expval(qml.PauliX(wires=i)) for i in range(self.n_qubits_)]
+
+            @qjit(autograph=True)
+            @qml.qnode(dev, **self.qnode_kwargs)
+            def circuitY(x):
+                embedding(x)
+                return [qml.expval(qml.PauliY(wires=i)) for i in range(self.n_qubits_)]
+
+            @qjit(autograph=True)
+            @qml.qnode(dev, **self.qnode_kwargs)
+            def circuitZ(x):
+                embedding(x)
+                return [qml.expval(qml.PauliZ(wires=i)) for i in range(self.n_qubits_)]
+
+            # @qjit(autograph=True)
+            def circuit_as_array(x):
+                xvals = jnp.array(circuitX(x))
+                yvals = jnp.array(circuitY(x))
+                zvals = jnp.array(circuitZ(x))
+                return jnp.vstack((xvals, yvals, zvals))
+
+        else:
+            @qml.qnode(dev, **self.qnode_kwargs)
+            def circuit(x):
+                embedding(x)
+                return (
+                        [qml.expval(qml.PauliX(wires=i)) for i in range(self.n_qubits_)]
+                        + [qml.expval(qml.PauliY(wires=i)) for i in range(self.n_qubits_)]
+                        + [qml.expval(qml.PauliZ(wires=i)) for i in range(self.n_qubits_)]
+                )
 
-        self.circuit = circuit
+            def circuit_as_array(x):
+                return jnp.array(circuit(x))
 
-        def circuit_as_array(x):
-            return jnp.array(circuit(x))
 
-        if self.jit:
+        if self.use_jax and self.jit:
             circuit_as_array = jax.jit(circuit_as_array)
-        circuit_as_array = jax.vmap(circuit_as_array, in_axes=(0))
-        circuit_as_array = chunk_vmapped_fn(circuit_as_array, 0, self.max_vmap)
+            circuit_as_array = jax.vmap(circuit_as_array, in_axes=(0))
+            circuit_as_array = chunk_vmapped_fn(circuit_as_array, 0, self.max_vmap)
+
+        elif "lightning" in self.dev_type and self.jit:
+            # circuit_as_array = qjit(circuit_as_array)
+            def batch_circuit_as_array(X):
+                return jnp.concatenate([jnp.array(circuitX(x)) for x in X])
+            circuit_as_array = batch_circuit_as_array
 
         return circuit_as_array
 
@@ -194,21 +235,48 @@ def precompute_kernel(self, X1, X2):
         valsZ_X1 = valsX1[:, 2]
         valsZ_X2 = valsX2[:, 2]
 
-        all_vals_X1 = np.reshape(np.concatenate((valsX_X1, valsY_X1, valsZ_X1)), -1)
+        all_vals_X1 = np.reshape(jnp.concatenate((valsX_X1, valsY_X1, valsZ_X1)), -1)
         default_gamma = 1 / np.var(all_vals_X1) / self.n_features_
 
-        # construct kernel following plots
-        kernel_matrix = np.zeros([dim1, dim2])
+        if self.use_jax:
+            #no actually using JAX here but it is part of the JAX pipeline
+            kernel_matrix = np.zeros([dim1, dim2])
+            for i in range(dim1):
+                for j in range(dim2):
+                    sumX = sum([(valsX_X1[i, q] - valsX_X2[j, q]) ** 2 for q in range(self.n_qubits_)])
+                    sumY = sum([(valsY_X1[i, q] - valsY_X2[j, q]) ** 2 for q in range(self.n_qubits_)])
+                    sumZ = sum([(valsZ_X1[i, q] - valsZ_X2[j, q]) ** 2 for q in range(self.n_qubits_)])
+
+                    kernel_matrix[i, j] = np.exp(
+                        -default_gamma * self.gamma_factor * (sumX + sumY + sumZ)
+                    )
 
-        for i in range(dim1):
-            for j in range(dim2):
-                sumX = sum([(valsX_X1[i, q] - valsX_X2[j, q]) ** 2 for q in range(self.n_qubits_)])
-                sumY = sum([(valsY_X1[i, q] - valsY_X2[j, q]) ** 2 for q in range(self.n_qubits_)])
-                sumZ = sum([(valsZ_X1[i, q] - valsZ_X2[j, q]) ** 2 for q in range(self.n_qubits_)])
+        else:
+            valsX_X1 = jnp.array(valsX_X1)
+            valsX_X2 = jnp.array(valsX_X2)
+            valsY_X1 = jnp.array(valsY_X1)
+            valsY_X2 = jnp.array(valsY_X2)
+            valsZ_X1 = jnp.array(valsZ_X1)
+            valsZ_X2 = jnp.array(valsZ_X2)
+
+            @qjit(autograph=True)
+            def construct_kernel(valsX_X1, valsX_X2, valsY_X1, valsY_X2, valsZ_X1, valsZ_X2,
+                                 dim1, dim2, default_gamma):
+                kernel_matrix = jnp.zeros([dim1, dim2])
+                for i in range(dim1):
+                    for j in range(dim2):
+                        sumX = sum([(valsX_X1[i, q] - valsX_X2[j, q]) ** 2 for q in range(self.n_qubits_)])
+                        sumY = sum([(valsY_X1[i, q] - valsY_X2[j, q]) ** 2 for q in range(self.n_qubits_)])
+                        sumZ = sum([(valsZ_X1[i, q] - valsZ_X2[j, q]) ** 2 for q in range(self.n_qubits_)])
+
+                        kernel_matrix = kernel_matrix.at[i,j].set(np.exp(
+                            -default_gamma * self.gamma_factor * (sumX + sumY + sumZ))
+                        )
+                return kernel_matrix
+
+            kernel_matrix = construct_kernel(valsX_X1, valsX_X2, valsY_X1, valsY_X2, valsZ_X1, valsZ_X2,
+                                 dim1, dim2, default_gamma)
 
-                kernel_matrix[i, j] = np.exp(
-                    -default_gamma * self.gamma_factor * (sumX + sumY + sumZ)
-                )
         return kernel_matrix
 
     def initialize(self, n_features, classes=None):
@@ -251,11 +319,12 @@ def fit(self, X, y):
         self.scaler = MinMaxScaler(feature_range=(-np.pi / 2, np.pi / 2))
         self.scaler.fit(X)
         X = self.transform(X)
+        X = jnp.array(X)
 
         self.params_ = {"X_train": X}
+        start = time.time()
         kernel_matrix = self.precompute_kernel(X, X)
 
-        start = time.time()
         self.svm.C = self.C
         self.svm.fit(kernel_matrix, y)
         end = time.time()

From 92b921da50fd769bc4f854b9c7b93438e8d91334 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Thu, 28 Mar 2024 15:08:52 +0100
Subject: [PATCH 011/100] add catalyst support

---
 .../models/dressed_quantum_circuit.py         | 57 ++++++++++++++-----
 1 file changed, 43 insertions(+), 14 deletions(-)

diff --git a/src/qml_benchmarks/models/dressed_quantum_circuit.py b/src/qml_benchmarks/models/dressed_quantum_circuit.py
index ed0c7fc8..354d22f7 100644
--- a/src/qml_benchmarks/models/dressed_quantum_circuit.py
+++ b/src/qml_benchmarks/models/dressed_quantum_circuit.py
@@ -13,6 +13,8 @@
 # limitations under the License.
 
 import pennylane as qml
+import catalyst
+from catalyst import qjit
 import numpy as np
 from sklearn.preprocessing import StandardScaler
 from sklearn.base import BaseEstimator, ClassifierMixin
@@ -26,6 +28,8 @@ def __init__(
         learning_rate=0.001,
         batch_size=32,
         max_vmap=None,
+        use_jax = False,
+        vmap = False,
         jit=True,
         max_steps=100000,
         convergence_interval=200,
@@ -66,6 +70,8 @@ def __init__(
         self.convergence_interval = convergence_interval
         self.dev_type = dev_type
         self.qnode_kwargs = qnode_kwargs
+        self.use_jax = use_jax
+        self.vmap = vmap
         self.jit = jit
         self.scaling = scaling
         self.random_state = random_state
@@ -130,9 +136,19 @@ def dressed_circuit(params, x):
             return x
 
         if self.jit:
-            dressed_circuit = jax.jit(dressed_circuit)
-        self.forward = jax.vmap(dressed_circuit, in_axes=(None, 0))
-        self.chunked_forward = chunk_vmapped_fn(self.forward, 1, self.max_vmap)
+            if self.use_jax:
+                dressed_circuit = jax.jit(dressed_circuit)
+            else:
+                dressed_circuit = qjit(dressed_circuit, autograph=True)
+
+        if self.vmap and self.use_jax:
+            self.forward = jax.vmap(dressed_circuit, in_axes=(None, 0))
+            self.chunked_forward = chunk_vmapped_fn(self.forward, 1, self.max_vmap)
+
+        else:
+            def forward(params,X):
+                return jnp.array([dressed_circuit(params,x) for x in X])
+            self.forward = forward
 
         return self.forward
 
@@ -199,16 +215,25 @@ def loss_fn(params, X, y):
             return jnp.mean(optax.softmax_cross_entropy(vals, labels))
 
         if self.jit:
-            loss_fn = jax.jit(loss_fn)
-        self.params_ = train(
-            self,
-            loss_fn,
-            optimizer,
-            X,
-            y,
-            self.generate_key,
-            convergence_interval=self.convergence_interval,
-        )
+            if self.use_jax:
+                loss_fn = jax.jit(loss_fn)
+            else:
+                loss_fn = qjit(loss_fn)
+
+        if self.use_jax:
+            self.params_ = train(
+                self,
+                loss_fn,
+                optimizer,
+                X,
+                y,
+                self.generate_key,
+                convergence_interval=self.convergence_interval,
+            )
+        else:
+            self.params_ = train_with_catalyst(
+                self, loss_fn, optimizer, X, y, self.generate_key,
+                convergence_interval=self.convergence_interval)
 
         return self
 
@@ -236,7 +261,11 @@ def predict_proba(self, X):
             (n_samples, n_classes)
         """
         X = self.transform(X)
-        return jax.nn.softmax(self.chunked_forward(self.params_, X))
+        if self.vmap:
+            out = self.chunked_forward(self.params_, X)
+        else:
+            out = self.forward(self.params_, X)
+        return jax.nn.softmax(out)
 
     def transform(self, X, preprocess=True):
         """

From 868a62c0141f7d1ced1a24e6587f0ffb2bb1e0e8 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Thu, 28 Mar 2024 15:24:46 +0100
Subject: [PATCH 012/100] attempt catalyst port

---
 src/qml_benchmarks/models/projected_quantum_kernel.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/qml_benchmarks/models/projected_quantum_kernel.py b/src/qml_benchmarks/models/projected_quantum_kernel.py
index 48658d9b..01c5abe4 100644
--- a/src/qml_benchmarks/models/projected_quantum_kernel.py
+++ b/src/qml_benchmarks/models/projected_quantum_kernel.py
@@ -174,12 +174,12 @@ def circuitZ(x):
                 embedding(x)
                 return [qml.expval(qml.PauliZ(wires=i)) for i in range(self.n_qubits_)]
 
-            # @qjit(autograph=True)
+            @qjit(autograph=True)
             def circuit_as_array(x):
                 xvals = jnp.array(circuitX(x))
                 yvals = jnp.array(circuitY(x))
                 zvals = jnp.array(circuitZ(x))
-                return jnp.vstack((xvals, yvals, zvals))
+                return jnp.concatenate((xvals, yvals, zvals))
 
         else:
             @qml.qnode(dev, **self.qnode_kwargs)
@@ -203,7 +203,7 @@ def circuit_as_array(x):
         elif "lightning" in self.dev_type and self.jit:
             # circuit_as_array = qjit(circuit_as_array)
             def batch_circuit_as_array(X):
-                return jnp.concatenate([jnp.array(circuitX(x)) for x in X])
+                return jnp.array([circuit_as_array(x) for x in X])
             circuit_as_array = batch_circuit_as_array
 
         return circuit_as_array

From af61f6534f43e3d71fe78c21cc6fa395e12c707b Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Tue, 9 Apr 2024 10:14:57 +0200
Subject: [PATCH 013/100] add code

---
 nersc/pm_podman.source | 45 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 45 insertions(+)
 create mode 100644 nersc/pm_podman.source

diff --git a/nersc/pm_podman.source b/nersc/pm_podman.source
new file mode 100644
index 00000000..8b3c4ee4
--- /dev/null
+++ b/nersc/pm_podman.source
@@ -0,0 +1,45 @@
+#!/bin/bash
+#  salloc -q interactive -C cpu -t 4:00:00 -A nstaff 
+
+IMG=balewski/ubu22-pennylane:p3
+#IMG=balewski/ubu22-qiskit1.0:p2b # + QML + IonQ +Braket + pytorch+PennyLane
+
+CFSH=/global/cfs/cdirs/mpccc/balewski/
+
+echo launch image $IMG
+echo you are launching Podman image ...  remeber to exit
+
+JNB_PORT=' '
+BASE_DIR=/quantumMind   # here git has home
+WORK_DIR=$BASE_DIR/PennyLane
+WORK_DIR2=/pennylane-qml-benchmarks
+DATA_VAULT=${CFSH}/quantDataVault2024
+DATA_DIR=/dataPennyLane_tmp
+
+echo "The number of arguments is: $#"
+#  encoded variables:    jnb
+for var in "$@"; do
+  echo "The length of argument '$var' is: ${#var}"
+  if [[ "jnb" ==  $var ]];  then
+      #JNB_PORT="    --publish 8833:8833 "
+      JNB_PORT="    -p 8833:8833 "
+     echo added  $JNB_PORT
+     echo "  cd notebooks;  jupyter notebook --ip 0.0.0.0 --no-browser --allow-root --port  8833 "
+  fi
+  # ... more ...
+done
+
+
+podman-hpc run -it \
+   --volume $HOME:/home \
+   --volume $CFSH/$BASE_DIR:/$BASE_DIR \
+   --volume $CFSH/$WORK_DIR:$WORK_DIR \
+   --volume $CFSH/$WORK_DIR2:$WORK_DIR2 \
+   --volume ${DATA_VAULT}:/dataVault \
+   --volume ${DATA_VAULT}/$DATA_DIR:/dataPennyLane_tmp \
+   -e HDF5_USE_FILE_LOCKING='FALSE' \
+   -e PennyLane_dataVault=/dataPennyLane_tmp \
+   --workdir $WORK_DIR  \
+   $IMG /bin/bash
+
+

From 43ddf7410005f8bf54b335dfad55b8c8b3390917 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <jbowles@x1004c0s0b0n1h0.chn.perlmutter.nersc.gov>
Date: Wed, 10 Apr 2024 03:08:51 -0700
Subject: [PATCH 014/100] test

---
 nersc/test.py | 14 ++++++++++++++
 1 file changed, 14 insertions(+)
 create mode 100644 nersc/test.py

diff --git a/nersc/test.py b/nersc/test.py
new file mode 100644
index 00000000..c5d2101e
--- /dev/null
+++ b/nersc/test.py
@@ -0,0 +1,14 @@
+import pennylane
+import qml_benchmarks
+from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier
+from qml_benchmarks.data.two_curves import generate_two_curves
+import numpy as np
+
+
+
+X,y = generate_two_curves(100,3, 2, 0.0, 0.1)
+model = IQPVariationalClassifier(use_jax=True, vmap=True)
+
+model.fit(X, y)
+print(model.loss_history_)
+np.savetxt('loss.txt', model.loss_history_)
\ No newline at end of file

From 56fc1539cb87b8693dd662420829ef5f0f6bc1bb Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Wed, 10 Apr 2024 12:22:29 +0200
Subject: [PATCH 015/100] update

---
 nersc/pm_podman.source | 21 ++-------------------
 1 file changed, 2 insertions(+), 19 deletions(-)

diff --git a/nersc/pm_podman.source b/nersc/pm_podman.source
index 8b3c4ee4..d13a528b 100644
--- a/nersc/pm_podman.source
+++ b/nersc/pm_podman.source
@@ -2,39 +2,22 @@
 #  salloc -q interactive -C cpu -t 4:00:00 -A nstaff 
 
 IMG=balewski/ubu22-pennylane:p3
-#IMG=balewski/ubu22-qiskit1.0:p2b # + QML + IonQ +Braket + pytorch+PennyLane
-
 CFSH=/global/cfs/cdirs/mpccc/balewski/
 
 echo launch image $IMG
 echo you are launching Podman image ...  remeber to exit
 
 JNB_PORT=' '
-BASE_DIR=/quantumMind   # here git has home
-WORK_DIR=$BASE_DIR/PennyLane
-WORK_DIR2=/pennylane-qml-benchmarks
+BASE_DIR=/qml-benchmarks   # here git has home
+WORK_DIR=$BASE_DIR/nersc
 DATA_VAULT=${CFSH}/quantDataVault2024
 DATA_DIR=/dataPennyLane_tmp
 
-echo "The number of arguments is: $#"
-#  encoded variables:    jnb
-for var in "$@"; do
-  echo "The length of argument '$var' is: ${#var}"
-  if [[ "jnb" ==  $var ]];  then
-      #JNB_PORT="    --publish 8833:8833 "
-      JNB_PORT="    -p 8833:8833 "
-     echo added  $JNB_PORT
-     echo "  cd notebooks;  jupyter notebook --ip 0.0.0.0 --no-browser --allow-root --port  8833 "
-  fi
-  # ... more ...
-done
-
 
 podman-hpc run -it \
    --volume $HOME:/home \
    --volume $CFSH/$BASE_DIR:/$BASE_DIR \
    --volume $CFSH/$WORK_DIR:$WORK_DIR \
-   --volume $CFSH/$WORK_DIR2:$WORK_DIR2 \
    --volume ${DATA_VAULT}:/dataVault \
    --volume ${DATA_VAULT}/$DATA_DIR:/dataPennyLane_tmp \
    -e HDF5_USE_FILE_LOCKING='FALSE' \

From 46e6eb58548720b7561653db2e09a98dfc0d2034 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Wed, 10 Apr 2024 12:28:45 +0200
Subject: [PATCH 016/100] update

---
 nersc/pm_podman.source | 7 +------
 1 file changed, 1 insertion(+), 6 deletions(-)

diff --git a/nersc/pm_podman.source b/nersc/pm_podman.source
index d13a528b..ee5a4dfc 100644
--- a/nersc/pm_podman.source
+++ b/nersc/pm_podman.source
@@ -2,7 +2,7 @@
 #  salloc -q interactive -C cpu -t 4:00:00 -A nstaff 
 
 IMG=balewski/ubu22-pennylane:p3
-CFSH=/global/cfs/cdirs/mpccc/balewski/
+CFSH=/global/cfs/cdirs/mpccc/jbowles/
 
 echo launch image $IMG
 echo you are launching Podman image ...  remeber to exit
@@ -10,18 +10,13 @@ echo you are launching Podman image ...  remeber to exit
 JNB_PORT=' '
 BASE_DIR=/qml-benchmarks   # here git has home
 WORK_DIR=$BASE_DIR/nersc
-DATA_VAULT=${CFSH}/quantDataVault2024
-DATA_DIR=/dataPennyLane_tmp
 
 
 podman-hpc run -it \
    --volume $HOME:/home \
    --volume $CFSH/$BASE_DIR:/$BASE_DIR \
    --volume $CFSH/$WORK_DIR:$WORK_DIR \
-   --volume ${DATA_VAULT}:/dataVault \
-   --volume ${DATA_VAULT}/$DATA_DIR:/dataPennyLane_tmp \
    -e HDF5_USE_FILE_LOCKING='FALSE' \
-   -e PennyLane_dataVault=/dataPennyLane_tmp \
    --workdir $WORK_DIR  \
    $IMG /bin/bash
 

From 81a614dfbf91f5196f6be6294afaa2f05e6ee004 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Wed, 10 Apr 2024 12:30:55 +0200
Subject: [PATCH 017/100] update

---
 nersc/pm_podman.source | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/nersc/pm_podman.source b/nersc/pm_podman.source
index ee5a4dfc..e57836d2 100644
--- a/nersc/pm_podman.source
+++ b/nersc/pm_podman.source
@@ -2,7 +2,7 @@
 #  salloc -q interactive -C cpu -t 4:00:00 -A nstaff 
 
 IMG=balewski/ubu22-pennylane:p3
-CFSH=/global/cfs/cdirs/mpccc/jbowles/
+CFSH=/global/homes/j/jbowles/qml-benchmarks
 
 echo launch image $IMG
 echo you are launching Podman image ...  remeber to exit

From 66169c780f1dac6259d1ab20302cdc6be29d3104 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Wed, 10 Apr 2024 12:32:39 +0200
Subject: [PATCH 018/100] update

---
 nersc/pm_podman.source | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/nersc/pm_podman.source b/nersc/pm_podman.source
index e57836d2..d011e337 100644
--- a/nersc/pm_podman.source
+++ b/nersc/pm_podman.source
@@ -2,7 +2,7 @@
 #  salloc -q interactive -C cpu -t 4:00:00 -A nstaff 
 
 IMG=balewski/ubu22-pennylane:p3
-CFSH=/global/homes/j/jbowles/qml-benchmarks
+CFSH=/global/homes/j/jbowles
 
 echo launch image $IMG
 echo you are launching Podman image ...  remeber to exit

From 5c01770d4aa734b3568e7be5f8f75f02142f1375 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Fri, 12 Apr 2024 16:12:15 +0200
Subject: [PATCH 019/100] v1

---
 .../hyperparam_settings.yaml                  |  4 +
 .../perf_ind_iqpvar_LR14_jax.py               | 78 +++++++++++++++++++
 2 files changed, 82 insertions(+)
 create mode 100644 nersc/performance_indicators/hyperparam_settings.yaml
 create mode 100644 nersc/performance_indicators/perf_ind_iqpvar_LR14_jax.py

diff --git a/nersc/performance_indicators/hyperparam_settings.yaml b/nersc/performance_indicators/hyperparam_settings.yaml
new file mode 100644
index 00000000..e51f5ec0
--- /dev/null
+++ b/nersc/performance_indicators/hyperparam_settings.yaml
@@ -0,0 +1,4 @@
+IQPVariationalClassifier:
+  learning_rate: 0.01
+  n_layers: 15
+  repeats: 10
\ No newline at end of file
diff --git a/nersc/performance_indicators/perf_ind_iqpvar_LR14_jax.py b/nersc/performance_indicators/perf_ind_iqpvar_LR14_jax.py
new file mode 100644
index 00000000..2ba55366
--- /dev/null
+++ b/nersc/performance_indicators/perf_ind_iqpvar_LR14_jax.py
@@ -0,0 +1,78 @@
+import qml_benchmarks
+import pennylane as qml
+import jax
+import jax.numpy as jnp
+import numpy as np
+import time
+import csv
+import os
+import yaml
+
+from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier
+from qml_benchmarks.hyperparam_search_utils import read_data
+
+with open('hyperparam_settings.yaml', "r") as file:
+    hp_settings = yaml.safe_load(file)
+
+hyperparams = {**hp_settings['IQPVariationalClassifier'], **{'use_jax':True, 'vmap':True, 'max_steps':100}}
+
+print(hyperparams)
+
+n_features = 2 #dataset dimension
+n_trials = 10 #number of trials to average over
+
+X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
+X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
+
+first_train_steps = []
+av_consec_train_steps = []
+predict_times = []
+
+for trial in range(n_trials):
+    jax.clear_caches()
+
+    model = IQPVariationalClassifier(**hyperparams)
+    model.fit(X_train, y_train)
+
+    #get step times from loss history data
+    step_times = np.array([model.loss_history_[1][i + 1] - model.loss_history_[1][i]
+                           for i in range(len(model.loss_history_[1]) - 1)])
+    step_times = np.insert(step_times, 0, [model.loss_history_[1][0]])
+
+    #first train step
+    first_train_steps.append(step_times[0])
+    #consecutive (average) train step
+    av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
+    #prediction time
+    time0 = time.time()
+    model.predict(X_test)
+    predict_times.append(time.time() - time0)
+
+
+#calculate mean and stds
+
+first_train_step = np.mean(first_train_steps)
+first_train_step_std = np.std(first_train_steps)
+
+consec_train_step = np.mean(av_consec_train_steps)
+consec_train_step_std = np.std(av_consec_train_steps)
+
+predict_time = np.mean(predict_times)
+predict_time_std = np.std(predict_times)
+
+#write to csv
+data = [first_train_step, first_train_step_std, consec_train_step, consec_train_step_std, predict_time, predict_time_std]
+
+model_name = model.__class__.__name__
+filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_JAX.csv"
+header = ['first_train_step', 'first_train_step_std', 'consec_train_step', 'consec_train_step_std', 'predict_time', 'predict_time_std']
+
+if not os.path.exists('JAX'):
+    # Create the directory
+    os.mkdir('JAX')
+
+with open('JAX/'+filename, mode="w", newline="") as file:
+    writer = csv.writer(file)
+    writer.writerow(header)
+    for row in [data]:
+        writer.writerow(row)

From 159666913a5ed5f48e5ebd97efb88605ab000bb5 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Fri, 12 Apr 2024 17:09:14 +0200
Subject: [PATCH 020/100] perf ind additions

---
 src/qml_benchmarks/model_utils.py | 21 ++++++++++++---------
 1 file changed, 12 insertions(+), 9 deletions(-)

diff --git a/src/qml_benchmarks/model_utils.py b/src/qml_benchmarks/model_utils.py
index 2645a128..93762bee 100644
--- a/src/qml_benchmarks/model_utils.py
+++ b/src/qml_benchmarks/model_utils.py
@@ -84,17 +84,17 @@ def update(params, opt_state, x, y):
         return params, opt_state, loss_val
 
     loss_history = []
+    steptimes = []
     converged = False
     start = time.time()
     for step in range(model.max_steps):
         key = random_key_generator()
         X_batch, y_batch = get_batch(X, y, key, batch_size=model.batch_size)
         params, opt_state, loss_val = update(params, opt_state, X_batch, y_batch)
+        steptimes.append(time.time() - start)
         loss_history.append(loss_val)
         logging.debug(f"{step} - loss: {loss_val}")
 
-        print(loss_val)
-
         if np.isnan(loss_val):
             logging.info(f"nan encountered. Training aborted.")
             break
@@ -112,15 +112,18 @@ def update(params, opt_state, x, y):
                 break
 
     end = time.time()
-    loss_history = np.array(loss_history)
-    model.loss_history_ = loss_history / np.max(np.abs(loss_history))
+    loss_history / np.max(np.abs(loss_history))
+    loss_history = np.array(np.vstack((loss_history,steptimes)))
+    model.loss_history_ = loss_history
     model.training_time_ = end - start
 
-    if not converged:
-        print("Loss did not converge:", loss_history)
-        raise ConvergenceWarning(
-            f"Model {model.__class__.__name__} has not converged after the maximum number of {model.max_steps} steps."
-        )
+    #removed for perforance profiling
+    #
+    # if not converged:
+    #     print("Loss did not converge:", loss_history)
+    #     raise ConvergenceWarning(
+    #         f"Model {model.__class__.__name__} has not converged after the maximum number of {model.max_steps} steps."
+    #     )
 
     return params
 

From 30ffcbbb5f66f6cf79d21a519da513e4f75a02e2 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Fri, 12 Apr 2024 17:22:56 +0200
Subject: [PATCH 021/100] v1

---
 .../save_last_job_info.sh                     | 21 +++++++++++++++++++
 1 file changed, 21 insertions(+)
 create mode 100644 nersc/performance_indicators/save_last_job_info.sh

diff --git a/nersc/performance_indicators/save_last_job_info.sh b/nersc/performance_indicators/save_last_job_info.sh
new file mode 100644
index 00000000..8a95c8e3
--- /dev/null
+++ b/nersc/performance_indicators/save_last_job_info.sh
@@ -0,0 +1,21 @@
+#!/bin/bash
+
+# Get the job ID of the last submitted job
+last_job_id=$(squeue -u $USER -t PD -h -o "%A" | tail -n 1)
+
+# Check if a job ID was found
+if [ -z "$last_job_id" ]; then
+    echo "No pending job found."
+    exit 1
+fi
+
+# Get the job information
+job_info=$(scontrol show job $last_job_id)
+
+# Determine the filename
+filename="${1:-job_info.txt}"
+
+# Save the job information to the specified file
+echo "$job_info" > "$filename"
+
+echo "Job information saved to $filename."
\ No newline at end of file

From d8bd09e7e3c98d53732b52a45b77cd4ecd8beae0 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Fri, 12 Apr 2024 17:24:49 +0200
Subject: [PATCH 022/100] v1

---
 nersc/performance_indicators/save_last_job_info.sh | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/nersc/performance_indicators/save_last_job_info.sh b/nersc/performance_indicators/save_last_job_info.sh
index 8a95c8e3..ab6d07e8 100644
--- a/nersc/performance_indicators/save_last_job_info.sh
+++ b/nersc/performance_indicators/save_last_job_info.sh
@@ -1,11 +1,11 @@
 #!/bin/bash
 
-# Get the job ID of the last submitted job
-last_job_id=$(squeue -u $USER -t PD -h -o "%A" | tail -n 1)
+# Get the job ID of the last completed job
+last_job_id=$(sacct --format=JobID -n -u $USER -S "$(date -d yesterday +"%Y-%m-%d")" -o JobID | tail -n 1)
 
 # Check if a job ID was found
 if [ -z "$last_job_id" ]; then
-    echo "No pending job found."
+    echo "No completed job found."
     exit 1
 fi
 
@@ -18,4 +18,4 @@ filename="${1:-job_info.txt}"
 # Save the job information to the specified file
 echo "$job_info" > "$filename"
 
-echo "Job information saved to $filename."
\ No newline at end of file
+echo "Job information saved to $filename."

From 7a65c8666ecb38024909743dbb83930b74cece76 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Mon, 15 Apr 2024 15:16:49 +0200
Subject: [PATCH 023/100] new code

---
 ...er_linearly_separable_14d_scontrol_JAX.csv |  0
 ...arly_separable_14d_scontrol_JAX_novmap.csv |  0
 ...eparable_2d_performance_indicators_JAX.csv |  2 +
 ...e_2d_performance_indicators_JAX_novmap.csv |  2 +
 .../perf_ind_iqpvar_LR14_jax.py               |  8 +-
 .../perf_ind_iqpvar_LR14_jax_novmap.py        | 80 +++++++++++++++++++
 .../perf_ind_iqpvar_LR15_jax.py               | 80 +++++++++++++++++++
 .../perf_ind_iqpvar_LR15_jax_novmap.py        | 80 +++++++++++++++++++
 .../perf_ind_iqpvar_LR16_jax.py               | 80 +++++++++++++++++++
 .../perf_ind_iqpvar_LR16_jax_novmap.py        | 80 +++++++++++++++++++
 .../perf_ind_iqpvar_LR17_jax.py               | 80 +++++++++++++++++++
 .../perf_ind_iqpvar_LR17_jax_novmap.py        | 80 +++++++++++++++++++
 .../perf_ind_iqpvar_LR18_jax.py               | 80 +++++++++++++++++++
 .../perf_ind_iqpvar_LR18_jax_novmap.py        | 80 +++++++++++++++++++
 .../perf_ind_iqpvar_LR19_jax.py               | 80 +++++++++++++++++++
 .../perf_ind_iqpvar_LR19_jax_novmap.py        | 80 +++++++++++++++++++
 .../perf_ind_iqpvar_LR20_jax.py               | 80 +++++++++++++++++++
 .../perf_ind_iqpvar_LR20_jax_novmap.py        | 80 +++++++++++++++++++
 18 files changed, 1049 insertions(+), 3 deletions(-)
 create mode 100644 nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_14d_scontrol_JAX.csv
 create mode 100644 nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_14d_scontrol_JAX_novmap.csv
 create mode 100644 nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_2d_performance_indicators_JAX.csv
 create mode 100644 nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_2d_performance_indicators_JAX_novmap.csv
 create mode 100644 nersc/performance_indicators/perf_ind_iqpvar_LR14_jax_novmap.py
 create mode 100644 nersc/performance_indicators/perf_ind_iqpvar_LR15_jax.py
 create mode 100644 nersc/performance_indicators/perf_ind_iqpvar_LR15_jax_novmap.py
 create mode 100644 nersc/performance_indicators/perf_ind_iqpvar_LR16_jax.py
 create mode 100644 nersc/performance_indicators/perf_ind_iqpvar_LR16_jax_novmap.py
 create mode 100644 nersc/performance_indicators/perf_ind_iqpvar_LR17_jax.py
 create mode 100644 nersc/performance_indicators/perf_ind_iqpvar_LR17_jax_novmap.py
 create mode 100644 nersc/performance_indicators/perf_ind_iqpvar_LR18_jax.py
 create mode 100644 nersc/performance_indicators/perf_ind_iqpvar_LR18_jax_novmap.py
 create mode 100644 nersc/performance_indicators/perf_ind_iqpvar_LR19_jax.py
 create mode 100644 nersc/performance_indicators/perf_ind_iqpvar_LR19_jax_novmap.py
 create mode 100644 nersc/performance_indicators/perf_ind_iqpvar_LR20_jax.py
 create mode 100644 nersc/performance_indicators/perf_ind_iqpvar_LR20_jax_novmap.py

diff --git a/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_14d_scontrol_JAX.csv b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_14d_scontrol_JAX.csv
new file mode 100644
index 00000000..e69de29b
diff --git a/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_14d_scontrol_JAX_novmap.csv b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_14d_scontrol_JAX_novmap.csv
new file mode 100644
index 00000000..e69de29b
diff --git a/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_2d_performance_indicators_JAX.csv b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_2d_performance_indicators_JAX.csv
new file mode 100644
index 00000000..8f6c7c7c
--- /dev/null
+++ b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_2d_performance_indicators_JAX.csv
@@ -0,0 +1,2 @@
+first_train_step,first_train_step_std,consec_train_step,consec_train_step_std,predict_time,predict_time_std,hyperparameters
+3.2171609163284303,0.10161538880525899,0.0010687878637602836,3.437768200541456e-05,1.3052996158599854,0.05948661322433076,"{'learning_rate': 0.01, 'n_layers': 15, 'repeats': 10, 'use_jax': True, 'vmap': True, 'max_steps': 100}"
diff --git a/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_2d_performance_indicators_JAX_novmap.csv b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_2d_performance_indicators_JAX_novmap.csv
new file mode 100644
index 00000000..dd23b67c
--- /dev/null
+++ b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_2d_performance_indicators_JAX_novmap.csv
@@ -0,0 +1,2 @@
+first_train_step,first_train_step_std,consec_train_step,consec_train_step_std,predict_time,predict_time_std,hyperparameters
+2.8686980962753297,0.05628433893041106,0.006390488749802715,0.0034330665556114876,0.5197975158691406,0.022069401914822987,"{'learning_rate': 0.01, 'n_layers': 15, 'repeats': 10, 'use_jax': True, 'vmap': True, 'max_vmap': 1, 'max_steps': 100}"
diff --git a/nersc/performance_indicators/perf_ind_iqpvar_LR14_jax.py b/nersc/performance_indicators/perf_ind_iqpvar_LR14_jax.py
index 2ba55366..8efa2976 100644
--- a/nersc/performance_indicators/perf_ind_iqpvar_LR14_jax.py
+++ b/nersc/performance_indicators/perf_ind_iqpvar_LR14_jax.py
@@ -18,7 +18,7 @@
 
 print(hyperparams)
 
-n_features = 2 #dataset dimension
+n_features = 14 #dataset dimension
 n_trials = 10 #number of trials to average over
 
 X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
@@ -61,11 +61,13 @@
 predict_time_std = np.std(predict_times)
 
 #write to csv
-data = [first_train_step, first_train_step_std, consec_train_step, consec_train_step_std, predict_time, predict_time_std]
+data = [first_train_step, first_train_step_std, consec_train_step, consec_train_step_std, predict_time,
+        predict_time_std, hyperparams]
 
 model_name = model.__class__.__name__
 filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_JAX.csv"
-header = ['first_train_step', 'first_train_step_std', 'consec_train_step', 'consec_train_step_std', 'predict_time', 'predict_time_std']
+header = ['first_train_step', 'first_train_step_std', 'consec_train_step', 'consec_train_step_std', 'predict_time',
+          'predict_time_std', 'hyperparameters']
 
 if not os.path.exists('JAX'):
     # Create the directory
diff --git a/nersc/performance_indicators/perf_ind_iqpvar_LR14_jax_novmap.py b/nersc/performance_indicators/perf_ind_iqpvar_LR14_jax_novmap.py
new file mode 100644
index 00000000..9bc9dd24
--- /dev/null
+++ b/nersc/performance_indicators/perf_ind_iqpvar_LR14_jax_novmap.py
@@ -0,0 +1,80 @@
+import qml_benchmarks
+import pennylane as qml
+import jax
+import jax.numpy as jnp
+import numpy as np
+import time
+import csv
+import os
+import yaml
+
+from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier
+from qml_benchmarks.hyperparam_search_utils import read_data
+
+with open('hyperparam_settings.yaml', "r") as file:
+    hp_settings = yaml.safe_load(file)
+
+hyperparams = {**hp_settings['IQPVariationalClassifier'], **{'use_jax':True, 'vmap':True, 'max_vmap':1, 'max_steps':100}}
+
+print(hyperparams)
+
+n_features = 14 #dataset dimension
+n_trials = 10 #number of trials to average over
+
+X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
+X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
+
+first_train_steps = []
+av_consec_train_steps = []
+predict_times = []
+
+for trial in range(n_trials):
+    jax.clear_caches()
+
+    model = IQPVariationalClassifier(**hyperparams)
+    model.fit(X_train, y_train)
+
+    #get step times from loss history data
+    step_times = np.array([model.loss_history_[1][i + 1] - model.loss_history_[1][i]
+                           for i in range(len(model.loss_history_[1]) - 1)])
+    step_times = np.insert(step_times, 0, [model.loss_history_[1][0]])
+
+    #first train step
+    first_train_steps.append(step_times[0])
+    #consecutive (average) train step
+    av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
+    #prediction time
+    time0 = time.time()
+    model.predict(X_test)
+    predict_times.append(time.time() - time0)
+
+
+#calculate mean and stds
+
+first_train_step = np.mean(first_train_steps)
+first_train_step_std = np.std(first_train_steps)
+
+consec_train_step = np.mean(av_consec_train_steps)
+consec_train_step_std = np.std(av_consec_train_steps)
+
+predict_time = np.mean(predict_times)
+predict_time_std = np.std(predict_times)
+
+#write to csv
+data = [first_train_step, first_train_step_std, consec_train_step, consec_train_step_std, predict_time,
+        predict_time_std, hyperparams]
+
+model_name = model.__class__.__name__
+filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_JAX_novmap.csv"
+header = ['first_train_step', 'first_train_step_std', 'consec_train_step', 'consec_train_step_std', 'predict_time',
+          'predict_time_std', 'hyperparameters']
+
+if not os.path.exists('JAX'):
+    # Create the directory
+    os.mkdir('JAX')
+
+with open('JAX/'+filename, mode="w", newline="") as file:
+    writer = csv.writer(file)
+    writer.writerow(header)
+    for row in [data]:
+        writer.writerow(row)
diff --git a/nersc/performance_indicators/perf_ind_iqpvar_LR15_jax.py b/nersc/performance_indicators/perf_ind_iqpvar_LR15_jax.py
new file mode 100644
index 00000000..20436274
--- /dev/null
+++ b/nersc/performance_indicators/perf_ind_iqpvar_LR15_jax.py
@@ -0,0 +1,80 @@
+import qml_benchmarks
+import pennylane as qml
+import jax
+import jax.numpy as jnp
+import numpy as np
+import time
+import csv
+import os
+import yaml
+
+from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier
+from qml_benchmarks.hyperparam_search_utils import read_data
+
+with open('hyperparam_settings.yaml', "r") as file:
+    hp_settings = yaml.safe_load(file)
+
+hyperparams = {**hp_settings['IQPVariationalClassifier'], **{'use_jax':True, 'vmap':True, 'max_steps':100}}
+
+print(hyperparams)
+
+n_features = 15 #dataset dimension
+n_trials = 10 #number of trials to average over
+
+X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
+X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
+
+first_train_steps = []
+av_consec_train_steps = []
+predict_times = []
+
+for trial in range(n_trials):
+    jax.clear_caches()
+
+    model = IQPVariationalClassifier(**hyperparams)
+    model.fit(X_train, y_train)
+
+    #get step times from loss history data
+    step_times = np.array([model.loss_history_[1][i + 1] - model.loss_history_[1][i]
+                           for i in range(len(model.loss_history_[1]) - 1)])
+    step_times = np.insert(step_times, 0, [model.loss_history_[1][0]])
+
+    #first train step
+    first_train_steps.append(step_times[0])
+    #consecutive (average) train step
+    av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
+    #prediction time
+    time0 = time.time()
+    model.predict(X_test)
+    predict_times.append(time.time() - time0)
+
+
+#calculate mean and stds
+
+first_train_step = np.mean(first_train_steps)
+first_train_step_std = np.std(first_train_steps)
+
+consec_train_step = np.mean(av_consec_train_steps)
+consec_train_step_std = np.std(av_consec_train_steps)
+
+predict_time = np.mean(predict_times)
+predict_time_std = np.std(predict_times)
+
+#write to csv
+data = [first_train_step, first_train_step_std, consec_train_step, consec_train_step_std, predict_time,
+        predict_time_std, hyperparams]
+
+model_name = model.__class__.__name__
+filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_JAX.csv"
+header = ['first_train_step', 'first_train_step_std', 'consec_train_step', 'consec_train_step_std', 'predict_time',
+          'predict_time_std', 'hyperparameters']
+
+if not os.path.exists('JAX'):
+    # Create the directory
+    os.mkdir('JAX')
+
+with open('JAX/'+filename, mode="w", newline="") as file:
+    writer = csv.writer(file)
+    writer.writerow(header)
+    for row in [data]:
+        writer.writerow(row)
diff --git a/nersc/performance_indicators/perf_ind_iqpvar_LR15_jax_novmap.py b/nersc/performance_indicators/perf_ind_iqpvar_LR15_jax_novmap.py
new file mode 100644
index 00000000..0c32497a
--- /dev/null
+++ b/nersc/performance_indicators/perf_ind_iqpvar_LR15_jax_novmap.py
@@ -0,0 +1,80 @@
+import qml_benchmarks
+import pennylane as qml
+import jax
+import jax.numpy as jnp
+import numpy as np
+import time
+import csv
+import os
+import yaml
+
+from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier
+from qml_benchmarks.hyperparam_search_utils import read_data
+
+with open('hyperparam_settings.yaml', "r") as file:
+    hp_settings = yaml.safe_load(file)
+
+hyperparams = {**hp_settings['IQPVariationalClassifier'], **{'use_jax':True, 'vmap':True, 'max_vmap':1, 'max_steps':100}}
+
+print(hyperparams)
+
+n_features = 15 #dataset dimension
+n_trials = 10 #number of trials to average over
+
+X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
+X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
+
+first_train_steps = []
+av_consec_train_steps = []
+predict_times = []
+
+for trial in range(n_trials):
+    jax.clear_caches()
+
+    model = IQPVariationalClassifier(**hyperparams)
+    model.fit(X_train, y_train)
+
+    #get step times from loss history data
+    step_times = np.array([model.loss_history_[1][i + 1] - model.loss_history_[1][i]
+                           for i in range(len(model.loss_history_[1]) - 1)])
+    step_times = np.insert(step_times, 0, [model.loss_history_[1][0]])
+
+    #first train step
+    first_train_steps.append(step_times[0])
+    #consecutive (average) train step
+    av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
+    #prediction time
+    time0 = time.time()
+    model.predict(X_test)
+    predict_times.append(time.time() - time0)
+
+
+#calculate mean and stds
+
+first_train_step = np.mean(first_train_steps)
+first_train_step_std = np.std(first_train_steps)
+
+consec_train_step = np.mean(av_consec_train_steps)
+consec_train_step_std = np.std(av_consec_train_steps)
+
+predict_time = np.mean(predict_times)
+predict_time_std = np.std(predict_times)
+
+#write to csv
+data = [first_train_step, first_train_step_std, consec_train_step, consec_train_step_std, predict_time,
+        predict_time_std, hyperparams]
+
+model_name = model.__class__.__name__
+filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_JAX_novmap.csv"
+header = ['first_train_step', 'first_train_step_std', 'consec_train_step', 'consec_train_step_std', 'predict_time',
+          'predict_time_std', 'hyperparameters']
+
+if not os.path.exists('JAX'):
+    # Create the directory
+    os.mkdir('JAX')
+
+with open('JAX/'+filename, mode="w", newline="") as file:
+    writer = csv.writer(file)
+    writer.writerow(header)
+    for row in [data]:
+        writer.writerow(row)
diff --git a/nersc/performance_indicators/perf_ind_iqpvar_LR16_jax.py b/nersc/performance_indicators/perf_ind_iqpvar_LR16_jax.py
new file mode 100644
index 00000000..190c7fd2
--- /dev/null
+++ b/nersc/performance_indicators/perf_ind_iqpvar_LR16_jax.py
@@ -0,0 +1,80 @@
+import qml_benchmarks
+import pennylane as qml
+import jax
+import jax.numpy as jnp
+import numpy as np
+import time
+import csv
+import os
+import yaml
+
+from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier
+from qml_benchmarks.hyperparam_search_utils import read_data
+
+with open('hyperparam_settings.yaml', "r") as file:
+    hp_settings = yaml.safe_load(file)
+
+hyperparams = {**hp_settings['IQPVariationalClassifier'], **{'use_jax':True, 'vmap':True, 'max_steps':100}}
+
+print(hyperparams)
+
+n_features = 16 #dataset dimension
+n_trials = 10 #number of trials to average over
+
+X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
+X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
+
+first_train_steps = []
+av_consec_train_steps = []
+predict_times = []
+
+for trial in range(n_trials):
+    jax.clear_caches()
+
+    model = IQPVariationalClassifier(**hyperparams)
+    model.fit(X_train, y_train)
+
+    #get step times from loss history data
+    step_times = np.array([model.loss_history_[1][i + 1] - model.loss_history_[1][i]
+                           for i in range(len(model.loss_history_[1]) - 1)])
+    step_times = np.insert(step_times, 0, [model.loss_history_[1][0]])
+
+    #first train step
+    first_train_steps.append(step_times[0])
+    #consecutive (average) train step
+    av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
+    #prediction time
+    time0 = time.time()
+    model.predict(X_test)
+    predict_times.append(time.time() - time0)
+
+
+#calculate mean and stds
+
+first_train_step = np.mean(first_train_steps)
+first_train_step_std = np.std(first_train_steps)
+
+consec_train_step = np.mean(av_consec_train_steps)
+consec_train_step_std = np.std(av_consec_train_steps)
+
+predict_time = np.mean(predict_times)
+predict_time_std = np.std(predict_times)
+
+#write to csv
+data = [first_train_step, first_train_step_std, consec_train_step, consec_train_step_std, predict_time,
+        predict_time_std, hyperparams]
+
+model_name = model.__class__.__name__
+filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_JAX.csv"
+header = ['first_train_step', 'first_train_step_std', 'consec_train_step', 'consec_train_step_std', 'predict_time',
+          'predict_time_std', 'hyperparameters']
+
+if not os.path.exists('JAX'):
+    # Create the directory
+    os.mkdir('JAX')
+
+with open('JAX/'+filename, mode="w", newline="") as file:
+    writer = csv.writer(file)
+    writer.writerow(header)
+    for row in [data]:
+        writer.writerow(row)
diff --git a/nersc/performance_indicators/perf_ind_iqpvar_LR16_jax_novmap.py b/nersc/performance_indicators/perf_ind_iqpvar_LR16_jax_novmap.py
new file mode 100644
index 00000000..427123a6
--- /dev/null
+++ b/nersc/performance_indicators/perf_ind_iqpvar_LR16_jax_novmap.py
@@ -0,0 +1,80 @@
+import qml_benchmarks
+import pennylane as qml
+import jax
+import jax.numpy as jnp
+import numpy as np
+import time
+import csv
+import os
+import yaml
+
+from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier
+from qml_benchmarks.hyperparam_search_utils import read_data
+
+with open('hyperparam_settings.yaml', "r") as file:
+    hp_settings = yaml.safe_load(file)
+
+hyperparams = {**hp_settings['IQPVariationalClassifier'], **{'use_jax':True, 'vmap':True, 'max_vmap':1, 'max_steps':100}}
+
+print(hyperparams)
+
+n_features = 16 #dataset dimension
+n_trials = 10 #number of trials to average over
+
+X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
+X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
+
+first_train_steps = []
+av_consec_train_steps = []
+predict_times = []
+
+for trial in range(n_trials):
+    jax.clear_caches()
+
+    model = IQPVariationalClassifier(**hyperparams)
+    model.fit(X_train, y_train)
+
+    #get step times from loss history data
+    step_times = np.array([model.loss_history_[1][i + 1] - model.loss_history_[1][i]
+                           for i in range(len(model.loss_history_[1]) - 1)])
+    step_times = np.insert(step_times, 0, [model.loss_history_[1][0]])
+
+    #first train step
+    first_train_steps.append(step_times[0])
+    #consecutive (average) train step
+    av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
+    #prediction time
+    time0 = time.time()
+    model.predict(X_test)
+    predict_times.append(time.time() - time0)
+
+
+#calculate mean and stds
+
+first_train_step = np.mean(first_train_steps)
+first_train_step_std = np.std(first_train_steps)
+
+consec_train_step = np.mean(av_consec_train_steps)
+consec_train_step_std = np.std(av_consec_train_steps)
+
+predict_time = np.mean(predict_times)
+predict_time_std = np.std(predict_times)
+
+#write to csv
+data = [first_train_step, first_train_step_std, consec_train_step, consec_train_step_std, predict_time,
+        predict_time_std, hyperparams]
+
+model_name = model.__class__.__name__
+filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_JAX_novmap.csv"
+header = ['first_train_step', 'first_train_step_std', 'consec_train_step', 'consec_train_step_std', 'predict_time',
+          'predict_time_std', 'hyperparameters']
+
+if not os.path.exists('JAX'):
+    # Create the directory
+    os.mkdir('JAX')
+
+with open('JAX/'+filename, mode="w", newline="") as file:
+    writer = csv.writer(file)
+    writer.writerow(header)
+    for row in [data]:
+        writer.writerow(row)
diff --git a/nersc/performance_indicators/perf_ind_iqpvar_LR17_jax.py b/nersc/performance_indicators/perf_ind_iqpvar_LR17_jax.py
new file mode 100644
index 00000000..910e1edf
--- /dev/null
+++ b/nersc/performance_indicators/perf_ind_iqpvar_LR17_jax.py
@@ -0,0 +1,80 @@
+import qml_benchmarks
+import pennylane as qml
+import jax
+import jax.numpy as jnp
+import numpy as np
+import time
+import csv
+import os
+import yaml
+
+from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier
+from qml_benchmarks.hyperparam_search_utils import read_data
+
+with open('hyperparam_settings.yaml', "r") as file:
+    hp_settings = yaml.safe_load(file)
+
+hyperparams = {**hp_settings['IQPVariationalClassifier'], **{'use_jax':True, 'vmap':True, 'max_steps':100}}
+
+print(hyperparams)
+
+n_features = 17 #dataset dimension
+n_trials = 10 #number of trials to average over
+
+X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
+X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
+
+first_train_steps = []
+av_consec_train_steps = []
+predict_times = []
+
+for trial in range(n_trials):
+    jax.clear_caches()
+
+    model = IQPVariationalClassifier(**hyperparams)
+    model.fit(X_train, y_train)
+
+    #get step times from loss history data
+    step_times = np.array([model.loss_history_[1][i + 1] - model.loss_history_[1][i]
+                           for i in range(len(model.loss_history_[1]) - 1)])
+    step_times = np.insert(step_times, 0, [model.loss_history_[1][0]])
+
+    #first train step
+    first_train_steps.append(step_times[0])
+    #consecutive (average) train step
+    av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
+    #prediction time
+    time0 = time.time()
+    model.predict(X_test)
+    predict_times.append(time.time() - time0)
+
+
+#calculate mean and stds
+
+first_train_step = np.mean(first_train_steps)
+first_train_step_std = np.std(first_train_steps)
+
+consec_train_step = np.mean(av_consec_train_steps)
+consec_train_step_std = np.std(av_consec_train_steps)
+
+predict_time = np.mean(predict_times)
+predict_time_std = np.std(predict_times)
+
+#write to csv
+data = [first_train_step, first_train_step_std, consec_train_step, consec_train_step_std, predict_time,
+        predict_time_std, hyperparams]
+
+model_name = model.__class__.__name__
+filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_JAX.csv"
+header = ['first_train_step', 'first_train_step_std', 'consec_train_step', 'consec_train_step_std', 'predict_time',
+          'predict_time_std', 'hyperparameters']
+
+if not os.path.exists('JAX'):
+    # Create the directory
+    os.mkdir('JAX')
+
+with open('JAX/'+filename, mode="w", newline="") as file:
+    writer = csv.writer(file)
+    writer.writerow(header)
+    for row in [data]:
+        writer.writerow(row)
diff --git a/nersc/performance_indicators/perf_ind_iqpvar_LR17_jax_novmap.py b/nersc/performance_indicators/perf_ind_iqpvar_LR17_jax_novmap.py
new file mode 100644
index 00000000..8e07e9b2
--- /dev/null
+++ b/nersc/performance_indicators/perf_ind_iqpvar_LR17_jax_novmap.py
@@ -0,0 +1,80 @@
+import qml_benchmarks
+import pennylane as qml
+import jax
+import jax.numpy as jnp
+import numpy as np
+import time
+import csv
+import os
+import yaml
+
+from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier
+from qml_benchmarks.hyperparam_search_utils import read_data
+
+with open('hyperparam_settings.yaml', "r") as file:
+    hp_settings = yaml.safe_load(file)
+
+hyperparams = {**hp_settings['IQPVariationalClassifier'], **{'use_jax':True, 'vmap':True, 'max_vmap':1, 'max_steps':100}}
+
+print(hyperparams)
+
+n_features = 17 #dataset dimension
+n_trials = 10 #number of trials to average over
+
+X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
+X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
+
+first_train_steps = []
+av_consec_train_steps = []
+predict_times = []
+
+for trial in range(n_trials):
+    jax.clear_caches()
+
+    model = IQPVariationalClassifier(**hyperparams)
+    model.fit(X_train, y_train)
+
+    #get step times from loss history data
+    step_times = np.array([model.loss_history_[1][i + 1] - model.loss_history_[1][i]
+                           for i in range(len(model.loss_history_[1]) - 1)])
+    step_times = np.insert(step_times, 0, [model.loss_history_[1][0]])
+
+    #first train step
+    first_train_steps.append(step_times[0])
+    #consecutive (average) train step
+    av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
+    #prediction time
+    time0 = time.time()
+    model.predict(X_test)
+    predict_times.append(time.time() - time0)
+
+
+#calculate mean and stds
+
+first_train_step = np.mean(first_train_steps)
+first_train_step_std = np.std(first_train_steps)
+
+consec_train_step = np.mean(av_consec_train_steps)
+consec_train_step_std = np.std(av_consec_train_steps)
+
+predict_time = np.mean(predict_times)
+predict_time_std = np.std(predict_times)
+
+#write to csv
+data = [first_train_step, first_train_step_std, consec_train_step, consec_train_step_std, predict_time,
+        predict_time_std, hyperparams]
+
+model_name = model.__class__.__name__
+filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_JAX_novmap.csv"
+header = ['first_train_step', 'first_train_step_std', 'consec_train_step', 'consec_train_step_std', 'predict_time',
+          'predict_time_std', 'hyperparameters']
+
+if not os.path.exists('JAX'):
+    # Create the directory
+    os.mkdir('JAX')
+
+with open('JAX/'+filename, mode="w", newline="") as file:
+    writer = csv.writer(file)
+    writer.writerow(header)
+    for row in [data]:
+        writer.writerow(row)
diff --git a/nersc/performance_indicators/perf_ind_iqpvar_LR18_jax.py b/nersc/performance_indicators/perf_ind_iqpvar_LR18_jax.py
new file mode 100644
index 00000000..569c8c70
--- /dev/null
+++ b/nersc/performance_indicators/perf_ind_iqpvar_LR18_jax.py
@@ -0,0 +1,80 @@
+import qml_benchmarks
+import pennylane as qml
+import jax
+import jax.numpy as jnp
+import numpy as np
+import time
+import csv
+import os
+import yaml
+
+from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier
+from qml_benchmarks.hyperparam_search_utils import read_data
+
+with open('hyperparam_settings.yaml', "r") as file:
+    hp_settings = yaml.safe_load(file)
+
+hyperparams = {**hp_settings['IQPVariationalClassifier'], **{'use_jax':True, 'vmap':True, 'max_steps':100}}
+
+print(hyperparams)
+
+n_features = 18 #dataset dimension
+n_trials = 10 #number of trials to average over
+
+X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
+X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
+
+first_train_steps = []
+av_consec_train_steps = []
+predict_times = []
+
+for trial in range(n_trials):
+    jax.clear_caches()
+
+    model = IQPVariationalClassifier(**hyperparams)
+    model.fit(X_train, y_train)
+
+    #get step times from loss history data
+    step_times = np.array([model.loss_history_[1][i + 1] - model.loss_history_[1][i]
+                           for i in range(len(model.loss_history_[1]) - 1)])
+    step_times = np.insert(step_times, 0, [model.loss_history_[1][0]])
+
+    #first train step
+    first_train_steps.append(step_times[0])
+    #consecutive (average) train step
+    av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
+    #prediction time
+    time0 = time.time()
+    model.predict(X_test)
+    predict_times.append(time.time() - time0)
+
+
+#calculate mean and stds
+
+first_train_step = np.mean(first_train_steps)
+first_train_step_std = np.std(first_train_steps)
+
+consec_train_step = np.mean(av_consec_train_steps)
+consec_train_step_std = np.std(av_consec_train_steps)
+
+predict_time = np.mean(predict_times)
+predict_time_std = np.std(predict_times)
+
+#write to csv
+data = [first_train_step, first_train_step_std, consec_train_step, consec_train_step_std, predict_time,
+        predict_time_std, hyperparams]
+
+model_name = model.__class__.__name__
+filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_JAX.csv"
+header = ['first_train_step', 'first_train_step_std', 'consec_train_step', 'consec_train_step_std', 'predict_time',
+          'predict_time_std', 'hyperparameters']
+
+if not os.path.exists('JAX'):
+    # Create the directory
+    os.mkdir('JAX')
+
+with open('JAX/'+filename, mode="w", newline="") as file:
+    writer = csv.writer(file)
+    writer.writerow(header)
+    for row in [data]:
+        writer.writerow(row)
diff --git a/nersc/performance_indicators/perf_ind_iqpvar_LR18_jax_novmap.py b/nersc/performance_indicators/perf_ind_iqpvar_LR18_jax_novmap.py
new file mode 100644
index 00000000..8d1efb3a
--- /dev/null
+++ b/nersc/performance_indicators/perf_ind_iqpvar_LR18_jax_novmap.py
@@ -0,0 +1,80 @@
+import qml_benchmarks
+import pennylane as qml
+import jax
+import jax.numpy as jnp
+import numpy as np
+import time
+import csv
+import os
+import yaml
+
+from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier
+from qml_benchmarks.hyperparam_search_utils import read_data
+
+with open('hyperparam_settings.yaml', "r") as file:
+    hp_settings = yaml.safe_load(file)
+
+hyperparams = {**hp_settings['IQPVariationalClassifier'], **{'use_jax':True, 'vmap':True, 'max_vmap':1, 'max_steps':100}}
+
+print(hyperparams)
+
+n_features = 18 #dataset dimension
+n_trials = 10 #number of trials to average over
+
+X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
+X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
+
+first_train_steps = []
+av_consec_train_steps = []
+predict_times = []
+
+for trial in range(n_trials):
+    jax.clear_caches()
+
+    model = IQPVariationalClassifier(**hyperparams)
+    model.fit(X_train, y_train)
+
+    #get step times from loss history data
+    step_times = np.array([model.loss_history_[1][i + 1] - model.loss_history_[1][i]
+                           for i in range(len(model.loss_history_[1]) - 1)])
+    step_times = np.insert(step_times, 0, [model.loss_history_[1][0]])
+
+    #first train step
+    first_train_steps.append(step_times[0])
+    #consecutive (average) train step
+    av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
+    #prediction time
+    time0 = time.time()
+    model.predict(X_test)
+    predict_times.append(time.time() - time0)
+
+
+#calculate mean and stds
+
+first_train_step = np.mean(first_train_steps)
+first_train_step_std = np.std(first_train_steps)
+
+consec_train_step = np.mean(av_consec_train_steps)
+consec_train_step_std = np.std(av_consec_train_steps)
+
+predict_time = np.mean(predict_times)
+predict_time_std = np.std(predict_times)
+
+#write to csv
+data = [first_train_step, first_train_step_std, consec_train_step, consec_train_step_std, predict_time,
+        predict_time_std, hyperparams]
+
+model_name = model.__class__.__name__
+filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_JAX_novmap.csv"
+header = ['first_train_step', 'first_train_step_std', 'consec_train_step', 'consec_train_step_std', 'predict_time',
+          'predict_time_std', 'hyperparameters']
+
+if not os.path.exists('JAX'):
+    # Create the directory
+    os.mkdir('JAX')
+
+with open('JAX/'+filename, mode="w", newline="") as file:
+    writer = csv.writer(file)
+    writer.writerow(header)
+    for row in [data]:
+        writer.writerow(row)
diff --git a/nersc/performance_indicators/perf_ind_iqpvar_LR19_jax.py b/nersc/performance_indicators/perf_ind_iqpvar_LR19_jax.py
new file mode 100644
index 00000000..d6798dea
--- /dev/null
+++ b/nersc/performance_indicators/perf_ind_iqpvar_LR19_jax.py
@@ -0,0 +1,80 @@
+import qml_benchmarks
+import pennylane as qml
+import jax
+import jax.numpy as jnp
+import numpy as np
+import time
+import csv
+import os
+import yaml
+
+from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier
+from qml_benchmarks.hyperparam_search_utils import read_data
+
+with open('hyperparam_settings.yaml', "r") as file:
+    hp_settings = yaml.safe_load(file)
+
+hyperparams = {**hp_settings['IQPVariationalClassifier'], **{'use_jax':True, 'vmap':True, 'max_steps':100}}
+
+print(hyperparams)
+
+n_features = 19 #dataset dimension
+n_trials = 10 #number of trials to average over
+
+X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
+X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
+
+first_train_steps = []
+av_consec_train_steps = []
+predict_times = []
+
+for trial in range(n_trials):
+    jax.clear_caches()
+
+    model = IQPVariationalClassifier(**hyperparams)
+    model.fit(X_train, y_train)
+
+    #get step times from loss history data
+    step_times = np.array([model.loss_history_[1][i + 1] - model.loss_history_[1][i]
+                           for i in range(len(model.loss_history_[1]) - 1)])
+    step_times = np.insert(step_times, 0, [model.loss_history_[1][0]])
+
+    #first train step
+    first_train_steps.append(step_times[0])
+    #consecutive (average) train step
+    av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
+    #prediction time
+    time0 = time.time()
+    model.predict(X_test)
+    predict_times.append(time.time() - time0)
+
+
+#calculate mean and stds
+
+first_train_step = np.mean(first_train_steps)
+first_train_step_std = np.std(first_train_steps)
+
+consec_train_step = np.mean(av_consec_train_steps)
+consec_train_step_std = np.std(av_consec_train_steps)
+
+predict_time = np.mean(predict_times)
+predict_time_std = np.std(predict_times)
+
+#write to csv
+data = [first_train_step, first_train_step_std, consec_train_step, consec_train_step_std, predict_time,
+        predict_time_std, hyperparams]
+
+model_name = model.__class__.__name__
+filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_JAX.csv"
+header = ['first_train_step', 'first_train_step_std', 'consec_train_step', 'consec_train_step_std', 'predict_time',
+          'predict_time_std', 'hyperparameters']
+
+if not os.path.exists('JAX'):
+    # Create the directory
+    os.mkdir('JAX')
+
+with open('JAX/'+filename, mode="w", newline="") as file:
+    writer = csv.writer(file)
+    writer.writerow(header)
+    for row in [data]:
+        writer.writerow(row)
diff --git a/nersc/performance_indicators/perf_ind_iqpvar_LR19_jax_novmap.py b/nersc/performance_indicators/perf_ind_iqpvar_LR19_jax_novmap.py
new file mode 100644
index 00000000..ccc8b142
--- /dev/null
+++ b/nersc/performance_indicators/perf_ind_iqpvar_LR19_jax_novmap.py
@@ -0,0 +1,80 @@
+import qml_benchmarks
+import pennylane as qml
+import jax
+import jax.numpy as jnp
+import numpy as np
+import time
+import csv
+import os
+import yaml
+
+from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier
+from qml_benchmarks.hyperparam_search_utils import read_data
+
+with open('hyperparam_settings.yaml', "r") as file:
+    hp_settings = yaml.safe_load(file)
+
+hyperparams = {**hp_settings['IQPVariationalClassifier'], **{'use_jax':True, 'vmap':True, 'max_vmap':1, 'max_steps':100}}
+
+print(hyperparams)
+
+n_features = 19 #dataset dimension
+n_trials = 10 #number of trials to average over
+
+X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
+X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
+
+first_train_steps = []
+av_consec_train_steps = []
+predict_times = []
+
+for trial in range(n_trials):
+    jax.clear_caches()
+
+    model = IQPVariationalClassifier(**hyperparams)
+    model.fit(X_train, y_train)
+
+    #get step times from loss history data
+    step_times = np.array([model.loss_history_[1][i + 1] - model.loss_history_[1][i]
+                           for i in range(len(model.loss_history_[1]) - 1)])
+    step_times = np.insert(step_times, 0, [model.loss_history_[1][0]])
+
+    #first train step
+    first_train_steps.append(step_times[0])
+    #consecutive (average) train step
+    av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
+    #prediction time
+    time0 = time.time()
+    model.predict(X_test)
+    predict_times.append(time.time() - time0)
+
+
+#calculate mean and stds
+
+first_train_step = np.mean(first_train_steps)
+first_train_step_std = np.std(first_train_steps)
+
+consec_train_step = np.mean(av_consec_train_steps)
+consec_train_step_std = np.std(av_consec_train_steps)
+
+predict_time = np.mean(predict_times)
+predict_time_std = np.std(predict_times)
+
+#write to csv
+data = [first_train_step, first_train_step_std, consec_train_step, consec_train_step_std, predict_time,
+        predict_time_std, hyperparams]
+
+model_name = model.__class__.__name__
+filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_JAX_novmap.csv"
+header = ['first_train_step', 'first_train_step_std', 'consec_train_step', 'consec_train_step_std', 'predict_time',
+          'predict_time_std', 'hyperparameters']
+
+if not os.path.exists('JAX'):
+    # Create the directory
+    os.mkdir('JAX')
+
+with open('JAX/'+filename, mode="w", newline="") as file:
+    writer = csv.writer(file)
+    writer.writerow(header)
+    for row in [data]:
+        writer.writerow(row)
diff --git a/nersc/performance_indicators/perf_ind_iqpvar_LR20_jax.py b/nersc/performance_indicators/perf_ind_iqpvar_LR20_jax.py
new file mode 100644
index 00000000..f21f8efe
--- /dev/null
+++ b/nersc/performance_indicators/perf_ind_iqpvar_LR20_jax.py
@@ -0,0 +1,80 @@
+import qml_benchmarks
+import pennylane as qml
+import jax
+import jax.numpy as jnp
+import numpy as np
+import time
+import csv
+import os
+import yaml
+
+from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier
+from qml_benchmarks.hyperparam_search_utils import read_data
+
+with open('hyperparam_settings.yaml', "r") as file:
+    hp_settings = yaml.safe_load(file)
+
+hyperparams = {**hp_settings['IQPVariationalClassifier'], **{'use_jax':True, 'vmap':True, 'max_steps':100}}
+
+print(hyperparams)
+
+n_features = 20 #dataset dimension
+n_trials = 10 #number of trials to average over
+
+X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
+X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
+
+first_train_steps = []
+av_consec_train_steps = []
+predict_times = []
+
+for trial in range(n_trials):
+    jax.clear_caches()
+
+    model = IQPVariationalClassifier(**hyperparams)
+    model.fit(X_train, y_train)
+
+    #get step times from loss history data
+    step_times = np.array([model.loss_history_[1][i + 1] - model.loss_history_[1][i]
+                           for i in range(len(model.loss_history_[1]) - 1)])
+    step_times = np.insert(step_times, 0, [model.loss_history_[1][0]])
+
+    #first train step
+    first_train_steps.append(step_times[0])
+    #consecutive (average) train step
+    av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
+    #prediction time
+    time0 = time.time()
+    model.predict(X_test)
+    predict_times.append(time.time() - time0)
+
+
+#calculate mean and stds
+
+first_train_step = np.mean(first_train_steps)
+first_train_step_std = np.std(first_train_steps)
+
+consec_train_step = np.mean(av_consec_train_steps)
+consec_train_step_std = np.std(av_consec_train_steps)
+
+predict_time = np.mean(predict_times)
+predict_time_std = np.std(predict_times)
+
+#write to csv
+data = [first_train_step, first_train_step_std, consec_train_step, consec_train_step_std, predict_time,
+        predict_time_std, hyperparams]
+
+model_name = model.__class__.__name__
+filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_JAX.csv"
+header = ['first_train_step', 'first_train_step_std', 'consec_train_step', 'consec_train_step_std', 'predict_time',
+          'predict_time_std', 'hyperparameters']
+
+if not os.path.exists('JAX'):
+    # Create the directory
+    os.mkdir('JAX')
+
+with open('JAX/'+filename, mode="w", newline="") as file:
+    writer = csv.writer(file)
+    writer.writerow(header)
+    for row in [data]:
+        writer.writerow(row)
diff --git a/nersc/performance_indicators/perf_ind_iqpvar_LR20_jax_novmap.py b/nersc/performance_indicators/perf_ind_iqpvar_LR20_jax_novmap.py
new file mode 100644
index 00000000..b0edd3f8
--- /dev/null
+++ b/nersc/performance_indicators/perf_ind_iqpvar_LR20_jax_novmap.py
@@ -0,0 +1,80 @@
+import qml_benchmarks
+import pennylane as qml
+import jax
+import jax.numpy as jnp
+import numpy as np
+import time
+import csv
+import os
+import yaml
+
+from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier
+from qml_benchmarks.hyperparam_search_utils import read_data
+
+with open('hyperparam_settings.yaml', "r") as file:
+    hp_settings = yaml.safe_load(file)
+
+hyperparams = {**hp_settings['IQPVariationalClassifier'], **{'use_jax':True, 'vmap':True, 'max_vmap':1, 'max_steps':100}}
+
+print(hyperparams)
+
+n_features = 20 #dataset dimension
+n_trials = 10 #number of trials to average over
+
+X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
+X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
+
+first_train_steps = []
+av_consec_train_steps = []
+predict_times = []
+
+for trial in range(n_trials):
+    jax.clear_caches()
+
+    model = IQPVariationalClassifier(**hyperparams)
+    model.fit(X_train, y_train)
+
+    #get step times from loss history data
+    step_times = np.array([model.loss_history_[1][i + 1] - model.loss_history_[1][i]
+                           for i in range(len(model.loss_history_[1]) - 1)])
+    step_times = np.insert(step_times, 0, [model.loss_history_[1][0]])
+
+    #first train step
+    first_train_steps.append(step_times[0])
+    #consecutive (average) train step
+    av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
+    #prediction time
+    time0 = time.time()
+    model.predict(X_test)
+    predict_times.append(time.time() - time0)
+
+
+#calculate mean and stds
+
+first_train_step = np.mean(first_train_steps)
+first_train_step_std = np.std(first_train_steps)
+
+consec_train_step = np.mean(av_consec_train_steps)
+consec_train_step_std = np.std(av_consec_train_steps)
+
+predict_time = np.mean(predict_times)
+predict_time_std = np.std(predict_times)
+
+#write to csv
+data = [first_train_step, first_train_step_std, consec_train_step, consec_train_step_std, predict_time,
+        predict_time_std, hyperparams]
+
+model_name = model.__class__.__name__
+filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_JAX_novmap.csv"
+header = ['first_train_step', 'first_train_step_std', 'consec_train_step', 'consec_train_step_std', 'predict_time',
+          'predict_time_std', 'hyperparameters']
+
+if not os.path.exists('JAX'):
+    # Create the directory
+    os.mkdir('JAX')
+
+with open('JAX/'+filename, mode="w", newline="") as file:
+    writer = csv.writer(file)
+    writer.writerow(header)
+    for row in [data]:
+        writer.writerow(row)

From eec5fa5f258b6fe827ed3d4ded6fe7cfc88e8e68 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Mon, 15 Apr 2024 15:37:06 +0200
Subject: [PATCH 024/100] add installs

---
 nersc/pm_podman.source | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/nersc/pm_podman.source b/nersc/pm_podman.source
index d011e337..ae99bfc8 100644
--- a/nersc/pm_podman.source
+++ b/nersc/pm_podman.source
@@ -5,13 +5,16 @@ IMG=balewski/ubu22-pennylane:p3
 CFSH=/global/homes/j/jbowles
 
 echo launch image $IMG
-echo you are launching Podman image ...  remeber to exit
+echo you are launching Podman image ...  remember to exit
+
+pip install ../
+pip install pennylane-catalyst
+pip install scipy==1.11.4
 
 JNB_PORT=' '
 BASE_DIR=/qml-benchmarks   # here git has home
 WORK_DIR=$BASE_DIR/nersc
 
-
 podman-hpc run -it \
    --volume $HOME:/home \
    --volume $CFSH/$BASE_DIR:/$BASE_DIR \

From c18820f3df37fa44e2d90080928b19ddaa42ebf7 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Mon, 15 Apr 2024 15:38:35 +0200
Subject: [PATCH 025/100] add installs

---
 nersc/installs.sh      | 5 +++++
 nersc/pm_podman.source | 4 ----
 2 files changed, 5 insertions(+), 4 deletions(-)
 create mode 100644 nersc/installs.sh

diff --git a/nersc/installs.sh b/nersc/installs.sh
new file mode 100644
index 00000000..ed740537
--- /dev/null
+++ b/nersc/installs.sh
@@ -0,0 +1,5 @@
+#!/bin/bash
+
+pip install ../
+pip install pennylane-catalyst
+pip install scipy==1.11.4
\ No newline at end of file
diff --git a/nersc/pm_podman.source b/nersc/pm_podman.source
index ae99bfc8..aa13f06e 100644
--- a/nersc/pm_podman.source
+++ b/nersc/pm_podman.source
@@ -7,10 +7,6 @@ CFSH=/global/homes/j/jbowles
 echo launch image $IMG
 echo you are launching Podman image ...  remember to exit
 
-pip install ../
-pip install pennylane-catalyst
-pip install scipy==1.11.4
-
 JNB_PORT=' '
 BASE_DIR=/qml-benchmarks   # here git has home
 WORK_DIR=$BASE_DIR/nersc

From 0f32a6ca422a917a1b6e86a9f8b6726e59842e73 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Mon, 15 Apr 2024 16:03:57 +0200
Subject: [PATCH 026/100] add qmetric

---
 ...er_linearly_separable_14d_scontrol_JAX.csv |  0
 ...arly_separable_14d_scontrol_JAX_novmap.csv |  0
 ...sifier_linearly_separable_15d_JAX_info.txt |  2 +
 ...sifier_linearly_separable_17d_JAX_info.txt |  2 +
 ...eparable_2d_performance_indicators_JAX.csv |  2 -
 ...e_2d_performance_indicators_JAX_novmap.csv |  2 -
 .../perf_ind_qmetric_LR14_jax.py              | 80 +++++++++++++++++++
 .../perf_ind_qmetric_LR15_jax.py              | 80 +++++++++++++++++++
 8 files changed, 164 insertions(+), 4 deletions(-)
 delete mode 100644 nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_14d_scontrol_JAX.csv
 delete mode 100644 nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_14d_scontrol_JAX_novmap.csv
 create mode 100644 nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_15d_JAX_info.txt
 create mode 100644 nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_17d_JAX_info.txt
 delete mode 100644 nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_2d_performance_indicators_JAX.csv
 delete mode 100644 nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_2d_performance_indicators_JAX_novmap.csv
 create mode 100644 nersc/performance_indicators/perf_ind_qmetric_LR14_jax.py
 create mode 100644 nersc/performance_indicators/perf_ind_qmetric_LR15_jax.py

diff --git a/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_14d_scontrol_JAX.csv b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_14d_scontrol_JAX.csv
deleted file mode 100644
index e69de29b..00000000
diff --git a/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_14d_scontrol_JAX_novmap.csv b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_14d_scontrol_JAX_novmap.csv
deleted file mode 100644
index e69de29b..00000000
diff --git a/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_15d_JAX_info.txt b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_15d_JAX_info.txt
new file mode 100644
index 00000000..74dede79
--- /dev/null
+++ b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_15d_JAX_info.txt
@@ -0,0 +1,2 @@
+salloc -q shared_interactive -C cpu -t 4:00:00 --cpus-per-task=32
+
diff --git a/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_17d_JAX_info.txt b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_17d_JAX_info.txt
new file mode 100644
index 00000000..74dede79
--- /dev/null
+++ b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_17d_JAX_info.txt
@@ -0,0 +1,2 @@
+salloc -q shared_interactive -C cpu -t 4:00:00 --cpus-per-task=32
+
diff --git a/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_2d_performance_indicators_JAX.csv b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_2d_performance_indicators_JAX.csv
deleted file mode 100644
index 8f6c7c7c..00000000
--- a/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_2d_performance_indicators_JAX.csv
+++ /dev/null
@@ -1,2 +0,0 @@
-first_train_step,first_train_step_std,consec_train_step,consec_train_step_std,predict_time,predict_time_std,hyperparameters
-3.2171609163284303,0.10161538880525899,0.0010687878637602836,3.437768200541456e-05,1.3052996158599854,0.05948661322433076,"{'learning_rate': 0.01, 'n_layers': 15, 'repeats': 10, 'use_jax': True, 'vmap': True, 'max_steps': 100}"
diff --git a/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_2d_performance_indicators_JAX_novmap.csv b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_2d_performance_indicators_JAX_novmap.csv
deleted file mode 100644
index dd23b67c..00000000
--- a/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_2d_performance_indicators_JAX_novmap.csv
+++ /dev/null
@@ -1,2 +0,0 @@
-first_train_step,first_train_step_std,consec_train_step,consec_train_step_std,predict_time,predict_time_std,hyperparameters
-2.8686980962753297,0.05628433893041106,0.006390488749802715,0.0034330665556114876,0.5197975158691406,0.022069401914822987,"{'learning_rate': 0.01, 'n_layers': 15, 'repeats': 10, 'use_jax': True, 'vmap': True, 'max_vmap': 1, 'max_steps': 100}"
diff --git a/nersc/performance_indicators/perf_ind_qmetric_LR14_jax.py b/nersc/performance_indicators/perf_ind_qmetric_LR14_jax.py
new file mode 100644
index 00000000..dec6996e
--- /dev/null
+++ b/nersc/performance_indicators/perf_ind_qmetric_LR14_jax.py
@@ -0,0 +1,80 @@
+import qml_benchmarks
+import pennylane as qml
+import jax
+import jax.numpy as jnp
+import numpy as np
+import time
+import csv
+import os
+import yaml
+
+from qml_benchmarks.models.quantum_metric_learning import QuantumMetricLearner
+from qml_benchmarks.hyperparam_search_utils import read_data
+
+with open('hyperparam_settings.yaml', "r") as file:
+    hp_settings = yaml.safe_load(file)
+
+hyperparams = {**hp_settings['IQPVariationalClassifier'], **{'use_jax':True, 'vmap':True, 'max_steps':100}}
+
+print(hyperparams)
+
+n_features = 14 #dataset dimension
+n_trials = 10 #number of trials to average over
+
+X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
+X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
+
+first_train_steps = []
+av_consec_train_steps = []
+predict_times = []
+
+for trial in range(n_trials):
+    jax.clear_caches()
+
+    model = QuantumMetricLearner(**hyperparams)
+    model.fit(X_train, y_train)
+
+    #get step times from loss history data
+    step_times = np.array([model.loss_history_[1][i + 1] - model.loss_history_[1][i]
+                           for i in range(len(model.loss_history_[1]) - 1)])
+    step_times = np.insert(step_times, 0, [model.loss_history_[1][0]])
+
+    #first train step
+    first_train_steps.append(step_times[0])
+    #consecutive (average) train step
+    av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
+    #prediction time
+    time0 = time.time()
+    model.predict(X_test)
+    predict_times.append(time.time() - time0)
+
+
+#calculate mean and stds
+
+first_train_step = np.mean(first_train_steps)
+first_train_step_std = np.std(first_train_steps)
+
+consec_train_step = np.mean(av_consec_train_steps)
+consec_train_step_std = np.std(av_consec_train_steps)
+
+predict_time = np.mean(predict_times)
+predict_time_std = np.std(predict_times)
+
+#write to csv
+data = [first_train_step, first_train_step_std, consec_train_step, consec_train_step_std, predict_time,
+        predict_time_std, hyperparams]
+
+model_name = model.__class__.__name__
+filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_JAX.csv"
+header = ['first_train_step', 'first_train_step_std', 'consec_train_step', 'consec_train_step_std', 'predict_time',
+          'predict_time_std', 'hyperparameters']
+
+if not os.path.exists('JAX'):
+    # Create the directory
+    os.mkdir('JAX')
+
+with open('JAX/'+filename, mode="w", newline="") as file:
+    writer = csv.writer(file)
+    writer.writerow(header)
+    for row in [data]:
+        writer.writerow(row)
diff --git a/nersc/performance_indicators/perf_ind_qmetric_LR15_jax.py b/nersc/performance_indicators/perf_ind_qmetric_LR15_jax.py
new file mode 100644
index 00000000..39b64de1
--- /dev/null
+++ b/nersc/performance_indicators/perf_ind_qmetric_LR15_jax.py
@@ -0,0 +1,80 @@
+import qml_benchmarks
+import pennylane as qml
+import jax
+import jax.numpy as jnp
+import numpy as np
+import time
+import csv
+import os
+import yaml
+
+from qml_benchmarks.models.quantum_metric_learning import QuantumMetricLearner
+from qml_benchmarks.hyperparam_search_utils import read_data
+
+with open('hyperparam_settings.yaml', "r") as file:
+    hp_settings = yaml.safe_load(file)
+
+hyperparams = {**hp_settings['IQPVariationalClassifier'], **{'use_jax':True, 'vmap':True, 'max_steps':100}}
+
+print(hyperparams)
+
+n_features = 15 #dataset dimension
+n_trials = 10 #number of trials to average over
+
+X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
+X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
+
+first_train_steps = []
+av_consec_train_steps = []
+predict_times = []
+
+for trial in range(n_trials):
+    jax.clear_caches()
+
+    model = QuantumMetricLearner(**hyperparams)
+    model.fit(X_train, y_train)
+
+    #get step times from loss history data
+    step_times = np.array([model.loss_history_[1][i + 1] - model.loss_history_[1][i]
+                           for i in range(len(model.loss_history_[1]) - 1)])
+    step_times = np.insert(step_times, 0, [model.loss_history_[1][0]])
+
+    #first train step
+    first_train_steps.append(step_times[0])
+    #consecutive (average) train step
+    av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
+    #prediction time
+    time0 = time.time()
+    model.predict(X_test)
+    predict_times.append(time.time() - time0)
+
+
+#calculate mean and stds
+
+first_train_step = np.mean(first_train_steps)
+first_train_step_std = np.std(first_train_steps)
+
+consec_train_step = np.mean(av_consec_train_steps)
+consec_train_step_std = np.std(av_consec_train_steps)
+
+predict_time = np.mean(predict_times)
+predict_time_std = np.std(predict_times)
+
+#write to csv
+data = [first_train_step, first_train_step_std, consec_train_step, consec_train_step_std, predict_time,
+        predict_time_std, hyperparams]
+
+model_name = model.__class__.__name__
+filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_JAX.csv"
+header = ['first_train_step', 'first_train_step_std', 'consec_train_step', 'consec_train_step_std', 'predict_time',
+          'predict_time_std', 'hyperparameters']
+
+if not os.path.exists('JAX'):
+    # Create the directory
+    os.mkdir('JAX')
+
+with open('JAX/'+filename, mode="w", newline="") as file:
+    writer = csv.writer(file)
+    writer.writerow(header)
+    for row in [data]:
+        writer.writerow(row)

From ae27b0e1ff51803139f5d1d6d79505871f65804b Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Mon, 15 Apr 2024 16:59:04 +0200
Subject: [PATCH 027/100] update

---
 ...earner_linearly_separable_14d_JAX_info.txt |  4 +
 .../hyperparam_settings.yaml                  | 14 +++-
 ...R14_jax.py => perf_ind_iqpker_LR14_jax.py} |  6 +-
 .../perf_ind_iqpvar_LR14_jax_novmap.py        | 80 -------------------
 .../perf_ind_iqpvar_LR15_jax.py               |  6 +-
 .../perf_ind_iqpvar_LR15_jax_novmap.py        | 80 -------------------
 .../perf_ind_iqpvar_LR16_jax_novmap.py        | 80 -------------------
 .../perf_ind_iqpvar_LR17_jax.py               |  6 +-
 .../perf_ind_iqpvar_LR17_jax_novmap.py        | 80 -------------------
 .../perf_ind_iqpvar_LR18_jax_novmap.py        | 80 -------------------
 .../perf_ind_iqpvar_LR19_jax_novmap.py        | 80 -------------------
 .../perf_ind_iqpvar_LR20_jax.py               |  6 +-
 .../perf_ind_iqpvar_LR20_jax_novmap.py        | 80 -------------------
 ...r_LR16_jax.py => perf_ind_pqk_LR15_jax.py} | 10 +--
 ...r_LR19_jax.py => perf_ind_pqk_LR17_jax.py} |  8 +-
 ...r_LR18_jax.py => perf_ind_pqk_LR20_jax.py} |  8 +-
 .../perf_ind_qmetric_LR14_jax.py              |  6 +-
 .../perf_ind_qmetric_LR15_jax.py              |  6 +-
 18 files changed, 48 insertions(+), 592 deletions(-)
 create mode 100644 nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_14d_JAX_info.txt
 rename nersc/performance_indicators/{perf_ind_iqpvar_LR14_jax.py => perf_ind_iqpker_LR14_jax.py} (90%)
 delete mode 100644 nersc/performance_indicators/perf_ind_iqpvar_LR14_jax_novmap.py
 delete mode 100644 nersc/performance_indicators/perf_ind_iqpvar_LR15_jax_novmap.py
 delete mode 100644 nersc/performance_indicators/perf_ind_iqpvar_LR16_jax_novmap.py
 delete mode 100644 nersc/performance_indicators/perf_ind_iqpvar_LR17_jax_novmap.py
 delete mode 100644 nersc/performance_indicators/perf_ind_iqpvar_LR18_jax_novmap.py
 delete mode 100644 nersc/performance_indicators/perf_ind_iqpvar_LR19_jax_novmap.py
 delete mode 100644 nersc/performance_indicators/perf_ind_iqpvar_LR20_jax_novmap.py
 rename nersc/performance_indicators/{perf_ind_iqpvar_LR16_jax.py => perf_ind_pqk_LR15_jax.py} (87%)
 rename nersc/performance_indicators/{perf_ind_iqpvar_LR19_jax.py => perf_ind_pqk_LR17_jax.py} (89%)
 rename nersc/performance_indicators/{perf_ind_iqpvar_LR18_jax.py => perf_ind_pqk_LR20_jax.py} (89%)

diff --git a/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_14d_JAX_info.txt b/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_14d_JAX_info.txt
new file mode 100644
index 00000000..6e40ef73
--- /dev/null
+++ b/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_14d_JAX_info.txt
@@ -0,0 +1,4 @@
+salloc -q shared_interactive -C cpu -t 4:00:00 --cpus-per-task=32
+
+
+
diff --git a/nersc/performance_indicators/hyperparam_settings.yaml b/nersc/performance_indicators/hyperparam_settings.yaml
index e51f5ec0..5a3e82c5 100644
--- a/nersc/performance_indicators/hyperparam_settings.yaml
+++ b/nersc/performance_indicators/hyperparam_settings.yaml
@@ -1,4 +1,16 @@
 IQPVariationalClassifier:
   learning_rate: 0.01
   n_layers: 15
-  repeats: 10
\ No newline at end of file
+  repeats: 10
+
+QuantumMetricLearner:
+  learning_rate: 0.01
+  n_layers: 4
+  batch_size: 16
+
+ProjectedQuantumKernel:
+  trotter_steps: 5
+
+IQPKernelClassifier:
+  repeats: 10
+
diff --git a/nersc/performance_indicators/perf_ind_iqpvar_LR14_jax.py b/nersc/performance_indicators/perf_ind_iqpker_LR14_jax.py
similarity index 90%
rename from nersc/performance_indicators/perf_ind_iqpvar_LR14_jax.py
rename to nersc/performance_indicators/perf_ind_iqpker_LR14_jax.py
index 8efa2976..9262450a 100644
--- a/nersc/performance_indicators/perf_ind_iqpvar_LR14_jax.py
+++ b/nersc/performance_indicators/perf_ind_iqpker_LR14_jax.py
@@ -8,13 +8,13 @@
 import os
 import yaml
 
-from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier
+from qml_benchmarks.models.projected_quantum_kernel import ProjectedQuantumKernel
 from qml_benchmarks.hyperparam_search_utils import read_data
 
 with open('hyperparam_settings.yaml', "r") as file:
     hp_settings = yaml.safe_load(file)
 
-hyperparams = {**hp_settings['IQPVariationalClassifier'], **{'use_jax':True, 'vmap':True, 'max_steps':100}}
+hyperparams = {**hp_settings['ProjectedQuantumKernel'], **{'use_jax':True, 'vmap':True, 'max_steps':100}}
 
 print(hyperparams)
 
@@ -31,7 +31,7 @@
 for trial in range(n_trials):
     jax.clear_caches()
 
-    model = IQPVariationalClassifier(**hyperparams)
+    model = QuantumMetricLearner(**hyperparams)
     model.fit(X_train, y_train)
 
     #get step times from loss history data
diff --git a/nersc/performance_indicators/perf_ind_iqpvar_LR14_jax_novmap.py b/nersc/performance_indicators/perf_ind_iqpvar_LR14_jax_novmap.py
deleted file mode 100644
index 9bc9dd24..00000000
--- a/nersc/performance_indicators/perf_ind_iqpvar_LR14_jax_novmap.py
+++ /dev/null
@@ -1,80 +0,0 @@
-import qml_benchmarks
-import pennylane as qml
-import jax
-import jax.numpy as jnp
-import numpy as np
-import time
-import csv
-import os
-import yaml
-
-from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier
-from qml_benchmarks.hyperparam_search_utils import read_data
-
-with open('hyperparam_settings.yaml', "r") as file:
-    hp_settings = yaml.safe_load(file)
-
-hyperparams = {**hp_settings['IQPVariationalClassifier'], **{'use_jax':True, 'vmap':True, 'max_vmap':1, 'max_steps':100}}
-
-print(hyperparams)
-
-n_features = 14 #dataset dimension
-n_trials = 10 #number of trials to average over
-
-X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
-X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
-
-first_train_steps = []
-av_consec_train_steps = []
-predict_times = []
-
-for trial in range(n_trials):
-    jax.clear_caches()
-
-    model = IQPVariationalClassifier(**hyperparams)
-    model.fit(X_train, y_train)
-
-    #get step times from loss history data
-    step_times = np.array([model.loss_history_[1][i + 1] - model.loss_history_[1][i]
-                           for i in range(len(model.loss_history_[1]) - 1)])
-    step_times = np.insert(step_times, 0, [model.loss_history_[1][0]])
-
-    #first train step
-    first_train_steps.append(step_times[0])
-    #consecutive (average) train step
-    av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
-    #prediction time
-    time0 = time.time()
-    model.predict(X_test)
-    predict_times.append(time.time() - time0)
-
-
-#calculate mean and stds
-
-first_train_step = np.mean(first_train_steps)
-first_train_step_std = np.std(first_train_steps)
-
-consec_train_step = np.mean(av_consec_train_steps)
-consec_train_step_std = np.std(av_consec_train_steps)
-
-predict_time = np.mean(predict_times)
-predict_time_std = np.std(predict_times)
-
-#write to csv
-data = [first_train_step, first_train_step_std, consec_train_step, consec_train_step_std, predict_time,
-        predict_time_std, hyperparams]
-
-model_name = model.__class__.__name__
-filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_JAX_novmap.csv"
-header = ['first_train_step', 'first_train_step_std', 'consec_train_step', 'consec_train_step_std', 'predict_time',
-          'predict_time_std', 'hyperparameters']
-
-if not os.path.exists('JAX'):
-    # Create the directory
-    os.mkdir('JAX')
-
-with open('JAX/'+filename, mode="w", newline="") as file:
-    writer = csv.writer(file)
-    writer.writerow(header)
-    for row in [data]:
-        writer.writerow(row)
diff --git a/nersc/performance_indicators/perf_ind_iqpvar_LR15_jax.py b/nersc/performance_indicators/perf_ind_iqpvar_LR15_jax.py
index 20436274..dece4d43 100644
--- a/nersc/performance_indicators/perf_ind_iqpvar_LR15_jax.py
+++ b/nersc/performance_indicators/perf_ind_iqpvar_LR15_jax.py
@@ -14,12 +14,12 @@
 with open('hyperparam_settings.yaml', "r") as file:
     hp_settings = yaml.safe_load(file)
 
-hyperparams = {**hp_settings['IQPVariationalClassifier'], **{'use_jax':True, 'vmap':True, 'max_steps':100}}
+hyperparams = {**hp_settings['IQPVariationalClassifier'], **{'use_jax':True, 'vmap':True, 'max_steps':10}}
 
 print(hyperparams)
 
 n_features = 15 #dataset dimension
-n_trials = 10 #number of trials to average over
+n_trials = 5 #number of trials to average over
 
 X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
 X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
@@ -45,7 +45,7 @@
     av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
     #prediction time
     time0 = time.time()
-    model.predict(X_test)
+    model.predict(X_test[:10])
     predict_times.append(time.time() - time0)
 
 
diff --git a/nersc/performance_indicators/perf_ind_iqpvar_LR15_jax_novmap.py b/nersc/performance_indicators/perf_ind_iqpvar_LR15_jax_novmap.py
deleted file mode 100644
index 0c32497a..00000000
--- a/nersc/performance_indicators/perf_ind_iqpvar_LR15_jax_novmap.py
+++ /dev/null
@@ -1,80 +0,0 @@
-import qml_benchmarks
-import pennylane as qml
-import jax
-import jax.numpy as jnp
-import numpy as np
-import time
-import csv
-import os
-import yaml
-
-from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier
-from qml_benchmarks.hyperparam_search_utils import read_data
-
-with open('hyperparam_settings.yaml', "r") as file:
-    hp_settings = yaml.safe_load(file)
-
-hyperparams = {**hp_settings['IQPVariationalClassifier'], **{'use_jax':True, 'vmap':True, 'max_vmap':1, 'max_steps':100}}
-
-print(hyperparams)
-
-n_features = 15 #dataset dimension
-n_trials = 10 #number of trials to average over
-
-X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
-X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
-
-first_train_steps = []
-av_consec_train_steps = []
-predict_times = []
-
-for trial in range(n_trials):
-    jax.clear_caches()
-
-    model = IQPVariationalClassifier(**hyperparams)
-    model.fit(X_train, y_train)
-
-    #get step times from loss history data
-    step_times = np.array([model.loss_history_[1][i + 1] - model.loss_history_[1][i]
-                           for i in range(len(model.loss_history_[1]) - 1)])
-    step_times = np.insert(step_times, 0, [model.loss_history_[1][0]])
-
-    #first train step
-    first_train_steps.append(step_times[0])
-    #consecutive (average) train step
-    av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
-    #prediction time
-    time0 = time.time()
-    model.predict(X_test)
-    predict_times.append(time.time() - time0)
-
-
-#calculate mean and stds
-
-first_train_step = np.mean(first_train_steps)
-first_train_step_std = np.std(first_train_steps)
-
-consec_train_step = np.mean(av_consec_train_steps)
-consec_train_step_std = np.std(av_consec_train_steps)
-
-predict_time = np.mean(predict_times)
-predict_time_std = np.std(predict_times)
-
-#write to csv
-data = [first_train_step, first_train_step_std, consec_train_step, consec_train_step_std, predict_time,
-        predict_time_std, hyperparams]
-
-model_name = model.__class__.__name__
-filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_JAX_novmap.csv"
-header = ['first_train_step', 'first_train_step_std', 'consec_train_step', 'consec_train_step_std', 'predict_time',
-          'predict_time_std', 'hyperparameters']
-
-if not os.path.exists('JAX'):
-    # Create the directory
-    os.mkdir('JAX')
-
-with open('JAX/'+filename, mode="w", newline="") as file:
-    writer = csv.writer(file)
-    writer.writerow(header)
-    for row in [data]:
-        writer.writerow(row)
diff --git a/nersc/performance_indicators/perf_ind_iqpvar_LR16_jax_novmap.py b/nersc/performance_indicators/perf_ind_iqpvar_LR16_jax_novmap.py
deleted file mode 100644
index 427123a6..00000000
--- a/nersc/performance_indicators/perf_ind_iqpvar_LR16_jax_novmap.py
+++ /dev/null
@@ -1,80 +0,0 @@
-import qml_benchmarks
-import pennylane as qml
-import jax
-import jax.numpy as jnp
-import numpy as np
-import time
-import csv
-import os
-import yaml
-
-from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier
-from qml_benchmarks.hyperparam_search_utils import read_data
-
-with open('hyperparam_settings.yaml', "r") as file:
-    hp_settings = yaml.safe_load(file)
-
-hyperparams = {**hp_settings['IQPVariationalClassifier'], **{'use_jax':True, 'vmap':True, 'max_vmap':1, 'max_steps':100}}
-
-print(hyperparams)
-
-n_features = 16 #dataset dimension
-n_trials = 10 #number of trials to average over
-
-X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
-X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
-
-first_train_steps = []
-av_consec_train_steps = []
-predict_times = []
-
-for trial in range(n_trials):
-    jax.clear_caches()
-
-    model = IQPVariationalClassifier(**hyperparams)
-    model.fit(X_train, y_train)
-
-    #get step times from loss history data
-    step_times = np.array([model.loss_history_[1][i + 1] - model.loss_history_[1][i]
-                           for i in range(len(model.loss_history_[1]) - 1)])
-    step_times = np.insert(step_times, 0, [model.loss_history_[1][0]])
-
-    #first train step
-    first_train_steps.append(step_times[0])
-    #consecutive (average) train step
-    av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
-    #prediction time
-    time0 = time.time()
-    model.predict(X_test)
-    predict_times.append(time.time() - time0)
-
-
-#calculate mean and stds
-
-first_train_step = np.mean(first_train_steps)
-first_train_step_std = np.std(first_train_steps)
-
-consec_train_step = np.mean(av_consec_train_steps)
-consec_train_step_std = np.std(av_consec_train_steps)
-
-predict_time = np.mean(predict_times)
-predict_time_std = np.std(predict_times)
-
-#write to csv
-data = [first_train_step, first_train_step_std, consec_train_step, consec_train_step_std, predict_time,
-        predict_time_std, hyperparams]
-
-model_name = model.__class__.__name__
-filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_JAX_novmap.csv"
-header = ['first_train_step', 'first_train_step_std', 'consec_train_step', 'consec_train_step_std', 'predict_time',
-          'predict_time_std', 'hyperparameters']
-
-if not os.path.exists('JAX'):
-    # Create the directory
-    os.mkdir('JAX')
-
-with open('JAX/'+filename, mode="w", newline="") as file:
-    writer = csv.writer(file)
-    writer.writerow(header)
-    for row in [data]:
-        writer.writerow(row)
diff --git a/nersc/performance_indicators/perf_ind_iqpvar_LR17_jax.py b/nersc/performance_indicators/perf_ind_iqpvar_LR17_jax.py
index 910e1edf..7dbc8e80 100644
--- a/nersc/performance_indicators/perf_ind_iqpvar_LR17_jax.py
+++ b/nersc/performance_indicators/perf_ind_iqpvar_LR17_jax.py
@@ -14,12 +14,12 @@
 with open('hyperparam_settings.yaml', "r") as file:
     hp_settings = yaml.safe_load(file)
 
-hyperparams = {**hp_settings['IQPVariationalClassifier'], **{'use_jax':True, 'vmap':True, 'max_steps':100}}
+hyperparams = {**hp_settings['IQPVariationalClassifier'], **{'use_jax':True, 'vmap':True, 'max_steps':10}}
 
 print(hyperparams)
 
 n_features = 17 #dataset dimension
-n_trials = 10 #number of trials to average over
+n_trials = 5 #number of trials to average over
 
 X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
 X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
@@ -45,7 +45,7 @@
     av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
     #prediction time
     time0 = time.time()
-    model.predict(X_test)
+    model.predict(X_test[:10])
     predict_times.append(time.time() - time0)
 
 
diff --git a/nersc/performance_indicators/perf_ind_iqpvar_LR17_jax_novmap.py b/nersc/performance_indicators/perf_ind_iqpvar_LR17_jax_novmap.py
deleted file mode 100644
index 8e07e9b2..00000000
--- a/nersc/performance_indicators/perf_ind_iqpvar_LR17_jax_novmap.py
+++ /dev/null
@@ -1,80 +0,0 @@
-import qml_benchmarks
-import pennylane as qml
-import jax
-import jax.numpy as jnp
-import numpy as np
-import time
-import csv
-import os
-import yaml
-
-from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier
-from qml_benchmarks.hyperparam_search_utils import read_data
-
-with open('hyperparam_settings.yaml', "r") as file:
-    hp_settings = yaml.safe_load(file)
-
-hyperparams = {**hp_settings['IQPVariationalClassifier'], **{'use_jax':True, 'vmap':True, 'max_vmap':1, 'max_steps':100}}
-
-print(hyperparams)
-
-n_features = 17 #dataset dimension
-n_trials = 10 #number of trials to average over
-
-X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
-X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
-
-first_train_steps = []
-av_consec_train_steps = []
-predict_times = []
-
-for trial in range(n_trials):
-    jax.clear_caches()
-
-    model = IQPVariationalClassifier(**hyperparams)
-    model.fit(X_train, y_train)
-
-    #get step times from loss history data
-    step_times = np.array([model.loss_history_[1][i + 1] - model.loss_history_[1][i]
-                           for i in range(len(model.loss_history_[1]) - 1)])
-    step_times = np.insert(step_times, 0, [model.loss_history_[1][0]])
-
-    #first train step
-    first_train_steps.append(step_times[0])
-    #consecutive (average) train step
-    av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
-    #prediction time
-    time0 = time.time()
-    model.predict(X_test)
-    predict_times.append(time.time() - time0)
-
-
-#calculate mean and stds
-
-first_train_step = np.mean(first_train_steps)
-first_train_step_std = np.std(first_train_steps)
-
-consec_train_step = np.mean(av_consec_train_steps)
-consec_train_step_std = np.std(av_consec_train_steps)
-
-predict_time = np.mean(predict_times)
-predict_time_std = np.std(predict_times)
-
-#write to csv
-data = [first_train_step, first_train_step_std, consec_train_step, consec_train_step_std, predict_time,
-        predict_time_std, hyperparams]
-
-model_name = model.__class__.__name__
-filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_JAX_novmap.csv"
-header = ['first_train_step', 'first_train_step_std', 'consec_train_step', 'consec_train_step_std', 'predict_time',
-          'predict_time_std', 'hyperparameters']
-
-if not os.path.exists('JAX'):
-    # Create the directory
-    os.mkdir('JAX')
-
-with open('JAX/'+filename, mode="w", newline="") as file:
-    writer = csv.writer(file)
-    writer.writerow(header)
-    for row in [data]:
-        writer.writerow(row)
diff --git a/nersc/performance_indicators/perf_ind_iqpvar_LR18_jax_novmap.py b/nersc/performance_indicators/perf_ind_iqpvar_LR18_jax_novmap.py
deleted file mode 100644
index 8d1efb3a..00000000
--- a/nersc/performance_indicators/perf_ind_iqpvar_LR18_jax_novmap.py
+++ /dev/null
@@ -1,80 +0,0 @@
-import qml_benchmarks
-import pennylane as qml
-import jax
-import jax.numpy as jnp
-import numpy as np
-import time
-import csv
-import os
-import yaml
-
-from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier
-from qml_benchmarks.hyperparam_search_utils import read_data
-
-with open('hyperparam_settings.yaml', "r") as file:
-    hp_settings = yaml.safe_load(file)
-
-hyperparams = {**hp_settings['IQPVariationalClassifier'], **{'use_jax':True, 'vmap':True, 'max_vmap':1, 'max_steps':100}}
-
-print(hyperparams)
-
-n_features = 18 #dataset dimension
-n_trials = 10 #number of trials to average over
-
-X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
-X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
-
-first_train_steps = []
-av_consec_train_steps = []
-predict_times = []
-
-for trial in range(n_trials):
-    jax.clear_caches()
-
-    model = IQPVariationalClassifier(**hyperparams)
-    model.fit(X_train, y_train)
-
-    #get step times from loss history data
-    step_times = np.array([model.loss_history_[1][i + 1] - model.loss_history_[1][i]
-                           for i in range(len(model.loss_history_[1]) - 1)])
-    step_times = np.insert(step_times, 0, [model.loss_history_[1][0]])
-
-    #first train step
-    first_train_steps.append(step_times[0])
-    #consecutive (average) train step
-    av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
-    #prediction time
-    time0 = time.time()
-    model.predict(X_test)
-    predict_times.append(time.time() - time0)
-
-
-#calculate mean and stds
-
-first_train_step = np.mean(first_train_steps)
-first_train_step_std = np.std(first_train_steps)
-
-consec_train_step = np.mean(av_consec_train_steps)
-consec_train_step_std = np.std(av_consec_train_steps)
-
-predict_time = np.mean(predict_times)
-predict_time_std = np.std(predict_times)
-
-#write to csv
-data = [first_train_step, first_train_step_std, consec_train_step, consec_train_step_std, predict_time,
-        predict_time_std, hyperparams]
-
-model_name = model.__class__.__name__
-filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_JAX_novmap.csv"
-header = ['first_train_step', 'first_train_step_std', 'consec_train_step', 'consec_train_step_std', 'predict_time',
-          'predict_time_std', 'hyperparameters']
-
-if not os.path.exists('JAX'):
-    # Create the directory
-    os.mkdir('JAX')
-
-with open('JAX/'+filename, mode="w", newline="") as file:
-    writer = csv.writer(file)
-    writer.writerow(header)
-    for row in [data]:
-        writer.writerow(row)
diff --git a/nersc/performance_indicators/perf_ind_iqpvar_LR19_jax_novmap.py b/nersc/performance_indicators/perf_ind_iqpvar_LR19_jax_novmap.py
deleted file mode 100644
index ccc8b142..00000000
--- a/nersc/performance_indicators/perf_ind_iqpvar_LR19_jax_novmap.py
+++ /dev/null
@@ -1,80 +0,0 @@
-import qml_benchmarks
-import pennylane as qml
-import jax
-import jax.numpy as jnp
-import numpy as np
-import time
-import csv
-import os
-import yaml
-
-from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier
-from qml_benchmarks.hyperparam_search_utils import read_data
-
-with open('hyperparam_settings.yaml', "r") as file:
-    hp_settings = yaml.safe_load(file)
-
-hyperparams = {**hp_settings['IQPVariationalClassifier'], **{'use_jax':True, 'vmap':True, 'max_vmap':1, 'max_steps':100}}
-
-print(hyperparams)
-
-n_features = 19 #dataset dimension
-n_trials = 10 #number of trials to average over
-
-X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
-X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
-
-first_train_steps = []
-av_consec_train_steps = []
-predict_times = []
-
-for trial in range(n_trials):
-    jax.clear_caches()
-
-    model = IQPVariationalClassifier(**hyperparams)
-    model.fit(X_train, y_train)
-
-    #get step times from loss history data
-    step_times = np.array([model.loss_history_[1][i + 1] - model.loss_history_[1][i]
-                           for i in range(len(model.loss_history_[1]) - 1)])
-    step_times = np.insert(step_times, 0, [model.loss_history_[1][0]])
-
-    #first train step
-    first_train_steps.append(step_times[0])
-    #consecutive (average) train step
-    av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
-    #prediction time
-    time0 = time.time()
-    model.predict(X_test)
-    predict_times.append(time.time() - time0)
-
-
-#calculate mean and stds
-
-first_train_step = np.mean(first_train_steps)
-first_train_step_std = np.std(first_train_steps)
-
-consec_train_step = np.mean(av_consec_train_steps)
-consec_train_step_std = np.std(av_consec_train_steps)
-
-predict_time = np.mean(predict_times)
-predict_time_std = np.std(predict_times)
-
-#write to csv
-data = [first_train_step, first_train_step_std, consec_train_step, consec_train_step_std, predict_time,
-        predict_time_std, hyperparams]
-
-model_name = model.__class__.__name__
-filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_JAX_novmap.csv"
-header = ['first_train_step', 'first_train_step_std', 'consec_train_step', 'consec_train_step_std', 'predict_time',
-          'predict_time_std', 'hyperparameters']
-
-if not os.path.exists('JAX'):
-    # Create the directory
-    os.mkdir('JAX')
-
-with open('JAX/'+filename, mode="w", newline="") as file:
-    writer = csv.writer(file)
-    writer.writerow(header)
-    for row in [data]:
-        writer.writerow(row)
diff --git a/nersc/performance_indicators/perf_ind_iqpvar_LR20_jax.py b/nersc/performance_indicators/perf_ind_iqpvar_LR20_jax.py
index f21f8efe..c6dc8770 100644
--- a/nersc/performance_indicators/perf_ind_iqpvar_LR20_jax.py
+++ b/nersc/performance_indicators/perf_ind_iqpvar_LR20_jax.py
@@ -14,12 +14,12 @@
 with open('hyperparam_settings.yaml', "r") as file:
     hp_settings = yaml.safe_load(file)
 
-hyperparams = {**hp_settings['IQPVariationalClassifier'], **{'use_jax':True, 'vmap':True, 'max_steps':100}}
+hyperparams = {**hp_settings['IQPVariationalClassifier'], **{'use_jax':True, 'vmap':True, 'max_steps':10}}
 
 print(hyperparams)
 
 n_features = 20 #dataset dimension
-n_trials = 10 #number of trials to average over
+n_trials = 5 #number of trials to average over
 
 X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
 X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
@@ -45,7 +45,7 @@
     av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
     #prediction time
     time0 = time.time()
-    model.predict(X_test)
+    model.predict(X_test[:10])
     predict_times.append(time.time() - time0)
 
 
diff --git a/nersc/performance_indicators/perf_ind_iqpvar_LR20_jax_novmap.py b/nersc/performance_indicators/perf_ind_iqpvar_LR20_jax_novmap.py
deleted file mode 100644
index b0edd3f8..00000000
--- a/nersc/performance_indicators/perf_ind_iqpvar_LR20_jax_novmap.py
+++ /dev/null
@@ -1,80 +0,0 @@
-import qml_benchmarks
-import pennylane as qml
-import jax
-import jax.numpy as jnp
-import numpy as np
-import time
-import csv
-import os
-import yaml
-
-from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier
-from qml_benchmarks.hyperparam_search_utils import read_data
-
-with open('hyperparam_settings.yaml', "r") as file:
-    hp_settings = yaml.safe_load(file)
-
-hyperparams = {**hp_settings['IQPVariationalClassifier'], **{'use_jax':True, 'vmap':True, 'max_vmap':1, 'max_steps':100}}
-
-print(hyperparams)
-
-n_features = 20 #dataset dimension
-n_trials = 10 #number of trials to average over
-
-X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
-X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
-
-first_train_steps = []
-av_consec_train_steps = []
-predict_times = []
-
-for trial in range(n_trials):
-    jax.clear_caches()
-
-    model = IQPVariationalClassifier(**hyperparams)
-    model.fit(X_train, y_train)
-
-    #get step times from loss history data
-    step_times = np.array([model.loss_history_[1][i + 1] - model.loss_history_[1][i]
-                           for i in range(len(model.loss_history_[1]) - 1)])
-    step_times = np.insert(step_times, 0, [model.loss_history_[1][0]])
-
-    #first train step
-    first_train_steps.append(step_times[0])
-    #consecutive (average) train step
-    av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
-    #prediction time
-    time0 = time.time()
-    model.predict(X_test)
-    predict_times.append(time.time() - time0)
-
-
-#calculate mean and stds
-
-first_train_step = np.mean(first_train_steps)
-first_train_step_std = np.std(first_train_steps)
-
-consec_train_step = np.mean(av_consec_train_steps)
-consec_train_step_std = np.std(av_consec_train_steps)
-
-predict_time = np.mean(predict_times)
-predict_time_std = np.std(predict_times)
-
-#write to csv
-data = [first_train_step, first_train_step_std, consec_train_step, consec_train_step_std, predict_time,
-        predict_time_std, hyperparams]
-
-model_name = model.__class__.__name__
-filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_JAX_novmap.csv"
-header = ['first_train_step', 'first_train_step_std', 'consec_train_step', 'consec_train_step_std', 'predict_time',
-          'predict_time_std', 'hyperparameters']
-
-if not os.path.exists('JAX'):
-    # Create the directory
-    os.mkdir('JAX')
-
-with open('JAX/'+filename, mode="w", newline="") as file:
-    writer = csv.writer(file)
-    writer.writerow(header)
-    for row in [data]:
-        writer.writerow(row)
diff --git a/nersc/performance_indicators/perf_ind_iqpvar_LR16_jax.py b/nersc/performance_indicators/perf_ind_pqk_LR15_jax.py
similarity index 87%
rename from nersc/performance_indicators/perf_ind_iqpvar_LR16_jax.py
rename to nersc/performance_indicators/perf_ind_pqk_LR15_jax.py
index 190c7fd2..7a32b087 100644
--- a/nersc/performance_indicators/perf_ind_iqpvar_LR16_jax.py
+++ b/nersc/performance_indicators/perf_ind_pqk_LR15_jax.py
@@ -8,18 +8,18 @@
 import os
 import yaml
 
-from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier
+from qml_benchmarks.models.projected_quantum_kernel import ProjectedQuantumKernel
 from qml_benchmarks.hyperparam_search_utils import read_data
 
 with open('hyperparam_settings.yaml', "r") as file:
     hp_settings = yaml.safe_load(file)
 
-hyperparams = {**hp_settings['IQPVariationalClassifier'], **{'use_jax':True, 'vmap':True, 'max_steps':100}}
+hyperparams = {**hp_settings['ProjectedQuantumKernel'], **{'use_jax':True, 'vmap':True, 'max_steps':100}}
 
 print(hyperparams)
 
-n_features = 16 #dataset dimension
-n_trials = 10 #number of trials to average over
+n_features = 3 #dataset dimension
+n_trials = 1 #number of trials to average over
 
 X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
 X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
@@ -31,7 +31,7 @@
 for trial in range(n_trials):
     jax.clear_caches()
 
-    model = IQPVariationalClassifier(**hyperparams)
+    model = QuantumMetricLearner(**hyperparams)
     model.fit(X_train, y_train)
 
     #get step times from loss history data
diff --git a/nersc/performance_indicators/perf_ind_iqpvar_LR19_jax.py b/nersc/performance_indicators/perf_ind_pqk_LR17_jax.py
similarity index 89%
rename from nersc/performance_indicators/perf_ind_iqpvar_LR19_jax.py
rename to nersc/performance_indicators/perf_ind_pqk_LR17_jax.py
index d6798dea..da5d6e8d 100644
--- a/nersc/performance_indicators/perf_ind_iqpvar_LR19_jax.py
+++ b/nersc/performance_indicators/perf_ind_pqk_LR17_jax.py
@@ -8,17 +8,17 @@
 import os
 import yaml
 
-from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier
+from qml_benchmarks.models.projected_quantum_kernel import ProjectedQuantumKernel
 from qml_benchmarks.hyperparam_search_utils import read_data
 
 with open('hyperparam_settings.yaml', "r") as file:
     hp_settings = yaml.safe_load(file)
 
-hyperparams = {**hp_settings['IQPVariationalClassifier'], **{'use_jax':True, 'vmap':True, 'max_steps':100}}
+hyperparams = {**hp_settings['ProjectedQuantumKernel'], **{'use_jax':True, 'vmap':True, 'max_steps':100}}
 
 print(hyperparams)
 
-n_features = 19 #dataset dimension
+n_features = 17 #dataset dimension
 n_trials = 10 #number of trials to average over
 
 X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
@@ -31,7 +31,7 @@
 for trial in range(n_trials):
     jax.clear_caches()
 
-    model = IQPVariationalClassifier(**hyperparams)
+    model = QuantumMetricLearner(**hyperparams)
     model.fit(X_train, y_train)
 
     #get step times from loss history data
diff --git a/nersc/performance_indicators/perf_ind_iqpvar_LR18_jax.py b/nersc/performance_indicators/perf_ind_pqk_LR20_jax.py
similarity index 89%
rename from nersc/performance_indicators/perf_ind_iqpvar_LR18_jax.py
rename to nersc/performance_indicators/perf_ind_pqk_LR20_jax.py
index 569c8c70..da5d6e8d 100644
--- a/nersc/performance_indicators/perf_ind_iqpvar_LR18_jax.py
+++ b/nersc/performance_indicators/perf_ind_pqk_LR20_jax.py
@@ -8,17 +8,17 @@
 import os
 import yaml
 
-from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier
+from qml_benchmarks.models.projected_quantum_kernel import ProjectedQuantumKernel
 from qml_benchmarks.hyperparam_search_utils import read_data
 
 with open('hyperparam_settings.yaml', "r") as file:
     hp_settings = yaml.safe_load(file)
 
-hyperparams = {**hp_settings['IQPVariationalClassifier'], **{'use_jax':True, 'vmap':True, 'max_steps':100}}
+hyperparams = {**hp_settings['ProjectedQuantumKernel'], **{'use_jax':True, 'vmap':True, 'max_steps':100}}
 
 print(hyperparams)
 
-n_features = 18 #dataset dimension
+n_features = 17 #dataset dimension
 n_trials = 10 #number of trials to average over
 
 X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
@@ -31,7 +31,7 @@
 for trial in range(n_trials):
     jax.clear_caches()
 
-    model = IQPVariationalClassifier(**hyperparams)
+    model = QuantumMetricLearner(**hyperparams)
     model.fit(X_train, y_train)
 
     #get step times from loss history data
diff --git a/nersc/performance_indicators/perf_ind_qmetric_LR14_jax.py b/nersc/performance_indicators/perf_ind_qmetric_LR14_jax.py
index dec6996e..cbb647c0 100644
--- a/nersc/performance_indicators/perf_ind_qmetric_LR14_jax.py
+++ b/nersc/performance_indicators/perf_ind_qmetric_LR14_jax.py
@@ -14,12 +14,12 @@
 with open('hyperparam_settings.yaml', "r") as file:
     hp_settings = yaml.safe_load(file)
 
-hyperparams = {**hp_settings['IQPVariationalClassifier'], **{'use_jax':True, 'vmap':True, 'max_steps':100}}
+hyperparams = {**hp_settings['QuantumMetricLearner'], **{'use_jax':True, 'vmap':True, 'max_steps':10}}
 
 print(hyperparams)
 
 n_features = 14 #dataset dimension
-n_trials = 10 #number of trials to average over
+n_trials = 5 #number of trials to average over
 
 X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
 X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
@@ -45,7 +45,7 @@
     av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
     #prediction time
     time0 = time.time()
-    model.predict(X_test)
+    model.predict(X_test[:10])
     predict_times.append(time.time() - time0)
 
 
diff --git a/nersc/performance_indicators/perf_ind_qmetric_LR15_jax.py b/nersc/performance_indicators/perf_ind_qmetric_LR15_jax.py
index 39b64de1..8a78431b 100644
--- a/nersc/performance_indicators/perf_ind_qmetric_LR15_jax.py
+++ b/nersc/performance_indicators/perf_ind_qmetric_LR15_jax.py
@@ -14,12 +14,12 @@
 with open('hyperparam_settings.yaml', "r") as file:
     hp_settings = yaml.safe_load(file)
 
-hyperparams = {**hp_settings['IQPVariationalClassifier'], **{'use_jax':True, 'vmap':True, 'max_steps':100}}
+hyperparams = {**hp_settings['QuantumMetricLearner'], **{'use_jax':True, 'vmap':True, 'max_steps':10}}
 
 print(hyperparams)
 
 n_features = 15 #dataset dimension
-n_trials = 10 #number of trials to average over
+n_trials = 5 #number of trials to average over
 
 X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
 X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
@@ -45,7 +45,7 @@
     av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
     #prediction time
     time0 = time.time()
-    model.predict(X_test)
+    model.predict(X_test[:10])
     predict_times.append(time.time() - time0)
 
 

From 2a254beadf8ac6325d81b627454b4f002158c271 Mon Sep 17 00:00:00 2001
From: root <bowles.physics@gmail.com>
Date: Mon, 15 Apr 2024 08:31:03 -0700
Subject: [PATCH 028/100] result

---
 ...sifier_linearly_separable_15d_performance_indicators_JAX.csv | 2 ++
 1 file changed, 2 insertions(+)
 create mode 100644 nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_15d_performance_indicators_JAX.csv

diff --git a/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_15d_performance_indicators_JAX.csv b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_15d_performance_indicators_JAX.csv
new file mode 100644
index 00000000..88c2dcb8
--- /dev/null
+++ b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_15d_performance_indicators_JAX.csv
@@ -0,0 +1,2 @@
+first_train_step,first_train_step_std,consec_train_step,consec_train_step_std,predict_time,predict_time_std,hyperparameters
+148.14196796417235,11.108692798642224,10.153086731168958,1.4567056279882924,49.51927604675293,2.8201209235990614,"{'learning_rate': 0.01, 'n_layers': 15, 'repeats': 10, 'use_jax': True, 'vmap': True, 'max_steps': 10}"

From 5cf8f2b8d41c44acfdc9c403a411973d3fc33aef Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Mon, 15 Apr 2024 17:37:29 +0200
Subject: [PATCH 029/100] add info

---
 ...sifier_linearly_separable_15d_JAX_info.txt | 28 ++++++++++++++++++-
 1 file changed, 27 insertions(+), 1 deletion(-)

diff --git a/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_15d_JAX_info.txt b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_15d_JAX_info.txt
index 74dede79..03554fb5 100644
--- a/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_15d_JAX_info.txt
+++ b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_15d_JAX_info.txt
@@ -1,2 +1,28 @@
-salloc -q shared_interactive -C cpu -t 4:00:00 --cpus-per-task=32
+salloc -q shared_interactive -C cpu -t 4:00:00 --cpus-per-task=16
 
+JobId=24420449 JobName=interactive
+UserId=jbowles(104993) GroupId=jbowles(104993) MCS_label=N/A
+Priority=73320 Nice=0 Account=m4139 QOS=shared_interactive
+JobState=COMPLETING Reason=NonZeroExitCode Dependency=(null)
+Requeue=0 Restarts=0 BatchFlag=0 Reboot=0 ExitCode=127:0
+RunTime=00:32:07 TimeLimit=04:00:00 TimeMin=N/A
+SubmitTime=2024-04-15T08:02:57 EligibleTime=2024-04-15T08:02:57
+AccrueTime=2024-04-15T08:03:06
+StartTime=2024-04-15T08:03:06 EndTime=2024-04-15T08:35:13 Deadline=N/A
+PreemptEligibleTime=2024-04-15T08:03:06 PreemptTime=None
+SuspendTime=None SecsPreSuspend=0 LastSchedEval=2024-04-15T08:03:06 Scheduler=Main
+Partition=shared_urgent_milan_ss11 AllocNode:Sid=login24:836159
+ReqNodeList=(null) ExcNodeList=(null)
+NodeList=nid004092
+BatchHost=nid004092
+NumNodes=1 NumCPUs=16 NumTasks=1 CPUs/Task=16 ReqB:S:C:T=0:0:*:*
+ReqTRES=cpu=16,mem=30624M,node=1,billing=16
+AllocTRES=cpu=16,mem=30624M,node=1,billing=16
+Socks/Node=* NtasksPerN:B:S:C=0:0:*:* CoreSpec=*
+MinCPUsNode=16 MinMemoryCPU=1914M MinTmpDiskNode=0
+Features=cpu DelayBoot=00:00:00
+OverSubscribe=OK Contiguous=0 Licenses=scratch:1 Network=(null)
+Command=/bin/sh
+WorkDir=/pscratch/sd/j/jbowles/qml-benchmarks/nersc
+AdminComment={"qos":"shared_interactive","arrayTaskId":4294967294,"resizing":0,"originCluster":"perlmutter","batchHost":"nid004092","jobExitCode":32512,"uid":104993,"cluster":"perlmutter","tresRequest":"1=16,2=30624,3=18446744073709551614,4=1,5=16","jobId":24420449,"gresRequest":"cpu=16,mem=30624M,node=1,billing=16","arrayJobId":0,"workingDirectory":"\/pscratch\/sd\/j\/jbowles\/qml-benchmarks\/nersc","licenses":"scratch:1","name":"interactive","partition":"shared_urgent_milan_ss11","packJobId":0,"reboot":0,"priority":73320,"features":"cpu","restartCnt":0,"argv":["\/bin\/sh","-c","srun --interactive --preserve-env --pty $SHELL"],"allocCpus":16,"startTime":1713193386,"packJobOffset":0,"jobAccount":"m4139","nodes":"nid004092","submitTime":1713193377,"timeLimit":240,"jobDerivedExitCode":0,"endTime":1713195313,"allocNodes":1}
+Power=
\ No newline at end of file

From 881dc94531364675f8cc4a0fc5407dfe27d7ed83 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Mon, 15 Apr 2024 17:40:25 +0200
Subject: [PATCH 030/100] add info

---
 ...sifier_linearly_separable_15d_JAX_info.txt | 239 +++++++++++++++++-
 1 file changed, 238 insertions(+), 1 deletion(-)

diff --git a/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_15d_JAX_info.txt b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_15d_JAX_info.txt
index 03554fb5..3fb0afc5 100644
--- a/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_15d_JAX_info.txt
+++ b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_15d_JAX_info.txt
@@ -25,4 +25,241 @@ OverSubscribe=OK Contiguous=0 Licenses=scratch:1 Network=(null)
 Command=/bin/sh
 WorkDir=/pscratch/sd/j/jbowles/qml-benchmarks/nersc
 AdminComment={"qos":"shared_interactive","arrayTaskId":4294967294,"resizing":0,"originCluster":"perlmutter","batchHost":"nid004092","jobExitCode":32512,"uid":104993,"cluster":"perlmutter","tresRequest":"1=16,2=30624,3=18446744073709551614,4=1,5=16","jobId":24420449,"gresRequest":"cpu=16,mem=30624M,node=1,billing=16","arrayJobId":0,"workingDirectory":"\/pscratch\/sd\/j\/jbowles\/qml-benchmarks\/nersc","licenses":"scratch:1","name":"interactive","partition":"shared_urgent_milan_ss11","packJobId":0,"reboot":0,"priority":73320,"features":"cpu","restartCnt":0,"argv":["\/bin\/sh","-c","srun --interactive --preserve-env --pty $SHELL"],"allocCpus":16,"startTime":1713193386,"packJobOffset":0,"jobAccount":"m4139","nodes":"nid004092","submitTime":1713193377,"timeLimit":240,"jobDerivedExitCode":0,"endTime":1713195313,"allocNodes":1}
-Power=
\ No newline at end of file
+Power=
+
+Package                   Version
+------------------------- --------------
+absl-py                   2.1.0
+alembic                   1.13.1
+annotated-types           0.6.0
+antlr4-python3-runtime    4.9.2
+anyio                     4.3.0
+appdirs                   1.4.4
+argon2-cffi               23.1.0
+argon2-cffi-bindings      21.2.0
+arrow                     1.3.0
+asttokens                 2.4.1
+astunparse                1.6.3
+async-lru                 2.0.4
+attrs                     23.2.0
+autograd                  1.6.2
+autoray                   0.6.9
+Babel                     2.14.0
+beautifulsoup4            4.12.3
+beniget                   0.4.1
+bitstring                 3.1.7
+bleach                    6.1.0
+cachetools                5.3.3
+certifi                   2024.2.2
+cffi                      1.16.0
+charset-normalizer        3.3.2
+chex                      0.1.86
+cirq-core                 1.3.0
+cirq-pasqal               1.3.0
+click                     8.1.7
+cloudpickle               3.0.0
+colorlog                  6.8.2
+comm                      0.2.2
+contourpy                 1.2.1
+cryptography              42.0.5
+cycler                    0.12.1
+dask                      2024.4.1
+dbus-python               1.2.18
+debugpy                   1.8.1
+decorator                 4.4.2
+defusedxml                0.7.1
+diastatic-malt            2.15.1
+dill                      0.3.8
+distro                    1.7.0
+duet                      0.2.9
+etils                     1.7.0
+exceptiongroup            1.2.0
+executing                 2.0.1
+fastdtw                   0.3.4
+fastjsonschema            2.19.1
+filelock                  3.9.0
+fire                      0.6.0
+flax                      0.8.2
+fonttools                 4.50.0
+fqdn                      1.5.1
+fsspec                    2024.3.1
+future                    1.0.0
+gast                      0.5.2
+greenlet                  3.0.3
+h11                       0.14.0
+h5py                      3.10.0
+httpcore                  1.0.5
+httpx                     0.27.0
+ibm-cloud-sdk-core        3.19.2
+ibm-platform-services     0.53.2
+idna                      3.6
+importlib_metadata        7.1.0
+importlib_resources       6.4.0
+ipykernel                 6.29.4
+ipython                   8.23.0
+ipywidgets                8.1.2
+isoduration               20.11.0
+jax                       0.4.23
+jaxlib                    0.4.23
+jaxopt                    0.8.3
+jedi                      0.19.1
+Jinja2                    3.1.3
+joblib                    1.3.2
+json5                     0.9.24
+jsonpointer               2.4
+jsonschema                4.21.1
+jsonschema-specifications 2023.12.1
+jupyter                   1.0.0
+jupyter_client            8.6.1
+jupyter-console           6.6.3
+jupyter_core              5.7.2
+jupyter-events            0.10.0
+jupyter-lsp               2.2.4
+jupyter_server            2.13.0
+jupyter_server_terminals  0.5.3
+jupyterlab                4.1.5
+jupyterlab_pygments       0.3.0
+jupyterlab_server         2.25.4
+jupyterlab_widgets        3.0.10
+kiwisolver                1.4.5
+lark-parser               0.12.0
+llvmlite                  0.42.0
+locket                    1.0.0
+Mako                      1.3.2
+markdown-it-py            3.0.0
+MarkupSafe                2.1.5
+matplotlib                3.8.4
+matplotlib-inline         0.1.6
+mdurl                     0.1.2
+mistune                   3.0.2
+ml-dtypes                 0.4.0
+mpmath                    1.3.0
+msgpack                   1.0.8
+nbclient                  0.10.0
+nbconvert                 7.16.3
+nbformat                  5.10.4
+nest-asyncio              1.6.0
+networkx                  3.2.1
+notebook                  7.1.2
+notebook_shim             0.2.4
+numba                     0.59.1
+numpy                     1.26.4
+olefile                   0.46
+opt-einsum                3.3.0
+optax                     0.2.2
+optuna                    3.6.1
+orbax-checkpoint          0.5.9
+overrides                 7.7.0
+packaging                 24.0
+pandas                    2.2.1
+pandocfilters             1.5.1
+parso                     0.8.3
+partd                     1.4.1
+patsy                     0.5.6
+pbr                       6.0.0
+PennyLane                 0.35.1
+PennyLane-Catalyst        0.5.0
+PennyLane-Cirq            0.34.0
+PennyLane_Lightning       0.35.1
+PennyLane_Lightning_GPU   0.35.1
+PennyLane-qiskit          0.35.1
+PennyLane-SF              0.29.0
+pexpect                   4.9.0
+pillow                    10.3.0
+pip                       24.0
+platformdirs              4.2.0
+ply                       3.11
+prometheus_client         0.20.0
+prompt-toolkit            3.0.43
+protobuf                  5.26.1
+psutil                    5.9.8
+ptyprocess                0.7.0
+pure-eval                 0.2.2
+pycparser                 2.22
+pydantic                  2.6.4
+pydantic_core             2.16.3
+pydantic-settings         2.2.1
+pydot                     2.0.0
+Pygments                  2.17.2
+PyGObject                 3.42.1
+PyJWT                     2.8.0
+pylatexenc                2.10
+pyparsing                 3.1.2
+pyspnego                  0.10.2
+python-dateutil           2.9.0.post0
+python-dotenv             1.0.1
+python-json-logger        2.0.7
+pythran                   0.10.0
+pytz                      2024.1
+PyYAML                    6.0.1
+pyzmq                     25.1.2
+qiskit                    1.0.2
+qiskit-aer                0.14.0.1
+qiskit-algorithms         0.3.0
+qiskit-ibm-experiment     0.4.7
+qiskit-ibm-provider       0.10.0
+qiskit-ibm-runtime        0.20.0
+qiskit-machine-learning   0.7.2
+qml_benchmarks            0.1
+qtconsole                 5.5.1
+QtPy                      2.4.1
+quantum-blackbird         0.5.0
+quantum-xir               0.2.2
+referencing               0.34.0
+requests                  2.31.0
+requests-ntlm             1.2.0
+rfc3339-validator         0.1.4
+rfc3986-validator         0.1.1
+rich                      13.7.1
+rpds-py                   0.18.0
+ruamel.yaml               0.18.6
+ruamel.yaml.clib          0.2.8
+rustworkx                 0.14.2
+scikit-learn              1.4.1.post1
+scipy                     1.11.4
+seaborn                   0.13.2
+semantic-version          2.10.0
+Send2Trash                1.8.2
+setuptools                59.6.0
+six                       1.16.0
+sniffio                   1.3.1
+sortedcontainers          2.4.0
+soupsieve                 2.5
+SQLAlchemy                2.0.29
+ssh-import-id             5.11
+stack-data                0.6.3
+statsmodels               0.14.1
+stevedore                 5.2.0
+StrawberryFields          0.23.0
+symengine                 0.11.0
+sympy                     1.12
+tensorstore               0.1.56
+termcolor                 2.4.0
+terminado                 0.18.1
+thewalrus                 0.21.0
+threadpoolctl             3.4.0
+tinycss2                  1.2.1
+toml                      0.10.2
+tomli                     2.0.1
+tomlkit                   0.12.4
+toolz                     0.12.1
+torch                     2.2.2+cpu
+torchaudio                2.2.2+cpu
+torchvision               0.17.2+cpu
+tornado                   6.4
+tqdm                      4.66.2
+traitlets                 5.14.2
+types-python-dateutil     2.9.0.20240316
+typing_extensions         4.10.0
+tzdata                    2024.1
+uri-template              1.3.0
+urllib3                   2.2.1
+wcwidth                   0.2.13
+webcolors                 1.13
+webencodings              0.5.1
+websocket-client          1.7.0
+websockets                12.0
+wheel                     0.37.1
+widgetsnbextension        4.0.10
+xanadu-cloud-client       0.3.2
+zipp                      3.18.1
\ No newline at end of file

From 7dfd26e532b3be9d0706935728043f3c95a6ec67 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Mon, 15 Apr 2024 17:47:37 +0200
Subject: [PATCH 031/100] update tests

---
 ...sifier_linearly_separable_17d_JAX_info.txt | 240 +++++++++++++++++-
 .../perf_ind_iqpker_LR14_jax.py               |  80 ------
 .../perf_ind_iqpker_LR15_jax.py               |  54 ++++
 .../perf_ind_pqk_LR15_jax.py                  |  54 +---
 .../perf_ind_pqk_LR17_jax.py                  |  52 +---
 .../perf_ind_pqk_LR20_jax.py                  |  54 +---
 6 files changed, 334 insertions(+), 200 deletions(-)
 delete mode 100644 nersc/performance_indicators/perf_ind_iqpker_LR14_jax.py
 create mode 100644 nersc/performance_indicators/perf_ind_iqpker_LR15_jax.py

diff --git a/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_17d_JAX_info.txt b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_17d_JAX_info.txt
index 74dede79..9e3568ca 100644
--- a/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_17d_JAX_info.txt
+++ b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_17d_JAX_info.txt
@@ -1,2 +1,240 @@
-salloc -q shared_interactive -C cpu -t 4:00:00 --cpus-per-task=32
+salloc -q shared_interactive -C cpu -t 4:00:00 --cpus-per-task=16
 
+ADD
+
+Package                   Version
+------------------------- --------------
+absl-py                   2.1.0
+alembic                   1.13.1
+annotated-types           0.6.0
+antlr4-python3-runtime    4.9.2
+anyio                     4.3.0
+appdirs                   1.4.4
+argon2-cffi               23.1.0
+argon2-cffi-bindings      21.2.0
+arrow                     1.3.0
+asttokens                 2.4.1
+astunparse                1.6.3
+async-lru                 2.0.4
+attrs                     23.2.0
+autograd                  1.6.2
+autoray                   0.6.9
+Babel                     2.14.0
+beautifulsoup4            4.12.3
+beniget                   0.4.1
+bitstring                 3.1.7
+bleach                    6.1.0
+cachetools                5.3.3
+certifi                   2024.2.2
+cffi                      1.16.0
+charset-normalizer        3.3.2
+chex                      0.1.86
+cirq-core                 1.3.0
+cirq-pasqal               1.3.0
+click                     8.1.7
+cloudpickle               3.0.0
+colorlog                  6.8.2
+comm                      0.2.2
+contourpy                 1.2.1
+cryptography              42.0.5
+cycler                    0.12.1
+dask                      2024.4.1
+dbus-python               1.2.18
+debugpy                   1.8.1
+decorator                 4.4.2
+defusedxml                0.7.1
+diastatic-malt            2.15.1
+dill                      0.3.8
+distro                    1.7.0
+duet                      0.2.9
+etils                     1.7.0
+exceptiongroup            1.2.0
+executing                 2.0.1
+fastdtw                   0.3.4
+fastjsonschema            2.19.1
+filelock                  3.9.0
+fire                      0.6.0
+flax                      0.8.2
+fonttools                 4.50.0
+fqdn                      1.5.1
+fsspec                    2024.3.1
+future                    1.0.0
+gast                      0.5.2
+greenlet                  3.0.3
+h11                       0.14.0
+h5py                      3.10.0
+httpcore                  1.0.5
+httpx                     0.27.0
+ibm-cloud-sdk-core        3.19.2
+ibm-platform-services     0.53.2
+idna                      3.6
+importlib_metadata        7.1.0
+importlib_resources       6.4.0
+ipykernel                 6.29.4
+ipython                   8.23.0
+ipywidgets                8.1.2
+isoduration               20.11.0
+jax                       0.4.23
+jaxlib                    0.4.23
+jaxopt                    0.8.3
+jedi                      0.19.1
+Jinja2                    3.1.3
+joblib                    1.3.2
+json5                     0.9.24
+jsonpointer               2.4
+jsonschema                4.21.1
+jsonschema-specifications 2023.12.1
+jupyter                   1.0.0
+jupyter_client            8.6.1
+jupyter-console           6.6.3
+jupyter_core              5.7.2
+jupyter-events            0.10.0
+jupyter-lsp               2.2.4
+jupyter_server            2.13.0
+jupyter_server_terminals  0.5.3
+jupyterlab                4.1.5
+jupyterlab_pygments       0.3.0
+jupyterlab_server         2.25.4
+jupyterlab_widgets        3.0.10
+kiwisolver                1.4.5
+lark-parser               0.12.0
+llvmlite                  0.42.0
+locket                    1.0.0
+Mako                      1.3.2
+markdown-it-py            3.0.0
+MarkupSafe                2.1.5
+matplotlib                3.8.4
+matplotlib-inline         0.1.6
+mdurl                     0.1.2
+mistune                   3.0.2
+ml-dtypes                 0.4.0
+mpmath                    1.3.0
+msgpack                   1.0.8
+nbclient                  0.10.0
+nbconvert                 7.16.3
+nbformat                  5.10.4
+nest-asyncio              1.6.0
+networkx                  3.2.1
+notebook                  7.1.2
+notebook_shim             0.2.4
+numba                     0.59.1
+numpy                     1.26.4
+olefile                   0.46
+opt-einsum                3.3.0
+optax                     0.2.2
+optuna                    3.6.1
+orbax-checkpoint          0.5.9
+overrides                 7.7.0
+packaging                 24.0
+pandas                    2.2.1
+pandocfilters             1.5.1
+parso                     0.8.3
+partd                     1.4.1
+patsy                     0.5.6
+pbr                       6.0.0
+PennyLane                 0.35.1
+PennyLane-Catalyst        0.5.0
+PennyLane-Cirq            0.34.0
+PennyLane_Lightning       0.35.1
+PennyLane_Lightning_GPU   0.35.1
+PennyLane-qiskit          0.35.1
+PennyLane-SF              0.29.0
+pexpect                   4.9.0
+pillow                    10.3.0
+pip                       24.0
+platformdirs              4.2.0
+ply                       3.11
+prometheus_client         0.20.0
+prompt-toolkit            3.0.43
+protobuf                  5.26.1
+psutil                    5.9.8
+ptyprocess                0.7.0
+pure-eval                 0.2.2
+pycparser                 2.22
+pydantic                  2.6.4
+pydantic_core             2.16.3
+pydantic-settings         2.2.1
+pydot                     2.0.0
+Pygments                  2.17.2
+PyGObject                 3.42.1
+PyJWT                     2.8.0
+pylatexenc                2.10
+pyparsing                 3.1.2
+pyspnego                  0.10.2
+python-dateutil           2.9.0.post0
+python-dotenv             1.0.1
+python-json-logger        2.0.7
+pythran                   0.10.0
+pytz                      2024.1
+PyYAML                    6.0.1
+pyzmq                     25.1.2
+qiskit                    1.0.2
+qiskit-aer                0.14.0.1
+qiskit-algorithms         0.3.0
+qiskit-ibm-experiment     0.4.7
+qiskit-ibm-provider       0.10.0
+qiskit-ibm-runtime        0.20.0
+qiskit-machine-learning   0.7.2
+qml_benchmarks            0.1
+qtconsole                 5.5.1
+QtPy                      2.4.1
+quantum-blackbird         0.5.0
+quantum-xir               0.2.2
+referencing               0.34.0
+requests                  2.31.0
+requests-ntlm             1.2.0
+rfc3339-validator         0.1.4
+rfc3986-validator         0.1.1
+rich                      13.7.1
+rpds-py                   0.18.0
+ruamel.yaml               0.18.6
+ruamel.yaml.clib          0.2.8
+rustworkx                 0.14.2
+scikit-learn              1.4.1.post1
+scipy                     1.11.4
+seaborn                   0.13.2
+semantic-version          2.10.0
+Send2Trash                1.8.2
+setuptools                59.6.0
+six                       1.16.0
+sniffio                   1.3.1
+sortedcontainers          2.4.0
+soupsieve                 2.5
+SQLAlchemy                2.0.29
+ssh-import-id             5.11
+stack-data                0.6.3
+statsmodels               0.14.1
+stevedore                 5.2.0
+StrawberryFields          0.23.0
+symengine                 0.11.0
+sympy                     1.12
+tensorstore               0.1.56
+termcolor                 2.4.0
+terminado                 0.18.1
+thewalrus                 0.21.0
+threadpoolctl             3.4.0
+tinycss2                  1.2.1
+toml                      0.10.2
+tomli                     2.0.1
+tomlkit                   0.12.4
+toolz                     0.12.1
+torch                     2.2.2+cpu
+torchaudio                2.2.2+cpu
+torchvision               0.17.2+cpu
+tornado                   6.4
+tqdm                      4.66.2
+traitlets                 5.14.2
+types-python-dateutil     2.9.0.20240316
+typing_extensions         4.10.0
+tzdata                    2024.1
+uri-template              1.3.0
+urllib3                   2.2.1
+wcwidth                   0.2.13
+webcolors                 1.13
+webencodings              0.5.1
+websocket-client          1.7.0
+websockets                12.0
+wheel                     0.37.1
+widgetsnbextension        4.0.10
+xanadu-cloud-client       0.3.2
+zipp                      3.18.1
\ No newline at end of file
diff --git a/nersc/performance_indicators/perf_ind_iqpker_LR14_jax.py b/nersc/performance_indicators/perf_ind_iqpker_LR14_jax.py
deleted file mode 100644
index 9262450a..00000000
--- a/nersc/performance_indicators/perf_ind_iqpker_LR14_jax.py
+++ /dev/null
@@ -1,80 +0,0 @@
-import qml_benchmarks
-import pennylane as qml
-import jax
-import jax.numpy as jnp
-import numpy as np
-import time
-import csv
-import os
-import yaml
-
-from qml_benchmarks.models.projected_quantum_kernel import ProjectedQuantumKernel
-from qml_benchmarks.hyperparam_search_utils import read_data
-
-with open('hyperparam_settings.yaml', "r") as file:
-    hp_settings = yaml.safe_load(file)
-
-hyperparams = {**hp_settings['ProjectedQuantumKernel'], **{'use_jax':True, 'vmap':True, 'max_steps':100}}
-
-print(hyperparams)
-
-n_features = 14 #dataset dimension
-n_trials = 10 #number of trials to average over
-
-X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
-X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
-
-first_train_steps = []
-av_consec_train_steps = []
-predict_times = []
-
-for trial in range(n_trials):
-    jax.clear_caches()
-
-    model = QuantumMetricLearner(**hyperparams)
-    model.fit(X_train, y_train)
-
-    #get step times from loss history data
-    step_times = np.array([model.loss_history_[1][i + 1] - model.loss_history_[1][i]
-                           for i in range(len(model.loss_history_[1]) - 1)])
-    step_times = np.insert(step_times, 0, [model.loss_history_[1][0]])
-
-    #first train step
-    first_train_steps.append(step_times[0])
-    #consecutive (average) train step
-    av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
-    #prediction time
-    time0 = time.time()
-    model.predict(X_test)
-    predict_times.append(time.time() - time0)
-
-
-#calculate mean and stds
-
-first_train_step = np.mean(first_train_steps)
-first_train_step_std = np.std(first_train_steps)
-
-consec_train_step = np.mean(av_consec_train_steps)
-consec_train_step_std = np.std(av_consec_train_steps)
-
-predict_time = np.mean(predict_times)
-predict_time_std = np.std(predict_times)
-
-#write to csv
-data = [first_train_step, first_train_step_std, consec_train_step, consec_train_step_std, predict_time,
-        predict_time_std, hyperparams]
-
-model_name = model.__class__.__name__
-filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_JAX.csv"
-header = ['first_train_step', 'first_train_step_std', 'consec_train_step', 'consec_train_step_std', 'predict_time',
-          'predict_time_std', 'hyperparameters']
-
-if not os.path.exists('JAX'):
-    # Create the directory
-    os.mkdir('JAX')
-
-with open('JAX/'+filename, mode="w", newline="") as file:
-    writer = csv.writer(file)
-    writer.writerow(header)
-    for row in [data]:
-        writer.writerow(row)
diff --git a/nersc/performance_indicators/perf_ind_iqpker_LR15_jax.py b/nersc/performance_indicators/perf_ind_iqpker_LR15_jax.py
new file mode 100644
index 00000000..0321dd17
--- /dev/null
+++ b/nersc/performance_indicators/perf_ind_iqpker_LR15_jax.py
@@ -0,0 +1,54 @@
+import qml_benchmarks
+import pennylane as qml
+import jax
+import jax.numpy as jnp
+import numpy as np
+import time
+import csv
+import os
+import yaml
+
+from qml_benchmarks.models.iqp_kernel import IQPKernelClassifier
+from qml_benchmarks.hyperparam_search_utils import read_data
+
+with open('hyperparam_settings.yaml', "r") as file:
+    hp_settings = yaml.safe_load(file)
+
+hyperparams = {**hp_settings['IQPKernelClassifier'], **{'use_jax':True, 'vmap':True}}
+
+print(hyperparams)
+
+n_features = 15 #dataset dimension
+
+X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
+X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
+
+model = IQPKernelClassifier(**hyperparams)
+model.fit(X_train[:50], y_train[:50])
+
+#kernel construction time
+construct_kernel_time = model.construct_kernel_time_
+#full training time
+training_time = model.training_time_
+#prediction time
+time0 = time.time()
+model.predict(X_test[:50])
+predict_time = time.time() - time0
+
+
+#write to csv
+data = [construct_kernel_time, training_time, predict_time, hyperparams]
+
+model_name = model.__class__.__name__
+filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_JAX.csv"
+header = ['construct_kernel_time', 'training_time', 'predict_time', 'hyperparameters']
+
+if not os.path.exists('JAX'):
+    # Create the directory
+    os.mkdir('JAX')
+
+with open('JAX/'+filename, mode="w", newline="") as file:
+    writer = csv.writer(file)
+    writer.writerow(header)
+    for row in [data]:
+        writer.writerow(row)
diff --git a/nersc/performance_indicators/perf_ind_pqk_LR15_jax.py b/nersc/performance_indicators/perf_ind_pqk_LR15_jax.py
index 7a32b087..5facd45f 100644
--- a/nersc/performance_indicators/perf_ind_pqk_LR15_jax.py
+++ b/nersc/performance_indicators/perf_ind_pqk_LR15_jax.py
@@ -14,60 +14,34 @@
 with open('hyperparam_settings.yaml', "r") as file:
     hp_settings = yaml.safe_load(file)
 
-hyperparams = {**hp_settings['ProjectedQuantumKernel'], **{'use_jax':True, 'vmap':True, 'max_steps':100}}
+hyperparams = {**hp_settings['ProjectedQuantumKernel'], **{'use_jax':True, 'vmap':True}}
 
 print(hyperparams)
 
-n_features = 3 #dataset dimension
-n_trials = 1 #number of trials to average over
+n_features = 15 #dataset dimension
 
 X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
 X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
 
-first_train_steps = []
-av_consec_train_steps = []
-predict_times = []
+model = ProjectedQuantumKernel(**hyperparams)
+model.fit(X_train[:50], y_train[:50])
 
-for trial in range(n_trials):
-    jax.clear_caches()
+#kernel construction time
+construct_kernel_time = model.construct_kernel_time_
+#full training time
+training_time = model.training_time_
+#prediction time
+time0 = time.time()
+model.predict(X_test[:50])
+predict_time = time.time() - time0
 
-    model = QuantumMetricLearner(**hyperparams)
-    model.fit(X_train, y_train)
-
-    #get step times from loss history data
-    step_times = np.array([model.loss_history_[1][i + 1] - model.loss_history_[1][i]
-                           for i in range(len(model.loss_history_[1]) - 1)])
-    step_times = np.insert(step_times, 0, [model.loss_history_[1][0]])
-
-    #first train step
-    first_train_steps.append(step_times[0])
-    #consecutive (average) train step
-    av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
-    #prediction time
-    time0 = time.time()
-    model.predict(X_test)
-    predict_times.append(time.time() - time0)
-
-
-#calculate mean and stds
-
-first_train_step = np.mean(first_train_steps)
-first_train_step_std = np.std(first_train_steps)
-
-consec_train_step = np.mean(av_consec_train_steps)
-consec_train_step_std = np.std(av_consec_train_steps)
-
-predict_time = np.mean(predict_times)
-predict_time_std = np.std(predict_times)
 
 #write to csv
-data = [first_train_step, first_train_step_std, consec_train_step, consec_train_step_std, predict_time,
-        predict_time_std, hyperparams]
+data = [construct_kernel_time, training_time, predict_time, hyperparams]
 
 model_name = model.__class__.__name__
 filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_JAX.csv"
-header = ['first_train_step', 'first_train_step_std', 'consec_train_step', 'consec_train_step_std', 'predict_time',
-          'predict_time_std', 'hyperparameters']
+header = ['construct_kernel_time', 'training_time', 'predict_time', 'hyperparameters']
 
 if not os.path.exists('JAX'):
     # Create the directory
diff --git a/nersc/performance_indicators/perf_ind_pqk_LR17_jax.py b/nersc/performance_indicators/perf_ind_pqk_LR17_jax.py
index da5d6e8d..ffe55626 100644
--- a/nersc/performance_indicators/perf_ind_pqk_LR17_jax.py
+++ b/nersc/performance_indicators/perf_ind_pqk_LR17_jax.py
@@ -14,60 +14,34 @@
 with open('hyperparam_settings.yaml', "r") as file:
     hp_settings = yaml.safe_load(file)
 
-hyperparams = {**hp_settings['ProjectedQuantumKernel'], **{'use_jax':True, 'vmap':True, 'max_steps':100}}
+hyperparams = {**hp_settings['ProjectedQuantumKernel'], **{'use_jax':True, 'vmap':True}}
 
 print(hyperparams)
 
 n_features = 17 #dataset dimension
-n_trials = 10 #number of trials to average over
 
 X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
 X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
 
-first_train_steps = []
-av_consec_train_steps = []
-predict_times = []
+model = ProjectedQuantumKernel(**hyperparams)
+model.fit(X_train[:50], y_train[:50])
 
-for trial in range(n_trials):
-    jax.clear_caches()
+#kernel construction time
+construct_kernel_time = model.construct_kernel_time_
+#full training time
+training_time = model.training_time_
+#prediction time
+time0 = time.time()
+model.predict(X_test[:50])
+predict_time = time.time() - time0
 
-    model = QuantumMetricLearner(**hyperparams)
-    model.fit(X_train, y_train)
-
-    #get step times from loss history data
-    step_times = np.array([model.loss_history_[1][i + 1] - model.loss_history_[1][i]
-                           for i in range(len(model.loss_history_[1]) - 1)])
-    step_times = np.insert(step_times, 0, [model.loss_history_[1][0]])
-
-    #first train step
-    first_train_steps.append(step_times[0])
-    #consecutive (average) train step
-    av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
-    #prediction time
-    time0 = time.time()
-    model.predict(X_test)
-    predict_times.append(time.time() - time0)
-
-
-#calculate mean and stds
-
-first_train_step = np.mean(first_train_steps)
-first_train_step_std = np.std(first_train_steps)
-
-consec_train_step = np.mean(av_consec_train_steps)
-consec_train_step_std = np.std(av_consec_train_steps)
-
-predict_time = np.mean(predict_times)
-predict_time_std = np.std(predict_times)
 
 #write to csv
-data = [first_train_step, first_train_step_std, consec_train_step, consec_train_step_std, predict_time,
-        predict_time_std, hyperparams]
+data = [construct_kernel_time, training_time, predict_time, hyperparams]
 
 model_name = model.__class__.__name__
 filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_JAX.csv"
-header = ['first_train_step', 'first_train_step_std', 'consec_train_step', 'consec_train_step_std', 'predict_time',
-          'predict_time_std', 'hyperparameters']
+header = ['construct_kernel_time', 'training_time', 'predict_time', 'hyperparameters']
 
 if not os.path.exists('JAX'):
     # Create the directory
diff --git a/nersc/performance_indicators/perf_ind_pqk_LR20_jax.py b/nersc/performance_indicators/perf_ind_pqk_LR20_jax.py
index da5d6e8d..df2ba712 100644
--- a/nersc/performance_indicators/perf_ind_pqk_LR20_jax.py
+++ b/nersc/performance_indicators/perf_ind_pqk_LR20_jax.py
@@ -14,60 +14,34 @@
 with open('hyperparam_settings.yaml', "r") as file:
     hp_settings = yaml.safe_load(file)
 
-hyperparams = {**hp_settings['ProjectedQuantumKernel'], **{'use_jax':True, 'vmap':True, 'max_steps':100}}
+hyperparams = {**hp_settings['ProjectedQuantumKernel'], **{'use_jax':True, 'vmap':True}}
 
 print(hyperparams)
 
-n_features = 17 #dataset dimension
-n_trials = 10 #number of trials to average over
+n_features = 20 #dataset dimension
 
 X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
 X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
 
-first_train_steps = []
-av_consec_train_steps = []
-predict_times = []
+model = ProjectedQuantumKernel(**hyperparams)
+model.fit(X_train[:50], y_train[:50])
 
-for trial in range(n_trials):
-    jax.clear_caches()
+#kernel construction time
+construct_kernel_time = model.construct_kernel_time_
+#full training time
+training_time = model.training_time_
+#prediction time
+time0 = time.time()
+model.predict(X_test[:50])
+predict_time = time.time() - time0
 
-    model = QuantumMetricLearner(**hyperparams)
-    model.fit(X_train, y_train)
-
-    #get step times from loss history data
-    step_times = np.array([model.loss_history_[1][i + 1] - model.loss_history_[1][i]
-                           for i in range(len(model.loss_history_[1]) - 1)])
-    step_times = np.insert(step_times, 0, [model.loss_history_[1][0]])
-
-    #first train step
-    first_train_steps.append(step_times[0])
-    #consecutive (average) train step
-    av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
-    #prediction time
-    time0 = time.time()
-    model.predict(X_test)
-    predict_times.append(time.time() - time0)
-
-
-#calculate mean and stds
-
-first_train_step = np.mean(first_train_steps)
-first_train_step_std = np.std(first_train_steps)
-
-consec_train_step = np.mean(av_consec_train_steps)
-consec_train_step_std = np.std(av_consec_train_steps)
-
-predict_time = np.mean(predict_times)
-predict_time_std = np.std(predict_times)
 
 #write to csv
-data = [first_train_step, first_train_step_std, consec_train_step, consec_train_step_std, predict_time,
-        predict_time_std, hyperparams]
+data = [construct_kernel_time, training_time, predict_time, hyperparams]
 
 model_name = model.__class__.__name__
 filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_JAX.csv"
-header = ['first_train_step', 'first_train_step_std', 'consec_train_step', 'consec_train_step_std', 'predict_time',
-          'predict_time_std', 'hyperparameters']
+header = ['construct_kernel_time', 'training_time', 'predict_time', 'hyperparameters']
 
 if not os.path.exists('JAX'):
     # Create the directory

From f452d2b20dbaf4539dfed9a33f9618175fcc92f3 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Tue, 16 Apr 2024 10:36:44 +0200
Subject: [PATCH 032/100] add slurm code

---
 nersc/performance_indicators/submit_job.slr | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)
 create mode 100644 nersc/performance_indicators/submit_job.slr

diff --git a/nersc/performance_indicators/submit_job.slr b/nersc/performance_indicators/submit_job.slr
new file mode 100644
index 00000000..a9bc6e1d
--- /dev/null
+++ b/nersc/performance_indicators/submit_job.slr
@@ -0,0 +1,16 @@
+#!/bin/bash
+#SBATCH -N 1
+#SBATCH -C cpu
+#SBATCH -q shared
+#SBATCH -J perf_ind_iqpvar_15
+#SBATCH --account=jbowles
+#SBATCH -t 12:00:00
+#SBATCH --ntasks-per-node=1
+#SBATCH --cpus-per-task=16
+#SBATCH --output out/%j.log
+#SBATCH  --licenses=scratch
+
+source ../pm_podman.source
+bash ../installs.sh
+
+time srun -n 1  python3 ../perf_ind_iqpvar_LR15_jax.py

From d08162630f949355b8592f9cf90642083eb370f0 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Tue, 16 Apr 2024 10:38:32 +0200
Subject: [PATCH 033/100] add slurm code

---
 nersc/performance_indicators/submit_job.slr | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/nersc/performance_indicators/submit_job.slr b/nersc/performance_indicators/submit_job.slr
index a9bc6e1d..df68bb15 100644
--- a/nersc/performance_indicators/submit_job.slr
+++ b/nersc/performance_indicators/submit_job.slr
@@ -3,7 +3,7 @@
 #SBATCH -C cpu
 #SBATCH -q shared
 #SBATCH -J perf_ind_iqpvar_15
-#SBATCH --account=jbowles
+# SBATCH --account=jbowles
 #SBATCH -t 12:00:00
 #SBATCH --ntasks-per-node=1
 #SBATCH --cpus-per-task=16

From f513c67757f615567d03089e5462aa7e66bb9c89 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Tue, 16 Apr 2024 10:40:28 +0200
Subject: [PATCH 034/100] add slurm code

---
 nersc/performance_indicators/submit_job.slr | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/nersc/performance_indicators/submit_job.slr b/nersc/performance_indicators/submit_job.slr
index df68bb15..4984709b 100644
--- a/nersc/performance_indicators/submit_job.slr
+++ b/nersc/performance_indicators/submit_job.slr
@@ -4,7 +4,7 @@
 #SBATCH -q shared
 #SBATCH -J perf_ind_iqpvar_15
 # SBATCH --account=jbowles
-#SBATCH -t 12:00:00
+#SBATCH -t 1:00:00
 #SBATCH --ntasks-per-node=1
 #SBATCH --cpus-per-task=16
 #SBATCH --output out/%j.log

From 74967ef4786454c370fdcfbe4edcdc6fac1ebc59 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Tue, 16 Apr 2024 10:43:24 +0200
Subject: [PATCH 035/100] add slurm code

---
 nersc/performance_indicators/submit_job.slr | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/nersc/performance_indicators/submit_job.slr b/nersc/performance_indicators/submit_job.slr
index 4984709b..a17ac2c8 100644
--- a/nersc/performance_indicators/submit_job.slr
+++ b/nersc/performance_indicators/submit_job.slr
@@ -3,7 +3,7 @@
 #SBATCH -C cpu
 #SBATCH -q shared
 #SBATCH -J perf_ind_iqpvar_15
-# SBATCH --account=jbowles
+#SBATCH --account=nstaff
 #SBATCH -t 1:00:00
 #SBATCH --ntasks-per-node=1
 #SBATCH --cpus-per-task=16

From 9462980b072729d1452bc5d326a6e4d382b850d1 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Tue, 16 Apr 2024 10:44:04 +0200
Subject: [PATCH 036/100] add slurm code

---
 nersc/performance_indicators/submit_job.slr | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/nersc/performance_indicators/submit_job.slr b/nersc/performance_indicators/submit_job.slr
index a17ac2c8..4984709b 100644
--- a/nersc/performance_indicators/submit_job.slr
+++ b/nersc/performance_indicators/submit_job.slr
@@ -3,7 +3,7 @@
 #SBATCH -C cpu
 #SBATCH -q shared
 #SBATCH -J perf_ind_iqpvar_15
-#SBATCH --account=nstaff
+# SBATCH --account=jbowles
 #SBATCH -t 1:00:00
 #SBATCH --ntasks-per-node=1
 #SBATCH --cpus-per-task=16

From 07b49a956833ded8af7bbc30ddb7f2c1027ffe41 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Tue, 16 Apr 2024 15:38:14 +0200
Subject: [PATCH 037/100] update code

---
 ...eparable_3d_performance_indicators_JAX.csv |   2 +
 ...3d_performance_indicators_JAX_packages.txt | 166 +++++++++++
 ...3d_performance_indicators_JAX_scontrol.txt |   0
 ...sifier_linearly_separable_15d_JAX_info.txt | 265 ------------------
 ...parable_15d_performance_indicators_JAX.csv |   2 -
 ...sifier_linearly_separable_17d_JAX_info.txt | 240 ----------------
 ...eparable_3d_performance_indicators_JAX.csv |   2 +
 ...3d_performance_indicators_JAX_packages.txt | 166 +++++++++++
 ...3d_performance_indicators_JAX_scontrol.txt |   0
 ...earner_linearly_separable_14d_JAX_info.txt |   4 -
 .../hyperparam_settings.yaml                  |   2 -
 .../perf_ind_iqpker_LR15_jax.py               |  54 +++-
 .../perf_ind_iqpvar_LR15_jax.py               |  55 +++-
 .../perf_ind_iqpvar_LR17_jax.py               |  80 ------
 .../perf_ind_iqpvar_LR20_jax.py               |  80 ------
 .../perf_ind_pqk_LR15_jax.py                  |  54 ----
 .../perf_ind_pqk_LR17_jax.py                  |  54 ----
 .../perf_ind_pqk_LR20_jax.py                  |  54 ----
 .../perf_ind_qmetric_LR14_jax.py              |  80 ------
 .../perf_ind_qmetric_LR15_jax.py              |  80 ------
 .../save_last_job_info.sh                     |  21 --
 nersc/performance_indicators/submit_job.slr   |   2 +-
 22 files changed, 418 insertions(+), 1045 deletions(-)
 create mode 100644 nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_3d_performance_indicators_JAX.csv
 create mode 100644 nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_3d_performance_indicators_JAX_packages.txt
 create mode 100644 nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_3d_performance_indicators_JAX_scontrol.txt
 delete mode 100644 nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_15d_JAX_info.txt
 delete mode 100644 nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_15d_performance_indicators_JAX.csv
 delete mode 100644 nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_17d_JAX_info.txt
 create mode 100644 nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_3d_performance_indicators_JAX.csv
 create mode 100644 nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_3d_performance_indicators_JAX_packages.txt
 create mode 100644 nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_3d_performance_indicators_JAX_scontrol.txt
 delete mode 100644 nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_14d_JAX_info.txt
 delete mode 100644 nersc/performance_indicators/perf_ind_iqpvar_LR17_jax.py
 delete mode 100644 nersc/performance_indicators/perf_ind_iqpvar_LR20_jax.py
 delete mode 100644 nersc/performance_indicators/perf_ind_pqk_LR15_jax.py
 delete mode 100644 nersc/performance_indicators/perf_ind_pqk_LR17_jax.py
 delete mode 100644 nersc/performance_indicators/perf_ind_pqk_LR20_jax.py
 delete mode 100644 nersc/performance_indicators/perf_ind_qmetric_LR14_jax.py
 delete mode 100644 nersc/performance_indicators/perf_ind_qmetric_LR15_jax.py
 delete mode 100644 nersc/performance_indicators/save_last_job_info.sh

diff --git a/nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_3d_performance_indicators_JAX.csv b/nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_3d_performance_indicators_JAX.csv
new file mode 100644
index 00000000..d0cc1ae5
--- /dev/null
+++ b/nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_3d_performance_indicators_JAX.csv
@@ -0,0 +1,2 @@
+construct_kernel_time,training_time,predict_time,hyperparameters
+1.1291069984436035,1.1299769878387451,1.092557668685913,"{'repeats': 10, 'use_jax': True, 'vmap': True, 'jit': True}"
diff --git a/nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_3d_performance_indicators_JAX_packages.txt b/nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_3d_performance_indicators_JAX_packages.txt
new file mode 100644
index 00000000..8859d638
--- /dev/null
+++ b/nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_3d_performance_indicators_JAX_packages.txt
@@ -0,0 +1,166 @@
+Package                   Version      Editable project location
+------------------------- ------------ ----------------------------------------------------------
+absl-py                   2.0.0
+anyio                     4.2.0
+appdirs                   1.4.4
+appnope                   0.1.3
+argon2-cffi               23.1.0
+argon2-cffi-bindings      21.2.0
+arrow                     1.3.0
+asttokens                 2.4.1
+astunparse                1.6.3
+async-lru                 2.0.4
+attrs                     23.1.0
+autograd                  1.6.2
+autoray                   0.6.7
+Babel                     2.14.0
+beautifulsoup4            4.12.2
+bleach                    6.1.0
+cachetools                5.3.2
+certifi                   2023.11.17
+cffi                      1.16.0
+charset-normalizer        3.3.2
+chex                      0.1.86
+comm                      0.2.0
+contourpy                 1.2.0
+cycler                    0.12.1
+debugpy                   1.8.0
+decorator                 5.1.1
+defusedxml                0.7.1
+diastatic-malt            2.15.1
+entrypoints               0.4
+etils                     1.5.2
+exceptiongroup            1.2.0
+executing                 2.0.1
+fastjsonschema            2.19.0
+filelock                  3.13.1
+flax                      0.8.0
+fonttools                 4.46.0
+fqdn                      1.5.1
+fsspec                    2023.12.2
+future                    0.18.3
+gast                      0.5.4
+idna                      3.6
+importlib-metadata        7.0.0
+importlib-resources       6.1.1
+ipykernel                 6.27.1
+ipython                   8.18.1
+ipython-genutils          0.2.0
+isoduration               20.11.0
+jax                       0.4.23
+jaxlib                    0.4.23
+jedi                      0.19.1
+Jinja2                    3.1.2
+joblib                    1.3.2
+json5                     0.9.14
+jsonpointer               2.4
+jsonschema                4.20.0
+jsonschema-specifications 2023.11.2
+jupyter_client            7.4.9
+jupyter_core              5.5.0
+jupyter-events            0.9.0
+jupyter-lsp               2.2.1
+jupyter_server            2.12.1
+jupyter_server_terminals  0.5.0
+jupyterlab                4.0.9
+jupyterlab_pygments       0.3.0
+jupyterlab_server         2.25.2
+kiwisolver                1.4.5
+llvmlite                  0.42.0
+markdown-it-py            3.0.0
+MarkupSafe                2.1.3
+matplotlib                3.8.2
+matplotlib-inline         0.1.6
+mdurl                     0.1.2
+mistune                   3.0.2
+ml-dtypes                 0.3.1
+mpmath                    1.3.0
+msgpack                   1.0.7
+multipledispatch          1.0.0
+nbclassic                 1.0.0
+nbclient                  0.9.0
+nbconvert                 7.12.0
+nbformat                  5.9.2
+nest-asyncio              1.5.8
+networkx                  3.2.1
+notebook                  6.5.6
+notebook_shim             0.2.3
+numba                     0.59.1
+numpy                     1.26.2
+numpyro                   0.14.0
+opt-einsum                3.3.0
+optax                     0.2.2
+orbax-checkpoint          0.4.8
+overrides                 7.4.0
+packaging                 23.2
+pandas                    2.1.4
+pandocfilters             1.5.0
+parso                     0.8.3
+PennyLane                 0.35.1
+PennyLane-Catalyst        0.5.0
+PennyLane_Lightning       0.35.1
+pexpect                   4.9.0
+Pillow                    10.1.0
+pip                       24.0
+platformdirs              4.1.0
+prometheus-client         0.19.0
+prompt-toolkit            3.0.43
+protobuf                  4.25.1
+psutil                    5.9.7
+ptyprocess                0.7.0
+pure-eval                 0.2.2
+pycparser                 2.21
+Pygments                  2.17.2
+pynndescent               0.5.12
+pyparsing                 3.1.1
+python-dateutil           2.8.2
+python-json-logger        2.0.7
+pytz                      2023.3.post1
+PyYAML                    6.0.1
+pyzmq                     24.0.1
+qgml                      2024.3.0
+qml-benchmarks            0.1          /Users/joseph/PycharmProjects/pennylane_jax/qml-benchmarks
+referencing               0.32.0
+requests                  2.31.0
+rfc3339-validator         0.1.4
+rfc3986-validator         0.1.1
+rich                      13.7.0
+rpds-py                   0.15.2
+rustworkx                 0.13.2
+scikit-learn              1.4.0
+scipy                     1.11.4
+seaborn                   0.13.0
+semantic-version          2.10.0
+Send2Trash                1.8.2
+setuptools                57.0.0
+six                       1.16.0
+sniffio                   1.3.0
+soupsieve                 2.5
+stack-data                0.6.3
+sympy                     1.12
+tensorstore               0.1.51
+termcolor                 2.4.0
+terminado                 0.18.0
+threadpoolctl             3.2.0
+tinycss2                  1.2.1
+toml                      0.10.2
+tomli                     2.0.1
+tomlkit                   0.12.3
+toolz                     0.12.0
+torch                     2.1.2
+torchvision               0.16.2
+tornado                   6.4
+tqdm                      4.66.2
+traitlets                 5.14.0
+types-python-dateutil     2.8.19.14
+typing_extensions         4.9.0
+tzdata                    2023.3
+umap-learn                0.5.6
+uri-template              1.3.0
+urllib3                   2.1.0
+wcwidth                   0.2.12
+webcolors                 1.13
+webencodings              0.5.1
+websocket-client          1.7.0
+wheel                     0.36.2
+zipp                      3.17.0
diff --git a/nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_3d_performance_indicators_JAX_scontrol.txt b/nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_3d_performance_indicators_JAX_scontrol.txt
new file mode 100644
index 00000000..e69de29b
diff --git a/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_15d_JAX_info.txt b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_15d_JAX_info.txt
deleted file mode 100644
index 3fb0afc5..00000000
--- a/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_15d_JAX_info.txt
+++ /dev/null
@@ -1,265 +0,0 @@
-salloc -q shared_interactive -C cpu -t 4:00:00 --cpus-per-task=16
-
-JobId=24420449 JobName=interactive
-UserId=jbowles(104993) GroupId=jbowles(104993) MCS_label=N/A
-Priority=73320 Nice=0 Account=m4139 QOS=shared_interactive
-JobState=COMPLETING Reason=NonZeroExitCode Dependency=(null)
-Requeue=0 Restarts=0 BatchFlag=0 Reboot=0 ExitCode=127:0
-RunTime=00:32:07 TimeLimit=04:00:00 TimeMin=N/A
-SubmitTime=2024-04-15T08:02:57 EligibleTime=2024-04-15T08:02:57
-AccrueTime=2024-04-15T08:03:06
-StartTime=2024-04-15T08:03:06 EndTime=2024-04-15T08:35:13 Deadline=N/A
-PreemptEligibleTime=2024-04-15T08:03:06 PreemptTime=None
-SuspendTime=None SecsPreSuspend=0 LastSchedEval=2024-04-15T08:03:06 Scheduler=Main
-Partition=shared_urgent_milan_ss11 AllocNode:Sid=login24:836159
-ReqNodeList=(null) ExcNodeList=(null)
-NodeList=nid004092
-BatchHost=nid004092
-NumNodes=1 NumCPUs=16 NumTasks=1 CPUs/Task=16 ReqB:S:C:T=0:0:*:*
-ReqTRES=cpu=16,mem=30624M,node=1,billing=16
-AllocTRES=cpu=16,mem=30624M,node=1,billing=16
-Socks/Node=* NtasksPerN:B:S:C=0:0:*:* CoreSpec=*
-MinCPUsNode=16 MinMemoryCPU=1914M MinTmpDiskNode=0
-Features=cpu DelayBoot=00:00:00
-OverSubscribe=OK Contiguous=0 Licenses=scratch:1 Network=(null)
-Command=/bin/sh
-WorkDir=/pscratch/sd/j/jbowles/qml-benchmarks/nersc
-AdminComment={"qos":"shared_interactive","arrayTaskId":4294967294,"resizing":0,"originCluster":"perlmutter","batchHost":"nid004092","jobExitCode":32512,"uid":104993,"cluster":"perlmutter","tresRequest":"1=16,2=30624,3=18446744073709551614,4=1,5=16","jobId":24420449,"gresRequest":"cpu=16,mem=30624M,node=1,billing=16","arrayJobId":0,"workingDirectory":"\/pscratch\/sd\/j\/jbowles\/qml-benchmarks\/nersc","licenses":"scratch:1","name":"interactive","partition":"shared_urgent_milan_ss11","packJobId":0,"reboot":0,"priority":73320,"features":"cpu","restartCnt":0,"argv":["\/bin\/sh","-c","srun --interactive --preserve-env --pty $SHELL"],"allocCpus":16,"startTime":1713193386,"packJobOffset":0,"jobAccount":"m4139","nodes":"nid004092","submitTime":1713193377,"timeLimit":240,"jobDerivedExitCode":0,"endTime":1713195313,"allocNodes":1}
-Power=
-
-Package                   Version
-------------------------- --------------
-absl-py                   2.1.0
-alembic                   1.13.1
-annotated-types           0.6.0
-antlr4-python3-runtime    4.9.2
-anyio                     4.3.0
-appdirs                   1.4.4
-argon2-cffi               23.1.0
-argon2-cffi-bindings      21.2.0
-arrow                     1.3.0
-asttokens                 2.4.1
-astunparse                1.6.3
-async-lru                 2.0.4
-attrs                     23.2.0
-autograd                  1.6.2
-autoray                   0.6.9
-Babel                     2.14.0
-beautifulsoup4            4.12.3
-beniget                   0.4.1
-bitstring                 3.1.7
-bleach                    6.1.0
-cachetools                5.3.3
-certifi                   2024.2.2
-cffi                      1.16.0
-charset-normalizer        3.3.2
-chex                      0.1.86
-cirq-core                 1.3.0
-cirq-pasqal               1.3.0
-click                     8.1.7
-cloudpickle               3.0.0
-colorlog                  6.8.2
-comm                      0.2.2
-contourpy                 1.2.1
-cryptography              42.0.5
-cycler                    0.12.1
-dask                      2024.4.1
-dbus-python               1.2.18
-debugpy                   1.8.1
-decorator                 4.4.2
-defusedxml                0.7.1
-diastatic-malt            2.15.1
-dill                      0.3.8
-distro                    1.7.0
-duet                      0.2.9
-etils                     1.7.0
-exceptiongroup            1.2.0
-executing                 2.0.1
-fastdtw                   0.3.4
-fastjsonschema            2.19.1
-filelock                  3.9.0
-fire                      0.6.0
-flax                      0.8.2
-fonttools                 4.50.0
-fqdn                      1.5.1
-fsspec                    2024.3.1
-future                    1.0.0
-gast                      0.5.2
-greenlet                  3.0.3
-h11                       0.14.0
-h5py                      3.10.0
-httpcore                  1.0.5
-httpx                     0.27.0
-ibm-cloud-sdk-core        3.19.2
-ibm-platform-services     0.53.2
-idna                      3.6
-importlib_metadata        7.1.0
-importlib_resources       6.4.0
-ipykernel                 6.29.4
-ipython                   8.23.0
-ipywidgets                8.1.2
-isoduration               20.11.0
-jax                       0.4.23
-jaxlib                    0.4.23
-jaxopt                    0.8.3
-jedi                      0.19.1
-Jinja2                    3.1.3
-joblib                    1.3.2
-json5                     0.9.24
-jsonpointer               2.4
-jsonschema                4.21.1
-jsonschema-specifications 2023.12.1
-jupyter                   1.0.0
-jupyter_client            8.6.1
-jupyter-console           6.6.3
-jupyter_core              5.7.2
-jupyter-events            0.10.0
-jupyter-lsp               2.2.4
-jupyter_server            2.13.0
-jupyter_server_terminals  0.5.3
-jupyterlab                4.1.5
-jupyterlab_pygments       0.3.0
-jupyterlab_server         2.25.4
-jupyterlab_widgets        3.0.10
-kiwisolver                1.4.5
-lark-parser               0.12.0
-llvmlite                  0.42.0
-locket                    1.0.0
-Mako                      1.3.2
-markdown-it-py            3.0.0
-MarkupSafe                2.1.5
-matplotlib                3.8.4
-matplotlib-inline         0.1.6
-mdurl                     0.1.2
-mistune                   3.0.2
-ml-dtypes                 0.4.0
-mpmath                    1.3.0
-msgpack                   1.0.8
-nbclient                  0.10.0
-nbconvert                 7.16.3
-nbformat                  5.10.4
-nest-asyncio              1.6.0
-networkx                  3.2.1
-notebook                  7.1.2
-notebook_shim             0.2.4
-numba                     0.59.1
-numpy                     1.26.4
-olefile                   0.46
-opt-einsum                3.3.0
-optax                     0.2.2
-optuna                    3.6.1
-orbax-checkpoint          0.5.9
-overrides                 7.7.0
-packaging                 24.0
-pandas                    2.2.1
-pandocfilters             1.5.1
-parso                     0.8.3
-partd                     1.4.1
-patsy                     0.5.6
-pbr                       6.0.0
-PennyLane                 0.35.1
-PennyLane-Catalyst        0.5.0
-PennyLane-Cirq            0.34.0
-PennyLane_Lightning       0.35.1
-PennyLane_Lightning_GPU   0.35.1
-PennyLane-qiskit          0.35.1
-PennyLane-SF              0.29.0
-pexpect                   4.9.0
-pillow                    10.3.0
-pip                       24.0
-platformdirs              4.2.0
-ply                       3.11
-prometheus_client         0.20.0
-prompt-toolkit            3.0.43
-protobuf                  5.26.1
-psutil                    5.9.8
-ptyprocess                0.7.0
-pure-eval                 0.2.2
-pycparser                 2.22
-pydantic                  2.6.4
-pydantic_core             2.16.3
-pydantic-settings         2.2.1
-pydot                     2.0.0
-Pygments                  2.17.2
-PyGObject                 3.42.1
-PyJWT                     2.8.0
-pylatexenc                2.10
-pyparsing                 3.1.2
-pyspnego                  0.10.2
-python-dateutil           2.9.0.post0
-python-dotenv             1.0.1
-python-json-logger        2.0.7
-pythran                   0.10.0
-pytz                      2024.1
-PyYAML                    6.0.1
-pyzmq                     25.1.2
-qiskit                    1.0.2
-qiskit-aer                0.14.0.1
-qiskit-algorithms         0.3.0
-qiskit-ibm-experiment     0.4.7
-qiskit-ibm-provider       0.10.0
-qiskit-ibm-runtime        0.20.0
-qiskit-machine-learning   0.7.2
-qml_benchmarks            0.1
-qtconsole                 5.5.1
-QtPy                      2.4.1
-quantum-blackbird         0.5.0
-quantum-xir               0.2.2
-referencing               0.34.0
-requests                  2.31.0
-requests-ntlm             1.2.0
-rfc3339-validator         0.1.4
-rfc3986-validator         0.1.1
-rich                      13.7.1
-rpds-py                   0.18.0
-ruamel.yaml               0.18.6
-ruamel.yaml.clib          0.2.8
-rustworkx                 0.14.2
-scikit-learn              1.4.1.post1
-scipy                     1.11.4
-seaborn                   0.13.2
-semantic-version          2.10.0
-Send2Trash                1.8.2
-setuptools                59.6.0
-six                       1.16.0
-sniffio                   1.3.1
-sortedcontainers          2.4.0
-soupsieve                 2.5
-SQLAlchemy                2.0.29
-ssh-import-id             5.11
-stack-data                0.6.3
-statsmodels               0.14.1
-stevedore                 5.2.0
-StrawberryFields          0.23.0
-symengine                 0.11.0
-sympy                     1.12
-tensorstore               0.1.56
-termcolor                 2.4.0
-terminado                 0.18.1
-thewalrus                 0.21.0
-threadpoolctl             3.4.0
-tinycss2                  1.2.1
-toml                      0.10.2
-tomli                     2.0.1
-tomlkit                   0.12.4
-toolz                     0.12.1
-torch                     2.2.2+cpu
-torchaudio                2.2.2+cpu
-torchvision               0.17.2+cpu
-tornado                   6.4
-tqdm                      4.66.2
-traitlets                 5.14.2
-types-python-dateutil     2.9.0.20240316
-typing_extensions         4.10.0
-tzdata                    2024.1
-uri-template              1.3.0
-urllib3                   2.2.1
-wcwidth                   0.2.13
-webcolors                 1.13
-webencodings              0.5.1
-websocket-client          1.7.0
-websockets                12.0
-wheel                     0.37.1
-widgetsnbextension        4.0.10
-xanadu-cloud-client       0.3.2
-zipp                      3.18.1
\ No newline at end of file
diff --git a/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_15d_performance_indicators_JAX.csv b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_15d_performance_indicators_JAX.csv
deleted file mode 100644
index 88c2dcb8..00000000
--- a/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_15d_performance_indicators_JAX.csv
+++ /dev/null
@@ -1,2 +0,0 @@
-first_train_step,first_train_step_std,consec_train_step,consec_train_step_std,predict_time,predict_time_std,hyperparameters
-148.14196796417235,11.108692798642224,10.153086731168958,1.4567056279882924,49.51927604675293,2.8201209235990614,"{'learning_rate': 0.01, 'n_layers': 15, 'repeats': 10, 'use_jax': True, 'vmap': True, 'max_steps': 10}"
diff --git a/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_17d_JAX_info.txt b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_17d_JAX_info.txt
deleted file mode 100644
index 9e3568ca..00000000
--- a/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_17d_JAX_info.txt
+++ /dev/null
@@ -1,240 +0,0 @@
-salloc -q shared_interactive -C cpu -t 4:00:00 --cpus-per-task=16
-
-ADD
-
-Package                   Version
-------------------------- --------------
-absl-py                   2.1.0
-alembic                   1.13.1
-annotated-types           0.6.0
-antlr4-python3-runtime    4.9.2
-anyio                     4.3.0
-appdirs                   1.4.4
-argon2-cffi               23.1.0
-argon2-cffi-bindings      21.2.0
-arrow                     1.3.0
-asttokens                 2.4.1
-astunparse                1.6.3
-async-lru                 2.0.4
-attrs                     23.2.0
-autograd                  1.6.2
-autoray                   0.6.9
-Babel                     2.14.0
-beautifulsoup4            4.12.3
-beniget                   0.4.1
-bitstring                 3.1.7
-bleach                    6.1.0
-cachetools                5.3.3
-certifi                   2024.2.2
-cffi                      1.16.0
-charset-normalizer        3.3.2
-chex                      0.1.86
-cirq-core                 1.3.0
-cirq-pasqal               1.3.0
-click                     8.1.7
-cloudpickle               3.0.0
-colorlog                  6.8.2
-comm                      0.2.2
-contourpy                 1.2.1
-cryptography              42.0.5
-cycler                    0.12.1
-dask                      2024.4.1
-dbus-python               1.2.18
-debugpy                   1.8.1
-decorator                 4.4.2
-defusedxml                0.7.1
-diastatic-malt            2.15.1
-dill                      0.3.8
-distro                    1.7.0
-duet                      0.2.9
-etils                     1.7.0
-exceptiongroup            1.2.0
-executing                 2.0.1
-fastdtw                   0.3.4
-fastjsonschema            2.19.1
-filelock                  3.9.0
-fire                      0.6.0
-flax                      0.8.2
-fonttools                 4.50.0
-fqdn                      1.5.1
-fsspec                    2024.3.1
-future                    1.0.0
-gast                      0.5.2
-greenlet                  3.0.3
-h11                       0.14.0
-h5py                      3.10.0
-httpcore                  1.0.5
-httpx                     0.27.0
-ibm-cloud-sdk-core        3.19.2
-ibm-platform-services     0.53.2
-idna                      3.6
-importlib_metadata        7.1.0
-importlib_resources       6.4.0
-ipykernel                 6.29.4
-ipython                   8.23.0
-ipywidgets                8.1.2
-isoduration               20.11.0
-jax                       0.4.23
-jaxlib                    0.4.23
-jaxopt                    0.8.3
-jedi                      0.19.1
-Jinja2                    3.1.3
-joblib                    1.3.2
-json5                     0.9.24
-jsonpointer               2.4
-jsonschema                4.21.1
-jsonschema-specifications 2023.12.1
-jupyter                   1.0.0
-jupyter_client            8.6.1
-jupyter-console           6.6.3
-jupyter_core              5.7.2
-jupyter-events            0.10.0
-jupyter-lsp               2.2.4
-jupyter_server            2.13.0
-jupyter_server_terminals  0.5.3
-jupyterlab                4.1.5
-jupyterlab_pygments       0.3.0
-jupyterlab_server         2.25.4
-jupyterlab_widgets        3.0.10
-kiwisolver                1.4.5
-lark-parser               0.12.0
-llvmlite                  0.42.0
-locket                    1.0.0
-Mako                      1.3.2
-markdown-it-py            3.0.0
-MarkupSafe                2.1.5
-matplotlib                3.8.4
-matplotlib-inline         0.1.6
-mdurl                     0.1.2
-mistune                   3.0.2
-ml-dtypes                 0.4.0
-mpmath                    1.3.0
-msgpack                   1.0.8
-nbclient                  0.10.0
-nbconvert                 7.16.3
-nbformat                  5.10.4
-nest-asyncio              1.6.0
-networkx                  3.2.1
-notebook                  7.1.2
-notebook_shim             0.2.4
-numba                     0.59.1
-numpy                     1.26.4
-olefile                   0.46
-opt-einsum                3.3.0
-optax                     0.2.2
-optuna                    3.6.1
-orbax-checkpoint          0.5.9
-overrides                 7.7.0
-packaging                 24.0
-pandas                    2.2.1
-pandocfilters             1.5.1
-parso                     0.8.3
-partd                     1.4.1
-patsy                     0.5.6
-pbr                       6.0.0
-PennyLane                 0.35.1
-PennyLane-Catalyst        0.5.0
-PennyLane-Cirq            0.34.0
-PennyLane_Lightning       0.35.1
-PennyLane_Lightning_GPU   0.35.1
-PennyLane-qiskit          0.35.1
-PennyLane-SF              0.29.0
-pexpect                   4.9.0
-pillow                    10.3.0
-pip                       24.0
-platformdirs              4.2.0
-ply                       3.11
-prometheus_client         0.20.0
-prompt-toolkit            3.0.43
-protobuf                  5.26.1
-psutil                    5.9.8
-ptyprocess                0.7.0
-pure-eval                 0.2.2
-pycparser                 2.22
-pydantic                  2.6.4
-pydantic_core             2.16.3
-pydantic-settings         2.2.1
-pydot                     2.0.0
-Pygments                  2.17.2
-PyGObject                 3.42.1
-PyJWT                     2.8.0
-pylatexenc                2.10
-pyparsing                 3.1.2
-pyspnego                  0.10.2
-python-dateutil           2.9.0.post0
-python-dotenv             1.0.1
-python-json-logger        2.0.7
-pythran                   0.10.0
-pytz                      2024.1
-PyYAML                    6.0.1
-pyzmq                     25.1.2
-qiskit                    1.0.2
-qiskit-aer                0.14.0.1
-qiskit-algorithms         0.3.0
-qiskit-ibm-experiment     0.4.7
-qiskit-ibm-provider       0.10.0
-qiskit-ibm-runtime        0.20.0
-qiskit-machine-learning   0.7.2
-qml_benchmarks            0.1
-qtconsole                 5.5.1
-QtPy                      2.4.1
-quantum-blackbird         0.5.0
-quantum-xir               0.2.2
-referencing               0.34.0
-requests                  2.31.0
-requests-ntlm             1.2.0
-rfc3339-validator         0.1.4
-rfc3986-validator         0.1.1
-rich                      13.7.1
-rpds-py                   0.18.0
-ruamel.yaml               0.18.6
-ruamel.yaml.clib          0.2.8
-rustworkx                 0.14.2
-scikit-learn              1.4.1.post1
-scipy                     1.11.4
-seaborn                   0.13.2
-semantic-version          2.10.0
-Send2Trash                1.8.2
-setuptools                59.6.0
-six                       1.16.0
-sniffio                   1.3.1
-sortedcontainers          2.4.0
-soupsieve                 2.5
-SQLAlchemy                2.0.29
-ssh-import-id             5.11
-stack-data                0.6.3
-statsmodels               0.14.1
-stevedore                 5.2.0
-StrawberryFields          0.23.0
-symengine                 0.11.0
-sympy                     1.12
-tensorstore               0.1.56
-termcolor                 2.4.0
-terminado                 0.18.1
-thewalrus                 0.21.0
-threadpoolctl             3.4.0
-tinycss2                  1.2.1
-toml                      0.10.2
-tomli                     2.0.1
-tomlkit                   0.12.4
-toolz                     0.12.1
-torch                     2.2.2+cpu
-torchaudio                2.2.2+cpu
-torchvision               0.17.2+cpu
-tornado                   6.4
-tqdm                      4.66.2
-traitlets                 5.14.2
-types-python-dateutil     2.9.0.20240316
-typing_extensions         4.10.0
-tzdata                    2024.1
-uri-template              1.3.0
-urllib3                   2.2.1
-wcwidth                   0.2.13
-webcolors                 1.13
-webencodings              0.5.1
-websocket-client          1.7.0
-websockets                12.0
-wheel                     0.37.1
-widgetsnbextension        4.0.10
-xanadu-cloud-client       0.3.2
-zipp                      3.18.1
\ No newline at end of file
diff --git a/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_3d_performance_indicators_JAX.csv b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_3d_performance_indicators_JAX.csv
new file mode 100644
index 00000000..7db3aa8d
--- /dev/null
+++ b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_3d_performance_indicators_JAX.csv
@@ -0,0 +1,2 @@
+first_train_step,first_train_step_std,consec_train_step,consec_train_step_std,predict_time,predict_time_std,hyperparameters
+7.222759056091308,0.302873401811286,0.002020430564880371,0.00016703330426884848,1.5872551918029785,0.07579271979049444,"{'learning_rate': 0.01, 'n_layers': 15, 'repeats': 10, 'use_jax': True, 'vmap': True, 'max_steps': 5, 'jit': True}"
diff --git a/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_3d_performance_indicators_JAX_packages.txt b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_3d_performance_indicators_JAX_packages.txt
new file mode 100644
index 00000000..8859d638
--- /dev/null
+++ b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_3d_performance_indicators_JAX_packages.txt
@@ -0,0 +1,166 @@
+Package                   Version      Editable project location
+------------------------- ------------ ----------------------------------------------------------
+absl-py                   2.0.0
+anyio                     4.2.0
+appdirs                   1.4.4
+appnope                   0.1.3
+argon2-cffi               23.1.0
+argon2-cffi-bindings      21.2.0
+arrow                     1.3.0
+asttokens                 2.4.1
+astunparse                1.6.3
+async-lru                 2.0.4
+attrs                     23.1.0
+autograd                  1.6.2
+autoray                   0.6.7
+Babel                     2.14.0
+beautifulsoup4            4.12.2
+bleach                    6.1.0
+cachetools                5.3.2
+certifi                   2023.11.17
+cffi                      1.16.0
+charset-normalizer        3.3.2
+chex                      0.1.86
+comm                      0.2.0
+contourpy                 1.2.0
+cycler                    0.12.1
+debugpy                   1.8.0
+decorator                 5.1.1
+defusedxml                0.7.1
+diastatic-malt            2.15.1
+entrypoints               0.4
+etils                     1.5.2
+exceptiongroup            1.2.0
+executing                 2.0.1
+fastjsonschema            2.19.0
+filelock                  3.13.1
+flax                      0.8.0
+fonttools                 4.46.0
+fqdn                      1.5.1
+fsspec                    2023.12.2
+future                    0.18.3
+gast                      0.5.4
+idna                      3.6
+importlib-metadata        7.0.0
+importlib-resources       6.1.1
+ipykernel                 6.27.1
+ipython                   8.18.1
+ipython-genutils          0.2.0
+isoduration               20.11.0
+jax                       0.4.23
+jaxlib                    0.4.23
+jedi                      0.19.1
+Jinja2                    3.1.2
+joblib                    1.3.2
+json5                     0.9.14
+jsonpointer               2.4
+jsonschema                4.20.0
+jsonschema-specifications 2023.11.2
+jupyter_client            7.4.9
+jupyter_core              5.5.0
+jupyter-events            0.9.0
+jupyter-lsp               2.2.1
+jupyter_server            2.12.1
+jupyter_server_terminals  0.5.0
+jupyterlab                4.0.9
+jupyterlab_pygments       0.3.0
+jupyterlab_server         2.25.2
+kiwisolver                1.4.5
+llvmlite                  0.42.0
+markdown-it-py            3.0.0
+MarkupSafe                2.1.3
+matplotlib                3.8.2
+matplotlib-inline         0.1.6
+mdurl                     0.1.2
+mistune                   3.0.2
+ml-dtypes                 0.3.1
+mpmath                    1.3.0
+msgpack                   1.0.7
+multipledispatch          1.0.0
+nbclassic                 1.0.0
+nbclient                  0.9.0
+nbconvert                 7.12.0
+nbformat                  5.9.2
+nest-asyncio              1.5.8
+networkx                  3.2.1
+notebook                  6.5.6
+notebook_shim             0.2.3
+numba                     0.59.1
+numpy                     1.26.2
+numpyro                   0.14.0
+opt-einsum                3.3.0
+optax                     0.2.2
+orbax-checkpoint          0.4.8
+overrides                 7.4.0
+packaging                 23.2
+pandas                    2.1.4
+pandocfilters             1.5.0
+parso                     0.8.3
+PennyLane                 0.35.1
+PennyLane-Catalyst        0.5.0
+PennyLane_Lightning       0.35.1
+pexpect                   4.9.0
+Pillow                    10.1.0
+pip                       24.0
+platformdirs              4.1.0
+prometheus-client         0.19.0
+prompt-toolkit            3.0.43
+protobuf                  4.25.1
+psutil                    5.9.7
+ptyprocess                0.7.0
+pure-eval                 0.2.2
+pycparser                 2.21
+Pygments                  2.17.2
+pynndescent               0.5.12
+pyparsing                 3.1.1
+python-dateutil           2.8.2
+python-json-logger        2.0.7
+pytz                      2023.3.post1
+PyYAML                    6.0.1
+pyzmq                     24.0.1
+qgml                      2024.3.0
+qml-benchmarks            0.1          /Users/joseph/PycharmProjects/pennylane_jax/qml-benchmarks
+referencing               0.32.0
+requests                  2.31.0
+rfc3339-validator         0.1.4
+rfc3986-validator         0.1.1
+rich                      13.7.0
+rpds-py                   0.15.2
+rustworkx                 0.13.2
+scikit-learn              1.4.0
+scipy                     1.11.4
+seaborn                   0.13.0
+semantic-version          2.10.0
+Send2Trash                1.8.2
+setuptools                57.0.0
+six                       1.16.0
+sniffio                   1.3.0
+soupsieve                 2.5
+stack-data                0.6.3
+sympy                     1.12
+tensorstore               0.1.51
+termcolor                 2.4.0
+terminado                 0.18.0
+threadpoolctl             3.2.0
+tinycss2                  1.2.1
+toml                      0.10.2
+tomli                     2.0.1
+tomlkit                   0.12.3
+toolz                     0.12.0
+torch                     2.1.2
+torchvision               0.16.2
+tornado                   6.4
+tqdm                      4.66.2
+traitlets                 5.14.0
+types-python-dateutil     2.8.19.14
+typing_extensions         4.9.0
+tzdata                    2023.3
+umap-learn                0.5.6
+uri-template              1.3.0
+urllib3                   2.1.0
+wcwidth                   0.2.12
+webcolors                 1.13
+webencodings              0.5.1
+websocket-client          1.7.0
+wheel                     0.36.2
+zipp                      3.17.0
diff --git a/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_3d_performance_indicators_JAX_scontrol.txt b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_3d_performance_indicators_JAX_scontrol.txt
new file mode 100644
index 00000000..e69de29b
diff --git a/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_14d_JAX_info.txt b/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_14d_JAX_info.txt
deleted file mode 100644
index 6e40ef73..00000000
--- a/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_14d_JAX_info.txt
+++ /dev/null
@@ -1,4 +0,0 @@
-salloc -q shared_interactive -C cpu -t 4:00:00 --cpus-per-task=32
-
-
-
diff --git a/nersc/performance_indicators/hyperparam_settings.yaml b/nersc/performance_indicators/hyperparam_settings.yaml
index 5a3e82c5..525220d4 100644
--- a/nersc/performance_indicators/hyperparam_settings.yaml
+++ b/nersc/performance_indicators/hyperparam_settings.yaml
@@ -1,10 +1,8 @@
 IQPVariationalClassifier:
-  learning_rate: 0.01
   n_layers: 15
   repeats: 10
 
 QuantumMetricLearner:
-  learning_rate: 0.01
   n_layers: 4
   batch_size: 16
 
diff --git a/nersc/performance_indicators/perf_ind_iqpker_LR15_jax.py b/nersc/performance_indicators/perf_ind_iqpker_LR15_jax.py
index 0321dd17..ef90813c 100644
--- a/nersc/performance_indicators/perf_ind_iqpker_LR15_jax.py
+++ b/nersc/performance_indicators/perf_ind_iqpker_LR15_jax.py
@@ -7,24 +7,38 @@
 import csv
 import os
 import yaml
-
-from qml_benchmarks.models.iqp_kernel import IQPKernelClassifier
+import subprocess
 from qml_benchmarks.hyperparam_search_utils import read_data
 
+#################################
+# settings for the performance indicator.
+# You only need to change this to make a different performance indicator
+
+from qml_benchmarks.models.iqp_kernel import IQPKernelClassifier as Model
+use_jax = True
+vmap = True
+jit = True
+perf_ind_name = 'JAX'  #a name for the performance indicator used for naming files
+n_features = 3 #dataset dimension
+n_train = 50 #number of data points to fit model
+n_test = 50 #number of datapoints for prediction
+
+#################################
+
+model_name = Model().__class__.__name__
+
+# get the 'worst case' hyperparameter settings for the model (those that require the most resources)
 with open('hyperparam_settings.yaml', "r") as file:
     hp_settings = yaml.safe_load(file)
 
-hyperparams = {**hp_settings['IQPKernelClassifier'], **{'use_jax':True, 'vmap':True}}
-
+hyperparams = {**hp_settings[model_name], **{'use_jax':use_jax, 'vmap':vmap, 'jit': jit}}
 print(hyperparams)
 
-n_features = 15 #dataset dimension
-
 X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
 X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
 
-model = IQPKernelClassifier(**hyperparams)
-model.fit(X_train[:50], y_train[:50])
+model = Model(**hyperparams)
+model.fit(X_train[:n_train], y_train[:n_train])
 
 #kernel construction time
 construct_kernel_time = model.construct_kernel_time_
@@ -32,23 +46,35 @@
 training_time = model.training_time_
 #prediction time
 time0 = time.time()
-model.predict(X_test[:50])
+model.predict(X_test[:n_test])
 predict_time = time.time() - time0
 
 
 #write to csv
 data = [construct_kernel_time, training_time, predict_time, hyperparams]
 
-model_name = model.__class__.__name__
-filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_JAX.csv"
+filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_{perf_ind_name}.csv"
+packages_filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_{perf_ind_name}_packages.txt"
+scontrol_filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_{perf_ind_name}_scontrol.txt"
+
 header = ['construct_kernel_time', 'training_time', 'predict_time', 'hyperparameters']
 
-if not os.path.exists('JAX'):
+if not os.path.exists(perf_ind_name):
     # Create the directory
-    os.mkdir('JAX')
+    os.mkdir(perf_ind_name)
 
-with open('JAX/'+filename, mode="w", newline="") as file:
+#write perf indicator data
+with open(perf_ind_name+'/'+filename, mode="w", newline="") as file:
     writer = csv.writer(file)
     writer.writerow(header)
     for row in [data]:
         writer.writerow(row)
+
+# get package list and write to file
+output = subprocess.check_output(['pip', 'list']).decode('utf-8')
+with open(perf_ind_name+'/'+packages_filename, 'w') as file:
+    file.write(output)
+
+# make an empty text file to store the scontrol data
+with open(perf_ind_name+'/'+scontrol_filename, 'w') as file:
+    pass
\ No newline at end of file
diff --git a/nersc/performance_indicators/perf_ind_iqpvar_LR15_jax.py b/nersc/performance_indicators/perf_ind_iqpvar_LR15_jax.py
index dece4d43..ad7465f2 100644
--- a/nersc/performance_indicators/perf_ind_iqpvar_LR15_jax.py
+++ b/nersc/performance_indicators/perf_ind_iqpvar_LR15_jax.py
@@ -7,20 +7,33 @@
 import csv
 import os
 import yaml
-
-from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier
+import subprocess
 from qml_benchmarks.hyperparam_search_utils import read_data
 
+#################################
+# settings for the performance indicator.
+# You only need to change this to make a different performance indicator
+
+from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier as Model
+use_jax = True
+vmap = True
+jit = True
+max_steps = 5 #the number of gradient descent steps to use to estimate the step time
+perf_ind_name = 'JAX'  #a name for the performance indicator used for naming files
+n_features = 3 #dataset dimension
+n_trials = 5 #number of trials to average over
+
+#################################
+
+model_name = Model().__class__.__name__
+
+# get the 'worst case' hyperparameter settings for the model (those that require the most resources)
 with open('hyperparam_settings.yaml', "r") as file:
     hp_settings = yaml.safe_load(file)
 
-hyperparams = {**hp_settings['IQPVariationalClassifier'], **{'use_jax':True, 'vmap':True, 'max_steps':10}}
-
+hyperparams = {**hp_settings[model_name], **{'use_jax':use_jax, 'vmap':vmap, 'max_steps':max_steps, 'jit': jit}}
 print(hyperparams)
 
-n_features = 15 #dataset dimension
-n_trials = 5 #number of trials to average over
-
 X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
 X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
 
@@ -31,7 +44,9 @@
 for trial in range(n_trials):
     jax.clear_caches()
 
-    model = IQPVariationalClassifier(**hyperparams)
+    model = Model(**hyperparams)
+
+    #note we train for max_steps only, so we won't reach convergence.
     model.fit(X_train, y_train)
 
     #get step times from loss history data
@@ -50,7 +65,6 @@
 
 
 #calculate mean and stds
-
 first_train_step = np.mean(first_train_steps)
 first_train_step_std = np.std(first_train_steps)
 
@@ -64,17 +78,30 @@
 data = [first_train_step, first_train_step_std, consec_train_step, consec_train_step_std, predict_time,
         predict_time_std, hyperparams]
 
-model_name = model.__class__.__name__
-filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_JAX.csv"
+filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_{perf_ind_name}.csv"
+packages_filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_{perf_ind_name}_packages.txt"
+scontrol_filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_{perf_ind_name}_scontrol.txt"
+
 header = ['first_train_step', 'first_train_step_std', 'consec_train_step', 'consec_train_step_std', 'predict_time',
           'predict_time_std', 'hyperparameters']
 
-if not os.path.exists('JAX'):
+if not os.path.exists(perf_ind_name):
     # Create the directory
-    os.mkdir('JAX')
+    os.mkdir(perf_ind_name)
 
-with open('JAX/'+filename, mode="w", newline="") as file:
+#write perf indicator data
+with open(perf_ind_name+'/'+filename, mode="w", newline="") as file:
     writer = csv.writer(file)
     writer.writerow(header)
     for row in [data]:
         writer.writerow(row)
+
+# get package list and write to file
+output = subprocess.check_output(['pip', 'list']).decode('utf-8')
+with open(perf_ind_name+'/'+packages_filename, 'w') as file:
+    file.write(output)
+
+# make an empty text file to store the scontrol data
+with open(perf_ind_name+'/'+scontrol_filename, 'w') as file:
+    pass
+
diff --git a/nersc/performance_indicators/perf_ind_iqpvar_LR17_jax.py b/nersc/performance_indicators/perf_ind_iqpvar_LR17_jax.py
deleted file mode 100644
index 7dbc8e80..00000000
--- a/nersc/performance_indicators/perf_ind_iqpvar_LR17_jax.py
+++ /dev/null
@@ -1,80 +0,0 @@
-import qml_benchmarks
-import pennylane as qml
-import jax
-import jax.numpy as jnp
-import numpy as np
-import time
-import csv
-import os
-import yaml
-
-from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier
-from qml_benchmarks.hyperparam_search_utils import read_data
-
-with open('hyperparam_settings.yaml', "r") as file:
-    hp_settings = yaml.safe_load(file)
-
-hyperparams = {**hp_settings['IQPVariationalClassifier'], **{'use_jax':True, 'vmap':True, 'max_steps':10}}
-
-print(hyperparams)
-
-n_features = 17 #dataset dimension
-n_trials = 5 #number of trials to average over
-
-X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
-X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
-
-first_train_steps = []
-av_consec_train_steps = []
-predict_times = []
-
-for trial in range(n_trials):
-    jax.clear_caches()
-
-    model = IQPVariationalClassifier(**hyperparams)
-    model.fit(X_train, y_train)
-
-    #get step times from loss history data
-    step_times = np.array([model.loss_history_[1][i + 1] - model.loss_history_[1][i]
-                           for i in range(len(model.loss_history_[1]) - 1)])
-    step_times = np.insert(step_times, 0, [model.loss_history_[1][0]])
-
-    #first train step
-    first_train_steps.append(step_times[0])
-    #consecutive (average) train step
-    av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
-    #prediction time
-    time0 = time.time()
-    model.predict(X_test[:10])
-    predict_times.append(time.time() - time0)
-
-
-#calculate mean and stds
-
-first_train_step = np.mean(first_train_steps)
-first_train_step_std = np.std(first_train_steps)
-
-consec_train_step = np.mean(av_consec_train_steps)
-consec_train_step_std = np.std(av_consec_train_steps)
-
-predict_time = np.mean(predict_times)
-predict_time_std = np.std(predict_times)
-
-#write to csv
-data = [first_train_step, first_train_step_std, consec_train_step, consec_train_step_std, predict_time,
-        predict_time_std, hyperparams]
-
-model_name = model.__class__.__name__
-filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_JAX.csv"
-header = ['first_train_step', 'first_train_step_std', 'consec_train_step', 'consec_train_step_std', 'predict_time',
-          'predict_time_std', 'hyperparameters']
-
-if not os.path.exists('JAX'):
-    # Create the directory
-    os.mkdir('JAX')
-
-with open('JAX/'+filename, mode="w", newline="") as file:
-    writer = csv.writer(file)
-    writer.writerow(header)
-    for row in [data]:
-        writer.writerow(row)
diff --git a/nersc/performance_indicators/perf_ind_iqpvar_LR20_jax.py b/nersc/performance_indicators/perf_ind_iqpvar_LR20_jax.py
deleted file mode 100644
index c6dc8770..00000000
--- a/nersc/performance_indicators/perf_ind_iqpvar_LR20_jax.py
+++ /dev/null
@@ -1,80 +0,0 @@
-import qml_benchmarks
-import pennylane as qml
-import jax
-import jax.numpy as jnp
-import numpy as np
-import time
-import csv
-import os
-import yaml
-
-from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier
-from qml_benchmarks.hyperparam_search_utils import read_data
-
-with open('hyperparam_settings.yaml', "r") as file:
-    hp_settings = yaml.safe_load(file)
-
-hyperparams = {**hp_settings['IQPVariationalClassifier'], **{'use_jax':True, 'vmap':True, 'max_steps':10}}
-
-print(hyperparams)
-
-n_features = 20 #dataset dimension
-n_trials = 5 #number of trials to average over
-
-X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
-X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
-
-first_train_steps = []
-av_consec_train_steps = []
-predict_times = []
-
-for trial in range(n_trials):
-    jax.clear_caches()
-
-    model = IQPVariationalClassifier(**hyperparams)
-    model.fit(X_train, y_train)
-
-    #get step times from loss history data
-    step_times = np.array([model.loss_history_[1][i + 1] - model.loss_history_[1][i]
-                           for i in range(len(model.loss_history_[1]) - 1)])
-    step_times = np.insert(step_times, 0, [model.loss_history_[1][0]])
-
-    #first train step
-    first_train_steps.append(step_times[0])
-    #consecutive (average) train step
-    av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
-    #prediction time
-    time0 = time.time()
-    model.predict(X_test[:10])
-    predict_times.append(time.time() - time0)
-
-
-#calculate mean and stds
-
-first_train_step = np.mean(first_train_steps)
-first_train_step_std = np.std(first_train_steps)
-
-consec_train_step = np.mean(av_consec_train_steps)
-consec_train_step_std = np.std(av_consec_train_steps)
-
-predict_time = np.mean(predict_times)
-predict_time_std = np.std(predict_times)
-
-#write to csv
-data = [first_train_step, first_train_step_std, consec_train_step, consec_train_step_std, predict_time,
-        predict_time_std, hyperparams]
-
-model_name = model.__class__.__name__
-filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_JAX.csv"
-header = ['first_train_step', 'first_train_step_std', 'consec_train_step', 'consec_train_step_std', 'predict_time',
-          'predict_time_std', 'hyperparameters']
-
-if not os.path.exists('JAX'):
-    # Create the directory
-    os.mkdir('JAX')
-
-with open('JAX/'+filename, mode="w", newline="") as file:
-    writer = csv.writer(file)
-    writer.writerow(header)
-    for row in [data]:
-        writer.writerow(row)
diff --git a/nersc/performance_indicators/perf_ind_pqk_LR15_jax.py b/nersc/performance_indicators/perf_ind_pqk_LR15_jax.py
deleted file mode 100644
index 5facd45f..00000000
--- a/nersc/performance_indicators/perf_ind_pqk_LR15_jax.py
+++ /dev/null
@@ -1,54 +0,0 @@
-import qml_benchmarks
-import pennylane as qml
-import jax
-import jax.numpy as jnp
-import numpy as np
-import time
-import csv
-import os
-import yaml
-
-from qml_benchmarks.models.projected_quantum_kernel import ProjectedQuantumKernel
-from qml_benchmarks.hyperparam_search_utils import read_data
-
-with open('hyperparam_settings.yaml', "r") as file:
-    hp_settings = yaml.safe_load(file)
-
-hyperparams = {**hp_settings['ProjectedQuantumKernel'], **{'use_jax':True, 'vmap':True}}
-
-print(hyperparams)
-
-n_features = 15 #dataset dimension
-
-X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
-X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
-
-model = ProjectedQuantumKernel(**hyperparams)
-model.fit(X_train[:50], y_train[:50])
-
-#kernel construction time
-construct_kernel_time = model.construct_kernel_time_
-#full training time
-training_time = model.training_time_
-#prediction time
-time0 = time.time()
-model.predict(X_test[:50])
-predict_time = time.time() - time0
-
-
-#write to csv
-data = [construct_kernel_time, training_time, predict_time, hyperparams]
-
-model_name = model.__class__.__name__
-filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_JAX.csv"
-header = ['construct_kernel_time', 'training_time', 'predict_time', 'hyperparameters']
-
-if not os.path.exists('JAX'):
-    # Create the directory
-    os.mkdir('JAX')
-
-with open('JAX/'+filename, mode="w", newline="") as file:
-    writer = csv.writer(file)
-    writer.writerow(header)
-    for row in [data]:
-        writer.writerow(row)
diff --git a/nersc/performance_indicators/perf_ind_pqk_LR17_jax.py b/nersc/performance_indicators/perf_ind_pqk_LR17_jax.py
deleted file mode 100644
index ffe55626..00000000
--- a/nersc/performance_indicators/perf_ind_pqk_LR17_jax.py
+++ /dev/null
@@ -1,54 +0,0 @@
-import qml_benchmarks
-import pennylane as qml
-import jax
-import jax.numpy as jnp
-import numpy as np
-import time
-import csv
-import os
-import yaml
-
-from qml_benchmarks.models.projected_quantum_kernel import ProjectedQuantumKernel
-from qml_benchmarks.hyperparam_search_utils import read_data
-
-with open('hyperparam_settings.yaml', "r") as file:
-    hp_settings = yaml.safe_load(file)
-
-hyperparams = {**hp_settings['ProjectedQuantumKernel'], **{'use_jax':True, 'vmap':True}}
-
-print(hyperparams)
-
-n_features = 17 #dataset dimension
-
-X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
-X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
-
-model = ProjectedQuantumKernel(**hyperparams)
-model.fit(X_train[:50], y_train[:50])
-
-#kernel construction time
-construct_kernel_time = model.construct_kernel_time_
-#full training time
-training_time = model.training_time_
-#prediction time
-time0 = time.time()
-model.predict(X_test[:50])
-predict_time = time.time() - time0
-
-
-#write to csv
-data = [construct_kernel_time, training_time, predict_time, hyperparams]
-
-model_name = model.__class__.__name__
-filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_JAX.csv"
-header = ['construct_kernel_time', 'training_time', 'predict_time', 'hyperparameters']
-
-if not os.path.exists('JAX'):
-    # Create the directory
-    os.mkdir('JAX')
-
-with open('JAX/'+filename, mode="w", newline="") as file:
-    writer = csv.writer(file)
-    writer.writerow(header)
-    for row in [data]:
-        writer.writerow(row)
diff --git a/nersc/performance_indicators/perf_ind_pqk_LR20_jax.py b/nersc/performance_indicators/perf_ind_pqk_LR20_jax.py
deleted file mode 100644
index df2ba712..00000000
--- a/nersc/performance_indicators/perf_ind_pqk_LR20_jax.py
+++ /dev/null
@@ -1,54 +0,0 @@
-import qml_benchmarks
-import pennylane as qml
-import jax
-import jax.numpy as jnp
-import numpy as np
-import time
-import csv
-import os
-import yaml
-
-from qml_benchmarks.models.projected_quantum_kernel import ProjectedQuantumKernel
-from qml_benchmarks.hyperparam_search_utils import read_data
-
-with open('hyperparam_settings.yaml', "r") as file:
-    hp_settings = yaml.safe_load(file)
-
-hyperparams = {**hp_settings['ProjectedQuantumKernel'], **{'use_jax':True, 'vmap':True}}
-
-print(hyperparams)
-
-n_features = 20 #dataset dimension
-
-X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
-X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
-
-model = ProjectedQuantumKernel(**hyperparams)
-model.fit(X_train[:50], y_train[:50])
-
-#kernel construction time
-construct_kernel_time = model.construct_kernel_time_
-#full training time
-training_time = model.training_time_
-#prediction time
-time0 = time.time()
-model.predict(X_test[:50])
-predict_time = time.time() - time0
-
-
-#write to csv
-data = [construct_kernel_time, training_time, predict_time, hyperparams]
-
-model_name = model.__class__.__name__
-filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_JAX.csv"
-header = ['construct_kernel_time', 'training_time', 'predict_time', 'hyperparameters']
-
-if not os.path.exists('JAX'):
-    # Create the directory
-    os.mkdir('JAX')
-
-with open('JAX/'+filename, mode="w", newline="") as file:
-    writer = csv.writer(file)
-    writer.writerow(header)
-    for row in [data]:
-        writer.writerow(row)
diff --git a/nersc/performance_indicators/perf_ind_qmetric_LR14_jax.py b/nersc/performance_indicators/perf_ind_qmetric_LR14_jax.py
deleted file mode 100644
index cbb647c0..00000000
--- a/nersc/performance_indicators/perf_ind_qmetric_LR14_jax.py
+++ /dev/null
@@ -1,80 +0,0 @@
-import qml_benchmarks
-import pennylane as qml
-import jax
-import jax.numpy as jnp
-import numpy as np
-import time
-import csv
-import os
-import yaml
-
-from qml_benchmarks.models.quantum_metric_learning import QuantumMetricLearner
-from qml_benchmarks.hyperparam_search_utils import read_data
-
-with open('hyperparam_settings.yaml', "r") as file:
-    hp_settings = yaml.safe_load(file)
-
-hyperparams = {**hp_settings['QuantumMetricLearner'], **{'use_jax':True, 'vmap':True, 'max_steps':10}}
-
-print(hyperparams)
-
-n_features = 14 #dataset dimension
-n_trials = 5 #number of trials to average over
-
-X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
-X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
-
-first_train_steps = []
-av_consec_train_steps = []
-predict_times = []
-
-for trial in range(n_trials):
-    jax.clear_caches()
-
-    model = QuantumMetricLearner(**hyperparams)
-    model.fit(X_train, y_train)
-
-    #get step times from loss history data
-    step_times = np.array([model.loss_history_[1][i + 1] - model.loss_history_[1][i]
-                           for i in range(len(model.loss_history_[1]) - 1)])
-    step_times = np.insert(step_times, 0, [model.loss_history_[1][0]])
-
-    #first train step
-    first_train_steps.append(step_times[0])
-    #consecutive (average) train step
-    av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
-    #prediction time
-    time0 = time.time()
-    model.predict(X_test[:10])
-    predict_times.append(time.time() - time0)
-
-
-#calculate mean and stds
-
-first_train_step = np.mean(first_train_steps)
-first_train_step_std = np.std(first_train_steps)
-
-consec_train_step = np.mean(av_consec_train_steps)
-consec_train_step_std = np.std(av_consec_train_steps)
-
-predict_time = np.mean(predict_times)
-predict_time_std = np.std(predict_times)
-
-#write to csv
-data = [first_train_step, first_train_step_std, consec_train_step, consec_train_step_std, predict_time,
-        predict_time_std, hyperparams]
-
-model_name = model.__class__.__name__
-filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_JAX.csv"
-header = ['first_train_step', 'first_train_step_std', 'consec_train_step', 'consec_train_step_std', 'predict_time',
-          'predict_time_std', 'hyperparameters']
-
-if not os.path.exists('JAX'):
-    # Create the directory
-    os.mkdir('JAX')
-
-with open('JAX/'+filename, mode="w", newline="") as file:
-    writer = csv.writer(file)
-    writer.writerow(header)
-    for row in [data]:
-        writer.writerow(row)
diff --git a/nersc/performance_indicators/perf_ind_qmetric_LR15_jax.py b/nersc/performance_indicators/perf_ind_qmetric_LR15_jax.py
deleted file mode 100644
index 8a78431b..00000000
--- a/nersc/performance_indicators/perf_ind_qmetric_LR15_jax.py
+++ /dev/null
@@ -1,80 +0,0 @@
-import qml_benchmarks
-import pennylane as qml
-import jax
-import jax.numpy as jnp
-import numpy as np
-import time
-import csv
-import os
-import yaml
-
-from qml_benchmarks.models.quantum_metric_learning import QuantumMetricLearner
-from qml_benchmarks.hyperparam_search_utils import read_data
-
-with open('hyperparam_settings.yaml', "r") as file:
-    hp_settings = yaml.safe_load(file)
-
-hyperparams = {**hp_settings['QuantumMetricLearner'], **{'use_jax':True, 'vmap':True, 'max_steps':10}}
-
-print(hyperparams)
-
-n_features = 15 #dataset dimension
-n_trials = 5 #number of trials to average over
-
-X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
-X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
-
-first_train_steps = []
-av_consec_train_steps = []
-predict_times = []
-
-for trial in range(n_trials):
-    jax.clear_caches()
-
-    model = QuantumMetricLearner(**hyperparams)
-    model.fit(X_train, y_train)
-
-    #get step times from loss history data
-    step_times = np.array([model.loss_history_[1][i + 1] - model.loss_history_[1][i]
-                           for i in range(len(model.loss_history_[1]) - 1)])
-    step_times = np.insert(step_times, 0, [model.loss_history_[1][0]])
-
-    #first train step
-    first_train_steps.append(step_times[0])
-    #consecutive (average) train step
-    av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
-    #prediction time
-    time0 = time.time()
-    model.predict(X_test[:10])
-    predict_times.append(time.time() - time0)
-
-
-#calculate mean and stds
-
-first_train_step = np.mean(first_train_steps)
-first_train_step_std = np.std(first_train_steps)
-
-consec_train_step = np.mean(av_consec_train_steps)
-consec_train_step_std = np.std(av_consec_train_steps)
-
-predict_time = np.mean(predict_times)
-predict_time_std = np.std(predict_times)
-
-#write to csv
-data = [first_train_step, first_train_step_std, consec_train_step, consec_train_step_std, predict_time,
-        predict_time_std, hyperparams]
-
-model_name = model.__class__.__name__
-filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_JAX.csv"
-header = ['first_train_step', 'first_train_step_std', 'consec_train_step', 'consec_train_step_std', 'predict_time',
-          'predict_time_std', 'hyperparameters']
-
-if not os.path.exists('JAX'):
-    # Create the directory
-    os.mkdir('JAX')
-
-with open('JAX/'+filename, mode="w", newline="") as file:
-    writer = csv.writer(file)
-    writer.writerow(header)
-    for row in [data]:
-        writer.writerow(row)
diff --git a/nersc/performance_indicators/save_last_job_info.sh b/nersc/performance_indicators/save_last_job_info.sh
deleted file mode 100644
index ab6d07e8..00000000
--- a/nersc/performance_indicators/save_last_job_info.sh
+++ /dev/null
@@ -1,21 +0,0 @@
-#!/bin/bash
-
-# Get the job ID of the last completed job
-last_job_id=$(sacct --format=JobID -n -u $USER -S "$(date -d yesterday +"%Y-%m-%d")" -o JobID | tail -n 1)
-
-# Check if a job ID was found
-if [ -z "$last_job_id" ]; then
-    echo "No completed job found."
-    exit 1
-fi
-
-# Get the job information
-job_info=$(scontrol show job $last_job_id)
-
-# Determine the filename
-filename="${1:-job_info.txt}"
-
-# Save the job information to the specified file
-echo "$job_info" > "$filename"
-
-echo "Job information saved to $filename."
diff --git a/nersc/performance_indicators/submit_job.slr b/nersc/performance_indicators/submit_job.slr
index 4984709b..90797939 100644
--- a/nersc/performance_indicators/submit_job.slr
+++ b/nersc/performance_indicators/submit_job.slr
@@ -4,7 +4,7 @@
 #SBATCH -q shared
 #SBATCH -J perf_ind_iqpvar_15
 # SBATCH --account=jbowles
-#SBATCH -t 1:00:00
+#SBATCH -t 8:00:00
 #SBATCH --ntasks-per-node=1
 #SBATCH --cpus-per-task=16
 #SBATCH --output out/%j.log

From ce6b0ffdfe45621030e21d6fd052442774ea506b Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Tue, 16 Apr 2024 15:42:09 +0200
Subject: [PATCH 038/100] delete

---
 nersc/test.py | 14 --------------
 1 file changed, 14 deletions(-)
 delete mode 100644 nersc/test.py

diff --git a/nersc/test.py b/nersc/test.py
deleted file mode 100644
index c5d2101e..00000000
--- a/nersc/test.py
+++ /dev/null
@@ -1,14 +0,0 @@
-import pennylane
-import qml_benchmarks
-from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier
-from qml_benchmarks.data.two_curves import generate_two_curves
-import numpy as np
-
-
-
-X,y = generate_two_curves(100,3, 2, 0.0, 0.1)
-model = IQPVariationalClassifier(use_jax=True, vmap=True)
-
-model.fit(X, y)
-print(model.loss_history_)
-np.savetxt('loss.txt', model.loss_history_)
\ No newline at end of file

From fe0a9ce9c7d46ae1abf09d0dc9aced6c7e44bff6 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Tue, 16 Apr 2024 15:45:33 +0200
Subject: [PATCH 039/100] update

---
 nersc/performance_indicators/perf_ind_iqpker_LR15_jax.py | 2 +-
 nersc/performance_indicators/perf_ind_iqpvar_LR15_jax.py | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/nersc/performance_indicators/perf_ind_iqpker_LR15_jax.py b/nersc/performance_indicators/perf_ind_iqpker_LR15_jax.py
index ef90813c..9627b08c 100644
--- a/nersc/performance_indicators/perf_ind_iqpker_LR15_jax.py
+++ b/nersc/performance_indicators/perf_ind_iqpker_LR15_jax.py
@@ -19,7 +19,7 @@
 vmap = True
 jit = True
 perf_ind_name = 'JAX'  #a name for the performance indicator used for naming files
-n_features = 3 #dataset dimension
+n_features = 15 #dataset dimension
 n_train = 50 #number of data points to fit model
 n_test = 50 #number of datapoints for prediction
 
diff --git a/nersc/performance_indicators/perf_ind_iqpvar_LR15_jax.py b/nersc/performance_indicators/perf_ind_iqpvar_LR15_jax.py
index ad7465f2..17f794c3 100644
--- a/nersc/performance_indicators/perf_ind_iqpvar_LR15_jax.py
+++ b/nersc/performance_indicators/perf_ind_iqpvar_LR15_jax.py
@@ -18,9 +18,9 @@
 use_jax = True
 vmap = True
 jit = True
-max_steps = 5 #the number of gradient descent steps to use to estimate the step time
+max_steps = 100 #the number of gradient descent steps to use to estimate the step time
 perf_ind_name = 'JAX'  #a name for the performance indicator used for naming files
-n_features = 3 #dataset dimension
+n_features = 15 #dataset dimension
 n_trials = 5 #number of trials to average over
 
 #################################

From b2bc7796d4368c8abc8180445c329b0cd89c442b Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Mon, 22 Apr 2024 15:09:39 +0200
Subject: [PATCH 040/100] update

---
 .../{perf_ind_iqpker_LR15_jax.py => perf_ind_kernel_jax.py}       | 0
 .../{perf_ind_iqpvar_LR15_jax.py => perf_ind_variational_jax.py}  | 0
 2 files changed, 0 insertions(+), 0 deletions(-)
 rename nersc/performance_indicators/{perf_ind_iqpker_LR15_jax.py => perf_ind_kernel_jax.py} (100%)
 rename nersc/performance_indicators/{perf_ind_iqpvar_LR15_jax.py => perf_ind_variational_jax.py} (100%)

diff --git a/nersc/performance_indicators/perf_ind_iqpker_LR15_jax.py b/nersc/performance_indicators/perf_ind_kernel_jax.py
similarity index 100%
rename from nersc/performance_indicators/perf_ind_iqpker_LR15_jax.py
rename to nersc/performance_indicators/perf_ind_kernel_jax.py
diff --git a/nersc/performance_indicators/perf_ind_iqpvar_LR15_jax.py b/nersc/performance_indicators/perf_ind_variational_jax.py
similarity index 100%
rename from nersc/performance_indicators/perf_ind_iqpvar_LR15_jax.py
rename to nersc/performance_indicators/perf_ind_variational_jax.py

From 595d3ef6a170829857c762ac73d55cca3fe9646d Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Mon, 22 Apr 2024 15:10:04 +0200
Subject: [PATCH 041/100] update

---
 nersc/pm_podman.source | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/nersc/pm_podman.source b/nersc/pm_podman.source
index aa13f06e..2b8000f7 100644
--- a/nersc/pm_podman.source
+++ b/nersc/pm_podman.source
@@ -2,7 +2,7 @@
 #  salloc -q interactive -C cpu -t 4:00:00 -A nstaff 
 
 IMG=balewski/ubu22-pennylane:p3
-CFSH=/global/homes/j/jbowles
+CFSH=/global/cfs/cdirs/m4139/
 
 echo launch image $IMG
 echo you are launching Podman image ...  remember to exit

From 62ac41eb2a2fd6d50d9ca60d82085131f188777c Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Mon, 22 Apr 2024 15:36:50 +0200
Subject: [PATCH 042/100] update

---
 nersc/performance_indicators/perf_ind_kernel_jax.py      | 6 ++----
 nersc/performance_indicators/perf_ind_variational_jax.py | 2 +-
 2 files changed, 3 insertions(+), 5 deletions(-)

diff --git a/nersc/performance_indicators/perf_ind_kernel_jax.py b/nersc/performance_indicators/perf_ind_kernel_jax.py
index 9627b08c..929aa826 100644
--- a/nersc/performance_indicators/perf_ind_kernel_jax.py
+++ b/nersc/performance_indicators/perf_ind_kernel_jax.py
@@ -20,8 +20,6 @@
 jit = True
 perf_ind_name = 'JAX'  #a name for the performance indicator used for naming files
 n_features = 15 #dataset dimension
-n_train = 50 #number of data points to fit model
-n_test = 50 #number of datapoints for prediction
 
 #################################
 
@@ -38,7 +36,7 @@
 X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
 
 model = Model(**hyperparams)
-model.fit(X_train[:n_train], y_train[:n_train])
+model.fit(X_train, y_train)
 
 #kernel construction time
 construct_kernel_time = model.construct_kernel_time_
@@ -46,7 +44,7 @@
 training_time = model.training_time_
 #prediction time
 time0 = time.time()
-model.predict(X_test[:n_test])
+model.predict(X_test)
 predict_time = time.time() - time0
 
 
diff --git a/nersc/performance_indicators/perf_ind_variational_jax.py b/nersc/performance_indicators/perf_ind_variational_jax.py
index 17f794c3..6e6790eb 100644
--- a/nersc/performance_indicators/perf_ind_variational_jax.py
+++ b/nersc/performance_indicators/perf_ind_variational_jax.py
@@ -60,7 +60,7 @@
     av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
     #prediction time
     time0 = time.time()
-    model.predict(X_test[:10])
+    model.predict(X_test)
     predict_times.append(time.time() - time0)
 
 

From a94578ec2f8f328164a7c47ac8adf89a8964af28 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Mon, 22 Apr 2024 15:38:23 +0200
Subject: [PATCH 043/100] update

---
 nersc/performance_indicators/submit_job.slr | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/nersc/performance_indicators/submit_job.slr b/nersc/performance_indicators/submit_job.slr
index 90797939..44cecf9e 100644
--- a/nersc/performance_indicators/submit_job.slr
+++ b/nersc/performance_indicators/submit_job.slr
@@ -4,13 +4,13 @@
 #SBATCH -q shared
 #SBATCH -J perf_ind_iqpvar_15
 # SBATCH --account=jbowles
-#SBATCH -t 8:00:00
+#SBATCH -t 24:00:00
 #SBATCH --ntasks-per-node=1
-#SBATCH --cpus-per-task=16
+#SBATCH --cpus-per-task=6
 #SBATCH --output out/%j.log
 #SBATCH  --licenses=scratch
 
 source ../pm_podman.source
 bash ../installs.sh
 
-time srun -n 1  python3 ../perf_ind_iqpvar_LR15_jax.py
+time srun -n 1  python3 ../perf_ind_variational_jax.py

From 4a78cf0d48b0f40ba894ccfa106ecda4750d0362 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Mon, 22 Apr 2024 16:29:01 +0200
Subject: [PATCH 044/100] update

---
 nersc/performance_indicators/README.md        | 65 +++++++++++++++++++
 .../perf_ind_kernel_jax.py                    |  4 ++
 .../perf_ind_variational_jax.py               |  6 +-
 nersc/performance_indicators/submit_job.slr   |  4 +-
 4 files changed, 76 insertions(+), 3 deletions(-)
 create mode 100644 nersc/performance_indicators/README.md

diff --git a/nersc/performance_indicators/README.md b/nersc/performance_indicators/README.md
new file mode 100644
index 00000000..44a37cdc
--- /dev/null
+++ b/nersc/performance_indicators/README.md
@@ -0,0 +1,65 @@
+# Instructions for running performance indicators
+
+To get the data for a given model we need to run the python script `perf_ind_variational.py` (for variational models) or 
+`perf_ind_kernel.py` (for kernel models). This gets the indicators 'first train step time', 'consec. train step time',
+ and 'prediction time'. 
+
+For variational models, there is an additional script `perf_ind_full_train.py` that attempts to run the model until 
+convergence in order to get the 'accuracy' and 'train steps for convergence' numbers. Since this can take a long time, it is a separate script. 
+
+The results are saved to a subdirectory specified at the
+start of the script (see e.g. JAX/). The script also saves the pip package list for
+future reference. 
+
+Within the scripts there are a number of settings that can be changed at the start. This should make it easy to reuse 
+the same code when we have updated models. **Please make sure you chose a unique choice of `perf_ind_name` when gathering results for a new
+workflow** (i.e. the name of the workflow) since results will be overwritten if the same name is used twice.
+
+
+The performance indicators use the hyperparameter settings specified in `hyperparam_settings.yaml`. These are chosen to 
+be those which require the most compute and shouldn't be changed. 
+
+## determinining the number of CPUs
+To avoid wasting resources, you should first determine how many CPUs are required. This can be done by launching an 
+interative node:
+
+`salloc -q shared_interactive -C cpu -t 4:00:00 --cpus-per-task=XX`
+
+where you control the number of cpus-per-task. Start with a high number and
+monitor the CPU usage with top by SSHing to the compute node in another terminal.
+
+On the interactive node, from the /nersc directory run 
+
+`source pm_podman.source`
+
+to load the container. Then 
+
+`bash installs.sh`
+
+To install the qml-benchmarks package (we also downgrade to a version of scipy that works with catalyst). Run the 
+performance indicator pythons script:
+
+`python3 perf_ind_variational.py`
+
+Note the CPU usage in top, and repeat this process (lowering the CPUs each time) until you determine a reasonable
+number of CPUs to use, and add this to the `CPUs` row of the google sheet. 
+
+## running a batch job
+
+Once the resources have been decided, edit the file `sbatch submit_job.slr` accordingly and launch a slurm job for 
+the performance indicator by running 
+
+`sbatch submit_job.slr`
+
+## recording memory usage
+Once the job has finished, run 
+
+`seff JOBID`
+
+where JOBID is the slurm job id. The value `Memory Utilized` is an estimate of the peak memory usage. 
+Add this to the sheet. 
+
+
+
+
+
diff --git a/nersc/performance_indicators/perf_ind_kernel_jax.py b/nersc/performance_indicators/perf_ind_kernel_jax.py
index 929aa826..494c1368 100644
--- a/nersc/performance_indicators/perf_ind_kernel_jax.py
+++ b/nersc/performance_indicators/perf_ind_kernel_jax.py
@@ -14,10 +14,14 @@
 # settings for the performance indicator.
 # You only need to change this to make a different performance indicator
 
+#define the model
 from qml_benchmarks.models.iqp_kernel import IQPKernelClassifier as Model
+
+#implementation attributes of model
 use_jax = True
 vmap = True
 jit = True
+
 perf_ind_name = 'JAX'  #a name for the performance indicator used for naming files
 n_features = 15 #dataset dimension
 
diff --git a/nersc/performance_indicators/perf_ind_variational_jax.py b/nersc/performance_indicators/perf_ind_variational_jax.py
index 6e6790eb..97175cd2 100644
--- a/nersc/performance_indicators/perf_ind_variational_jax.py
+++ b/nersc/performance_indicators/perf_ind_variational_jax.py
@@ -14,13 +14,17 @@
 # settings for the performance indicator.
 # You only need to change this to make a different performance indicator
 
+#define model
 from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier as Model
+
+#implementation attributes of model
 use_jax = True
 vmap = True
 jit = True
+
 max_steps = 100 #the number of gradient descent steps to use to estimate the step time
 perf_ind_name = 'JAX'  #a name for the performance indicator used for naming files
-n_features = 15 #dataset dimension
+n_features = 17 #dataset dimension
 n_trials = 5 #number of trials to average over
 
 #################################
diff --git a/nersc/performance_indicators/submit_job.slr b/nersc/performance_indicators/submit_job.slr
index 44cecf9e..cf170f80 100644
--- a/nersc/performance_indicators/submit_job.slr
+++ b/nersc/performance_indicators/submit_job.slr
@@ -2,11 +2,11 @@
 #SBATCH -N 1
 #SBATCH -C cpu
 #SBATCH -q shared
-#SBATCH -J perf_ind_iqpvar_15
+#SBATCH -J perf_ind_iqpvar_17
 # SBATCH --account=jbowles
 #SBATCH -t 24:00:00
 #SBATCH --ntasks-per-node=1
-#SBATCH --cpus-per-task=6
+#SBATCH --cpus-per-task=10
 #SBATCH --output out/%j.log
 #SBATCH  --licenses=scratch
 

From 2698fc5ceff8031475bff9d196dc9f7cada73672 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Mon, 22 Apr 2024 16:31:31 +0200
Subject: [PATCH 045/100] update

---
 nersc/performance_indicators/perf_ind_variational_jax.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/nersc/performance_indicators/perf_ind_variational_jax.py b/nersc/performance_indicators/perf_ind_variational_jax.py
index 97175cd2..5e64ad2d 100644
--- a/nersc/performance_indicators/perf_ind_variational_jax.py
+++ b/nersc/performance_indicators/perf_ind_variational_jax.py
@@ -24,7 +24,7 @@
 
 max_steps = 100 #the number of gradient descent steps to use to estimate the step time
 perf_ind_name = 'JAX'  #a name for the performance indicator used for naming files
-n_features = 17 #dataset dimension
+n_features = 20 #dataset dimension
 n_trials = 5 #number of trials to average over
 
 #################################

From 2690831363ba23ed2c40a07530782eaa80f22035 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Mon, 22 Apr 2024 16:52:46 +0200
Subject: [PATCH 046/100] update

---
 nersc/performance_indicators/submit_job.slr | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/nersc/performance_indicators/submit_job.slr b/nersc/performance_indicators/submit_job.slr
index cf170f80..8f1ad1e9 100644
--- a/nersc/performance_indicators/submit_job.slr
+++ b/nersc/performance_indicators/submit_job.slr
@@ -2,11 +2,11 @@
 #SBATCH -N 1
 #SBATCH -C cpu
 #SBATCH -q shared
-#SBATCH -J perf_ind_iqpvar_17
+#SBATCH -J perf_ind_iqpvar_20
 # SBATCH --account=jbowles
 #SBATCH -t 24:00:00
 #SBATCH --ntasks-per-node=1
-#SBATCH --cpus-per-task=10
+#SBATCH --cpus-per-task=24
 #SBATCH --output out/%j.log
 #SBATCH  --licenses=scratch
 

From 8513a930260e4f0aa46cc27c863f04658baa79c4 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Mon, 22 Apr 2024 16:56:27 +0200
Subject: [PATCH 047/100] .

---
 nersc/performance_indicators/perf_ind_variational_jax.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/nersc/performance_indicators/perf_ind_variational_jax.py b/nersc/performance_indicators/perf_ind_variational_jax.py
index 5e64ad2d..69f590ac 100644
--- a/nersc/performance_indicators/perf_ind_variational_jax.py
+++ b/nersc/performance_indicators/perf_ind_variational_jax.py
@@ -15,7 +15,7 @@
 # You only need to change this to make a different performance indicator
 
 #define model
-from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier as Model
+from qml_benchmarks.models.quantum_metric_learning import QuantumMetricLearner as Model
 
 #implementation attributes of model
 use_jax = True
@@ -24,7 +24,7 @@
 
 max_steps = 100 #the number of gradient descent steps to use to estimate the step time
 perf_ind_name = 'JAX'  #a name for the performance indicator used for naming files
-n_features = 20 #dataset dimension
+n_features = 13 #dataset dimension
 n_trials = 5 #number of trials to average over
 
 #################################

From e813a7faa0373ab9c82122d6afc84afe79a3e6c3 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Mon, 22 Apr 2024 17:00:52 +0200
Subject: [PATCH 048/100] add use jax

---
 src/qml_benchmarks/models/quantum_metric_learning.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/qml_benchmarks/models/quantum_metric_learning.py b/src/qml_benchmarks/models/quantum_metric_learning.py
index dcca5191..3fedbf3d 100644
--- a/src/qml_benchmarks/models/quantum_metric_learning.py
+++ b/src/qml_benchmarks/models/quantum_metric_learning.py
@@ -50,6 +50,8 @@ def __init__(
         max_steps=50000,
         learning_rate=0.01,
         batch_size=32,
+        use_jax = True,
+        vmap = True,
         max_vmap=4,
         jit=True,
         random_state=42,
@@ -100,6 +102,8 @@ def __init__(
         self.learning_rate = learning_rate
         self.batch_size = batch_size
         self.jit = jit
+        self.use_jax = use_jax
+        self.vmap = vmap
         self.dev_type = dev_type
         self.qnode_kwargs = qnode_kwargs
         self.scaling = scaling

From 282d90bc0f4c39d7207a4e1753d9ee7aa747df19 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Wed, 24 Apr 2024 11:16:38 +0200
Subject: [PATCH 049/100] .

---
 nersc/performance_indicators/check.py | 6 ++++++
 1 file changed, 6 insertions(+)
 create mode 100644 nersc/performance_indicators/check.py

diff --git a/nersc/performance_indicators/check.py b/nersc/performance_indicators/check.py
new file mode 100644
index 00000000..a0308044
--- /dev/null
+++ b/nersc/performance_indicators/check.py
@@ -0,0 +1,6 @@
+import numpy as np
+
+array = np.random.rand(10,10)
+
+np.savetxt("check1.txt", array)
+np.savetxt("WORKDIR/check2.txt", array)

From 2cba02236e37e820f8b3f9d3395c69dde6f60d19 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <jbowles@login24.chn.perlmutter.nersc.gov>
Date: Wed, 24 Apr 2024 03:16:29 -0700
Subject: [PATCH 050/100] update

---
 ...eparable_2d_performance_indicators_JAX.csv |   2 +
 ...2d_performance_indicators_JAX_packages.txt | 236 ++++++++++++++++++
 ...2d_performance_indicators_JAX_scontrol.txt |   0
 nersc/performance_indicators/check.py         |   1 +
 .../perf_ind_variational_jax.py               |   4 +-
 nersc/performance_indicators/submit_job.slr   |   9 +-
 nersc/pm_podman.source                        |   6 +-
 7 files changed, 249 insertions(+), 9 deletions(-)
 create mode 100644 nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_2d_performance_indicators_JAX.csv
 create mode 100644 nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_2d_performance_indicators_JAX_packages.txt
 create mode 100644 nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_2d_performance_indicators_JAX_scontrol.txt

diff --git a/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_2d_performance_indicators_JAX.csv b/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_2d_performance_indicators_JAX.csv
new file mode 100644
index 00000000..f2297385
--- /dev/null
+++ b/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_2d_performance_indicators_JAX.csv
@@ -0,0 +1,2 @@
+first_train_step,first_train_step_std,consec_train_step,consec_train_step_std,predict_time,predict_time_std,hyperparameters
+16.04134237766266,0.1555722951889038,0.005458146634728018,1.619921790228969e-05,4.981820344924927,0.05548095703125,"{'n_layers': 4, 'batch_size': 16, 'use_jax': True, 'vmap': True, 'max_steps': 100, 'jit': True}"
diff --git a/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_2d_performance_indicators_JAX_packages.txt b/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_2d_performance_indicators_JAX_packages.txt
new file mode 100644
index 00000000..0083ba07
--- /dev/null
+++ b/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_2d_performance_indicators_JAX_packages.txt
@@ -0,0 +1,236 @@
+Package                   Version
+------------------------- --------------
+absl-py                   2.1.0
+alembic                   1.13.1
+annotated-types           0.6.0
+antlr4-python3-runtime    4.9.2
+anyio                     4.3.0
+appdirs                   1.4.4
+argon2-cffi               23.1.0
+argon2-cffi-bindings      21.2.0
+arrow                     1.3.0
+asttokens                 2.4.1
+astunparse                1.6.3
+async-lru                 2.0.4
+attrs                     23.2.0
+autograd                  1.6.2
+autoray                   0.6.9
+Babel                     2.14.0
+beautifulsoup4            4.12.3
+beniget                   0.4.1
+bitstring                 3.1.7
+bleach                    6.1.0
+cachetools                5.3.3
+certifi                   2024.2.2
+cffi                      1.16.0
+charset-normalizer        3.3.2
+chex                      0.1.86
+cirq-core                 1.3.0
+cirq-pasqal               1.3.0
+click                     8.1.7
+cloudpickle               3.0.0
+colorlog                  6.8.2
+comm                      0.2.2
+contourpy                 1.2.1
+cryptography              42.0.5
+cycler                    0.12.1
+dask                      2024.4.1
+dbus-python               1.2.18
+debugpy                   1.8.1
+decorator                 4.4.2
+defusedxml                0.7.1
+diastatic-malt            2.15.1
+dill                      0.3.8
+distro                    1.7.0
+duet                      0.2.9
+etils                     1.7.0
+exceptiongroup            1.2.0
+executing                 2.0.1
+fastdtw                   0.3.4
+fastjsonschema            2.19.1
+filelock                  3.9.0
+fire                      0.6.0
+flax                      0.8.2
+fonttools                 4.50.0
+fqdn                      1.5.1
+fsspec                    2024.3.1
+future                    1.0.0
+gast                      0.5.2
+greenlet                  3.0.3
+h11                       0.14.0
+h5py                      3.10.0
+httpcore                  1.0.5
+httpx                     0.27.0
+ibm-cloud-sdk-core        3.19.2
+ibm-platform-services     0.53.2
+idna                      3.6
+importlib_metadata        7.1.0
+importlib_resources       6.4.0
+ipykernel                 6.29.4
+ipython                   8.23.0
+ipywidgets                8.1.2
+isoduration               20.11.0
+jax                       0.4.23
+jaxlib                    0.4.23
+jaxopt                    0.8.3
+jedi                      0.19.1
+Jinja2                    3.1.3
+joblib                    1.3.2
+json5                     0.9.24
+jsonpointer               2.4
+jsonschema                4.21.1
+jsonschema-specifications 2023.12.1
+jupyter                   1.0.0
+jupyter_client            8.6.1
+jupyter-console           6.6.3
+jupyter_core              5.7.2
+jupyter-events            0.10.0
+jupyter-lsp               2.2.4
+jupyter_server            2.13.0
+jupyter_server_terminals  0.5.3
+jupyterlab                4.1.5
+jupyterlab_pygments       0.3.0
+jupyterlab_server         2.25.4
+jupyterlab_widgets        3.0.10
+kiwisolver                1.4.5
+lark-parser               0.12.0
+llvmlite                  0.42.0
+locket                    1.0.0
+Mako                      1.3.2
+markdown-it-py            3.0.0
+MarkupSafe                2.1.5
+matplotlib                3.8.4
+matplotlib-inline         0.1.6
+mdurl                     0.1.2
+mistune                   3.0.2
+ml-dtypes                 0.4.0
+mpmath                    1.3.0
+msgpack                   1.0.8
+nbclient                  0.10.0
+nbconvert                 7.16.3
+nbformat                  5.10.4
+nest-asyncio              1.6.0
+networkx                  3.2.1
+notebook                  7.1.2
+notebook_shim             0.2.4
+numba                     0.59.1
+numpy                     1.26.4
+olefile                   0.46
+opt-einsum                3.3.0
+optax                     0.2.2
+optuna                    3.6.1
+orbax-checkpoint          0.5.10
+overrides                 7.7.0
+packaging                 24.0
+pandas                    2.2.1
+pandocfilters             1.5.1
+parso                     0.8.3
+partd                     1.4.1
+patsy                     0.5.6
+pbr                       6.0.0
+PennyLane                 0.35.1
+PennyLane-Catalyst        0.5.0
+PennyLane-Cirq            0.34.0
+PennyLane_Lightning       0.35.1
+PennyLane_Lightning_GPU   0.35.1
+PennyLane-qiskit          0.35.1
+PennyLane-SF              0.29.0
+pexpect                   4.9.0
+pillow                    10.3.0
+pip                       24.0
+platformdirs              4.2.0
+ply                       3.11
+prometheus_client         0.20.0
+prompt-toolkit            3.0.43
+protobuf                  5.26.1
+psutil                    5.9.8
+ptyprocess                0.7.0
+pure-eval                 0.2.2
+pycparser                 2.22
+pydantic                  2.6.4
+pydantic_core             2.16.3
+pydantic-settings         2.2.1
+pydot                     2.0.0
+Pygments                  2.17.2
+PyGObject                 3.42.1
+PyJWT                     2.8.0
+pylatexenc                2.10
+pyparsing                 3.1.2
+pyspnego                  0.10.2
+python-dateutil           2.9.0.post0
+python-dotenv             1.0.1
+python-json-logger        2.0.7
+pythran                   0.10.0
+pytz                      2024.1
+PyYAML                    6.0.1
+pyzmq                     25.1.2
+qiskit                    1.0.2
+qiskit-aer                0.14.0.1
+qiskit-algorithms         0.3.0
+qiskit-ibm-experiment     0.4.7
+qiskit-ibm-provider       0.10.0
+qiskit-ibm-runtime        0.20.0
+qiskit-machine-learning   0.7.2
+qml_benchmarks            0.1
+qtconsole                 5.5.1
+QtPy                      2.4.1
+quantum-blackbird         0.5.0
+quantum-xir               0.2.2
+referencing               0.34.0
+requests                  2.31.0
+requests-ntlm             1.2.0
+rfc3339-validator         0.1.4
+rfc3986-validator         0.1.1
+rich                      13.7.1
+rpds-py                   0.18.0
+ruamel.yaml               0.18.6
+ruamel.yaml.clib          0.2.8
+rustworkx                 0.14.2
+scikit-learn              1.4.1.post1
+scipy                     1.11.4
+seaborn                   0.13.2
+semantic-version          2.10.0
+Send2Trash                1.8.2
+setuptools                59.6.0
+six                       1.16.0
+sniffio                   1.3.1
+sortedcontainers          2.4.0
+soupsieve                 2.5
+SQLAlchemy                2.0.29
+ssh-import-id             5.11
+stack-data                0.6.3
+statsmodels               0.14.1
+stevedore                 5.2.0
+StrawberryFields          0.23.0
+symengine                 0.11.0
+sympy                     1.12
+tensorstore               0.1.56
+termcolor                 2.4.0
+terminado                 0.18.1
+thewalrus                 0.21.0
+threadpoolctl             3.4.0
+tinycss2                  1.2.1
+toml                      0.10.2
+tomli                     2.0.1
+tomlkit                   0.12.4
+toolz                     0.12.1
+torch                     2.2.2+cpu
+torchaudio                2.2.2+cpu
+torchvision               0.17.2+cpu
+tornado                   6.4
+tqdm                      4.66.2
+traitlets                 5.14.2
+types-python-dateutil     2.9.0.20240316
+typing_extensions         4.10.0
+tzdata                    2024.1
+uri-template              1.3.0
+urllib3                   2.2.1
+wcwidth                   0.2.13
+webcolors                 1.13
+webencodings              0.5.1
+websocket-client          1.7.0
+websockets                12.0
+wheel                     0.37.1
+widgetsnbextension        4.0.10
+xanadu-cloud-client       0.3.2
+zipp                      3.18.1
diff --git a/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_2d_performance_indicators_JAX_scontrol.txt b/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_2d_performance_indicators_JAX_scontrol.txt
new file mode 100644
index 00000000..e69de29b
diff --git a/nersc/performance_indicators/check.py b/nersc/performance_indicators/check.py
index a0308044..c27dfdf3 100644
--- a/nersc/performance_indicators/check.py
+++ b/nersc/performance_indicators/check.py
@@ -2,5 +2,6 @@
 
 array = np.random.rand(10,10)
 
+print('testing')
 np.savetxt("check1.txt", array)
 np.savetxt("WORKDIR/check2.txt", array)
diff --git a/nersc/performance_indicators/perf_ind_variational_jax.py b/nersc/performance_indicators/perf_ind_variational_jax.py
index 69f590ac..6bd94e52 100644
--- a/nersc/performance_indicators/perf_ind_variational_jax.py
+++ b/nersc/performance_indicators/perf_ind_variational_jax.py
@@ -15,7 +15,7 @@
 # You only need to change this to make a different performance indicator
 
 #define model
-from qml_benchmarks.models.quantum_metric_learning import QuantumMetricLearner as Model
+from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier as Model
 
 #implementation attributes of model
 use_jax = True
@@ -24,7 +24,7 @@
 
 max_steps = 100 #the number of gradient descent steps to use to estimate the step time
 perf_ind_name = 'JAX'  #a name for the performance indicator used for naming files
-n_features = 13 #dataset dimension
+n_features = 15 #dataset dimension
 n_trials = 5 #number of trials to average over
 
 #################################
diff --git a/nersc/performance_indicators/submit_job.slr b/nersc/performance_indicators/submit_job.slr
index 8f1ad1e9..d9924d7b 100644
--- a/nersc/performance_indicators/submit_job.slr
+++ b/nersc/performance_indicators/submit_job.slr
@@ -2,15 +2,16 @@
 #SBATCH -N 1
 #SBATCH -C cpu
 #SBATCH -q shared
-#SBATCH -J perf_ind_iqpvar_20
+#SBATCH -J perf_ind_iqp15
 # SBATCH --account=jbowles
-#SBATCH -t 24:00:00
+#SBATCH -t 01:00:00
 #SBATCH --ntasks-per-node=1
-#SBATCH --cpus-per-task=24
+#SBATCH --cpus-per-task=6
 #SBATCH --output out/%j.log
 #SBATCH  --licenses=scratch
 
 source ../pm_podman.source
 bash ../installs.sh
 
-time srun -n 1  python3 ../perf_ind_variational_jax.py
+time srun -n 1  python3 perf_ind_variational_jax.py
+
diff --git a/nersc/pm_podman.source b/nersc/pm_podman.source
index 2b8000f7..5d69d8e5 100644
--- a/nersc/pm_podman.source
+++ b/nersc/pm_podman.source
@@ -2,14 +2,14 @@
 #  salloc -q interactive -C cpu -t 4:00:00 -A nstaff 
 
 IMG=balewski/ubu22-pennylane:p3
-CFSH=/global/cfs/cdirs/m4139/
+CFSH=/global/cfs/cdirs/m4139
 
 echo launch image $IMG
 echo you are launching Podman image ...  remember to exit
 
 JNB_PORT=' '
-BASE_DIR=/qml-benchmarks   # here git has home
-WORK_DIR=$BASE_DIR/nersc
+BASE_DIR=qml-benchmarks   # here git has home
+WORK_DIR=/$BASE_DIR/nersc
 
 podman-hpc run -it \
    --volume $HOME:/home \

From 45c4d726a36b981b536ef963924f1acb9423b711 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <jbowles@x1004c0s5b0n0h0.chn.perlmutter.nersc.gov>
Date: Wed, 24 Apr 2024 06:38:24 -0700
Subject: [PATCH 051/100] update

---
 nersc/oom                                           | 0
 nersc/performance_indicators/perf_ind_kernel_jax.py | 4 ++--
 nersc/performance_indicators/submit_job.slr         | 6 ++++--
 nersc/pm_podman.source                              | 2 ++
 4 files changed, 8 insertions(+), 4 deletions(-)
 create mode 100644 nersc/oom

diff --git a/nersc/oom b/nersc/oom
new file mode 100644
index 00000000..e69de29b
diff --git a/nersc/performance_indicators/perf_ind_kernel_jax.py b/nersc/performance_indicators/perf_ind_kernel_jax.py
index 494c1368..0984dd77 100644
--- a/nersc/performance_indicators/perf_ind_kernel_jax.py
+++ b/nersc/performance_indicators/perf_ind_kernel_jax.py
@@ -23,7 +23,7 @@
 jit = True
 
 perf_ind_name = 'JAX'  #a name for the performance indicator used for naming files
-n_features = 15 #dataset dimension
+n_features = 16 #dataset dimension
 
 #################################
 
@@ -79,4 +79,4 @@
 
 # make an empty text file to store the scontrol data
 with open(perf_ind_name+'/'+scontrol_filename, 'w') as file:
-    pass
\ No newline at end of file
+    pass
diff --git a/nersc/performance_indicators/submit_job.slr b/nersc/performance_indicators/submit_job.slr
index d9924d7b..d8f4ee53 100644
--- a/nersc/performance_indicators/submit_job.slr
+++ b/nersc/performance_indicators/submit_job.slr
@@ -10,8 +10,10 @@
 #SBATCH --output out/%j.log
 #SBATCH  --licenses=scratch
 
+
+
 source ../pm_podman.source
-bash ../installs.sh
+bash installs.sh
 
-time srun -n 1  python3 perf_ind_variational_jax.py
 
+time srun -n 1  python3 performance_indicators/perf_ind_variational_jax.py
diff --git a/nersc/pm_podman.source b/nersc/pm_podman.source
index 5d69d8e5..f267304b 100644
--- a/nersc/pm_podman.source
+++ b/nersc/pm_podman.source
@@ -4,6 +4,8 @@
 IMG=balewski/ubu22-pennylane:p3
 CFSH=/global/cfs/cdirs/m4139
 
+# export PODMANHPC_ADDITIONAL_STORES=/global/cfs/cdirs/nstaff/balewski/podman_common
+
 echo launch image $IMG
 echo you are launching Podman image ...  remember to exit
 

From 2cb5917b58928a7062297dbcf15b91275b6e1f0a Mon Sep 17 00:00:00 2001
From: Jan Balewski <janstar1122@gmail.com>
Date: Thu, 25 Apr 2024 14:08:40 -0700
Subject: [PATCH 052/100] prototype of Slurm job w/ Podman works

---
 nersc/performance_indicators/submit_job.slr  | 19 --------
 nersc/performance_indicators/submit_job1.slr | 46 ++++++++++++++++++++
 nersc/performance_indicators/wrap_podman.sh  | 23 ++++++++++
 3 files changed, 69 insertions(+), 19 deletions(-)
 delete mode 100644 nersc/performance_indicators/submit_job.slr
 create mode 100755 nersc/performance_indicators/submit_job1.slr
 create mode 100755 nersc/performance_indicators/wrap_podman.sh

diff --git a/nersc/performance_indicators/submit_job.slr b/nersc/performance_indicators/submit_job.slr
deleted file mode 100644
index d8f4ee53..00000000
--- a/nersc/performance_indicators/submit_job.slr
+++ /dev/null
@@ -1,19 +0,0 @@
-#!/bin/bash
-#SBATCH -N 1
-#SBATCH -C cpu
-#SBATCH -q shared
-#SBATCH -J perf_ind_iqp15
-# SBATCH --account=jbowles
-#SBATCH -t 01:00:00
-#SBATCH --ntasks-per-node=1
-#SBATCH --cpus-per-task=6
-#SBATCH --output out/%j.log
-#SBATCH  --licenses=scratch
-
-
-
-source ../pm_podman.source
-bash installs.sh
-
-
-time srun -n 1  python3 performance_indicators/perf_ind_variational_jax.py
diff --git a/nersc/performance_indicators/submit_job1.slr b/nersc/performance_indicators/submit_job1.slr
new file mode 100755
index 00000000..0c0aaa50
--- /dev/null
+++ b/nersc/performance_indicators/submit_job1.slr
@@ -0,0 +1,46 @@
+#!/bin/bash
+#SBATCH -N 1
+#SBATCH -C cpu
+#SBATCH -q shared -t 01:00:00
+# SBATCH -q debug -t 5:00  # charged for full node
+#SBATCH -J perf_ind_iqp15
+# SBATCH --account=jbowles
+#SBATCH --ntasks-per-node=1
+#SBATCH --cpus-per-task=6
+#SBATCH --output out1/%j.log
+#SBATCH  --licenses=scratch
+
+
+echo 'S:start'
+
+# .... container
+# salloc -q interactive -C cpu -t 4:00:00 -A nstaff
+
+IMG=balewski/ubu22-pennylane:p4
+echo use image $IMG
+POD_PUB=/dvs_ro/cfs/cdirs/nstaff/balewski/podman_common/  
+export PODMANHPC_ADDITIONAL_STORES=$POD_PUB
+
+#.... output sandbox 
+CODE_DIR=`pwd`
+myJob=${SLURM_JOBID}
+outPath=$SCRATCH/tmp_pennylane_jobs/$myJob
+echo outPath=$outPath
+mkdir -p $outPath
+cp -rp *.py wrap_podman.sh   $outPath
+
+#... location of input on CSF
+CFSH=/global/cfs/cdirs/m4139
+BASE_DIR=qml-benchmarks   # here git has home
+WORK_DIR=/$BASE_DIR/nersc
+
+#... the task to be executed
+CMD=" python -u  perf_ind_variational_jax.py "
+
+cd $outPath
+G=1
+echo 'S:ready to run '; pwd
+srun -n $G  ./wrap_podman.sh $IMG " $CMD "  $outPath  $CFSH $BASE_DIR $WORK_DIR
+
+echo 'S:end '; date
+
diff --git a/nersc/performance_indicators/wrap_podman.sh b/nersc/performance_indicators/wrap_podman.sh
new file mode 100755
index 00000000..f601c873
--- /dev/null
+++ b/nersc/performance_indicators/wrap_podman.sh
@@ -0,0 +1,23 @@
+#!/bin/bash
+echo W:myRank is $SLURM_PROCID
+IMG=$1
+CMD=$2
+outPath=$3
+CFSH=$4
+BASE_DIR=$5
+WORK_DIR=$6
+
+if [ $SLURM_PROCID -eq 0 ] ; then 
+   echo W:IMG=$IMG 
+   echo W:CMD=$CMD
+   #echo Q:fire $
+fi
+echo 'W:start'
+time podman-hpc run -it \
+	   --volume $HOME:/home \
+	   --volume $CFSH/$BASE_DIR:/$BASE_DIR \
+	   --volume $CFSH/$WORK_DIR:$WORK_DIR \
+	   -e HDF5_USE_FILE_LOCKING='FALSE' \
+	   --workdir $WORK_DIR  \
+	   $IMG $CMD 
+echo 'W:done'

From feaf6679c52687f705f5e1ee21c9471c917fc1b2 Mon Sep 17 00:00:00 2001
From: Jan Balewski <janstar1122@gmail.com>
Date: Fri, 26 Apr 2024 13:29:59 -0700
Subject: [PATCH 053/100] double dream works

---
 .gitignore                                    |   2 +-
 nersc/container/Dockerfile.ubu22-PennyLane    |  51 ++++++++
 nersc/container/Readme                        | 114 ++++++++++++++++++
 nersc/installs.sh                             |   5 -
 .../perf_ind_variational_jax.py               |   4 +-
 nersc/submit_job2.slr                         |  51 ++++++++
 nersc/wrap_podman2.sh                         |  34 ++++++
 7 files changed, 253 insertions(+), 8 deletions(-)
 create mode 100644 nersc/container/Dockerfile.ubu22-PennyLane
 create mode 100644 nersc/container/Readme
 delete mode 100644 nersc/installs.sh
 create mode 100755 nersc/submit_job2.slr
 create mode 100755 nersc/wrap_podman2.sh

diff --git a/.gitignore b/.gitignore
index a2b2f4d7..48fb0467 100644
--- a/.gitignore
+++ b/.gitignore
@@ -179,5 +179,5 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 .idea/
-
+old/
 
diff --git a/nersc/container/Dockerfile.ubu22-PennyLane b/nersc/container/Dockerfile.ubu22-PennyLane
new file mode 100644
index 00000000..ec5bd051
--- /dev/null
+++ b/nersc/container/Dockerfile.ubu22-PennyLane
@@ -0,0 +1,51 @@
+FROM ubuntu:22.04
+
+# time podman-hpc build -f Dockerfile.ubu22-PennyLane -t balewski/ubu22-pennylane:p5   .
+#>>> real	6m47.959s
+
+# podman commit  -a "Jan Balewski"  c3d48cf15876XXX  balewski/XXXkit-qml:p2ch
+# on laptop:
+#  podman run -it balewski/ubu22-pennylane:p0 bash
+# time python3 -c 'import pennylane'
+
+# is needed by  tzdada which sets sth for one of libs in section 2
+ARG DEBIAN_FRONTEND=noninteractive
+ENV TZ=America/Los_Angeles
+
+RUN echo 1a-AAAAAAAAAAAAAAAAAAAAAAAAAAAAA  OS update   && \
+    apt-get update && \
+    apt-get install -y  locales autoconf automake gcc g++ make vim wget  ssh  openssh-server sudo   git emacs aptitude build-essential xterm python3-pip   python3-tk python3-scipy python3-dev iputils-ping net-tools screen feh hdf5-tools python3-bitstring mlocate graphviz  tzdata x11-apps && \
+    apt-get clean all
+
+# Forbid installing qiskit-terra  for qiskit 1.0 
+RUN echo 2b-AAAAAAAAAAAAAAAAAAAAAAAAAAAAA   Qiskit 1.0 libs  && \
+    pip3 install -c https://qisk.it/1-0-constraints qiskit[visualization] qiskit-ibm-runtime qiskit-machine-learning qiskit_ibm_experiment qiskit-aer
+
+
+RUN echo 2c-AAAAAAAAAAAAAAAAAAAAAAAAAAAAA   ML+QML  libs  && \
+    pip3 install scikit-learn pandas seaborn[stats]  networkx[default] optuna
+
+
+RUN echo 2d-AAAAAAAAAAAAAAAAAAAAAAAAAAAAA   python libs  && \
+    pip3 install --upgrade pip && \
+    pip3 install matplotlib h5py ruamel.yaml  scipy    jupyter notebook  bitstring
+
+# notes: python3-tk instals TK for matplotlib to display graphic
+RUN echo 2e-AAAAAAAAAAAAAAAAAAAAAAAAAAAAA   python libs  && \
+    pip3 install  --upgrade Pillow
+
+
+# based on  https://pennylane.ai/install/
+# based on  https://pytorch.org/get-started/locally/
+
+RUN echo 3a-AAAAAAAAAAAAAAAAAAAAAAAAAAAAA   pennylane  && \
+    pip3 install pennylane --upgrade && \
+    pip3 install pennylane-lightning pennylane-lightning[gpu] pennylane-sf pennylane-qiskit pennylane-cirq pennylane-catalyst scipy==1.11.4  && \
+    pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
+
+RUN echo 3b-AAAAAAAAAAAAAAAAAAAAAAAAAAAAA   JAX  && \
+    pip3 install -U "jax[cpu]" && \
+    pip3 install jaxopt optax
+
+
+
diff --git a/nersc/container/Readme b/nersc/container/Readme
new file mode 100644
index 00000000..797be716
--- /dev/null
+++ b/nersc/container/Readme
@@ -0,0 +1,114 @@
+Usage of Podman container for PennyLane
+
+details are here
+https://docs.google.com/document/d/1LHYlTXtOyA5vZSrJF0IdZf3L88DcguMWXyHAPLV-MvU/edit?usp=sharing
+
+
+Summary:
+This is an instruction on how to build and use a Podman image.
+The example consists of three parts:
+A) Building the image, which I've already completed on PM.
+B) Running the image on PM in three modes: simple interactive, advanced interactive, and a Slurm job that executes your code inside a podman image.
+C) Running the same image on your laptop. You'll need to build it as described in A), but then running it is straightforward.
+
+
+
+= = = 
+A = = = = = = = = = =  building container (one time) = = = = = = = = =
+= = = 
+Dockerfile: Dockerfile.ubu22-PennyLane
+Laptop: 
+    podman build -f Dockerfile.ubu22-PennyLane -t balewski/ubu22-pennylane:p5   .
+Perlmutter
+   podman-hpc build -f Dockerfile.ubu22-PennyLane -t balewski/ubu22-pennylane:p5   .
+   POD_PUB=/cfs/cdirs/nstaff/balewski/podman_common/
+   podman-hpc --squash-dir /global/$POD_PUB migrate  balewski/ubu22-pennylane:p5 
+   Note,  the instruction above is for user:balewski whohas access to CFS/nstaff - change the strings as needed
+
+
+
+= = = 
+B = = = = = = = = = =   Perlmutter @ NERSC = = = = = = = = =
+= = = 
+
+B.1  - - - - - -  simple interactive use  on a worker node - - - - - - - 
+
+ssh perlmutter
+export PS1="\h:\W> "  # to reduce too long names of directoriesexport PS1="\h:\W> "  # to reduce too long names of directories
+
+POD_PUB=/cfs/cdirs/nstaff/balewski/podman_common/
+export PODMANHPC_ADDITIONAL_STORES=/dvs_ro/$POD_PUB
+cd  $SCRATCH/tmp
+salloc -q interactive -C cpu -t 4:00:00 -A nstaff
+podman-hpc run -it  balewski/ubu22-pennylane:p5 bash
+ python3 -c 'import pennylane'
+exit
+
+
+B.2  - - - - - -  advanced interactive use  on a worker node - - - - - - - 
+copy & customize  https://github.com/balewski/quantumMind/blob/main/PennyLane/pm_podman.source   
+to configure shortcuts to your liking
+
+ssh perlmutter
+salloc -q interactive -C cpu -t 4:00:00 -A nstaff
+source pm_podman.source 
+
+[ do your work ]
+
+exit  #from podman
+exit # from salloc
+
+
+
+B.3  - - - - - -  advanced : Slurm job with multiple task  using Podman image- - - - - - - 
+  copy and customize https://github.com/balewski/quantumMind/tree/main/PennyLane/qml_intro/
+batchPodman.slr
+wrap_podman.sh
+ 
+You will edit code stored  in CFS area but running job will read/write from/to SCRATCH using its copy.  Do NOT start Podman image manually
+
+ssh perlmutter
+ssh CFS  (to your area )
+sbatch batchPodman.slr
+
+
+B.4  - - - - - -  advanced :  install project spceciffic software, one time
+NOT possible now - working on it
+
+
+B.5  - - - - - -  advanced : jupiter notebook on the laptop powered by Podman PennyLane image running on PM - - - - - - - 
+  Easy to do 
+
+
+
+
+= = = 
+C = = = = = = = = = =   Laptop = = = = = = = = =
+= = = 
+
+C.1  - - - - - -  simple interactive use  on a laptop - - - - - - - 
+$ podman run -it  balewski/ubu22-pennylane:p3 bash
+ python3 -c 'import pennylane'
+exit
+
+
+C.2  - - - - - -  advanced w/ predefined volume mounts on a laptop  - - - - - - - 
+copy & customize  https://github.com/balewski/quantumMind/blob/main/PennyLane/laptop_podman.source 
+to configure shortcuts to your liking
+
+source laptop_podman.source 
+
+[ do your work ]
+
+exit  #from podman
+
+If you need to reset podman, do :   source restart_podman.source
+
+
+C.4  - - - - - -  advanced : jupiter notebook on the laptop powered by Podman PennyLane image running on the laptop - - - - - - - 
+
+source laptop_podman.source jnb 
+(exec inside running image)  jupyter notebook --ip 0.0.0.0 --no-browser --allow-root --port  8833 
+copy http  string and paste into local browser , e.g:   http://127.0.0.1:8833/tree?token=7c5cdf5e5c4f1a9d2a616a739988132d59f1a7ca3f4c0779
+
+Troubleshooting: if JNB remains blank change port ID in  laptop_podman.source  and restart the image
diff --git a/nersc/installs.sh b/nersc/installs.sh
deleted file mode 100644
index ed740537..00000000
--- a/nersc/installs.sh
+++ /dev/null
@@ -1,5 +0,0 @@
-#!/bin/bash
-
-pip install ../
-pip install pennylane-catalyst
-pip install scipy==1.11.4
\ No newline at end of file
diff --git a/nersc/performance_indicators/perf_ind_variational_jax.py b/nersc/performance_indicators/perf_ind_variational_jax.py
index 6bd94e52..f36a21fd 100644
--- a/nersc/performance_indicators/perf_ind_variational_jax.py
+++ b/nersc/performance_indicators/perf_ind_variational_jax.py
@@ -24,8 +24,8 @@
 
 max_steps = 100 #the number of gradient descent steps to use to estimate the step time
 perf_ind_name = 'JAX'  #a name for the performance indicator used for naming files
-n_features = 15 #dataset dimension
-n_trials = 5 #number of trials to average over
+n_features = 2 #dataset dimension
+n_trials = 2 #number of trials to average over
 
 #################################
 
diff --git a/nersc/submit_job2.slr b/nersc/submit_job2.slr
new file mode 100755
index 00000000..ba912a42
--- /dev/null
+++ b/nersc/submit_job2.slr
@@ -0,0 +1,51 @@
+#!/bin/bash
+#SBATCH -N 1
+#SBATCH -C cpu
+# SBATCH -q shared -t 01:00:00
+#SBATCH -q debug -t 5:00  # charged for full node
+#SBATCH -J perf_ind_iqp15
+# SBATCH --account=jbowles
+#SBATCH --ntasks-per-node=1
+#SBATCH --cpus-per-task=6
+#SBATCH --output out1/%j.log
+#SBATCH  --licenses=scratch
+
+
+echo 'S:start'
+
+# .... container
+# salloc -q interactive -C cpu -t 4:00:00 -A nstaff
+
+IMG=balewski/ubu22-pennylane:p5
+echo use image $IMG
+POD_PUB=/dvs_ro/cfs/cdirs/nstaff/balewski/podman_common/  
+export PODMANHPC_ADDITIONAL_STORES=$POD_PUB
+
+#.... output sandbox 
+CODE_DIR=`pwd`
+myJob=${SLURM_JOBID}
+outPath=$SCRATCH/tmp_pennylane_jobs/$myJob
+echo outPath=$outPath
+mkdir -p $outPath
+cp -rp performance_indicators *.slr *.sh   $outPath
+
+#... location of input on CSF
+CFSH=/global/cfs/cdirs/m4139  # for Joseph
+#CFSH=/pscratch/sd/b/balewski  # for Jan, do not remove
+BASE_DIR=/qml-benchmarks   # here git has home
+WORK_DIR=$BASE_DIR/nersc
+
+#... the task to be executed
+CMD=" python3 -u  performance_indicators/perf_ind_variational_jax.py "
+#CMD=" python3 -u -c \"import pennylane \"; echo done1  "
+#CMD=" python3 -u -c \"import qml_benchmarks \";  echo done2 host  "
+
+
+cd $outPath
+G=1
+echo 'S:ready to run '; pwd
+srun -n $G  ./wrap_podman2.sh $IMG " $CMD "  $outPath  $CFSH $BASE_DIR $WORK_DIR
+
+sleep 1
+echo 'S:end '; date
+
diff --git a/nersc/wrap_podman2.sh b/nersc/wrap_podman2.sh
new file mode 100755
index 00000000..333112c8
--- /dev/null
+++ b/nersc/wrap_podman2.sh
@@ -0,0 +1,34 @@
+#!/bin/bash
+echo W:myRank is $SLURM_PROCID
+IMG=$1
+CMD=$2
+outPath=$3
+CFSH=$4
+BASE_DIR=$5
+WORK_DIR=$6
+
+if [ $SLURM_PROCID -eq 0 ] ; then 
+   echo W:IMG=$IMG 
+   echo W:CMD=$CMD
+   #echo Q:fire $
+fi
+
+echo W:BASE_DIR=$BASE_DIR
+echo 'W:start podman'
+podman-hpc run -it \
+	   --volume $HOME:/home \
+	   --volume $CFSH/$BASE_DIR:/$BASE_DIR \
+	   --volume $CFSH/$WORK_DIR:$WORK_DIR \
+	   -e HDF5_USE_FILE_LOCKING='FALSE' \
+	   --workdir $WORK_DIR \
+	   $IMG <<EOF 
+cd $BASE_DIR
+echo P:pwd; pwd
+pip3 install .
+cd -
+$CMD
+echo P:done
+exit
+EOF
+  
+echo 'W:done podman'

From 8aad113f9a664788e542bb3770703e654ed7c25f Mon Sep 17 00:00:00 2001
From: Jan Balewski <janstar1122@gmail.com>
Date: Fri, 26 Apr 2024 13:33:53 -0700
Subject: [PATCH 054/100] ok

---
 nersc/submit_job2.slr | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/nersc/submit_job2.slr b/nersc/submit_job2.slr
index ba912a42..b51165b4 100755
--- a/nersc/submit_job2.slr
+++ b/nersc/submit_job2.slr
@@ -1,10 +1,9 @@
 #!/bin/bash
 #SBATCH -N 1
 #SBATCH -C cpu
-# SBATCH -q shared -t 01:00:00
-#SBATCH -q debug -t 5:00  # charged for full node
+#SBATCH -q shared -t 01:00:00
+# SBATCH -q debug -t 5:00  # charged for full node
 #SBATCH -J perf_ind_iqp15
-# SBATCH --account=jbowles
 #SBATCH --ntasks-per-node=1
 #SBATCH --cpus-per-task=6
 #SBATCH --output out1/%j.log

From 63ce6dd7a5146a4de7eb602a33af66a08371e923 Mon Sep 17 00:00:00 2001
From: josephbowles <bowles.physics@gmail.com>
Date: Mon, 29 Apr 2024 08:18:25 -0700
Subject: [PATCH 055/100] working sbatch script

---
 nersc/submit_job2.slr | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/nersc/submit_job2.slr b/nersc/submit_job2.slr
index b51165b4..30783d88 100755
--- a/nersc/submit_job2.slr
+++ b/nersc/submit_job2.slr
@@ -15,9 +15,10 @@ echo 'S:start'
 # .... container
 # salloc -q interactive -C cpu -t 4:00:00 -A nstaff
 
-IMG=balewski/ubu22-pennylane:p5
+IMG=jbowles/ubu22-pennylane:p5
 echo use image $IMG
-POD_PUB=/dvs_ro/cfs/cdirs/nstaff/balewski/podman_common/  
+POD_PUB=/global/cfs/cdirs/m4139/qml-benchmarks/nersc/podman/
+# POD_PUB=/dvs_ro/cfs/cdirs/nstaff/balewski/podman_common/  
 export PODMANHPC_ADDITIONAL_STORES=$POD_PUB
 
 #.... output sandbox 

From 375b141f45c5ff23af045b19c4949b86775b1f49 Mon Sep 17 00:00:00 2001
From: josephbowles <bowles.physics@gmail.com>
Date: Tue, 30 Apr 2024 01:51:13 -0700
Subject: [PATCH 056/100] update jax version

---
 nersc/Dockerfile.ubu22-PennyLane | 51 ++++++++++++++++++++++++++++++++
 1 file changed, 51 insertions(+)
 create mode 100644 nersc/Dockerfile.ubu22-PennyLane

diff --git a/nersc/Dockerfile.ubu22-PennyLane b/nersc/Dockerfile.ubu22-PennyLane
new file mode 100644
index 00000000..56b6e431
--- /dev/null
+++ b/nersc/Dockerfile.ubu22-PennyLane
@@ -0,0 +1,51 @@
+FROM ubuntu:22.04
+
+# time podman-hpc build -f Dockerfile.ubu22-PennyLane -t balewski/ubu22-pennylane:p5   .
+#>>> real	6m47.959s
+
+# podman commit  -a "Jan Balewski"  c3d48cf15876XXX  balewski/XXXkit-qml:p2ch
+# on laptop:
+#  podman run -it balewski/ubu22-pennylane:p0 bash
+# time python3 -c 'import pennylane'
+
+# is needed by  tzdada which sets sth for one of libs in section 2
+ARG DEBIAN_FRONTEND=noninteractive
+ENV TZ=America/Los_Angeles
+
+RUN echo 1a-AAAAAAAAAAAAAAAAAAAAAAAAAAAAA  OS update   && \
+    apt-get update && \
+    apt-get install -y  locales autoconf automake gcc g++ make vim wget  ssh  openssh-server sudo   git emacs aptitude build-essential xterm python3-pip   python3-tk python3-scipy python3-dev iputils-ping net-tools screen feh hdf5-tools python3-bitstring mlocate graphviz  tzdata x11-apps && \
+    apt-get clean all
+
+# Forbid installing qiskit-terra  for qiskit 1.0 
+RUN echo 2b-AAAAAAAAAAAAAAAAAAAAAAAAAAAAA   Qiskit 1.0 libs  && \
+    pip3 install -c https://qisk.it/1-0-constraints qiskit[visualization] qiskit-ibm-runtime qiskit-machine-learning qiskit_ibm_experiment qiskit-aer
+
+
+RUN echo 2c-AAAAAAAAAAAAAAAAAAAAAAAAAAAAA   ML+QML  libs  && \
+    pip3 install scikit-learn pandas seaborn[stats]  networkx[default] optuna
+
+
+RUN echo 2d-AAAAAAAAAAAAAAAAAAAAAAAAAAAAA   python libs  && \
+    pip3 install --upgrade pip && \
+    pip3 install matplotlib h5py ruamel.yaml  scipy    jupyter notebook  bitstring
+
+# notes: python3-tk instals TK for matplotlib to display graphic
+RUN echo 2e-AAAAAAAAAAAAAAAAAAAAAAAAAAAAA   python libs  && \
+    pip3 install  --upgrade Pillow
+
+
+# based on  https://pennylane.ai/install/
+# based on  https://pytorch.org/get-started/locally/
+
+RUN echo 3a-AAAAAAAAAAAAAAAAAAAAAAAAAAAAA   pennylane  && \
+    pip3 install pennylane --upgrade && \
+    pip3 install pennylane-lightning pennylane-lightning[gpu] pennylane-sf pennylane-qiskit pennylane-cirq pennylane-catalyst scipy==1.11.4  && \
+    pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
+
+RUN echo 3b-AAAAAAAAAAAAAAAAAAAAAAAAAAAAA   JAX  && \
+    pip3 install "jax[cpu]"==0.4.23
+    pip3 install jaxopt optax
+
+
+

From 7b40ba11b3291b72d0ce49df96bd79af258538f5 Mon Sep 17 00:00:00 2001
From: josephbowles <bowles.physics@gmail.com>
Date: Tue, 30 Apr 2024 03:59:36 -0700
Subject: [PATCH 057/100] working sbatch

---
 nersc/Dockerfile.ubu22-PennyLane              |   2 +-
 ...eparable_3d_performance_indicators_JAX.csv |   2 -
 ...3d_performance_indicators_JAX_packages.txt | 166 ------------
 ...parable_15d_performance_indicators_JAX.csv |   2 +
 ...d_performance_indicators_JAX_packages.txt} |  80 +++---
 ...d_performance_indicators_JAX_scontrol.txt} |   0
 ...eparable_2d_performance_indicators_JAX.csv |   2 +
 ...2d_performance_indicators_JAX_packages.txt | 236 ++++++++++++++++++
 ...d_performance_indicators_JAX_scontrol.txt} |   0
 ...eparable_3d_performance_indicators_JAX.csv |   2 -
 ...3d_performance_indicators_JAX_packages.txt | 166 ------------
 ...eparable_2d_performance_indicators_JAX.csv |   2 -
 ...2d_performance_indicators_JAX_scontrol.txt |   0
 .../perf_ind_kernel_jax.py                    |   5 +-
 .../perf_ind_variational_jax.py               |   5 +-
 nersc/submit_job2.slr                         |   6 +-
 16 files changed, 292 insertions(+), 384 deletions(-)
 delete mode 100644 nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_3d_performance_indicators_JAX.csv
 delete mode 100644 nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_3d_performance_indicators_JAX_packages.txt
 create mode 100644 nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_15d_performance_indicators_JAX.csv
 rename nersc/performance_indicators/JAX/{QuantumMetricLearner_linearly_separable_2d_performance_indicators_JAX_packages.txt => IQPVariationalClassifier_linearly_separable_15d_performance_indicators_JAX_packages.txt} (82%)
 rename nersc/performance_indicators/JAX/{IQPKernelClassifier_linearly_separable_3d_performance_indicators_JAX_scontrol.txt => IQPVariationalClassifier_linearly_separable_15d_performance_indicators_JAX_scontrol.txt} (100%)
 create mode 100644 nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_2d_performance_indicators_JAX.csv
 create mode 100644 nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_2d_performance_indicators_JAX_packages.txt
 rename nersc/performance_indicators/JAX/{IQPVariationalClassifier_linearly_separable_3d_performance_indicators_JAX_scontrol.txt => IQPVariationalClassifier_linearly_separable_2d_performance_indicators_JAX_scontrol.txt} (100%)
 delete mode 100644 nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_3d_performance_indicators_JAX.csv
 delete mode 100644 nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_3d_performance_indicators_JAX_packages.txt
 delete mode 100644 nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_2d_performance_indicators_JAX.csv
 delete mode 100644 nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_2d_performance_indicators_JAX_scontrol.txt

diff --git a/nersc/Dockerfile.ubu22-PennyLane b/nersc/Dockerfile.ubu22-PennyLane
index 56b6e431..d4849da5 100644
--- a/nersc/Dockerfile.ubu22-PennyLane
+++ b/nersc/Dockerfile.ubu22-PennyLane
@@ -44,7 +44,7 @@ RUN echo 3a-AAAAAAAAAAAAAAAAAAAAAAAAAAAAA   pennylane  && \
     pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
 
 RUN echo 3b-AAAAAAAAAAAAAAAAAAAAAAAAAAAAA   JAX  && \
-    pip3 install "jax[cpu]"==0.4.23
+    pip3 install "jax[cpu]"==0.4.23 && \
     pip3 install jaxopt optax
 
 
diff --git a/nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_3d_performance_indicators_JAX.csv b/nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_3d_performance_indicators_JAX.csv
deleted file mode 100644
index d0cc1ae5..00000000
--- a/nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_3d_performance_indicators_JAX.csv
+++ /dev/null
@@ -1,2 +0,0 @@
-construct_kernel_time,training_time,predict_time,hyperparameters
-1.1291069984436035,1.1299769878387451,1.092557668685913,"{'repeats': 10, 'use_jax': True, 'vmap': True, 'jit': True}"
diff --git a/nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_3d_performance_indicators_JAX_packages.txt b/nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_3d_performance_indicators_JAX_packages.txt
deleted file mode 100644
index 8859d638..00000000
--- a/nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_3d_performance_indicators_JAX_packages.txt
+++ /dev/null
@@ -1,166 +0,0 @@
-Package                   Version      Editable project location
-------------------------- ------------ ----------------------------------------------------------
-absl-py                   2.0.0
-anyio                     4.2.0
-appdirs                   1.4.4
-appnope                   0.1.3
-argon2-cffi               23.1.0
-argon2-cffi-bindings      21.2.0
-arrow                     1.3.0
-asttokens                 2.4.1
-astunparse                1.6.3
-async-lru                 2.0.4
-attrs                     23.1.0
-autograd                  1.6.2
-autoray                   0.6.7
-Babel                     2.14.0
-beautifulsoup4            4.12.2
-bleach                    6.1.0
-cachetools                5.3.2
-certifi                   2023.11.17
-cffi                      1.16.0
-charset-normalizer        3.3.2
-chex                      0.1.86
-comm                      0.2.0
-contourpy                 1.2.0
-cycler                    0.12.1
-debugpy                   1.8.0
-decorator                 5.1.1
-defusedxml                0.7.1
-diastatic-malt            2.15.1
-entrypoints               0.4
-etils                     1.5.2
-exceptiongroup            1.2.0
-executing                 2.0.1
-fastjsonschema            2.19.0
-filelock                  3.13.1
-flax                      0.8.0
-fonttools                 4.46.0
-fqdn                      1.5.1
-fsspec                    2023.12.2
-future                    0.18.3
-gast                      0.5.4
-idna                      3.6
-importlib-metadata        7.0.0
-importlib-resources       6.1.1
-ipykernel                 6.27.1
-ipython                   8.18.1
-ipython-genutils          0.2.0
-isoduration               20.11.0
-jax                       0.4.23
-jaxlib                    0.4.23
-jedi                      0.19.1
-Jinja2                    3.1.2
-joblib                    1.3.2
-json5                     0.9.14
-jsonpointer               2.4
-jsonschema                4.20.0
-jsonschema-specifications 2023.11.2
-jupyter_client            7.4.9
-jupyter_core              5.5.0
-jupyter-events            0.9.0
-jupyter-lsp               2.2.1
-jupyter_server            2.12.1
-jupyter_server_terminals  0.5.0
-jupyterlab                4.0.9
-jupyterlab_pygments       0.3.0
-jupyterlab_server         2.25.2
-kiwisolver                1.4.5
-llvmlite                  0.42.0
-markdown-it-py            3.0.0
-MarkupSafe                2.1.3
-matplotlib                3.8.2
-matplotlib-inline         0.1.6
-mdurl                     0.1.2
-mistune                   3.0.2
-ml-dtypes                 0.3.1
-mpmath                    1.3.0
-msgpack                   1.0.7
-multipledispatch          1.0.0
-nbclassic                 1.0.0
-nbclient                  0.9.0
-nbconvert                 7.12.0
-nbformat                  5.9.2
-nest-asyncio              1.5.8
-networkx                  3.2.1
-notebook                  6.5.6
-notebook_shim             0.2.3
-numba                     0.59.1
-numpy                     1.26.2
-numpyro                   0.14.0
-opt-einsum                3.3.0
-optax                     0.2.2
-orbax-checkpoint          0.4.8
-overrides                 7.4.0
-packaging                 23.2
-pandas                    2.1.4
-pandocfilters             1.5.0
-parso                     0.8.3
-PennyLane                 0.35.1
-PennyLane-Catalyst        0.5.0
-PennyLane_Lightning       0.35.1
-pexpect                   4.9.0
-Pillow                    10.1.0
-pip                       24.0
-platformdirs              4.1.0
-prometheus-client         0.19.0
-prompt-toolkit            3.0.43
-protobuf                  4.25.1
-psutil                    5.9.7
-ptyprocess                0.7.0
-pure-eval                 0.2.2
-pycparser                 2.21
-Pygments                  2.17.2
-pynndescent               0.5.12
-pyparsing                 3.1.1
-python-dateutil           2.8.2
-python-json-logger        2.0.7
-pytz                      2023.3.post1
-PyYAML                    6.0.1
-pyzmq                     24.0.1
-qgml                      2024.3.0
-qml-benchmarks            0.1          /Users/joseph/PycharmProjects/pennylane_jax/qml-benchmarks
-referencing               0.32.0
-requests                  2.31.0
-rfc3339-validator         0.1.4
-rfc3986-validator         0.1.1
-rich                      13.7.0
-rpds-py                   0.15.2
-rustworkx                 0.13.2
-scikit-learn              1.4.0
-scipy                     1.11.4
-seaborn                   0.13.0
-semantic-version          2.10.0
-Send2Trash                1.8.2
-setuptools                57.0.0
-six                       1.16.0
-sniffio                   1.3.0
-soupsieve                 2.5
-stack-data                0.6.3
-sympy                     1.12
-tensorstore               0.1.51
-termcolor                 2.4.0
-terminado                 0.18.0
-threadpoolctl             3.2.0
-tinycss2                  1.2.1
-toml                      0.10.2
-tomli                     2.0.1
-tomlkit                   0.12.3
-toolz                     0.12.0
-torch                     2.1.2
-torchvision               0.16.2
-tornado                   6.4
-tqdm                      4.66.2
-traitlets                 5.14.0
-types-python-dateutil     2.8.19.14
-typing_extensions         4.9.0
-tzdata                    2023.3
-umap-learn                0.5.6
-uri-template              1.3.0
-urllib3                   2.1.0
-wcwidth                   0.2.12
-webcolors                 1.13
-webencodings              0.5.1
-websocket-client          1.7.0
-wheel                     0.36.2
-zipp                      3.17.0
diff --git a/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_15d_performance_indicators_JAX.csv b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_15d_performance_indicators_JAX.csv
new file mode 100644
index 00000000..9e27eb16
--- /dev/null
+++ b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_15d_performance_indicators_JAX.csv
@@ -0,0 +1,2 @@
+first_train_step,first_train_step_std,consec_train_step,consec_train_step_std,predict_time,predict_time_std,hyperparameters
+182.53770942687987,9.854155175126856,24.854176211597945,0.899221214100941,191.73146567344665,13.744482148495923,"{'n_layers': 15, 'repeats': 10, 'use_jax': True, 'vmap': True, 'max_steps': 100, 'jit': True}"
diff --git a/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_2d_performance_indicators_JAX_packages.txt b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_15d_performance_indicators_JAX_packages.txt
similarity index 82%
rename from nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_2d_performance_indicators_JAX_packages.txt
rename to nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_15d_performance_indicators_JAX_packages.txt
index 0083ba07..bd314f52 100644
--- a/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_2d_performance_indicators_JAX_packages.txt
+++ b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_15d_performance_indicators_JAX_packages.txt
@@ -34,7 +34,7 @@ comm                      0.2.2
 contourpy                 1.2.1
 cryptography              42.0.5
 cycler                    0.12.1
-dask                      2024.4.1
+dask                      2024.4.2
 dbus-python               1.2.18
 debugpy                   1.8.1
 decorator                 4.4.2
@@ -44,39 +44,39 @@ dill                      0.3.8
 distro                    1.7.0
 duet                      0.2.9
 etils                     1.7.0
-exceptiongroup            1.2.0
+exceptiongroup            1.2.1
 executing                 2.0.1
 fastdtw                   0.3.4
 fastjsonschema            2.19.1
-filelock                  3.9.0
+filelock                  3.13.1
 fire                      0.6.0
 flax                      0.8.2
-fonttools                 4.50.0
+fonttools                 4.51.0
 fqdn                      1.5.1
 fsspec                    2024.3.1
 future                    1.0.0
 gast                      0.5.2
 greenlet                  3.0.3
 h11                       0.14.0
-h5py                      3.10.0
+h5py                      3.11.0
 httpcore                  1.0.5
 httpx                     0.27.0
-ibm-cloud-sdk-core        3.19.2
-ibm-platform-services     0.53.2
-idna                      3.6
+ibm-cloud-sdk-core        3.20.0
+ibm-platform-services     0.53.6
+idna                      3.7
 importlib_metadata        7.1.0
 importlib_resources       6.4.0
 ipykernel                 6.29.4
-ipython                   8.23.0
+ipython                   8.24.0
 ipywidgets                8.1.2
 isoduration               20.11.0
-jax                       0.4.23
-jaxlib                    0.4.23
+jax                       0.4.26
+jaxlib                    0.4.26
 jaxopt                    0.8.3
 jedi                      0.19.1
 Jinja2                    3.1.3
-joblib                    1.3.2
-json5                     0.9.24
+joblib                    1.4.0
+json5                     0.9.25
 jsonpointer               2.4
 jsonschema                4.21.1
 jsonschema-specifications 2023.12.1
@@ -85,22 +85,22 @@ jupyter_client            8.6.1
 jupyter-console           6.6.3
 jupyter_core              5.7.2
 jupyter-events            0.10.0
-jupyter-lsp               2.2.4
-jupyter_server            2.13.0
+jupyter-lsp               2.2.5
+jupyter_server            2.14.0
 jupyter_server_terminals  0.5.3
-jupyterlab                4.1.5
+jupyterlab                4.1.8
 jupyterlab_pygments       0.3.0
-jupyterlab_server         2.25.4
+jupyterlab_server         2.27.1
 jupyterlab_widgets        3.0.10
 kiwisolver                1.4.5
 lark-parser               0.12.0
 llvmlite                  0.42.0
 locket                    1.0.0
-Mako                      1.3.2
+Mako                      1.3.3
 markdown-it-py            3.0.0
 MarkupSafe                2.1.5
 matplotlib                3.8.4
-matplotlib-inline         0.1.6
+matplotlib-inline         0.1.7
 mdurl                     0.1.2
 mistune                   3.0.2
 ml-dtypes                 0.4.0
@@ -110,8 +110,8 @@ nbclient                  0.10.0
 nbconvert                 7.16.3
 nbformat                  5.10.4
 nest-asyncio              1.6.0
-networkx                  3.2.1
-notebook                  7.1.2
+networkx                  3.3
+notebook                  7.1.3
 notebook_shim             0.2.4
 numba                     0.59.1
 numpy                     1.26.4
@@ -122,9 +122,9 @@ optuna                    3.6.1
 orbax-checkpoint          0.5.10
 overrides                 7.7.0
 packaging                 24.0
-pandas                    2.2.1
+pandas                    2.2.2
 pandocfilters             1.5.1
-parso                     0.8.3
+parso                     0.8.4
 partd                     1.4.1
 patsy                     0.5.6
 pbr                       6.0.0
@@ -138,7 +138,7 @@ PennyLane-SF              0.29.0
 pexpect                   4.9.0
 pillow                    10.3.0
 pip                       24.0
-platformdirs              4.2.0
+platformdirs              4.2.1
 ply                       3.11
 prometheus_client         0.20.0
 prompt-toolkit            3.0.43
@@ -147,8 +147,8 @@ psutil                    5.9.8
 ptyprocess                0.7.0
 pure-eval                 0.2.2
 pycparser                 2.22
-pydantic                  2.6.4
-pydantic_core             2.16.3
+pydantic                  2.7.1
+pydantic_core             2.18.2
 pydantic-settings         2.2.1
 pydot                     2.0.0
 Pygments                  2.17.2
@@ -163,12 +163,12 @@ python-json-logger        2.0.7
 pythran                   0.10.0
 pytz                      2024.1
 PyYAML                    6.0.1
-pyzmq                     25.1.2
+pyzmq                     26.0.2
 qiskit                    1.0.2
 qiskit-aer                0.14.0.1
 qiskit-algorithms         0.3.0
 qiskit-ibm-experiment     0.4.7
-qiskit-ibm-provider       0.10.0
+qiskit-ibm-provider       0.11.0
 qiskit-ibm-runtime        0.20.0
 qiskit-machine-learning   0.7.2
 qml_benchmarks            0.1
@@ -176,7 +176,7 @@ qtconsole                 5.5.1
 QtPy                      2.4.1
 quantum-blackbird         0.5.0
 quantum-xir               0.2.2
-referencing               0.34.0
+referencing               0.35.0
 requests                  2.31.0
 requests-ntlm             1.2.0
 rfc3339-validator         0.1.4
@@ -186,11 +186,11 @@ rpds-py                   0.18.0
 ruamel.yaml               0.18.6
 ruamel.yaml.clib          0.2.8
 rustworkx                 0.14.2
-scikit-learn              1.4.1.post1
+scikit-learn              1.4.2
 scipy                     1.11.4
 seaborn                   0.13.2
 semantic-version          2.10.0
-Send2Trash                1.8.2
+Send2Trash                1.8.3
 setuptools                59.6.0
 six                       1.16.0
 sniffio                   1.3.1
@@ -199,36 +199,36 @@ soupsieve                 2.5
 SQLAlchemy                2.0.29
 ssh-import-id             5.11
 stack-data                0.6.3
-statsmodels               0.14.1
+statsmodels               0.14.2
 stevedore                 5.2.0
 StrawberryFields          0.23.0
 symengine                 0.11.0
 sympy                     1.12
-tensorstore               0.1.56
+tensorstore               0.1.58
 termcolor                 2.4.0
 terminado                 0.18.1
 thewalrus                 0.21.0
 threadpoolctl             3.4.0
-tinycss2                  1.2.1
+tinycss2                  1.3.0
 toml                      0.10.2
 tomli                     2.0.1
 tomlkit                   0.12.4
 toolz                     0.12.1
-torch                     2.2.2+cpu
-torchaudio                2.2.2+cpu
-torchvision               0.17.2+cpu
+torch                     2.3.0+cpu
+torchaudio                2.3.0+cpu
+torchvision               0.18.0+cpu
 tornado                   6.4
 tqdm                      4.66.2
-traitlets                 5.14.2
+traitlets                 5.14.3
 types-python-dateutil     2.9.0.20240316
-typing_extensions         4.10.0
+typing_extensions         4.11.0
 tzdata                    2024.1
 uri-template              1.3.0
 urllib3                   2.2.1
 wcwidth                   0.2.13
 webcolors                 1.13
 webencodings              0.5.1
-websocket-client          1.7.0
+websocket-client          1.8.0
 websockets                12.0
 wheel                     0.37.1
 widgetsnbextension        4.0.10
diff --git a/nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_3d_performance_indicators_JAX_scontrol.txt b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_15d_performance_indicators_JAX_scontrol.txt
similarity index 100%
rename from nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_3d_performance_indicators_JAX_scontrol.txt
rename to nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_15d_performance_indicators_JAX_scontrol.txt
diff --git a/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_2d_performance_indicators_JAX.csv b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_2d_performance_indicators_JAX.csv
new file mode 100644
index 00000000..257a5e6e
--- /dev/null
+++ b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_2d_performance_indicators_JAX.csv
@@ -0,0 +1,2 @@
+first_train_step,first_train_step_std,consec_train_step,consec_train_step_std,predict_time,predict_time_std,hyperparameters
+6.239000558853149,0.13502860069274902,0.0019811849401454736,9.583583985916306e-05,2.7535842657089233,0.0826483964920044,"{'n_layers': 15, 'repeats': 10, 'use_jax': True, 'vmap': True, 'max_steps': 100, 'jit': True}"
diff --git a/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_2d_performance_indicators_JAX_packages.txt b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_2d_performance_indicators_JAX_packages.txt
new file mode 100644
index 00000000..bd314f52
--- /dev/null
+++ b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_2d_performance_indicators_JAX_packages.txt
@@ -0,0 +1,236 @@
+Package                   Version
+------------------------- --------------
+absl-py                   2.1.0
+alembic                   1.13.1
+annotated-types           0.6.0
+antlr4-python3-runtime    4.9.2
+anyio                     4.3.0
+appdirs                   1.4.4
+argon2-cffi               23.1.0
+argon2-cffi-bindings      21.2.0
+arrow                     1.3.0
+asttokens                 2.4.1
+astunparse                1.6.3
+async-lru                 2.0.4
+attrs                     23.2.0
+autograd                  1.6.2
+autoray                   0.6.9
+Babel                     2.14.0
+beautifulsoup4            4.12.3
+beniget                   0.4.1
+bitstring                 3.1.7
+bleach                    6.1.0
+cachetools                5.3.3
+certifi                   2024.2.2
+cffi                      1.16.0
+charset-normalizer        3.3.2
+chex                      0.1.86
+cirq-core                 1.3.0
+cirq-pasqal               1.3.0
+click                     8.1.7
+cloudpickle               3.0.0
+colorlog                  6.8.2
+comm                      0.2.2
+contourpy                 1.2.1
+cryptography              42.0.5
+cycler                    0.12.1
+dask                      2024.4.2
+dbus-python               1.2.18
+debugpy                   1.8.1
+decorator                 4.4.2
+defusedxml                0.7.1
+diastatic-malt            2.15.1
+dill                      0.3.8
+distro                    1.7.0
+duet                      0.2.9
+etils                     1.7.0
+exceptiongroup            1.2.1
+executing                 2.0.1
+fastdtw                   0.3.4
+fastjsonschema            2.19.1
+filelock                  3.13.1
+fire                      0.6.0
+flax                      0.8.2
+fonttools                 4.51.0
+fqdn                      1.5.1
+fsspec                    2024.3.1
+future                    1.0.0
+gast                      0.5.2
+greenlet                  3.0.3
+h11                       0.14.0
+h5py                      3.11.0
+httpcore                  1.0.5
+httpx                     0.27.0
+ibm-cloud-sdk-core        3.20.0
+ibm-platform-services     0.53.6
+idna                      3.7
+importlib_metadata        7.1.0
+importlib_resources       6.4.0
+ipykernel                 6.29.4
+ipython                   8.24.0
+ipywidgets                8.1.2
+isoduration               20.11.0
+jax                       0.4.26
+jaxlib                    0.4.26
+jaxopt                    0.8.3
+jedi                      0.19.1
+Jinja2                    3.1.3
+joblib                    1.4.0
+json5                     0.9.25
+jsonpointer               2.4
+jsonschema                4.21.1
+jsonschema-specifications 2023.12.1
+jupyter                   1.0.0
+jupyter_client            8.6.1
+jupyter-console           6.6.3
+jupyter_core              5.7.2
+jupyter-events            0.10.0
+jupyter-lsp               2.2.5
+jupyter_server            2.14.0
+jupyter_server_terminals  0.5.3
+jupyterlab                4.1.8
+jupyterlab_pygments       0.3.0
+jupyterlab_server         2.27.1
+jupyterlab_widgets        3.0.10
+kiwisolver                1.4.5
+lark-parser               0.12.0
+llvmlite                  0.42.0
+locket                    1.0.0
+Mako                      1.3.3
+markdown-it-py            3.0.0
+MarkupSafe                2.1.5
+matplotlib                3.8.4
+matplotlib-inline         0.1.7
+mdurl                     0.1.2
+mistune                   3.0.2
+ml-dtypes                 0.4.0
+mpmath                    1.3.0
+msgpack                   1.0.8
+nbclient                  0.10.0
+nbconvert                 7.16.3
+nbformat                  5.10.4
+nest-asyncio              1.6.0
+networkx                  3.3
+notebook                  7.1.3
+notebook_shim             0.2.4
+numba                     0.59.1
+numpy                     1.26.4
+olefile                   0.46
+opt-einsum                3.3.0
+optax                     0.2.2
+optuna                    3.6.1
+orbax-checkpoint          0.5.10
+overrides                 7.7.0
+packaging                 24.0
+pandas                    2.2.2
+pandocfilters             1.5.1
+parso                     0.8.4
+partd                     1.4.1
+patsy                     0.5.6
+pbr                       6.0.0
+PennyLane                 0.35.1
+PennyLane-Catalyst        0.5.0
+PennyLane-Cirq            0.34.0
+PennyLane_Lightning       0.35.1
+PennyLane_Lightning_GPU   0.35.1
+PennyLane-qiskit          0.35.1
+PennyLane-SF              0.29.0
+pexpect                   4.9.0
+pillow                    10.3.0
+pip                       24.0
+platformdirs              4.2.1
+ply                       3.11
+prometheus_client         0.20.0
+prompt-toolkit            3.0.43
+protobuf                  5.26.1
+psutil                    5.9.8
+ptyprocess                0.7.0
+pure-eval                 0.2.2
+pycparser                 2.22
+pydantic                  2.7.1
+pydantic_core             2.18.2
+pydantic-settings         2.2.1
+pydot                     2.0.0
+Pygments                  2.17.2
+PyGObject                 3.42.1
+PyJWT                     2.8.0
+pylatexenc                2.10
+pyparsing                 3.1.2
+pyspnego                  0.10.2
+python-dateutil           2.9.0.post0
+python-dotenv             1.0.1
+python-json-logger        2.0.7
+pythran                   0.10.0
+pytz                      2024.1
+PyYAML                    6.0.1
+pyzmq                     26.0.2
+qiskit                    1.0.2
+qiskit-aer                0.14.0.1
+qiskit-algorithms         0.3.0
+qiskit-ibm-experiment     0.4.7
+qiskit-ibm-provider       0.11.0
+qiskit-ibm-runtime        0.20.0
+qiskit-machine-learning   0.7.2
+qml_benchmarks            0.1
+qtconsole                 5.5.1
+QtPy                      2.4.1
+quantum-blackbird         0.5.0
+quantum-xir               0.2.2
+referencing               0.35.0
+requests                  2.31.0
+requests-ntlm             1.2.0
+rfc3339-validator         0.1.4
+rfc3986-validator         0.1.1
+rich                      13.7.1
+rpds-py                   0.18.0
+ruamel.yaml               0.18.6
+ruamel.yaml.clib          0.2.8
+rustworkx                 0.14.2
+scikit-learn              1.4.2
+scipy                     1.11.4
+seaborn                   0.13.2
+semantic-version          2.10.0
+Send2Trash                1.8.3
+setuptools                59.6.0
+six                       1.16.0
+sniffio                   1.3.1
+sortedcontainers          2.4.0
+soupsieve                 2.5
+SQLAlchemy                2.0.29
+ssh-import-id             5.11
+stack-data                0.6.3
+statsmodels               0.14.2
+stevedore                 5.2.0
+StrawberryFields          0.23.0
+symengine                 0.11.0
+sympy                     1.12
+tensorstore               0.1.58
+termcolor                 2.4.0
+terminado                 0.18.1
+thewalrus                 0.21.0
+threadpoolctl             3.4.0
+tinycss2                  1.3.0
+toml                      0.10.2
+tomli                     2.0.1
+tomlkit                   0.12.4
+toolz                     0.12.1
+torch                     2.3.0+cpu
+torchaudio                2.3.0+cpu
+torchvision               0.18.0+cpu
+tornado                   6.4
+tqdm                      4.66.2
+traitlets                 5.14.3
+types-python-dateutil     2.9.0.20240316
+typing_extensions         4.11.0
+tzdata                    2024.1
+uri-template              1.3.0
+urllib3                   2.2.1
+wcwidth                   0.2.13
+webcolors                 1.13
+webencodings              0.5.1
+websocket-client          1.8.0
+websockets                12.0
+wheel                     0.37.1
+widgetsnbextension        4.0.10
+xanadu-cloud-client       0.3.2
+zipp                      3.18.1
diff --git a/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_3d_performance_indicators_JAX_scontrol.txt b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_2d_performance_indicators_JAX_scontrol.txt
similarity index 100%
rename from nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_3d_performance_indicators_JAX_scontrol.txt
rename to nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_2d_performance_indicators_JAX_scontrol.txt
diff --git a/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_3d_performance_indicators_JAX.csv b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_3d_performance_indicators_JAX.csv
deleted file mode 100644
index 7db3aa8d..00000000
--- a/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_3d_performance_indicators_JAX.csv
+++ /dev/null
@@ -1,2 +0,0 @@
-first_train_step,first_train_step_std,consec_train_step,consec_train_step_std,predict_time,predict_time_std,hyperparameters
-7.222759056091308,0.302873401811286,0.002020430564880371,0.00016703330426884848,1.5872551918029785,0.07579271979049444,"{'learning_rate': 0.01, 'n_layers': 15, 'repeats': 10, 'use_jax': True, 'vmap': True, 'max_steps': 5, 'jit': True}"
diff --git a/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_3d_performance_indicators_JAX_packages.txt b/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_3d_performance_indicators_JAX_packages.txt
deleted file mode 100644
index 8859d638..00000000
--- a/nersc/performance_indicators/JAX/IQPVariationalClassifier_linearly_separable_3d_performance_indicators_JAX_packages.txt
+++ /dev/null
@@ -1,166 +0,0 @@
-Package                   Version      Editable project location
-------------------------- ------------ ----------------------------------------------------------
-absl-py                   2.0.0
-anyio                     4.2.0
-appdirs                   1.4.4
-appnope                   0.1.3
-argon2-cffi               23.1.0
-argon2-cffi-bindings      21.2.0
-arrow                     1.3.0
-asttokens                 2.4.1
-astunparse                1.6.3
-async-lru                 2.0.4
-attrs                     23.1.0
-autograd                  1.6.2
-autoray                   0.6.7
-Babel                     2.14.0
-beautifulsoup4            4.12.2
-bleach                    6.1.0
-cachetools                5.3.2
-certifi                   2023.11.17
-cffi                      1.16.0
-charset-normalizer        3.3.2
-chex                      0.1.86
-comm                      0.2.0
-contourpy                 1.2.0
-cycler                    0.12.1
-debugpy                   1.8.0
-decorator                 5.1.1
-defusedxml                0.7.1
-diastatic-malt            2.15.1
-entrypoints               0.4
-etils                     1.5.2
-exceptiongroup            1.2.0
-executing                 2.0.1
-fastjsonschema            2.19.0
-filelock                  3.13.1
-flax                      0.8.0
-fonttools                 4.46.0
-fqdn                      1.5.1
-fsspec                    2023.12.2
-future                    0.18.3
-gast                      0.5.4
-idna                      3.6
-importlib-metadata        7.0.0
-importlib-resources       6.1.1
-ipykernel                 6.27.1
-ipython                   8.18.1
-ipython-genutils          0.2.0
-isoduration               20.11.0
-jax                       0.4.23
-jaxlib                    0.4.23
-jedi                      0.19.1
-Jinja2                    3.1.2
-joblib                    1.3.2
-json5                     0.9.14
-jsonpointer               2.4
-jsonschema                4.20.0
-jsonschema-specifications 2023.11.2
-jupyter_client            7.4.9
-jupyter_core              5.5.0
-jupyter-events            0.9.0
-jupyter-lsp               2.2.1
-jupyter_server            2.12.1
-jupyter_server_terminals  0.5.0
-jupyterlab                4.0.9
-jupyterlab_pygments       0.3.0
-jupyterlab_server         2.25.2
-kiwisolver                1.4.5
-llvmlite                  0.42.0
-markdown-it-py            3.0.0
-MarkupSafe                2.1.3
-matplotlib                3.8.2
-matplotlib-inline         0.1.6
-mdurl                     0.1.2
-mistune                   3.0.2
-ml-dtypes                 0.3.1
-mpmath                    1.3.0
-msgpack                   1.0.7
-multipledispatch          1.0.0
-nbclassic                 1.0.0
-nbclient                  0.9.0
-nbconvert                 7.12.0
-nbformat                  5.9.2
-nest-asyncio              1.5.8
-networkx                  3.2.1
-notebook                  6.5.6
-notebook_shim             0.2.3
-numba                     0.59.1
-numpy                     1.26.2
-numpyro                   0.14.0
-opt-einsum                3.3.0
-optax                     0.2.2
-orbax-checkpoint          0.4.8
-overrides                 7.4.0
-packaging                 23.2
-pandas                    2.1.4
-pandocfilters             1.5.0
-parso                     0.8.3
-PennyLane                 0.35.1
-PennyLane-Catalyst        0.5.0
-PennyLane_Lightning       0.35.1
-pexpect                   4.9.0
-Pillow                    10.1.0
-pip                       24.0
-platformdirs              4.1.0
-prometheus-client         0.19.0
-prompt-toolkit            3.0.43
-protobuf                  4.25.1
-psutil                    5.9.7
-ptyprocess                0.7.0
-pure-eval                 0.2.2
-pycparser                 2.21
-Pygments                  2.17.2
-pynndescent               0.5.12
-pyparsing                 3.1.1
-python-dateutil           2.8.2
-python-json-logger        2.0.7
-pytz                      2023.3.post1
-PyYAML                    6.0.1
-pyzmq                     24.0.1
-qgml                      2024.3.0
-qml-benchmarks            0.1          /Users/joseph/PycharmProjects/pennylane_jax/qml-benchmarks
-referencing               0.32.0
-requests                  2.31.0
-rfc3339-validator         0.1.4
-rfc3986-validator         0.1.1
-rich                      13.7.0
-rpds-py                   0.15.2
-rustworkx                 0.13.2
-scikit-learn              1.4.0
-scipy                     1.11.4
-seaborn                   0.13.0
-semantic-version          2.10.0
-Send2Trash                1.8.2
-setuptools                57.0.0
-six                       1.16.0
-sniffio                   1.3.0
-soupsieve                 2.5
-stack-data                0.6.3
-sympy                     1.12
-tensorstore               0.1.51
-termcolor                 2.4.0
-terminado                 0.18.0
-threadpoolctl             3.2.0
-tinycss2                  1.2.1
-toml                      0.10.2
-tomli                     2.0.1
-tomlkit                   0.12.3
-toolz                     0.12.0
-torch                     2.1.2
-torchvision               0.16.2
-tornado                   6.4
-tqdm                      4.66.2
-traitlets                 5.14.0
-types-python-dateutil     2.8.19.14
-typing_extensions         4.9.0
-tzdata                    2023.3
-umap-learn                0.5.6
-uri-template              1.3.0
-urllib3                   2.1.0
-wcwidth                   0.2.12
-webcolors                 1.13
-webencodings              0.5.1
-websocket-client          1.7.0
-wheel                     0.36.2
-zipp                      3.17.0
diff --git a/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_2d_performance_indicators_JAX.csv b/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_2d_performance_indicators_JAX.csv
deleted file mode 100644
index f2297385..00000000
--- a/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_2d_performance_indicators_JAX.csv
+++ /dev/null
@@ -1,2 +0,0 @@
-first_train_step,first_train_step_std,consec_train_step,consec_train_step_std,predict_time,predict_time_std,hyperparameters
-16.04134237766266,0.1555722951889038,0.005458146634728018,1.619921790228969e-05,4.981820344924927,0.05548095703125,"{'n_layers': 4, 'batch_size': 16, 'use_jax': True, 'vmap': True, 'max_steps': 100, 'jit': True}"
diff --git a/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_2d_performance_indicators_JAX_scontrol.txt b/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_2d_performance_indicators_JAX_scontrol.txt
deleted file mode 100644
index e69de29b..00000000
diff --git a/nersc/performance_indicators/perf_ind_kernel_jax.py b/nersc/performance_indicators/perf_ind_kernel_jax.py
index 0984dd77..474a6533 100644
--- a/nersc/performance_indicators/perf_ind_kernel_jax.py
+++ b/nersc/performance_indicators/perf_ind_kernel_jax.py
@@ -10,6 +10,9 @@
 import subprocess
 from qml_benchmarks.hyperparam_search_utils import read_data
 
+print(os.getcwd())
+os.chdir('performance_indicators')
+
 #################################
 # settings for the performance indicator.
 # You only need to change this to make a different performance indicator
@@ -23,7 +26,7 @@
 jit = True
 
 perf_ind_name = 'JAX'  #a name for the performance indicator used for naming files
-n_features = 16 #dataset dimension
+n_features = 2  #dataset dimension
 
 #################################
 
diff --git a/nersc/performance_indicators/perf_ind_variational_jax.py b/nersc/performance_indicators/perf_ind_variational_jax.py
index f36a21fd..b1b6298f 100644
--- a/nersc/performance_indicators/perf_ind_variational_jax.py
+++ b/nersc/performance_indicators/perf_ind_variational_jax.py
@@ -10,12 +10,15 @@
 import subprocess
 from qml_benchmarks.hyperparam_search_utils import read_data
 
+print(os.getcwd())
+os.chdir('performance_indicators')
+
 #################################
 # settings for the performance indicator.
 # You only need to change this to make a different performance indicator
 
 #define model
-from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier as Model
+from qml_benchmarks.models.quantum_metric_learning import QuantumMetricLearner as Model
 
 #implementation attributes of model
 use_jax = True
diff --git a/nersc/submit_job2.slr b/nersc/submit_job2.slr
index 30783d88..a51e976d 100755
--- a/nersc/submit_job2.slr
+++ b/nersc/submit_job2.slr
@@ -1,9 +1,9 @@
 #!/bin/bash
 #SBATCH -N 1
 #SBATCH -C cpu
-#SBATCH -q shared -t 01:00:00
+#SBATCH -q shared -t 10:00:00
 # SBATCH -q debug -t 5:00  # charged for full node
-#SBATCH -J perf_ind_iqp15
+#SBATCH -J perf_ind_qkernel2
 #SBATCH --ntasks-per-node=1
 #SBATCH --cpus-per-task=6
 #SBATCH --output out1/%j.log
@@ -36,7 +36,7 @@ BASE_DIR=/qml-benchmarks   # here git has home
 WORK_DIR=$BASE_DIR/nersc
 
 #... the task to be executed
-CMD=" python3 -u  performance_indicators/perf_ind_variational_jax.py "
+CMD=" python3 -u  performance_indicators/perf_ind_kernel_jax.py "
 #CMD=" python3 -u -c \"import pennylane \"; echo done1  "
 #CMD=" python3 -u -c \"import qml_benchmarks \";  echo done2 host  "
 

From 5c3ccb2ec2424b621718889b760023d54732a0ce Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Tue, 30 Apr 2024 13:03:25 +0200
Subject: [PATCH 058/100] add perf attributes

---
 src/qml_benchmarks/models/iqp_kernel.py               | 6 +++---
 src/qml_benchmarks/models/projected_quantum_kernel.py | 4 ++--
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/src/qml_benchmarks/models/iqp_kernel.py b/src/qml_benchmarks/models/iqp_kernel.py
index 3c83d744..8c8bf74f 100644
--- a/src/qml_benchmarks/models/iqp_kernel.py
+++ b/src/qml_benchmarks/models/iqp_kernel.py
@@ -221,15 +221,15 @@ def fit(self, X, y):
         X = self.transform(X)
 
         self.params_ = {"x_train": X}
+        start = time.time()
         kernel_matrix = self.precompute_kernel(X, X)
+        self.construct_kernel_time_ = time.time() - start
 
-        start = time.time()
         # we are updating this value here, in case it was
         # changed after initialising the model
         self.svm.C = self.C
         self.svm.fit(kernel_matrix, y)
-        end = time.time()
-        self.training_time_ = end - start
+        self.training_time_ = time.time() - start
 
         return self
 
diff --git a/src/qml_benchmarks/models/projected_quantum_kernel.py b/src/qml_benchmarks/models/projected_quantum_kernel.py
index 01c5abe4..f01dca4f 100644
--- a/src/qml_benchmarks/models/projected_quantum_kernel.py
+++ b/src/qml_benchmarks/models/projected_quantum_kernel.py
@@ -324,11 +324,11 @@ def fit(self, X, y):
         self.params_ = {"X_train": X}
         start = time.time()
         kernel_matrix = self.precompute_kernel(X, X)
+        self.construct_kernel_time_ = time.time() - start
 
         self.svm.C = self.C
         self.svm.fit(kernel_matrix, y)
-        end = time.time()
-        self.training_time_ = end - start
+        self.training_time_ = time.time() - start
 
         return self
 

From c380d43abf865ee13cb3b54551ad1e38e667c87f Mon Sep 17 00:00:00 2001
From: Jan Balewski <janstar1122@gmail.com>
Date: Tue, 30 Apr 2024 15:29:24 -0700
Subject: [PATCH 059/100] reads data but crash in Python scalars

---
 .gitignore                                    |   5 +
 nersc/container/Readme                        |  29 ++-
 nersc/oom                                     |   0
 .../perf_ind_kernel_jax.py                    | 175 ++++++++++--------
 nersc/submit_job2.slr                         |  31 +++-
 nersc/wrap_podman2.sh                         |  22 ++-
 6 files changed, 166 insertions(+), 96 deletions(-)
 delete mode 100644 nersc/oom

diff --git a/.gitignore b/.gitignore
index 48fb0467..0eb00f6e 100644
--- a/.gitignore
+++ b/.gitignore
@@ -180,4 +180,9 @@ cython_debug/
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 .idea/
 old/
+.bash_history
+.local/
+.viminfo
+hyperplanes_diff/
+linearly_separable/
 
diff --git a/nersc/container/Readme b/nersc/container/Readme
index 797be716..626b9a2d 100644
--- a/nersc/container/Readme
+++ b/nersc/container/Readme
@@ -22,8 +22,10 @@ Laptop:
 Perlmutter
    podman-hpc build -f Dockerfile.ubu22-PennyLane -t balewski/ubu22-pennylane:p5   .
    POD_PUB=/cfs/cdirs/nstaff/balewski/podman_common/
-   podman-hpc --squash-dir /global/$POD_PUB migrate  balewski/ubu22-pennylane:p5 
-   Note,  the instruction above is for user:balewski whohas access to CFS/nstaff - change the strings as needed
+   podman-hpc --squash-dir /global/$POD_PUB migrate  balewski/ubu22-pennylane:p5
+   chmod -R a+rx  /global/$POD_PUB   # to allow anyone to use this image
+   
+   Note,  the instruction above is for user:balewski who has access to CFS/nstaff - change the strings as needed
 
 
 
@@ -72,9 +74,28 @@ ssh CFS  (to your area )
 sbatch batchPodman.slr
 
 
-B.4  - - - - - -  advanced :  install project spceciffic software, one time
-NOT possible now - working on it
+B.4  - - - - - -  advanced :  install project spceciffic software, NOT recommended but possible
 
+ssh pm
+
+IMG=jbowles/ubu22-pennylane:p5
+POD_PUB=/dvs_ro/cfs/cdirs/m4139/qml-benchmarks/nersc/podman/
+export PODMANHPC_ADDITIONAL_STORES=$POD_PUB
+CFSH=$SCRATCH/pennylane_wrk  # for Jan
+#CFSH=/global/cfs/cdirs/m4139  # for Joseph
+podman-hpc run -it --volume  $CFSH/qml-benchmarks:/root --workdir /root  $IMG  bash
+pip3 install    --user .       # inside image, create private ./local
+
+exit
+
+Testing:
+podman-hpc run -it  …
+python3 -u -c "import qml_benchmarks "
+exit
+
+AND when lanuching podman image you must have this volume mount, on the top of all other volumes you need:
+
+ --volume  $CFSH/qml-benchmarks:/root 
 
 B.5  - - - - - -  advanced : jupiter notebook on the laptop powered by Podman PennyLane image running on PM - - - - - - - 
   Easy to do 
diff --git a/nersc/oom b/nersc/oom
deleted file mode 100644
index e69de29b..00000000
diff --git a/nersc/performance_indicators/perf_ind_kernel_jax.py b/nersc/performance_indicators/perf_ind_kernel_jax.py
index 474a6533..4d6c0d8a 100644
--- a/nersc/performance_indicators/perf_ind_kernel_jax.py
+++ b/nersc/performance_indicators/perf_ind_kernel_jax.py
@@ -10,76 +10,105 @@
 import subprocess
 from qml_benchmarks.hyperparam_search_utils import read_data
 
-print(os.getcwd())
-os.chdir('performance_indicators')
-
-#################################
-# settings for the performance indicator.
-# You only need to change this to make a different performance indicator
-
-#define the model
-from qml_benchmarks.models.iqp_kernel import IQPKernelClassifier as Model
-
-#implementation attributes of model
-use_jax = True
-vmap = True
-jit = True
-
-perf_ind_name = 'JAX'  #a name for the performance indicator used for naming files
-n_features = 2  #dataset dimension
-
-#################################
-
-model_name = Model().__class__.__name__
-
-# get the 'worst case' hyperparameter settings for the model (those that require the most resources)
-with open('hyperparam_settings.yaml', "r") as file:
-    hp_settings = yaml.safe_load(file)
-
-hyperparams = {**hp_settings[model_name], **{'use_jax':use_jax, 'vmap':vmap, 'jit': jit}}
-print(hyperparams)
-
-X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
-X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
-
-model = Model(**hyperparams)
-model.fit(X_train, y_train)
-
-#kernel construction time
-construct_kernel_time = model.construct_kernel_time_
-#full training time
-training_time = model.training_time_
-#prediction time
-time0 = time.time()
-model.predict(X_test)
-predict_time = time.time() - time0
-
-
-#write to csv
-data = [construct_kernel_time, training_time, predict_time, hyperparams]
-
-filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_{perf_ind_name}.csv"
-packages_filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_{perf_ind_name}_packages.txt"
-scontrol_filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_{perf_ind_name}_scontrol.txt"
-
-header = ['construct_kernel_time', 'training_time', 'predict_time', 'hyperparameters']
-
-if not os.path.exists(perf_ind_name):
-    # Create the directory
-    os.mkdir(perf_ind_name)
-
-#write perf indicator data
-with open(perf_ind_name+'/'+filename, mode="w", newline="") as file:
-    writer = csv.writer(file)
-    writer.writerow(header)
-    for row in [data]:
-        writer.writerow(row)
-
-# get package list and write to file
-output = subprocess.check_output(['pip', 'list']).decode('utf-8')
-with open(perf_ind_name+'/'+packages_filename, 'w') as file:
-    file.write(output)
-
-# make an empty text file to store the scontrol data
-with open(perf_ind_name+'/'+scontrol_filename, 'w') as file:
-    pass
+import argparse
+def get_parser():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-v","--verbosity",type=int,choices=[0, 1, 2,3,4],  help="increase output verbosity", default=1, dest='verb')
+    parser.add_argument("--inputPath",default='../linearly_separable',help='input data location')
+    parser.add_argument('-n','--numFeatures',type=int,default=2, help="dataset dimension ")
+
+    args = parser.parse_args()
+    
+    print( 'myArg-program:',parser.prog)
+    for arg in vars(args):  print( 'myArg:',arg, getattr(args, arg))
+
+    #assert os.path.exists(args.outPath)
+    return args
+
+#=================================
+#=================================
+#  M A I N 
+#=================================
+#=================================
+if __name__=="__main__":
+    args=get_parser()
+
+    print(os.getcwd())
+    os.chdir('performance_indicators')
+
+    #################################
+    # settings for the performance indicator.
+    # You only need to change this to make a different performance indicator
+
+    #define the model
+    from qml_benchmarks.models.iqp_kernel import IQPKernelClassifier as Model
+
+    #implementation attributes of model
+    use_jax = True
+    vmap = True
+    jit = True
+
+    perf_ind_name = 'JAX'  #a name for the performance indicator used for naming files
+    n_features = args.numFeatures  #dataset dimension
+
+    #################################
+
+    model_name = Model().__class__.__name__
+
+    # get the 'worst case' hyperparameter settings for the model (those that require the most resources)
+    with open('hyperparam_settings.yaml', "r") as file:
+        hp_settings = yaml.safe_load(file)
+
+    hyperparams = {**hp_settings[model_name], **{'use_jax':use_jax, 'vmap':vmap, 'jit': jit}}
+    print(hyperparams)
+    assert os.path.exists(args.inputPath)
+    #inpF1=f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv'
+    inpF1=os.path.join(args.inputPath,'linearly_separable_%dd_train.csv'%(n_features))
+    inpF2=inpF1.replace('train','test')
+    print('M:inpF1',inpF1)
+    X_train,y_train = read_data(inpF1)
+    print('M:inpF2',inpF2)
+    X_test,y_test = read_data(inpF2)
+
+    model = Model(**hyperparams)
+    model.fit(X_train, y_train)
+
+    #kernel construction time
+    construct_kernel_time = model.construct_kernel_time_
+    #full training time
+    training_time = model.training_time_
+    #prediction time
+    time0 = time.time()
+    model.predict(X_test)
+    predict_time = time.time() - time0
+
+
+    #write to csv
+    data = [construct_kernel_time, training_time, predict_time, hyperparams]
+
+    filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_{perf_ind_name}.csv"
+    packages_filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_{perf_ind_name}_packages.txt"
+    scontrol_filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_{perf_ind_name}_scontrol.txt"
+
+    header = ['construct_kernel_time', 'training_time', 'predict_time', 'hyperparameters']
+
+    if not os.path.exists(perf_ind_name):
+        # Create the directory
+        os.mkdir(perf_ind_name)
+
+    #write perf indicator data
+    with open(perf_ind_name+'/'+filename, mode="w", newline="") as file:
+        writer = csv.writer(file)
+        writer.writerow(header)
+        for row in [data]:
+            writer.writerow(row)
+
+    # get package list and write to file
+    output = subprocess.check_output(['pip', 'list']).decode('utf-8')
+    with open(perf_ind_name+'/'+packages_filename, 'w') as file:
+        file.write(output)
+
+    # make an empty text file to store the scontrol data
+    with open(perf_ind_name+'/'+scontrol_filename, 'w') as file:
+        pass
+    print('M:done')
diff --git a/nersc/submit_job2.slr b/nersc/submit_job2.slr
index a51e976d..2dce16db 100755
--- a/nersc/submit_job2.slr
+++ b/nersc/submit_job2.slr
@@ -1,24 +1,29 @@
 #!/bin/bash
 #SBATCH -N 1
 #SBATCH -C cpu
-#SBATCH -q shared -t 10:00:00
-# SBATCH -q debug -t 5:00  # charged for full node
+# SBATCH -q shared -t 10:00:00
+#SBATCH -q debug -t 5:00  # charged for full node
 #SBATCH -J perf_ind_qkernel2
 #SBATCH --ntasks-per-node=1
 #SBATCH --cpus-per-task=6
 #SBATCH --output out1/%j.log
 #SBATCH  --licenses=scratch
-
+# - - - E N D    O F    SLURM    C O M M A N D S
+set -u ;  # exit  if you try to use an uninitialized variable
 
 echo 'S:start'
 
+numFeatures=3
+
 # .... container
 # salloc -q interactive -C cpu -t 4:00:00 -A nstaff
 
-IMG=jbowles/ubu22-pennylane:p5
+IMG=balewski/ubu22-pennylane:p5  # old
+#IMG=jbowles/ubu22-pennylane:p5
 echo use image $IMG
-POD_PUB=/global/cfs/cdirs/m4139/qml-benchmarks/nersc/podman/
-# POD_PUB=/dvs_ro/cfs/cdirs/nstaff/balewski/podman_common/  
+
+#POD_PUB=//dvs_ro/cfs/cdirs/m4139/qml-benchmarks/nersc/podman/
+POD_PUB=/dvs_ro/cfs/cdirs/nstaff/balewski/podman_common/  #old
 export PODMANHPC_ADDITIONAL_STORES=$POD_PUB
 
 #.... output sandbox 
@@ -29,19 +34,25 @@ echo outPath=$outPath
 mkdir -p $outPath
 cp -rp performance_indicators *.slr *.sh   $outPath
 
+
 #... location of input on CSF
-CFSH=/global/cfs/cdirs/m4139  # for Joseph
-#CFSH=/pscratch/sd/b/balewski  # for Jan, do not remove
+#CFSH=/global/cfs/cdirs/m4139  # for Joseph
+CFSH=/pscratch/sd/b/balewski/pennylane_wrk  # for Jan
 BASE_DIR=/qml-benchmarks   # here git has home
 WORK_DIR=$BASE_DIR/nersc
 
 #... the task to be executed
-CMD=" python3 -u  performance_indicators/perf_ind_kernel_jax.py "
+CMD=" python3 -u  performance_indicators/perf_ind_kernel_jax.py --numFeatures $numFeatures --inputPath ../../../linearly_separable "
 #CMD=" python3 -u -c \"import pennylane \"; echo done1  "
 #CMD=" python3 -u -c \"import qml_benchmarks \";  echo done2 host  "
 
-
 cd $outPath
+# this is hack Jan  does not like  but 'paper' dir is not part of pip-install
+# qml-benchmarks> cp -rp paper /pscratch/sd/b/balewski/tmp_pennylane_jobs
+
+#ls -l $CFSH/$BASE_DIR/paper
+#ln -s $BASE_DIR/paper ..
+
 G=1
 echo 'S:ready to run '; pwd
 srun -n $G  ./wrap_podman2.sh $IMG " $CMD "  $outPath  $CFSH $BASE_DIR $WORK_DIR
diff --git a/nersc/wrap_podman2.sh b/nersc/wrap_podman2.sh
index 333112c8..a6b263c0 100755
--- a/nersc/wrap_podman2.sh
+++ b/nersc/wrap_podman2.sh
@@ -16,19 +16,23 @@ fi
 echo W:BASE_DIR=$BASE_DIR
 echo 'W:start podman'
 podman-hpc run -it \
-	   --volume $HOME:/home \
-	   --volume $CFSH/$BASE_DIR:/$BASE_DIR \
-	   --volume $CFSH/$WORK_DIR:$WORK_DIR \
-	   -e HDF5_USE_FILE_LOCKING='FALSE' \
-	   --workdir $WORK_DIR \
-	   $IMG <<EOF 
-cd $BASE_DIR
+    --volume $CFSH/$BASE_DIR:/root \
+    --volume $CFSH/$BASE_DIR:$BASE_DIR \
+    --volume $CFSH/$BASE_DIR/linearly_separable:/linearly_separable \
+    --volume $CFSH/$WORK_DIR:$WORK_DIR \
+    -e HDF5_USE_FILE_LOCKING='FALSE' \
+    --workdir $WORK_DIR \
+    $IMG <<EOF 
 echo P:pwd; pwd
-pip3 install .
-cd -
+echo P:all
+ls -l ../..
 $CMD
 echo P:done
 exit
 EOF
   
 echo 'W:done podman'
+
+# spare
+# 	   --volume $HOME:/home \
+    # --volume  $CFSH/qml-benchmarks:/root 

From 1de6e38469b632ed8912248f2871425ab64473a8 Mon Sep 17 00:00:00 2001
From: Jan Balewski <janstar1122@gmail.com>
Date: Wed, 1 May 2024 07:17:30 -0700
Subject: [PATCH 060/100] cleanup

---
 nersc/toy_slurm/Readme                | 33 ++++++++++++++++
 nersc/toy_slurm/input.json            |  4 ++
 nersc/toy_slurm/run_cpu_task.py       | 43 +++++++++++++++++++++
 nersc/toy_slurm/run_gpu_task.py       | 55 +++++++++++++++++++++++++++
 nersc/toy_slurm/submit_shared_cpu.slr | 35 +++++++++++++++++
 nersc/toy_slurm/submit_shared_gpu.slr | 40 +++++++++++++++++++
 6 files changed, 210 insertions(+)
 create mode 100644 nersc/toy_slurm/Readme
 create mode 100644 nersc/toy_slurm/input.json
 create mode 100755 nersc/toy_slurm/run_cpu_task.py
 create mode 100755 nersc/toy_slurm/run_gpu_task.py
 create mode 100755 nersc/toy_slurm/submit_shared_cpu.slr
 create mode 100755 nersc/toy_slurm/submit_shared_gpu.slr

diff --git a/nersc/toy_slurm/Readme b/nersc/toy_slurm/Readme
new file mode 100644
index 00000000..472577d2
--- /dev/null
+++ b/nersc/toy_slurm/Readme
@@ -0,0 +1,33 @@
+Example of Slurm scripts using fraction of a node (aka shared queue)
+
+Jan's tricks:
+* write Slurm output to out/<JobId>.log , merge stdout+sderr
+* disable selcted Slurm command with '-', e.g.: this is how to change queue by just moving '-'
+#SBATCH -q shared
+#-SBATCH -q debug
+* run task under time to see the burden on the system
+* not hardcode paths, use $SCRATCH
+* $ sqs  command at NERSC is your friend
+* alasy ask for 30% more time , 30% more CPUs, 5GB more RAM than you think you need
+
+Preparation:
+* alwasy start Slurm job from CFS (Community File System) so the Slurm-output is written here
+cd /global/cfs/cdirs/mXXX/yyy
+mkdir out
+
+= = = = = =   CPU shared node = = = = =
+Inspect: submit_shared_cpu.slr
+- reads config from CFS , it controlls how long code runs
+- writes output to SCRATCH to sub-dir whos name is jobId
+- python task is rank-aware and each runk writes different output
+
+$ sbatch submit_shared_cpu.slr
+
+$ scontrol show job <JobID>
+
+
+= = = = = =   GPU shared  node = = = = =
+
+$ sbatch submit_shared_gpu.slr
+- similar workflow, but now computation is done on GPU
+- 
diff --git a/nersc/toy_slurm/input.json b/nersc/toy_slurm/input.json
new file mode 100644
index 00000000..35163625
--- /dev/null
+++ b/nersc/toy_slurm/input.json
@@ -0,0 +1,4 @@
+{"nrows": 25,
+"ncols": 55,
+"iterations": 60,
+"pause": 1.1}
diff --git a/nersc/toy_slurm/run_cpu_task.py b/nersc/toy_slurm/run_cpu_task.py
new file mode 100755
index 00000000..a0083758
--- /dev/null
+++ b/nersc/toy_slurm/run_cpu_task.py
@@ -0,0 +1,43 @@
+#!/usr/bin/env python3
+import numpy as np
+import json,os 
+import time
+import argparse
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    
+    parser.add_argument("--input", type=str, help="input file")
+    parser.add_argument("--output", type=str, help="output file")
+    args = parser.parse_args()
+
+    return args
+
+
+
+def main():
+    args = parse_args(); 
+    for arg in vars(args):
+        print( 'myArgs:',arg, getattr(args, arg))
+    myRank=os.environ['SLURM_PROCID']
+    print('I am rank:',myRank)
+    with open(args.input, 'r') as file:
+        info = json.load(file)
+    matrix = np.random.random((info['nrows'],info['ncols']))
+    for i in range(info['iterations']):
+        matrix += np.random.random((info['nrows'],info['ncols']))
+        time.sleep(info['pause'])
+
+    output=[]
+    for i in range(info['nrows']):
+        output.append(np.sum(matrix[i,:]))
+
+    outF=args.output+myRank
+    with open(outF, 'w') as file:
+        for i in range(len(output)):
+            file.write(str(output[i])+"\n")
+        file.write("---------------------")
+
+if __name__ == "__main__":
+    main()
diff --git a/nersc/toy_slurm/run_gpu_task.py b/nersc/toy_slurm/run_gpu_task.py
new file mode 100755
index 00000000..9bb540ab
--- /dev/null
+++ b/nersc/toy_slurm/run_gpu_task.py
@@ -0,0 +1,55 @@
+#!/usr/bin/env python3
+import numpy as np
+import cupy as cp
+import json,os 
+import time
+import argparse
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    
+    parser.add_argument("--input", type=str, help="input file")
+    parser.add_argument("--output", type=str, help="output file")
+    args = parser.parse_args()
+
+    return args
+
+def main():
+    args = parse_args(); 
+    myRank=os.environ['SLURM_PROCID']
+    print('I am rank:',myRank)
+
+    # Get the number of available CUDA devices
+    num_gpus = cp.cuda.runtime.getDeviceCount()
+    print("Number of GPUs available: %d" % num_gpus)
+
+
+    # Access and print some GPU attributes
+    for i in range(num_gpus):
+        with cp.cuda.Device(i):
+            compute_capability = cp.cuda.Device().compute_capability
+            total_memory = cp.cuda.Device().mem_info[1]
+            print("GPU %d: Compute Capability: %s, Total Memory: %d bytes" % (i, compute_capability, total_memory))
+
+    assert num_gpus ==1  # make sure  there GPU-rank mapping is unique
+    
+    with open(args.input, 'r') as file:
+        info = json.load(file)
+    matrix = cp.random.random((info['nrows'],info['ncols']), np.float64)
+    for i in range(info['iterations']):
+        matrix += cp.random.random((info['nrows'],info['ncols']), np.float64)
+        time.sleep(info['pause'])
+
+    output=[]
+    for i in range(info['nrows']):
+        output.append(np.sum(matrix[i,:]))
+
+    outF=args.output+myRank
+    with open(outF, 'w') as file:
+        for i in range(len(output)):
+            file.write(str(output[i])+"\n")
+        file.write("---------------------")
+
+if __name__ == "__main__":
+    main()
diff --git a/nersc/toy_slurm/submit_shared_cpu.slr b/nersc/toy_slurm/submit_shared_cpu.slr
new file mode 100755
index 00000000..8c5dee07
--- /dev/null
+++ b/nersc/toy_slurm/submit_shared_cpu.slr
@@ -0,0 +1,35 @@
+#!/bin/bash
+#SBATCH -N 1
+#SBATCH -C cpu
+#SBATCH -q shared
+#-SBATCH -q debug  # charged for full node
+#SBATCH -J test_cpu12
+#SBATCH --account=nstaff
+#SBATCH -t 00:08:00
+#SBATCH --ntasks-per-node=2
+#SBATCH --cpus-per-task=4
+#-SBATCH --mem=25GB    # optional, use only if needed
+#SBATCH --output out/%j.log
+#SBATCH  --licenses=scratch
+# - - - E N D    O F    SLURM    C O M M A N D S
+
+nprocspn=${SLURM_NTASKS_PER_NODE}
+#nprocspn=1  # special case for partial use of full node
+
+N=${SLURM_NNODES}
+G=$[ $N * $nprocspn ]
+jobId=${SLURM_JOBID}
+echo S:  G=$G  N=$N
+
+OUT_DIR=$SCRATCH/penny_tmp/$jobId
+mkdir -p $OUT_DIR
+
+CFSH=/global/cfs/cdirs/mpccc/balewski/pennylane_nesap/toy_slurm
+
+#  .... starting job proper
+module load pytorch
+time srun -n $G  python ./run_cpu_task.py --input $CFSH/input.json --output $OUT_DIR/cpu_output.txt
+
+echo S:train-done
+date
+
diff --git a/nersc/toy_slurm/submit_shared_gpu.slr b/nersc/toy_slurm/submit_shared_gpu.slr
new file mode 100755
index 00000000..5abd33e7
--- /dev/null
+++ b/nersc/toy_slurm/submit_shared_gpu.slr
@@ -0,0 +1,40 @@
+#!/bin/bash
+#SBATCH -N 1
+#SBATCH -C gpu
+#SBATCH -q shared
+#-SBATCH -q debug  # charged for full node
+#SBATCH -J test_gpu12
+#SBATCH --account=nstaff
+#SBATCH -t 00:08:00
+#SBATCH --gpus 2  # total
+#SBATCH --ntasks-per-node=2  # must match # GPUs/node if full node is used
+#SBATCH --gpus-per-task=1  # assures a different GPU is given to each task
+#SBATCH --output out/%j.log
+#SBATCH  --licenses=scratch
+# - - - E N D    O F    SLURM    C O M M A N D S
+
+#env|grpe SLURM
+
+nprocspn=${SLURM_NTASKS}
+#nprocspn=1  # special case for partial use of full node
+
+N=${SLURM_NNODES}
+G=$[ $N * $nprocspn ]
+jobId=${SLURM_JOBID}
+echo S:  G=$G  N=$N
+
+OUT_DIR=$SCRATCH/penny_tmp/$jobId
+mkdir -p $OUT_DIR
+
+CFSH=/global/cfs/cdirs/mpccc/balewski/pennylane_nesap/toy_slurm
+
+echo "S: start nvidia-smi"
+nvidia-smi
+module load python
+
+#  .... starting job proper
+time srun -n $G --gpus-per-task=1  python ./run_gpu_task.py --input $CFSH/input.json --output $OUT_DIR/gpu_output.txt
+
+echo S:train-done
+date
+

From a76e44aaa4181e3a612ad5885a15c15dfc9707fb Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Thu, 2 May 2024 15:57:04 +0200
Subject: [PATCH 061/100] data gen file

---
 .../generate_linearly_separable.py            | 40 +++++++++++++++++++
 1 file changed, 40 insertions(+)
 create mode 100644 nersc/performance_indicators/generate_linearly_separable.py

diff --git a/nersc/performance_indicators/generate_linearly_separable.py b/nersc/performance_indicators/generate_linearly_separable.py
new file mode 100644
index 00000000..477f3273
--- /dev/null
+++ b/nersc/performance_indicators/generate_linearly_separable.py
@@ -0,0 +1,40 @@
+# Copyright 2024 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Generate datasets for the  LINEARLY SEPARABLE benchmark."""
+
+import os
+import numpy as np
+from sklearn.model_selection import train_test_split
+from qml_benchmarks.data import generate_linearly_separable
+
+np.random.seed(42)
+
+os.makedirs("linearly_separable", exist_ok=True)
+
+n_samples = 300
+
+for n_features in range(2, 21):
+    margin = 0.02 * n_features
+
+    X, y = generate_linearly_separable(n_samples, n_features, margin)
+    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)
+
+    name_train = f"linearly_separable/linearly_separable_{n_features}d_train.csv"
+    data_train = np.c_[X_train, y_train]
+    np.savetxt(name_train, data_train, delimiter=",")
+
+    name_test = f"linearly_separable/linearly_separable_{n_features}d_test.csv"
+    data_test = np.c_[X_test, y_test]
+    np.savetxt(name_test, data_test, delimiter=",")

From 15f499881f857f3b14bde9a0f7203097de8ed258 Mon Sep 17 00:00:00 2001
From: josephbowles <bowles.physics@gmail.com>
Date: Thu, 2 May 2024 07:23:42 -0700
Subject: [PATCH 062/100] results

---
 ...parable_13d_performance_indicators_JAX.csv |   2 +
 ...3d_performance_indicators_JAX_packages.txt | 236 ++++++++++++++++++
 ...3d_performance_indicators_JAX_scontrol.txt |   0
 ...eparable_2d_performance_indicators_JAX.csv |   2 +
 ...2d_performance_indicators_JAX_packages.txt | 236 ++++++++++++++++++
 ...2d_performance_indicators_JAX_scontrol.txt |   0
 ...parable_15d_performance_indicators_JAX.csv |   2 +
 ...5d_performance_indicators_JAX_packages.txt | 236 ++++++++++++++++++
 ...5d_performance_indicators_JAX_scontrol.txt |   0
 ...parable_17d_performance_indicators_JAX.csv |   2 +
 ...7d_performance_indicators_JAX_packages.txt | 236 ++++++++++++++++++
 ...7d_performance_indicators_JAX_scontrol.txt |   0
 ...parable_20d_performance_indicators_JAX.csv |   2 +
 ...0d_performance_indicators_JAX_packages.txt | 236 ++++++++++++++++++
 ...0d_performance_indicators_JAX_scontrol.txt |   0
 ...parable_13d_performance_indicators_JAX.csv |   2 +
 ...3d_performance_indicators_JAX_packages.txt | 236 ++++++++++++++++++
 ...3d_performance_indicators_JAX_scontrol.txt |   0
 ...parable_16d_performance_indicators_JAX.csv |   2 +
 ...6d_performance_indicators_JAX_packages.txt | 236 ++++++++++++++++++
 ...6d_performance_indicators_JAX_scontrol.txt |   0
 ...eparable_2d_performance_indicators_JAX.csv |   2 +
 ...2d_performance_indicators_JAX_packages.txt | 236 ++++++++++++++++++
 ...2d_performance_indicators_JAX_scontrol.txt |   0
 24 files changed, 1904 insertions(+)
 create mode 100644 nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_13d_performance_indicators_JAX.csv
 create mode 100644 nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_13d_performance_indicators_JAX_packages.txt
 create mode 100644 nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_13d_performance_indicators_JAX_scontrol.txt
 create mode 100644 nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_2d_performance_indicators_JAX.csv
 create mode 100644 nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_2d_performance_indicators_JAX_packages.txt
 create mode 100644 nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_2d_performance_indicators_JAX_scontrol.txt
 create mode 100644 nersc/performance_indicators/JAX/ProjectedQuantumKernel_linearly_separable_15d_performance_indicators_JAX.csv
 create mode 100644 nersc/performance_indicators/JAX/ProjectedQuantumKernel_linearly_separable_15d_performance_indicators_JAX_packages.txt
 create mode 100644 nersc/performance_indicators/JAX/ProjectedQuantumKernel_linearly_separable_15d_performance_indicators_JAX_scontrol.txt
 create mode 100644 nersc/performance_indicators/JAX/ProjectedQuantumKernel_linearly_separable_17d_performance_indicators_JAX.csv
 create mode 100644 nersc/performance_indicators/JAX/ProjectedQuantumKernel_linearly_separable_17d_performance_indicators_JAX_packages.txt
 create mode 100644 nersc/performance_indicators/JAX/ProjectedQuantumKernel_linearly_separable_17d_performance_indicators_JAX_scontrol.txt
 create mode 100644 nersc/performance_indicators/JAX/ProjectedQuantumKernel_linearly_separable_20d_performance_indicators_JAX.csv
 create mode 100644 nersc/performance_indicators/JAX/ProjectedQuantumKernel_linearly_separable_20d_performance_indicators_JAX_packages.txt
 create mode 100644 nersc/performance_indicators/JAX/ProjectedQuantumKernel_linearly_separable_20d_performance_indicators_JAX_scontrol.txt
 create mode 100644 nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_13d_performance_indicators_JAX.csv
 create mode 100644 nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_13d_performance_indicators_JAX_packages.txt
 create mode 100644 nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_13d_performance_indicators_JAX_scontrol.txt
 create mode 100644 nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_16d_performance_indicators_JAX.csv
 create mode 100644 nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_16d_performance_indicators_JAX_packages.txt
 create mode 100644 nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_16d_performance_indicators_JAX_scontrol.txt
 create mode 100644 nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_2d_performance_indicators_JAX.csv
 create mode 100644 nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_2d_performance_indicators_JAX_packages.txt
 create mode 100644 nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_2d_performance_indicators_JAX_scontrol.txt

diff --git a/nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_13d_performance_indicators_JAX.csv b/nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_13d_performance_indicators_JAX.csv
new file mode 100644
index 00000000..4633b5f3
--- /dev/null
+++ b/nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_13d_performance_indicators_JAX.csv
@@ -0,0 +1,2 @@
+construct_kernel_time,training_time,predict_time,hyperparameters
+3860.4390046596527,3860.4435765743256,3755.371912240982,"{'repeats': 10, 'use_jax': True, 'vmap': True, 'jit': True}"
diff --git a/nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_13d_performance_indicators_JAX_packages.txt b/nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_13d_performance_indicators_JAX_packages.txt
new file mode 100644
index 00000000..f5d4f0ca
--- /dev/null
+++ b/nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_13d_performance_indicators_JAX_packages.txt
@@ -0,0 +1,236 @@
+Package                   Version
+------------------------- --------------
+absl-py                   2.1.0
+alembic                   1.13.1
+annotated-types           0.6.0
+antlr4-python3-runtime    4.9.2
+anyio                     4.3.0
+appdirs                   1.4.4
+argon2-cffi               23.1.0
+argon2-cffi-bindings      21.2.0
+arrow                     1.3.0
+asttokens                 2.4.1
+astunparse                1.6.3
+async-lru                 2.0.4
+attrs                     23.2.0
+autograd                  1.6.2
+autoray                   0.6.9
+Babel                     2.14.0
+beautifulsoup4            4.12.3
+beniget                   0.4.1
+bitstring                 3.1.7
+bleach                    6.1.0
+cachetools                5.3.3
+certifi                   2024.2.2
+cffi                      1.16.0
+charset-normalizer        3.3.2
+chex                      0.1.86
+cirq-core                 1.3.0
+cirq-pasqal               1.3.0
+click                     8.1.7
+cloudpickle               3.0.0
+colorlog                  6.8.2
+comm                      0.2.2
+contourpy                 1.2.1
+cryptography              42.0.5
+cycler                    0.12.1
+dask                      2024.4.2
+dbus-python               1.2.18
+debugpy                   1.8.1
+decorator                 4.4.2
+defusedxml                0.7.1
+diastatic-malt            2.15.1
+dill                      0.3.8
+distro                    1.7.0
+duet                      0.2.9
+etils                     1.7.0
+exceptiongroup            1.2.1
+executing                 2.0.1
+fastdtw                   0.3.4
+fastjsonschema            2.19.1
+filelock                  3.13.1
+fire                      0.6.0
+flax                      0.8.3
+fonttools                 4.51.0
+fqdn                      1.5.1
+fsspec                    2024.3.1
+future                    1.0.0
+gast                      0.5.2
+greenlet                  3.0.3
+h11                       0.14.0
+h5py                      3.11.0
+httpcore                  1.0.5
+httpx                     0.27.0
+ibm-cloud-sdk-core        3.20.0
+ibm-platform-services     0.53.6
+idna                      3.7
+importlib_metadata        7.1.0
+importlib_resources       6.4.0
+ipykernel                 6.29.4
+ipython                   8.24.0
+ipywidgets                8.1.2
+isoduration               20.11.0
+jax                       0.4.23
+jaxlib                    0.4.23
+jaxopt                    0.8.3
+jedi                      0.19.1
+Jinja2                    3.1.3
+joblib                    1.4.0
+json5                     0.9.25
+jsonpointer               2.4
+jsonschema                4.21.1
+jsonschema-specifications 2023.12.1
+jupyter                   1.0.0
+jupyter_client            8.6.1
+jupyter-console           6.6.3
+jupyter_core              5.7.2
+jupyter-events            0.10.0
+jupyter-lsp               2.2.5
+jupyter_server            2.14.0
+jupyter_server_terminals  0.5.3
+jupyterlab                4.1.8
+jupyterlab_pygments       0.3.0
+jupyterlab_server         2.27.1
+jupyterlab_widgets        3.0.10
+kiwisolver                1.4.5
+lark-parser               0.12.0
+llvmlite                  0.42.0
+locket                    1.0.0
+Mako                      1.3.3
+markdown-it-py            3.0.0
+MarkupSafe                2.1.5
+matplotlib                3.8.4
+matplotlib-inline         0.1.7
+mdurl                     0.1.2
+mistune                   3.0.2
+ml-dtypes                 0.4.0
+mpmath                    1.3.0
+msgpack                   1.0.8
+nbclient                  0.10.0
+nbconvert                 7.16.4
+nbformat                  5.10.4
+nest-asyncio              1.6.0
+networkx                  3.3
+notebook                  7.1.3
+notebook_shim             0.2.4
+numba                     0.59.1
+numpy                     1.26.4
+olefile                   0.46
+opt-einsum                3.3.0
+optax                     0.2.2
+optuna                    3.6.1
+orbax-checkpoint          0.5.10
+overrides                 7.7.0
+packaging                 24.0
+pandas                    2.2.2
+pandocfilters             1.5.1
+parso                     0.8.4
+partd                     1.4.1
+patsy                     0.5.6
+pbr                       6.0.0
+PennyLane                 0.35.1
+PennyLane-Catalyst        0.5.0
+PennyLane-Cirq            0.34.0
+PennyLane_Lightning       0.35.1
+PennyLane_Lightning_GPU   0.35.1
+PennyLane-qiskit          0.35.1
+PennyLane-SF              0.29.0
+pexpect                   4.9.0
+pillow                    10.3.0
+pip                       24.0
+platformdirs              4.2.1
+ply                       3.11
+prometheus_client         0.20.0
+prompt-toolkit            3.0.43
+protobuf                  5.26.1
+psutil                    5.9.8
+ptyprocess                0.7.0
+pure-eval                 0.2.2
+pycparser                 2.22
+pydantic                  2.7.1
+pydantic_core             2.18.2
+pydantic-settings         2.2.1
+pydot                     2.0.0
+Pygments                  2.17.2
+PyGObject                 3.42.1
+PyJWT                     2.8.0
+pylatexenc                2.10
+pyparsing                 3.1.2
+pyspnego                  0.10.2
+python-dateutil           2.9.0.post0
+python-dotenv             1.0.1
+python-json-logger        2.0.7
+pythran                   0.10.0
+pytz                      2024.1
+PyYAML                    6.0.1
+pyzmq                     26.0.2
+qiskit                    1.0.2
+qiskit-aer                0.14.1
+qiskit-algorithms         0.3.0
+qiskit-ibm-experiment     0.4.7
+qiskit-ibm-provider       0.11.0
+qiskit-ibm-runtime        0.20.0
+qiskit-machine-learning   0.7.2
+qml_benchmarks            0.1
+qtconsole                 5.5.1
+QtPy                      2.4.1
+quantum-blackbird         0.5.0
+quantum-xir               0.2.2
+referencing               0.35.0
+requests                  2.31.0
+requests-ntlm             1.2.0
+rfc3339-validator         0.1.4
+rfc3986-validator         0.1.1
+rich                      13.7.1
+rpds-py                   0.18.0
+ruamel.yaml               0.18.6
+ruamel.yaml.clib          0.2.8
+rustworkx                 0.14.2
+scikit-learn              1.4.2
+scipy                     1.11.4
+seaborn                   0.13.2
+semantic-version          2.10.0
+Send2Trash                1.8.3
+setuptools                59.6.0
+six                       1.16.0
+sniffio                   1.3.1
+sortedcontainers          2.4.0
+soupsieve                 2.5
+SQLAlchemy                2.0.29
+ssh-import-id             5.11
+stack-data                0.6.3
+statsmodels               0.14.2
+stevedore                 5.2.0
+StrawberryFields          0.23.0
+symengine                 0.11.0
+sympy                     1.12
+tensorstore               0.1.58
+termcolor                 2.4.0
+terminado                 0.18.1
+thewalrus                 0.21.0
+threadpoolctl             3.5.0
+tinycss2                  1.3.0
+toml                      0.10.2
+tomli                     2.0.1
+tomlkit                   0.12.4
+toolz                     0.12.1
+torch                     2.3.0+cpu
+torchaudio                2.3.0+cpu
+torchvision               0.18.0+cpu
+tornado                   6.4
+tqdm                      4.66.2
+traitlets                 5.14.3
+types-python-dateutil     2.9.0.20240316
+typing_extensions         4.11.0
+tzdata                    2024.1
+uri-template              1.3.0
+urllib3                   2.2.1
+wcwidth                   0.2.13
+webcolors                 1.13
+webencodings              0.5.1
+websocket-client          1.8.0
+websockets                12.0
+wheel                     0.37.1
+widgetsnbextension        4.0.10
+xanadu-cloud-client       0.3.2
+zipp                      3.18.1
diff --git a/nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_13d_performance_indicators_JAX_scontrol.txt b/nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_13d_performance_indicators_JAX_scontrol.txt
new file mode 100644
index 00000000..e69de29b
diff --git a/nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_2d_performance_indicators_JAX.csv b/nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_2d_performance_indicators_JAX.csv
new file mode 100644
index 00000000..f7aec58e
--- /dev/null
+++ b/nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_2d_performance_indicators_JAX.csv
@@ -0,0 +1,2 @@
+construct_kernel_time,training_time,predict_time,hyperparameters
+2.4044647216796875,2.408006429672241,1.8492097854614258,"{'repeats': 10, 'use_jax': True, 'vmap': True, 'jit': True}"
diff --git a/nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_2d_performance_indicators_JAX_packages.txt b/nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_2d_performance_indicators_JAX_packages.txt
new file mode 100644
index 00000000..f5d4f0ca
--- /dev/null
+++ b/nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_2d_performance_indicators_JAX_packages.txt
@@ -0,0 +1,236 @@
+Package                   Version
+------------------------- --------------
+absl-py                   2.1.0
+alembic                   1.13.1
+annotated-types           0.6.0
+antlr4-python3-runtime    4.9.2
+anyio                     4.3.0
+appdirs                   1.4.4
+argon2-cffi               23.1.0
+argon2-cffi-bindings      21.2.0
+arrow                     1.3.0
+asttokens                 2.4.1
+astunparse                1.6.3
+async-lru                 2.0.4
+attrs                     23.2.0
+autograd                  1.6.2
+autoray                   0.6.9
+Babel                     2.14.0
+beautifulsoup4            4.12.3
+beniget                   0.4.1
+bitstring                 3.1.7
+bleach                    6.1.0
+cachetools                5.3.3
+certifi                   2024.2.2
+cffi                      1.16.0
+charset-normalizer        3.3.2
+chex                      0.1.86
+cirq-core                 1.3.0
+cirq-pasqal               1.3.0
+click                     8.1.7
+cloudpickle               3.0.0
+colorlog                  6.8.2
+comm                      0.2.2
+contourpy                 1.2.1
+cryptography              42.0.5
+cycler                    0.12.1
+dask                      2024.4.2
+dbus-python               1.2.18
+debugpy                   1.8.1
+decorator                 4.4.2
+defusedxml                0.7.1
+diastatic-malt            2.15.1
+dill                      0.3.8
+distro                    1.7.0
+duet                      0.2.9
+etils                     1.7.0
+exceptiongroup            1.2.1
+executing                 2.0.1
+fastdtw                   0.3.4
+fastjsonschema            2.19.1
+filelock                  3.13.1
+fire                      0.6.0
+flax                      0.8.3
+fonttools                 4.51.0
+fqdn                      1.5.1
+fsspec                    2024.3.1
+future                    1.0.0
+gast                      0.5.2
+greenlet                  3.0.3
+h11                       0.14.0
+h5py                      3.11.0
+httpcore                  1.0.5
+httpx                     0.27.0
+ibm-cloud-sdk-core        3.20.0
+ibm-platform-services     0.53.6
+idna                      3.7
+importlib_metadata        7.1.0
+importlib_resources       6.4.0
+ipykernel                 6.29.4
+ipython                   8.24.0
+ipywidgets                8.1.2
+isoduration               20.11.0
+jax                       0.4.23
+jaxlib                    0.4.23
+jaxopt                    0.8.3
+jedi                      0.19.1
+Jinja2                    3.1.3
+joblib                    1.4.0
+json5                     0.9.25
+jsonpointer               2.4
+jsonschema                4.21.1
+jsonschema-specifications 2023.12.1
+jupyter                   1.0.0
+jupyter_client            8.6.1
+jupyter-console           6.6.3
+jupyter_core              5.7.2
+jupyter-events            0.10.0
+jupyter-lsp               2.2.5
+jupyter_server            2.14.0
+jupyter_server_terminals  0.5.3
+jupyterlab                4.1.8
+jupyterlab_pygments       0.3.0
+jupyterlab_server         2.27.1
+jupyterlab_widgets        3.0.10
+kiwisolver                1.4.5
+lark-parser               0.12.0
+llvmlite                  0.42.0
+locket                    1.0.0
+Mako                      1.3.3
+markdown-it-py            3.0.0
+MarkupSafe                2.1.5
+matplotlib                3.8.4
+matplotlib-inline         0.1.7
+mdurl                     0.1.2
+mistune                   3.0.2
+ml-dtypes                 0.4.0
+mpmath                    1.3.0
+msgpack                   1.0.8
+nbclient                  0.10.0
+nbconvert                 7.16.4
+nbformat                  5.10.4
+nest-asyncio              1.6.0
+networkx                  3.3
+notebook                  7.1.3
+notebook_shim             0.2.4
+numba                     0.59.1
+numpy                     1.26.4
+olefile                   0.46
+opt-einsum                3.3.0
+optax                     0.2.2
+optuna                    3.6.1
+orbax-checkpoint          0.5.10
+overrides                 7.7.0
+packaging                 24.0
+pandas                    2.2.2
+pandocfilters             1.5.1
+parso                     0.8.4
+partd                     1.4.1
+patsy                     0.5.6
+pbr                       6.0.0
+PennyLane                 0.35.1
+PennyLane-Catalyst        0.5.0
+PennyLane-Cirq            0.34.0
+PennyLane_Lightning       0.35.1
+PennyLane_Lightning_GPU   0.35.1
+PennyLane-qiskit          0.35.1
+PennyLane-SF              0.29.0
+pexpect                   4.9.0
+pillow                    10.3.0
+pip                       24.0
+platformdirs              4.2.1
+ply                       3.11
+prometheus_client         0.20.0
+prompt-toolkit            3.0.43
+protobuf                  5.26.1
+psutil                    5.9.8
+ptyprocess                0.7.0
+pure-eval                 0.2.2
+pycparser                 2.22
+pydantic                  2.7.1
+pydantic_core             2.18.2
+pydantic-settings         2.2.1
+pydot                     2.0.0
+Pygments                  2.17.2
+PyGObject                 3.42.1
+PyJWT                     2.8.0
+pylatexenc                2.10
+pyparsing                 3.1.2
+pyspnego                  0.10.2
+python-dateutil           2.9.0.post0
+python-dotenv             1.0.1
+python-json-logger        2.0.7
+pythran                   0.10.0
+pytz                      2024.1
+PyYAML                    6.0.1
+pyzmq                     26.0.2
+qiskit                    1.0.2
+qiskit-aer                0.14.1
+qiskit-algorithms         0.3.0
+qiskit-ibm-experiment     0.4.7
+qiskit-ibm-provider       0.11.0
+qiskit-ibm-runtime        0.20.0
+qiskit-machine-learning   0.7.2
+qml_benchmarks            0.1
+qtconsole                 5.5.1
+QtPy                      2.4.1
+quantum-blackbird         0.5.0
+quantum-xir               0.2.2
+referencing               0.35.0
+requests                  2.31.0
+requests-ntlm             1.2.0
+rfc3339-validator         0.1.4
+rfc3986-validator         0.1.1
+rich                      13.7.1
+rpds-py                   0.18.0
+ruamel.yaml               0.18.6
+ruamel.yaml.clib          0.2.8
+rustworkx                 0.14.2
+scikit-learn              1.4.2
+scipy                     1.11.4
+seaborn                   0.13.2
+semantic-version          2.10.0
+Send2Trash                1.8.3
+setuptools                59.6.0
+six                       1.16.0
+sniffio                   1.3.1
+sortedcontainers          2.4.0
+soupsieve                 2.5
+SQLAlchemy                2.0.29
+ssh-import-id             5.11
+stack-data                0.6.3
+statsmodels               0.14.2
+stevedore                 5.2.0
+StrawberryFields          0.23.0
+symengine                 0.11.0
+sympy                     1.12
+tensorstore               0.1.58
+termcolor                 2.4.0
+terminado                 0.18.1
+thewalrus                 0.21.0
+threadpoolctl             3.5.0
+tinycss2                  1.3.0
+toml                      0.10.2
+tomli                     2.0.1
+tomlkit                   0.12.4
+toolz                     0.12.1
+torch                     2.3.0+cpu
+torchaudio                2.3.0+cpu
+torchvision               0.18.0+cpu
+tornado                   6.4
+tqdm                      4.66.2
+traitlets                 5.14.3
+types-python-dateutil     2.9.0.20240316
+typing_extensions         4.11.0
+tzdata                    2024.1
+uri-template              1.3.0
+urllib3                   2.2.1
+wcwidth                   0.2.13
+webcolors                 1.13
+webencodings              0.5.1
+websocket-client          1.8.0
+websockets                12.0
+wheel                     0.37.1
+widgetsnbextension        4.0.10
+xanadu-cloud-client       0.3.2
+zipp                      3.18.1
diff --git a/nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_2d_performance_indicators_JAX_scontrol.txt b/nersc/performance_indicators/JAX/IQPKernelClassifier_linearly_separable_2d_performance_indicators_JAX_scontrol.txt
new file mode 100644
index 00000000..e69de29b
diff --git a/nersc/performance_indicators/JAX/ProjectedQuantumKernel_linearly_separable_15d_performance_indicators_JAX.csv b/nersc/performance_indicators/JAX/ProjectedQuantumKernel_linearly_separable_15d_performance_indicators_JAX.csv
new file mode 100644
index 00000000..7d9877ce
--- /dev/null
+++ b/nersc/performance_indicators/JAX/ProjectedQuantumKernel_linearly_separable_15d_performance_indicators_JAX.csv
@@ -0,0 +1,2 @@
+construct_kernel_time,training_time,predict_time,hyperparameters
+68.69405555725098,68.69814586639404,70.27257776260376,"{'trotter_steps': 5, 'use_jax': True, 'vmap': True, 'jit': True}"
diff --git a/nersc/performance_indicators/JAX/ProjectedQuantumKernel_linearly_separable_15d_performance_indicators_JAX_packages.txt b/nersc/performance_indicators/JAX/ProjectedQuantumKernel_linearly_separable_15d_performance_indicators_JAX_packages.txt
new file mode 100644
index 00000000..f5d4f0ca
--- /dev/null
+++ b/nersc/performance_indicators/JAX/ProjectedQuantumKernel_linearly_separable_15d_performance_indicators_JAX_packages.txt
@@ -0,0 +1,236 @@
+Package                   Version
+------------------------- --------------
+absl-py                   2.1.0
+alembic                   1.13.1
+annotated-types           0.6.0
+antlr4-python3-runtime    4.9.2
+anyio                     4.3.0
+appdirs                   1.4.4
+argon2-cffi               23.1.0
+argon2-cffi-bindings      21.2.0
+arrow                     1.3.0
+asttokens                 2.4.1
+astunparse                1.6.3
+async-lru                 2.0.4
+attrs                     23.2.0
+autograd                  1.6.2
+autoray                   0.6.9
+Babel                     2.14.0
+beautifulsoup4            4.12.3
+beniget                   0.4.1
+bitstring                 3.1.7
+bleach                    6.1.0
+cachetools                5.3.3
+certifi                   2024.2.2
+cffi                      1.16.0
+charset-normalizer        3.3.2
+chex                      0.1.86
+cirq-core                 1.3.0
+cirq-pasqal               1.3.0
+click                     8.1.7
+cloudpickle               3.0.0
+colorlog                  6.8.2
+comm                      0.2.2
+contourpy                 1.2.1
+cryptography              42.0.5
+cycler                    0.12.1
+dask                      2024.4.2
+dbus-python               1.2.18
+debugpy                   1.8.1
+decorator                 4.4.2
+defusedxml                0.7.1
+diastatic-malt            2.15.1
+dill                      0.3.8
+distro                    1.7.0
+duet                      0.2.9
+etils                     1.7.0
+exceptiongroup            1.2.1
+executing                 2.0.1
+fastdtw                   0.3.4
+fastjsonschema            2.19.1
+filelock                  3.13.1
+fire                      0.6.0
+flax                      0.8.3
+fonttools                 4.51.0
+fqdn                      1.5.1
+fsspec                    2024.3.1
+future                    1.0.0
+gast                      0.5.2
+greenlet                  3.0.3
+h11                       0.14.0
+h5py                      3.11.0
+httpcore                  1.0.5
+httpx                     0.27.0
+ibm-cloud-sdk-core        3.20.0
+ibm-platform-services     0.53.6
+idna                      3.7
+importlib_metadata        7.1.0
+importlib_resources       6.4.0
+ipykernel                 6.29.4
+ipython                   8.24.0
+ipywidgets                8.1.2
+isoduration               20.11.0
+jax                       0.4.23
+jaxlib                    0.4.23
+jaxopt                    0.8.3
+jedi                      0.19.1
+Jinja2                    3.1.3
+joblib                    1.4.0
+json5                     0.9.25
+jsonpointer               2.4
+jsonschema                4.21.1
+jsonschema-specifications 2023.12.1
+jupyter                   1.0.0
+jupyter_client            8.6.1
+jupyter-console           6.6.3
+jupyter_core              5.7.2
+jupyter-events            0.10.0
+jupyter-lsp               2.2.5
+jupyter_server            2.14.0
+jupyter_server_terminals  0.5.3
+jupyterlab                4.1.8
+jupyterlab_pygments       0.3.0
+jupyterlab_server         2.27.1
+jupyterlab_widgets        3.0.10
+kiwisolver                1.4.5
+lark-parser               0.12.0
+llvmlite                  0.42.0
+locket                    1.0.0
+Mako                      1.3.3
+markdown-it-py            3.0.0
+MarkupSafe                2.1.5
+matplotlib                3.8.4
+matplotlib-inline         0.1.7
+mdurl                     0.1.2
+mistune                   3.0.2
+ml-dtypes                 0.4.0
+mpmath                    1.3.0
+msgpack                   1.0.8
+nbclient                  0.10.0
+nbconvert                 7.16.4
+nbformat                  5.10.4
+nest-asyncio              1.6.0
+networkx                  3.3
+notebook                  7.1.3
+notebook_shim             0.2.4
+numba                     0.59.1
+numpy                     1.26.4
+olefile                   0.46
+opt-einsum                3.3.0
+optax                     0.2.2
+optuna                    3.6.1
+orbax-checkpoint          0.5.10
+overrides                 7.7.0
+packaging                 24.0
+pandas                    2.2.2
+pandocfilters             1.5.1
+parso                     0.8.4
+partd                     1.4.1
+patsy                     0.5.6
+pbr                       6.0.0
+PennyLane                 0.35.1
+PennyLane-Catalyst        0.5.0
+PennyLane-Cirq            0.34.0
+PennyLane_Lightning       0.35.1
+PennyLane_Lightning_GPU   0.35.1
+PennyLane-qiskit          0.35.1
+PennyLane-SF              0.29.0
+pexpect                   4.9.0
+pillow                    10.3.0
+pip                       24.0
+platformdirs              4.2.1
+ply                       3.11
+prometheus_client         0.20.0
+prompt-toolkit            3.0.43
+protobuf                  5.26.1
+psutil                    5.9.8
+ptyprocess                0.7.0
+pure-eval                 0.2.2
+pycparser                 2.22
+pydantic                  2.7.1
+pydantic_core             2.18.2
+pydantic-settings         2.2.1
+pydot                     2.0.0
+Pygments                  2.17.2
+PyGObject                 3.42.1
+PyJWT                     2.8.0
+pylatexenc                2.10
+pyparsing                 3.1.2
+pyspnego                  0.10.2
+python-dateutil           2.9.0.post0
+python-dotenv             1.0.1
+python-json-logger        2.0.7
+pythran                   0.10.0
+pytz                      2024.1
+PyYAML                    6.0.1
+pyzmq                     26.0.2
+qiskit                    1.0.2
+qiskit-aer                0.14.1
+qiskit-algorithms         0.3.0
+qiskit-ibm-experiment     0.4.7
+qiskit-ibm-provider       0.11.0
+qiskit-ibm-runtime        0.20.0
+qiskit-machine-learning   0.7.2
+qml_benchmarks            0.1
+qtconsole                 5.5.1
+QtPy                      2.4.1
+quantum-blackbird         0.5.0
+quantum-xir               0.2.2
+referencing               0.35.0
+requests                  2.31.0
+requests-ntlm             1.2.0
+rfc3339-validator         0.1.4
+rfc3986-validator         0.1.1
+rich                      13.7.1
+rpds-py                   0.18.0
+ruamel.yaml               0.18.6
+ruamel.yaml.clib          0.2.8
+rustworkx                 0.14.2
+scikit-learn              1.4.2
+scipy                     1.11.4
+seaborn                   0.13.2
+semantic-version          2.10.0
+Send2Trash                1.8.3
+setuptools                59.6.0
+six                       1.16.0
+sniffio                   1.3.1
+sortedcontainers          2.4.0
+soupsieve                 2.5
+SQLAlchemy                2.0.29
+ssh-import-id             5.11
+stack-data                0.6.3
+statsmodels               0.14.2
+stevedore                 5.2.0
+StrawberryFields          0.23.0
+symengine                 0.11.0
+sympy                     1.12
+tensorstore               0.1.58
+termcolor                 2.4.0
+terminado                 0.18.1
+thewalrus                 0.21.0
+threadpoolctl             3.5.0
+tinycss2                  1.3.0
+toml                      0.10.2
+tomli                     2.0.1
+tomlkit                   0.12.4
+toolz                     0.12.1
+torch                     2.3.0+cpu
+torchaudio                2.3.0+cpu
+torchvision               0.18.0+cpu
+tornado                   6.4
+tqdm                      4.66.2
+traitlets                 5.14.3
+types-python-dateutil     2.9.0.20240316
+typing_extensions         4.11.0
+tzdata                    2024.1
+uri-template              1.3.0
+urllib3                   2.2.1
+wcwidth                   0.2.13
+webcolors                 1.13
+webencodings              0.5.1
+websocket-client          1.8.0
+websockets                12.0
+wheel                     0.37.1
+widgetsnbextension        4.0.10
+xanadu-cloud-client       0.3.2
+zipp                      3.18.1
diff --git a/nersc/performance_indicators/JAX/ProjectedQuantumKernel_linearly_separable_15d_performance_indicators_JAX_scontrol.txt b/nersc/performance_indicators/JAX/ProjectedQuantumKernel_linearly_separable_15d_performance_indicators_JAX_scontrol.txt
new file mode 100644
index 00000000..e69de29b
diff --git a/nersc/performance_indicators/JAX/ProjectedQuantumKernel_linearly_separable_17d_performance_indicators_JAX.csv b/nersc/performance_indicators/JAX/ProjectedQuantumKernel_linearly_separable_17d_performance_indicators_JAX.csv
new file mode 100644
index 00000000..90cfb263
--- /dev/null
+++ b/nersc/performance_indicators/JAX/ProjectedQuantumKernel_linearly_separable_17d_performance_indicators_JAX.csv
@@ -0,0 +1,2 @@
+construct_kernel_time,training_time,predict_time,hyperparameters
+143.43669366836548,143.44123721122742,146.31366968154907,"{'trotter_steps': 5, 'use_jax': True, 'vmap': True, 'jit': True}"
diff --git a/nersc/performance_indicators/JAX/ProjectedQuantumKernel_linearly_separable_17d_performance_indicators_JAX_packages.txt b/nersc/performance_indicators/JAX/ProjectedQuantumKernel_linearly_separable_17d_performance_indicators_JAX_packages.txt
new file mode 100644
index 00000000..f5d4f0ca
--- /dev/null
+++ b/nersc/performance_indicators/JAX/ProjectedQuantumKernel_linearly_separable_17d_performance_indicators_JAX_packages.txt
@@ -0,0 +1,236 @@
+Package                   Version
+------------------------- --------------
+absl-py                   2.1.0
+alembic                   1.13.1
+annotated-types           0.6.0
+antlr4-python3-runtime    4.9.2
+anyio                     4.3.0
+appdirs                   1.4.4
+argon2-cffi               23.1.0
+argon2-cffi-bindings      21.2.0
+arrow                     1.3.0
+asttokens                 2.4.1
+astunparse                1.6.3
+async-lru                 2.0.4
+attrs                     23.2.0
+autograd                  1.6.2
+autoray                   0.6.9
+Babel                     2.14.0
+beautifulsoup4            4.12.3
+beniget                   0.4.1
+bitstring                 3.1.7
+bleach                    6.1.0
+cachetools                5.3.3
+certifi                   2024.2.2
+cffi                      1.16.0
+charset-normalizer        3.3.2
+chex                      0.1.86
+cirq-core                 1.3.0
+cirq-pasqal               1.3.0
+click                     8.1.7
+cloudpickle               3.0.0
+colorlog                  6.8.2
+comm                      0.2.2
+contourpy                 1.2.1
+cryptography              42.0.5
+cycler                    0.12.1
+dask                      2024.4.2
+dbus-python               1.2.18
+debugpy                   1.8.1
+decorator                 4.4.2
+defusedxml                0.7.1
+diastatic-malt            2.15.1
+dill                      0.3.8
+distro                    1.7.0
+duet                      0.2.9
+etils                     1.7.0
+exceptiongroup            1.2.1
+executing                 2.0.1
+fastdtw                   0.3.4
+fastjsonschema            2.19.1
+filelock                  3.13.1
+fire                      0.6.0
+flax                      0.8.3
+fonttools                 4.51.0
+fqdn                      1.5.1
+fsspec                    2024.3.1
+future                    1.0.0
+gast                      0.5.2
+greenlet                  3.0.3
+h11                       0.14.0
+h5py                      3.11.0
+httpcore                  1.0.5
+httpx                     0.27.0
+ibm-cloud-sdk-core        3.20.0
+ibm-platform-services     0.53.6
+idna                      3.7
+importlib_metadata        7.1.0
+importlib_resources       6.4.0
+ipykernel                 6.29.4
+ipython                   8.24.0
+ipywidgets                8.1.2
+isoduration               20.11.0
+jax                       0.4.23
+jaxlib                    0.4.23
+jaxopt                    0.8.3
+jedi                      0.19.1
+Jinja2                    3.1.3
+joblib                    1.4.0
+json5                     0.9.25
+jsonpointer               2.4
+jsonschema                4.21.1
+jsonschema-specifications 2023.12.1
+jupyter                   1.0.0
+jupyter_client            8.6.1
+jupyter-console           6.6.3
+jupyter_core              5.7.2
+jupyter-events            0.10.0
+jupyter-lsp               2.2.5
+jupyter_server            2.14.0
+jupyter_server_terminals  0.5.3
+jupyterlab                4.1.8
+jupyterlab_pygments       0.3.0
+jupyterlab_server         2.27.1
+jupyterlab_widgets        3.0.10
+kiwisolver                1.4.5
+lark-parser               0.12.0
+llvmlite                  0.42.0
+locket                    1.0.0
+Mako                      1.3.3
+markdown-it-py            3.0.0
+MarkupSafe                2.1.5
+matplotlib                3.8.4
+matplotlib-inline         0.1.7
+mdurl                     0.1.2
+mistune                   3.0.2
+ml-dtypes                 0.4.0
+mpmath                    1.3.0
+msgpack                   1.0.8
+nbclient                  0.10.0
+nbconvert                 7.16.4
+nbformat                  5.10.4
+nest-asyncio              1.6.0
+networkx                  3.3
+notebook                  7.1.3
+notebook_shim             0.2.4
+numba                     0.59.1
+numpy                     1.26.4
+olefile                   0.46
+opt-einsum                3.3.0
+optax                     0.2.2
+optuna                    3.6.1
+orbax-checkpoint          0.5.10
+overrides                 7.7.0
+packaging                 24.0
+pandas                    2.2.2
+pandocfilters             1.5.1
+parso                     0.8.4
+partd                     1.4.1
+patsy                     0.5.6
+pbr                       6.0.0
+PennyLane                 0.35.1
+PennyLane-Catalyst        0.5.0
+PennyLane-Cirq            0.34.0
+PennyLane_Lightning       0.35.1
+PennyLane_Lightning_GPU   0.35.1
+PennyLane-qiskit          0.35.1
+PennyLane-SF              0.29.0
+pexpect                   4.9.0
+pillow                    10.3.0
+pip                       24.0
+platformdirs              4.2.1
+ply                       3.11
+prometheus_client         0.20.0
+prompt-toolkit            3.0.43
+protobuf                  5.26.1
+psutil                    5.9.8
+ptyprocess                0.7.0
+pure-eval                 0.2.2
+pycparser                 2.22
+pydantic                  2.7.1
+pydantic_core             2.18.2
+pydantic-settings         2.2.1
+pydot                     2.0.0
+Pygments                  2.17.2
+PyGObject                 3.42.1
+PyJWT                     2.8.0
+pylatexenc                2.10
+pyparsing                 3.1.2
+pyspnego                  0.10.2
+python-dateutil           2.9.0.post0
+python-dotenv             1.0.1
+python-json-logger        2.0.7
+pythran                   0.10.0
+pytz                      2024.1
+PyYAML                    6.0.1
+pyzmq                     26.0.2
+qiskit                    1.0.2
+qiskit-aer                0.14.1
+qiskit-algorithms         0.3.0
+qiskit-ibm-experiment     0.4.7
+qiskit-ibm-provider       0.11.0
+qiskit-ibm-runtime        0.20.0
+qiskit-machine-learning   0.7.2
+qml_benchmarks            0.1
+qtconsole                 5.5.1
+QtPy                      2.4.1
+quantum-blackbird         0.5.0
+quantum-xir               0.2.2
+referencing               0.35.0
+requests                  2.31.0
+requests-ntlm             1.2.0
+rfc3339-validator         0.1.4
+rfc3986-validator         0.1.1
+rich                      13.7.1
+rpds-py                   0.18.0
+ruamel.yaml               0.18.6
+ruamel.yaml.clib          0.2.8
+rustworkx                 0.14.2
+scikit-learn              1.4.2
+scipy                     1.11.4
+seaborn                   0.13.2
+semantic-version          2.10.0
+Send2Trash                1.8.3
+setuptools                59.6.0
+six                       1.16.0
+sniffio                   1.3.1
+sortedcontainers          2.4.0
+soupsieve                 2.5
+SQLAlchemy                2.0.29
+ssh-import-id             5.11
+stack-data                0.6.3
+statsmodels               0.14.2
+stevedore                 5.2.0
+StrawberryFields          0.23.0
+symengine                 0.11.0
+sympy                     1.12
+tensorstore               0.1.58
+termcolor                 2.4.0
+terminado                 0.18.1
+thewalrus                 0.21.0
+threadpoolctl             3.5.0
+tinycss2                  1.3.0
+toml                      0.10.2
+tomli                     2.0.1
+tomlkit                   0.12.4
+toolz                     0.12.1
+torch                     2.3.0+cpu
+torchaudio                2.3.0+cpu
+torchvision               0.18.0+cpu
+tornado                   6.4
+tqdm                      4.66.2
+traitlets                 5.14.3
+types-python-dateutil     2.9.0.20240316
+typing_extensions         4.11.0
+tzdata                    2024.1
+uri-template              1.3.0
+urllib3                   2.2.1
+wcwidth                   0.2.13
+webcolors                 1.13
+webencodings              0.5.1
+websocket-client          1.8.0
+websockets                12.0
+wheel                     0.37.1
+widgetsnbextension        4.0.10
+xanadu-cloud-client       0.3.2
+zipp                      3.18.1
diff --git a/nersc/performance_indicators/JAX/ProjectedQuantumKernel_linearly_separable_17d_performance_indicators_JAX_scontrol.txt b/nersc/performance_indicators/JAX/ProjectedQuantumKernel_linearly_separable_17d_performance_indicators_JAX_scontrol.txt
new file mode 100644
index 00000000..e69de29b
diff --git a/nersc/performance_indicators/JAX/ProjectedQuantumKernel_linearly_separable_20d_performance_indicators_JAX.csv b/nersc/performance_indicators/JAX/ProjectedQuantumKernel_linearly_separable_20d_performance_indicators_JAX.csv
new file mode 100644
index 00000000..10f48e95
--- /dev/null
+++ b/nersc/performance_indicators/JAX/ProjectedQuantumKernel_linearly_separable_20d_performance_indicators_JAX.csv
@@ -0,0 +1,2 @@
+construct_kernel_time,training_time,predict_time,hyperparameters
+987.022438287735,987.0266752243042,905.2389245033264,"{'trotter_steps': 5, 'use_jax': True, 'vmap': True, 'jit': True}"
diff --git a/nersc/performance_indicators/JAX/ProjectedQuantumKernel_linearly_separable_20d_performance_indicators_JAX_packages.txt b/nersc/performance_indicators/JAX/ProjectedQuantumKernel_linearly_separable_20d_performance_indicators_JAX_packages.txt
new file mode 100644
index 00000000..f5d4f0ca
--- /dev/null
+++ b/nersc/performance_indicators/JAX/ProjectedQuantumKernel_linearly_separable_20d_performance_indicators_JAX_packages.txt
@@ -0,0 +1,236 @@
+Package                   Version
+------------------------- --------------
+absl-py                   2.1.0
+alembic                   1.13.1
+annotated-types           0.6.0
+antlr4-python3-runtime    4.9.2
+anyio                     4.3.0
+appdirs                   1.4.4
+argon2-cffi               23.1.0
+argon2-cffi-bindings      21.2.0
+arrow                     1.3.0
+asttokens                 2.4.1
+astunparse                1.6.3
+async-lru                 2.0.4
+attrs                     23.2.0
+autograd                  1.6.2
+autoray                   0.6.9
+Babel                     2.14.0
+beautifulsoup4            4.12.3
+beniget                   0.4.1
+bitstring                 3.1.7
+bleach                    6.1.0
+cachetools                5.3.3
+certifi                   2024.2.2
+cffi                      1.16.0
+charset-normalizer        3.3.2
+chex                      0.1.86
+cirq-core                 1.3.0
+cirq-pasqal               1.3.0
+click                     8.1.7
+cloudpickle               3.0.0
+colorlog                  6.8.2
+comm                      0.2.2
+contourpy                 1.2.1
+cryptography              42.0.5
+cycler                    0.12.1
+dask                      2024.4.2
+dbus-python               1.2.18
+debugpy                   1.8.1
+decorator                 4.4.2
+defusedxml                0.7.1
+diastatic-malt            2.15.1
+dill                      0.3.8
+distro                    1.7.0
+duet                      0.2.9
+etils                     1.7.0
+exceptiongroup            1.2.1
+executing                 2.0.1
+fastdtw                   0.3.4
+fastjsonschema            2.19.1
+filelock                  3.13.1
+fire                      0.6.0
+flax                      0.8.3
+fonttools                 4.51.0
+fqdn                      1.5.1
+fsspec                    2024.3.1
+future                    1.0.0
+gast                      0.5.2
+greenlet                  3.0.3
+h11                       0.14.0
+h5py                      3.11.0
+httpcore                  1.0.5
+httpx                     0.27.0
+ibm-cloud-sdk-core        3.20.0
+ibm-platform-services     0.53.6
+idna                      3.7
+importlib_metadata        7.1.0
+importlib_resources       6.4.0
+ipykernel                 6.29.4
+ipython                   8.24.0
+ipywidgets                8.1.2
+isoduration               20.11.0
+jax                       0.4.23
+jaxlib                    0.4.23
+jaxopt                    0.8.3
+jedi                      0.19.1
+Jinja2                    3.1.3
+joblib                    1.4.0
+json5                     0.9.25
+jsonpointer               2.4
+jsonschema                4.21.1
+jsonschema-specifications 2023.12.1
+jupyter                   1.0.0
+jupyter_client            8.6.1
+jupyter-console           6.6.3
+jupyter_core              5.7.2
+jupyter-events            0.10.0
+jupyter-lsp               2.2.5
+jupyter_server            2.14.0
+jupyter_server_terminals  0.5.3
+jupyterlab                4.1.8
+jupyterlab_pygments       0.3.0
+jupyterlab_server         2.27.1
+jupyterlab_widgets        3.0.10
+kiwisolver                1.4.5
+lark-parser               0.12.0
+llvmlite                  0.42.0
+locket                    1.0.0
+Mako                      1.3.3
+markdown-it-py            3.0.0
+MarkupSafe                2.1.5
+matplotlib                3.8.4
+matplotlib-inline         0.1.7
+mdurl                     0.1.2
+mistune                   3.0.2
+ml-dtypes                 0.4.0
+mpmath                    1.3.0
+msgpack                   1.0.8
+nbclient                  0.10.0
+nbconvert                 7.16.4
+nbformat                  5.10.4
+nest-asyncio              1.6.0
+networkx                  3.3
+notebook                  7.1.3
+notebook_shim             0.2.4
+numba                     0.59.1
+numpy                     1.26.4
+olefile                   0.46
+opt-einsum                3.3.0
+optax                     0.2.2
+optuna                    3.6.1
+orbax-checkpoint          0.5.10
+overrides                 7.7.0
+packaging                 24.0
+pandas                    2.2.2
+pandocfilters             1.5.1
+parso                     0.8.4
+partd                     1.4.1
+patsy                     0.5.6
+pbr                       6.0.0
+PennyLane                 0.35.1
+PennyLane-Catalyst        0.5.0
+PennyLane-Cirq            0.34.0
+PennyLane_Lightning       0.35.1
+PennyLane_Lightning_GPU   0.35.1
+PennyLane-qiskit          0.35.1
+PennyLane-SF              0.29.0
+pexpect                   4.9.0
+pillow                    10.3.0
+pip                       24.0
+platformdirs              4.2.1
+ply                       3.11
+prometheus_client         0.20.0
+prompt-toolkit            3.0.43
+protobuf                  5.26.1
+psutil                    5.9.8
+ptyprocess                0.7.0
+pure-eval                 0.2.2
+pycparser                 2.22
+pydantic                  2.7.1
+pydantic_core             2.18.2
+pydantic-settings         2.2.1
+pydot                     2.0.0
+Pygments                  2.17.2
+PyGObject                 3.42.1
+PyJWT                     2.8.0
+pylatexenc                2.10
+pyparsing                 3.1.2
+pyspnego                  0.10.2
+python-dateutil           2.9.0.post0
+python-dotenv             1.0.1
+python-json-logger        2.0.7
+pythran                   0.10.0
+pytz                      2024.1
+PyYAML                    6.0.1
+pyzmq                     26.0.2
+qiskit                    1.0.2
+qiskit-aer                0.14.1
+qiskit-algorithms         0.3.0
+qiskit-ibm-experiment     0.4.7
+qiskit-ibm-provider       0.11.0
+qiskit-ibm-runtime        0.20.0
+qiskit-machine-learning   0.7.2
+qml_benchmarks            0.1
+qtconsole                 5.5.1
+QtPy                      2.4.1
+quantum-blackbird         0.5.0
+quantum-xir               0.2.2
+referencing               0.35.0
+requests                  2.31.0
+requests-ntlm             1.2.0
+rfc3339-validator         0.1.4
+rfc3986-validator         0.1.1
+rich                      13.7.1
+rpds-py                   0.18.0
+ruamel.yaml               0.18.6
+ruamel.yaml.clib          0.2.8
+rustworkx                 0.14.2
+scikit-learn              1.4.2
+scipy                     1.11.4
+seaborn                   0.13.2
+semantic-version          2.10.0
+Send2Trash                1.8.3
+setuptools                59.6.0
+six                       1.16.0
+sniffio                   1.3.1
+sortedcontainers          2.4.0
+soupsieve                 2.5
+SQLAlchemy                2.0.29
+ssh-import-id             5.11
+stack-data                0.6.3
+statsmodels               0.14.2
+stevedore                 5.2.0
+StrawberryFields          0.23.0
+symengine                 0.11.0
+sympy                     1.12
+tensorstore               0.1.58
+termcolor                 2.4.0
+terminado                 0.18.1
+thewalrus                 0.21.0
+threadpoolctl             3.5.0
+tinycss2                  1.3.0
+toml                      0.10.2
+tomli                     2.0.1
+tomlkit                   0.12.4
+toolz                     0.12.1
+torch                     2.3.0+cpu
+torchaudio                2.3.0+cpu
+torchvision               0.18.0+cpu
+tornado                   6.4
+tqdm                      4.66.2
+traitlets                 5.14.3
+types-python-dateutil     2.9.0.20240316
+typing_extensions         4.11.0
+tzdata                    2024.1
+uri-template              1.3.0
+urllib3                   2.2.1
+wcwidth                   0.2.13
+webcolors                 1.13
+webencodings              0.5.1
+websocket-client          1.8.0
+websockets                12.0
+wheel                     0.37.1
+widgetsnbextension        4.0.10
+xanadu-cloud-client       0.3.2
+zipp                      3.18.1
diff --git a/nersc/performance_indicators/JAX/ProjectedQuantumKernel_linearly_separable_20d_performance_indicators_JAX_scontrol.txt b/nersc/performance_indicators/JAX/ProjectedQuantumKernel_linearly_separable_20d_performance_indicators_JAX_scontrol.txt
new file mode 100644
index 00000000..e69de29b
diff --git a/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_13d_performance_indicators_JAX.csv b/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_13d_performance_indicators_JAX.csv
new file mode 100644
index 00000000..bb1d1bf0
--- /dev/null
+++ b/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_13d_performance_indicators_JAX.csv
@@ -0,0 +1,2 @@
+first_train_step,first_train_step_std,consec_train_step,consec_train_step_std,predict_time,predict_time_std,hyperparameters
+118.64471855163575,13.176545198064549,5.997364590866397,0.3148851718071143,603.2157266139984,30.745927128215357,"{'n_layers': 4, 'batch_size': 16, 'use_jax': True, 'vmap': True, 'max_steps': 100, 'jit': True}"
diff --git a/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_13d_performance_indicators_JAX_packages.txt b/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_13d_performance_indicators_JAX_packages.txt
new file mode 100644
index 00000000..f5d4f0ca
--- /dev/null
+++ b/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_13d_performance_indicators_JAX_packages.txt
@@ -0,0 +1,236 @@
+Package                   Version
+------------------------- --------------
+absl-py                   2.1.0
+alembic                   1.13.1
+annotated-types           0.6.0
+antlr4-python3-runtime    4.9.2
+anyio                     4.3.0
+appdirs                   1.4.4
+argon2-cffi               23.1.0
+argon2-cffi-bindings      21.2.0
+arrow                     1.3.0
+asttokens                 2.4.1
+astunparse                1.6.3
+async-lru                 2.0.4
+attrs                     23.2.0
+autograd                  1.6.2
+autoray                   0.6.9
+Babel                     2.14.0
+beautifulsoup4            4.12.3
+beniget                   0.4.1
+bitstring                 3.1.7
+bleach                    6.1.0
+cachetools                5.3.3
+certifi                   2024.2.2
+cffi                      1.16.0
+charset-normalizer        3.3.2
+chex                      0.1.86
+cirq-core                 1.3.0
+cirq-pasqal               1.3.0
+click                     8.1.7
+cloudpickle               3.0.0
+colorlog                  6.8.2
+comm                      0.2.2
+contourpy                 1.2.1
+cryptography              42.0.5
+cycler                    0.12.1
+dask                      2024.4.2
+dbus-python               1.2.18
+debugpy                   1.8.1
+decorator                 4.4.2
+defusedxml                0.7.1
+diastatic-malt            2.15.1
+dill                      0.3.8
+distro                    1.7.0
+duet                      0.2.9
+etils                     1.7.0
+exceptiongroup            1.2.1
+executing                 2.0.1
+fastdtw                   0.3.4
+fastjsonschema            2.19.1
+filelock                  3.13.1
+fire                      0.6.0
+flax                      0.8.3
+fonttools                 4.51.0
+fqdn                      1.5.1
+fsspec                    2024.3.1
+future                    1.0.0
+gast                      0.5.2
+greenlet                  3.0.3
+h11                       0.14.0
+h5py                      3.11.0
+httpcore                  1.0.5
+httpx                     0.27.0
+ibm-cloud-sdk-core        3.20.0
+ibm-platform-services     0.53.6
+idna                      3.7
+importlib_metadata        7.1.0
+importlib_resources       6.4.0
+ipykernel                 6.29.4
+ipython                   8.24.0
+ipywidgets                8.1.2
+isoduration               20.11.0
+jax                       0.4.23
+jaxlib                    0.4.23
+jaxopt                    0.8.3
+jedi                      0.19.1
+Jinja2                    3.1.3
+joblib                    1.4.0
+json5                     0.9.25
+jsonpointer               2.4
+jsonschema                4.21.1
+jsonschema-specifications 2023.12.1
+jupyter                   1.0.0
+jupyter_client            8.6.1
+jupyter-console           6.6.3
+jupyter_core              5.7.2
+jupyter-events            0.10.0
+jupyter-lsp               2.2.5
+jupyter_server            2.14.0
+jupyter_server_terminals  0.5.3
+jupyterlab                4.1.8
+jupyterlab_pygments       0.3.0
+jupyterlab_server         2.27.1
+jupyterlab_widgets        3.0.10
+kiwisolver                1.4.5
+lark-parser               0.12.0
+llvmlite                  0.42.0
+locket                    1.0.0
+Mako                      1.3.3
+markdown-it-py            3.0.0
+MarkupSafe                2.1.5
+matplotlib                3.8.4
+matplotlib-inline         0.1.7
+mdurl                     0.1.2
+mistune                   3.0.2
+ml-dtypes                 0.4.0
+mpmath                    1.3.0
+msgpack                   1.0.8
+nbclient                  0.10.0
+nbconvert                 7.16.4
+nbformat                  5.10.4
+nest-asyncio              1.6.0
+networkx                  3.3
+notebook                  7.1.3
+notebook_shim             0.2.4
+numba                     0.59.1
+numpy                     1.26.4
+olefile                   0.46
+opt-einsum                3.3.0
+optax                     0.2.2
+optuna                    3.6.1
+orbax-checkpoint          0.5.10
+overrides                 7.7.0
+packaging                 24.0
+pandas                    2.2.2
+pandocfilters             1.5.1
+parso                     0.8.4
+partd                     1.4.1
+patsy                     0.5.6
+pbr                       6.0.0
+PennyLane                 0.35.1
+PennyLane-Catalyst        0.5.0
+PennyLane-Cirq            0.34.0
+PennyLane_Lightning       0.35.1
+PennyLane_Lightning_GPU   0.35.1
+PennyLane-qiskit          0.35.1
+PennyLane-SF              0.29.0
+pexpect                   4.9.0
+pillow                    10.3.0
+pip                       24.0
+platformdirs              4.2.1
+ply                       3.11
+prometheus_client         0.20.0
+prompt-toolkit            3.0.43
+protobuf                  5.26.1
+psutil                    5.9.8
+ptyprocess                0.7.0
+pure-eval                 0.2.2
+pycparser                 2.22
+pydantic                  2.7.1
+pydantic_core             2.18.2
+pydantic-settings         2.2.1
+pydot                     2.0.0
+Pygments                  2.17.2
+PyGObject                 3.42.1
+PyJWT                     2.8.0
+pylatexenc                2.10
+pyparsing                 3.1.2
+pyspnego                  0.10.2
+python-dateutil           2.9.0.post0
+python-dotenv             1.0.1
+python-json-logger        2.0.7
+pythran                   0.10.0
+pytz                      2024.1
+PyYAML                    6.0.1
+pyzmq                     26.0.2
+qiskit                    1.0.2
+qiskit-aer                0.14.1
+qiskit-algorithms         0.3.0
+qiskit-ibm-experiment     0.4.7
+qiskit-ibm-provider       0.11.0
+qiskit-ibm-runtime        0.20.0
+qiskit-machine-learning   0.7.2
+qml_benchmarks            0.1
+qtconsole                 5.5.1
+QtPy                      2.4.1
+quantum-blackbird         0.5.0
+quantum-xir               0.2.2
+referencing               0.35.0
+requests                  2.31.0
+requests-ntlm             1.2.0
+rfc3339-validator         0.1.4
+rfc3986-validator         0.1.1
+rich                      13.7.1
+rpds-py                   0.18.0
+ruamel.yaml               0.18.6
+ruamel.yaml.clib          0.2.8
+rustworkx                 0.14.2
+scikit-learn              1.4.2
+scipy                     1.11.4
+seaborn                   0.13.2
+semantic-version          2.10.0
+Send2Trash                1.8.3
+setuptools                59.6.0
+six                       1.16.0
+sniffio                   1.3.1
+sortedcontainers          2.4.0
+soupsieve                 2.5
+SQLAlchemy                2.0.29
+ssh-import-id             5.11
+stack-data                0.6.3
+statsmodels               0.14.2
+stevedore                 5.2.0
+StrawberryFields          0.23.0
+symengine                 0.11.0
+sympy                     1.12
+tensorstore               0.1.58
+termcolor                 2.4.0
+terminado                 0.18.1
+thewalrus                 0.21.0
+threadpoolctl             3.5.0
+tinycss2                  1.3.0
+toml                      0.10.2
+tomli                     2.0.1
+tomlkit                   0.12.4
+toolz                     0.12.1
+torch                     2.3.0+cpu
+torchaudio                2.3.0+cpu
+torchvision               0.18.0+cpu
+tornado                   6.4
+tqdm                      4.66.2
+traitlets                 5.14.3
+types-python-dateutil     2.9.0.20240316
+typing_extensions         4.11.0
+tzdata                    2024.1
+uri-template              1.3.0
+urllib3                   2.2.1
+wcwidth                   0.2.13
+webcolors                 1.13
+webencodings              0.5.1
+websocket-client          1.8.0
+websockets                12.0
+wheel                     0.37.1
+widgetsnbextension        4.0.10
+xanadu-cloud-client       0.3.2
+zipp                      3.18.1
diff --git a/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_13d_performance_indicators_JAX_scontrol.txt b/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_13d_performance_indicators_JAX_scontrol.txt
new file mode 100644
index 00000000..e69de29b
diff --git a/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_16d_performance_indicators_JAX.csv b/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_16d_performance_indicators_JAX.csv
new file mode 100644
index 00000000..13b06b9e
--- /dev/null
+++ b/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_16d_performance_indicators_JAX.csv
@@ -0,0 +1,2 @@
+first_train_step,first_train_step_std,consec_train_step,consec_train_step_std,predict_time,predict_time_std,hyperparameters
+156.70789346694946,10.264342281303652,18.11830727981799,0.7520945854665012,846.8818587303161,21.645421239376127,"{'n_layers': 4, 'batch_size': 16, 'use_jax': True, 'vmap': True, 'max_steps': 100, 'jit': True}"
diff --git a/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_16d_performance_indicators_JAX_packages.txt b/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_16d_performance_indicators_JAX_packages.txt
new file mode 100644
index 00000000..f5d4f0ca
--- /dev/null
+++ b/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_16d_performance_indicators_JAX_packages.txt
@@ -0,0 +1,236 @@
+Package                   Version
+------------------------- --------------
+absl-py                   2.1.0
+alembic                   1.13.1
+annotated-types           0.6.0
+antlr4-python3-runtime    4.9.2
+anyio                     4.3.0
+appdirs                   1.4.4
+argon2-cffi               23.1.0
+argon2-cffi-bindings      21.2.0
+arrow                     1.3.0
+asttokens                 2.4.1
+astunparse                1.6.3
+async-lru                 2.0.4
+attrs                     23.2.0
+autograd                  1.6.2
+autoray                   0.6.9
+Babel                     2.14.0
+beautifulsoup4            4.12.3
+beniget                   0.4.1
+bitstring                 3.1.7
+bleach                    6.1.0
+cachetools                5.3.3
+certifi                   2024.2.2
+cffi                      1.16.0
+charset-normalizer        3.3.2
+chex                      0.1.86
+cirq-core                 1.3.0
+cirq-pasqal               1.3.0
+click                     8.1.7
+cloudpickle               3.0.0
+colorlog                  6.8.2
+comm                      0.2.2
+contourpy                 1.2.1
+cryptography              42.0.5
+cycler                    0.12.1
+dask                      2024.4.2
+dbus-python               1.2.18
+debugpy                   1.8.1
+decorator                 4.4.2
+defusedxml                0.7.1
+diastatic-malt            2.15.1
+dill                      0.3.8
+distro                    1.7.0
+duet                      0.2.9
+etils                     1.7.0
+exceptiongroup            1.2.1
+executing                 2.0.1
+fastdtw                   0.3.4
+fastjsonschema            2.19.1
+filelock                  3.13.1
+fire                      0.6.0
+flax                      0.8.3
+fonttools                 4.51.0
+fqdn                      1.5.1
+fsspec                    2024.3.1
+future                    1.0.0
+gast                      0.5.2
+greenlet                  3.0.3
+h11                       0.14.0
+h5py                      3.11.0
+httpcore                  1.0.5
+httpx                     0.27.0
+ibm-cloud-sdk-core        3.20.0
+ibm-platform-services     0.53.6
+idna                      3.7
+importlib_metadata        7.1.0
+importlib_resources       6.4.0
+ipykernel                 6.29.4
+ipython                   8.24.0
+ipywidgets                8.1.2
+isoduration               20.11.0
+jax                       0.4.23
+jaxlib                    0.4.23
+jaxopt                    0.8.3
+jedi                      0.19.1
+Jinja2                    3.1.3
+joblib                    1.4.0
+json5                     0.9.25
+jsonpointer               2.4
+jsonschema                4.21.1
+jsonschema-specifications 2023.12.1
+jupyter                   1.0.0
+jupyter_client            8.6.1
+jupyter-console           6.6.3
+jupyter_core              5.7.2
+jupyter-events            0.10.0
+jupyter-lsp               2.2.5
+jupyter_server            2.14.0
+jupyter_server_terminals  0.5.3
+jupyterlab                4.1.8
+jupyterlab_pygments       0.3.0
+jupyterlab_server         2.27.1
+jupyterlab_widgets        3.0.10
+kiwisolver                1.4.5
+lark-parser               0.12.0
+llvmlite                  0.42.0
+locket                    1.0.0
+Mako                      1.3.3
+markdown-it-py            3.0.0
+MarkupSafe                2.1.5
+matplotlib                3.8.4
+matplotlib-inline         0.1.7
+mdurl                     0.1.2
+mistune                   3.0.2
+ml-dtypes                 0.4.0
+mpmath                    1.3.0
+msgpack                   1.0.8
+nbclient                  0.10.0
+nbconvert                 7.16.4
+nbformat                  5.10.4
+nest-asyncio              1.6.0
+networkx                  3.3
+notebook                  7.1.3
+notebook_shim             0.2.4
+numba                     0.59.1
+numpy                     1.26.4
+olefile                   0.46
+opt-einsum                3.3.0
+optax                     0.2.2
+optuna                    3.6.1
+orbax-checkpoint          0.5.10
+overrides                 7.7.0
+packaging                 24.0
+pandas                    2.2.2
+pandocfilters             1.5.1
+parso                     0.8.4
+partd                     1.4.1
+patsy                     0.5.6
+pbr                       6.0.0
+PennyLane                 0.35.1
+PennyLane-Catalyst        0.5.0
+PennyLane-Cirq            0.34.0
+PennyLane_Lightning       0.35.1
+PennyLane_Lightning_GPU   0.35.1
+PennyLane-qiskit          0.35.1
+PennyLane-SF              0.29.0
+pexpect                   4.9.0
+pillow                    10.3.0
+pip                       24.0
+platformdirs              4.2.1
+ply                       3.11
+prometheus_client         0.20.0
+prompt-toolkit            3.0.43
+protobuf                  5.26.1
+psutil                    5.9.8
+ptyprocess                0.7.0
+pure-eval                 0.2.2
+pycparser                 2.22
+pydantic                  2.7.1
+pydantic_core             2.18.2
+pydantic-settings         2.2.1
+pydot                     2.0.0
+Pygments                  2.17.2
+PyGObject                 3.42.1
+PyJWT                     2.8.0
+pylatexenc                2.10
+pyparsing                 3.1.2
+pyspnego                  0.10.2
+python-dateutil           2.9.0.post0
+python-dotenv             1.0.1
+python-json-logger        2.0.7
+pythran                   0.10.0
+pytz                      2024.1
+PyYAML                    6.0.1
+pyzmq                     26.0.2
+qiskit                    1.0.2
+qiskit-aer                0.14.1
+qiskit-algorithms         0.3.0
+qiskit-ibm-experiment     0.4.7
+qiskit-ibm-provider       0.11.0
+qiskit-ibm-runtime        0.20.0
+qiskit-machine-learning   0.7.2
+qml_benchmarks            0.1
+qtconsole                 5.5.1
+QtPy                      2.4.1
+quantum-blackbird         0.5.0
+quantum-xir               0.2.2
+referencing               0.35.0
+requests                  2.31.0
+requests-ntlm             1.2.0
+rfc3339-validator         0.1.4
+rfc3986-validator         0.1.1
+rich                      13.7.1
+rpds-py                   0.18.0
+ruamel.yaml               0.18.6
+ruamel.yaml.clib          0.2.8
+rustworkx                 0.14.2
+scikit-learn              1.4.2
+scipy                     1.11.4
+seaborn                   0.13.2
+semantic-version          2.10.0
+Send2Trash                1.8.3
+setuptools                59.6.0
+six                       1.16.0
+sniffio                   1.3.1
+sortedcontainers          2.4.0
+soupsieve                 2.5
+SQLAlchemy                2.0.29
+ssh-import-id             5.11
+stack-data                0.6.3
+statsmodels               0.14.2
+stevedore                 5.2.0
+StrawberryFields          0.23.0
+symengine                 0.11.0
+sympy                     1.12
+tensorstore               0.1.58
+termcolor                 2.4.0
+terminado                 0.18.1
+thewalrus                 0.21.0
+threadpoolctl             3.5.0
+tinycss2                  1.3.0
+toml                      0.10.2
+tomli                     2.0.1
+tomlkit                   0.12.4
+toolz                     0.12.1
+torch                     2.3.0+cpu
+torchaudio                2.3.0+cpu
+torchvision               0.18.0+cpu
+tornado                   6.4
+tqdm                      4.66.2
+traitlets                 5.14.3
+types-python-dateutil     2.9.0.20240316
+typing_extensions         4.11.0
+tzdata                    2024.1
+uri-template              1.3.0
+urllib3                   2.2.1
+wcwidth                   0.2.13
+webcolors                 1.13
+webencodings              0.5.1
+websocket-client          1.8.0
+websockets                12.0
+wheel                     0.37.1
+widgetsnbextension        4.0.10
+xanadu-cloud-client       0.3.2
+zipp                      3.18.1
diff --git a/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_16d_performance_indicators_JAX_scontrol.txt b/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_16d_performance_indicators_JAX_scontrol.txt
new file mode 100644
index 00000000..e69de29b
diff --git a/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_2d_performance_indicators_JAX.csv b/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_2d_performance_indicators_JAX.csv
new file mode 100644
index 00000000..0f549554
--- /dev/null
+++ b/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_2d_performance_indicators_JAX.csv
@@ -0,0 +1,2 @@
+first_train_step,first_train_step_std,consec_train_step,consec_train_step_std,predict_time,predict_time_std,hyperparameters
+18.878392934799194,0.00267791748046875,0.0069307666836362905,3.6954879760742188e-06,5.774362087249756,0.06407976150512695,"{'n_layers': 4, 'batch_size': 16, 'use_jax': True, 'vmap': True, 'max_steps': 100, 'jit': True}"
diff --git a/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_2d_performance_indicators_JAX_packages.txt b/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_2d_performance_indicators_JAX_packages.txt
new file mode 100644
index 00000000..f5d4f0ca
--- /dev/null
+++ b/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_2d_performance_indicators_JAX_packages.txt
@@ -0,0 +1,236 @@
+Package                   Version
+------------------------- --------------
+absl-py                   2.1.0
+alembic                   1.13.1
+annotated-types           0.6.0
+antlr4-python3-runtime    4.9.2
+anyio                     4.3.0
+appdirs                   1.4.4
+argon2-cffi               23.1.0
+argon2-cffi-bindings      21.2.0
+arrow                     1.3.0
+asttokens                 2.4.1
+astunparse                1.6.3
+async-lru                 2.0.4
+attrs                     23.2.0
+autograd                  1.6.2
+autoray                   0.6.9
+Babel                     2.14.0
+beautifulsoup4            4.12.3
+beniget                   0.4.1
+bitstring                 3.1.7
+bleach                    6.1.0
+cachetools                5.3.3
+certifi                   2024.2.2
+cffi                      1.16.0
+charset-normalizer        3.3.2
+chex                      0.1.86
+cirq-core                 1.3.0
+cirq-pasqal               1.3.0
+click                     8.1.7
+cloudpickle               3.0.0
+colorlog                  6.8.2
+comm                      0.2.2
+contourpy                 1.2.1
+cryptography              42.0.5
+cycler                    0.12.1
+dask                      2024.4.2
+dbus-python               1.2.18
+debugpy                   1.8.1
+decorator                 4.4.2
+defusedxml                0.7.1
+diastatic-malt            2.15.1
+dill                      0.3.8
+distro                    1.7.0
+duet                      0.2.9
+etils                     1.7.0
+exceptiongroup            1.2.1
+executing                 2.0.1
+fastdtw                   0.3.4
+fastjsonschema            2.19.1
+filelock                  3.13.1
+fire                      0.6.0
+flax                      0.8.3
+fonttools                 4.51.0
+fqdn                      1.5.1
+fsspec                    2024.3.1
+future                    1.0.0
+gast                      0.5.2
+greenlet                  3.0.3
+h11                       0.14.0
+h5py                      3.11.0
+httpcore                  1.0.5
+httpx                     0.27.0
+ibm-cloud-sdk-core        3.20.0
+ibm-platform-services     0.53.6
+idna                      3.7
+importlib_metadata        7.1.0
+importlib_resources       6.4.0
+ipykernel                 6.29.4
+ipython                   8.24.0
+ipywidgets                8.1.2
+isoduration               20.11.0
+jax                       0.4.23
+jaxlib                    0.4.23
+jaxopt                    0.8.3
+jedi                      0.19.1
+Jinja2                    3.1.3
+joblib                    1.4.0
+json5                     0.9.25
+jsonpointer               2.4
+jsonschema                4.21.1
+jsonschema-specifications 2023.12.1
+jupyter                   1.0.0
+jupyter_client            8.6.1
+jupyter-console           6.6.3
+jupyter_core              5.7.2
+jupyter-events            0.10.0
+jupyter-lsp               2.2.5
+jupyter_server            2.14.0
+jupyter_server_terminals  0.5.3
+jupyterlab                4.1.8
+jupyterlab_pygments       0.3.0
+jupyterlab_server         2.27.1
+jupyterlab_widgets        3.0.10
+kiwisolver                1.4.5
+lark-parser               0.12.0
+llvmlite                  0.42.0
+locket                    1.0.0
+Mako                      1.3.3
+markdown-it-py            3.0.0
+MarkupSafe                2.1.5
+matplotlib                3.8.4
+matplotlib-inline         0.1.7
+mdurl                     0.1.2
+mistune                   3.0.2
+ml-dtypes                 0.4.0
+mpmath                    1.3.0
+msgpack                   1.0.8
+nbclient                  0.10.0
+nbconvert                 7.16.4
+nbformat                  5.10.4
+nest-asyncio              1.6.0
+networkx                  3.3
+notebook                  7.1.3
+notebook_shim             0.2.4
+numba                     0.59.1
+numpy                     1.26.4
+olefile                   0.46
+opt-einsum                3.3.0
+optax                     0.2.2
+optuna                    3.6.1
+orbax-checkpoint          0.5.10
+overrides                 7.7.0
+packaging                 24.0
+pandas                    2.2.2
+pandocfilters             1.5.1
+parso                     0.8.4
+partd                     1.4.1
+patsy                     0.5.6
+pbr                       6.0.0
+PennyLane                 0.35.1
+PennyLane-Catalyst        0.5.0
+PennyLane-Cirq            0.34.0
+PennyLane_Lightning       0.35.1
+PennyLane_Lightning_GPU   0.35.1
+PennyLane-qiskit          0.35.1
+PennyLane-SF              0.29.0
+pexpect                   4.9.0
+pillow                    10.3.0
+pip                       24.0
+platformdirs              4.2.1
+ply                       3.11
+prometheus_client         0.20.0
+prompt-toolkit            3.0.43
+protobuf                  5.26.1
+psutil                    5.9.8
+ptyprocess                0.7.0
+pure-eval                 0.2.2
+pycparser                 2.22
+pydantic                  2.7.1
+pydantic_core             2.18.2
+pydantic-settings         2.2.1
+pydot                     2.0.0
+Pygments                  2.17.2
+PyGObject                 3.42.1
+PyJWT                     2.8.0
+pylatexenc                2.10
+pyparsing                 3.1.2
+pyspnego                  0.10.2
+python-dateutil           2.9.0.post0
+python-dotenv             1.0.1
+python-json-logger        2.0.7
+pythran                   0.10.0
+pytz                      2024.1
+PyYAML                    6.0.1
+pyzmq                     26.0.2
+qiskit                    1.0.2
+qiskit-aer                0.14.1
+qiskit-algorithms         0.3.0
+qiskit-ibm-experiment     0.4.7
+qiskit-ibm-provider       0.11.0
+qiskit-ibm-runtime        0.20.0
+qiskit-machine-learning   0.7.2
+qml_benchmarks            0.1
+qtconsole                 5.5.1
+QtPy                      2.4.1
+quantum-blackbird         0.5.0
+quantum-xir               0.2.2
+referencing               0.35.0
+requests                  2.31.0
+requests-ntlm             1.2.0
+rfc3339-validator         0.1.4
+rfc3986-validator         0.1.1
+rich                      13.7.1
+rpds-py                   0.18.0
+ruamel.yaml               0.18.6
+ruamel.yaml.clib          0.2.8
+rustworkx                 0.14.2
+scikit-learn              1.4.2
+scipy                     1.11.4
+seaborn                   0.13.2
+semantic-version          2.10.0
+Send2Trash                1.8.3
+setuptools                59.6.0
+six                       1.16.0
+sniffio                   1.3.1
+sortedcontainers          2.4.0
+soupsieve                 2.5
+SQLAlchemy                2.0.29
+ssh-import-id             5.11
+stack-data                0.6.3
+statsmodels               0.14.2
+stevedore                 5.2.0
+StrawberryFields          0.23.0
+symengine                 0.11.0
+sympy                     1.12
+tensorstore               0.1.58
+termcolor                 2.4.0
+terminado                 0.18.1
+thewalrus                 0.21.0
+threadpoolctl             3.5.0
+tinycss2                  1.3.0
+toml                      0.10.2
+tomli                     2.0.1
+tomlkit                   0.12.4
+toolz                     0.12.1
+torch                     2.3.0+cpu
+torchaudio                2.3.0+cpu
+torchvision               0.18.0+cpu
+tornado                   6.4
+tqdm                      4.66.2
+traitlets                 5.14.3
+types-python-dateutil     2.9.0.20240316
+typing_extensions         4.11.0
+tzdata                    2024.1
+uri-template              1.3.0
+urllib3                   2.2.1
+wcwidth                   0.2.13
+webcolors                 1.13
+webencodings              0.5.1
+websocket-client          1.8.0
+websockets                12.0
+wheel                     0.37.1
+widgetsnbextension        4.0.10
+xanadu-cloud-client       0.3.2
+zipp                      3.18.1
diff --git a/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_2d_performance_indicators_JAX_scontrol.txt b/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_2d_performance_indicators_JAX_scontrol.txt
new file mode 100644
index 00000000..e69de29b

From 34ca1cf7be2be6cdc58f0e4b710bc780d0579414 Mon Sep 17 00:00:00 2001
From: josephbowles <bowles.physics@gmail.com>
Date: Thu, 2 May 2024 07:25:09 -0700
Subject: [PATCH 063/100] update

---
 nersc/performance_indicators/perf_ind_kernel_jax.py      | 6 +++---
 nersc/performance_indicators/perf_ind_variational_jax.py | 7 +++----
 2 files changed, 6 insertions(+), 7 deletions(-)

diff --git a/nersc/performance_indicators/perf_ind_kernel_jax.py b/nersc/performance_indicators/perf_ind_kernel_jax.py
index 4d6c0d8a..fa61a73a 100644
--- a/nersc/performance_indicators/perf_ind_kernel_jax.py
+++ b/nersc/performance_indicators/perf_ind_kernel_jax.py
@@ -14,7 +14,7 @@
 def get_parser():
     parser = argparse.ArgumentParser()
     parser.add_argument("-v","--verbosity",type=int,choices=[0, 1, 2,3,4],  help="increase output verbosity", default=1, dest='verb')
-    parser.add_argument("--inputPath",default='../linearly_separable',help='input data location')
+    parser.add_argument("--inputPath",default='linearly_separable/',help='input data location')
     parser.add_argument('-n','--numFeatures',type=int,default=2, help="dataset dimension ")
 
     args = parser.parse_args()
@@ -33,8 +33,8 @@ def get_parser():
 if __name__=="__main__":
     args=get_parser()
 
-    print(os.getcwd())
-    os.chdir('performance_indicators')
+    #print(os.getcwd())
+    #os.chdir('performance_indicators')
 
     #################################
     # settings for the performance indicator.
diff --git a/nersc/performance_indicators/perf_ind_variational_jax.py b/nersc/performance_indicators/perf_ind_variational_jax.py
index b1b6298f..a62b2af5 100644
--- a/nersc/performance_indicators/perf_ind_variational_jax.py
+++ b/nersc/performance_indicators/perf_ind_variational_jax.py
@@ -11,7 +11,7 @@
 from qml_benchmarks.hyperparam_search_utils import read_data
 
 print(os.getcwd())
-os.chdir('performance_indicators')
+# os.chdir('nersc/performance_indicators')
 
 #################################
 # settings for the performance indicator.
@@ -41,8 +41,8 @@
 hyperparams = {**hp_settings[model_name], **{'use_jax':use_jax, 'vmap':vmap, 'max_steps':max_steps, 'jit': jit}}
 print(hyperparams)
 
-X_train,y_train = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
-X_test,y_test = read_data(f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv')
+X_train,y_train = read_data(f'linearly_separable/linearly_separable_{n_features}d_train.csv')
+X_test,y_test = read_data(f'linearly_separable/linearly_separable_{n_features}d_train.csv')
 
 first_train_steps = []
 av_consec_train_steps = []
@@ -111,4 +111,3 @@
 # make an empty text file to store the scontrol data
 with open(perf_ind_name+'/'+scontrol_filename, 'w') as file:
     pass
-

From 7f52b53d9bb9aa799848c55a904971149b1e1d10 Mon Sep 17 00:00:00 2001
From: josephbowles <bowles.physics@gmail.com>
Date: Mon, 6 May 2024 04:01:59 -0700
Subject: [PATCH 064/100] result

---
 ...earner_linearly_separable_16d_performance_indicators_JAX.csv | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_16d_performance_indicators_JAX.csv b/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_16d_performance_indicators_JAX.csv
index 13b06b9e..6ee18e68 100644
--- a/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_16d_performance_indicators_JAX.csv
+++ b/nersc/performance_indicators/JAX/QuantumMetricLearner_linearly_separable_16d_performance_indicators_JAX.csv
@@ -1,2 +1,2 @@
 first_train_step,first_train_step_std,consec_train_step,consec_train_step_std,predict_time,predict_time_std,hyperparameters
-156.70789346694946,10.264342281303652,18.11830727981799,0.7520945854665012,846.8818587303161,21.645421239376127,"{'n_layers': 4, 'batch_size': 16, 'use_jax': True, 'vmap': True, 'max_steps': 100, 'jit': True}"
+186.66433362960817,26.55748553966961,28.853529106005272,2.0727284764603633,2363.9554340362547,173.4738127352816,"{'n_layers': 4, 'batch_size': 16, 'use_jax': True, 'vmap': True, 'max_steps': 100, 'jit': True}"

From 07b91a9192ab51e0432465c76693a588525ae56d Mon Sep 17 00:00:00 2001
From: josephbowles <bowles.physics@gmail.com>
Date: Mon, 6 May 2024 05:57:11 -0700
Subject: [PATCH 065/100] working sbatch

---
 .../perf_ind_kernel_jax.py                     |  8 ++++----
 nersc/submit_job2.slr                          | 18 +++++++++---------
 nersc/wrap_podman2.sh                          |  2 +-
 3 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/nersc/performance_indicators/perf_ind_kernel_jax.py b/nersc/performance_indicators/perf_ind_kernel_jax.py
index fa61a73a..d982f005 100644
--- a/nersc/performance_indicators/perf_ind_kernel_jax.py
+++ b/nersc/performance_indicators/perf_ind_kernel_jax.py
@@ -56,7 +56,7 @@ def get_parser():
     model_name = Model().__class__.__name__
 
     # get the 'worst case' hyperparameter settings for the model (those that require the most resources)
-    with open('hyperparam_settings.yaml', "r") as file:
+    with open('performance_indicators/hyperparam_settings.yaml', "r") as file:
         hp_settings = yaml.safe_load(file)
 
     hyperparams = {**hp_settings[model_name], **{'use_jax':use_jax, 'vmap':vmap, 'jit': jit}}
@@ -97,7 +97,7 @@ def get_parser():
         os.mkdir(perf_ind_name)
 
     #write perf indicator data
-    with open(perf_ind_name+'/'+filename, mode="w", newline="") as file:
+    with open('performance_indicators/'+perf_ind_name+'/'+filename, mode="w", newline="") as file:
         writer = csv.writer(file)
         writer.writerow(header)
         for row in [data]:
@@ -105,10 +105,10 @@ def get_parser():
 
     # get package list and write to file
     output = subprocess.check_output(['pip', 'list']).decode('utf-8')
-    with open(perf_ind_name+'/'+packages_filename, 'w') as file:
+    with open('performance_indicators/'+perf_ind_name+'/'+packages_filename, 'w') as file:
         file.write(output)
 
     # make an empty text file to store the scontrol data
-    with open(perf_ind_name+'/'+scontrol_filename, 'w') as file:
+    with open('performance_indicators/'+perf_ind_name+'/'+scontrol_filename, 'w') as file:
         pass
     print('M:done')
diff --git a/nersc/submit_job2.slr b/nersc/submit_job2.slr
index 2dce16db..b17925ec 100755
--- a/nersc/submit_job2.slr
+++ b/nersc/submit_job2.slr
@@ -1,8 +1,8 @@
 #!/bin/bash
 #SBATCH -N 1
 #SBATCH -C cpu
-# SBATCH -q shared -t 10:00:00
-#SBATCH -q debug -t 5:00  # charged for full node
+#SBATCH -q shared -t 1:00:00
+# SBATCH -q debug -t 5:00  # charged for full node
 #SBATCH -J perf_ind_qkernel2
 #SBATCH --ntasks-per-node=1
 #SBATCH --cpus-per-task=6
@@ -18,12 +18,12 @@ numFeatures=3
 # .... container
 # salloc -q interactive -C cpu -t 4:00:00 -A nstaff
 
-IMG=balewski/ubu22-pennylane:p5  # old
-#IMG=jbowles/ubu22-pennylane:p5
+#IMG=balewski/ubu22-pennylane:p5  # old
+IMG=jbowles/ubu22-pennylane:p5
 echo use image $IMG
 
-#POD_PUB=//dvs_ro/cfs/cdirs/m4139/qml-benchmarks/nersc/podman/
-POD_PUB=/dvs_ro/cfs/cdirs/nstaff/balewski/podman_common/  #old
+POD_PUB=//dvs_ro/cfs/cdirs/m4139/qml-benchmarks/nersc/podman/
+#POD_PUB=/dvs_ro/cfs/cdirs/nstaff/balewski/podman_common/  #old
 export PODMANHPC_ADDITIONAL_STORES=$POD_PUB
 
 #.... output sandbox 
@@ -36,13 +36,13 @@ cp -rp performance_indicators *.slr *.sh   $outPath
 
 
 #... location of input on CSF
-#CFSH=/global/cfs/cdirs/m4139  # for Joseph
-CFSH=/pscratch/sd/b/balewski/pennylane_wrk  # for Jan
+CFSH=/global/cfs/cdirs/m4139  # for Joseph
+# CFSH=/pscratch/sd/b/balewski/pennylane_wrk  # for Jan
 BASE_DIR=/qml-benchmarks   # here git has home
 WORK_DIR=$BASE_DIR/nersc
 
 #... the task to be executed
-CMD=" python3 -u  performance_indicators/perf_ind_kernel_jax.py --numFeatures $numFeatures --inputPath ../../../linearly_separable "
+CMD=" python3 -u  performance_indicators/perf_ind_kernel_jax.py --numFeatures $numFeatures --inputPath performance_indicators/linearly_separable/ "
 #CMD=" python3 -u -c \"import pennylane \"; echo done1  "
 #CMD=" python3 -u -c \"import qml_benchmarks \";  echo done2 host  "
 
diff --git a/nersc/wrap_podman2.sh b/nersc/wrap_podman2.sh
index a6b263c0..ee8ec5d0 100755
--- a/nersc/wrap_podman2.sh
+++ b/nersc/wrap_podman2.sh
@@ -18,7 +18,7 @@ echo 'W:start podman'
 podman-hpc run -it \
     --volume $CFSH/$BASE_DIR:/root \
     --volume $CFSH/$BASE_DIR:$BASE_DIR \
-    --volume $CFSH/$BASE_DIR/linearly_separable:/linearly_separable \
+    --volume $CFSH/$BASE_DIR/nersc/performance_indicators/linearly_separable:/linearly_separable \
     --volume $CFSH/$WORK_DIR:$WORK_DIR \
     -e HDF5_USE_FILE_LOCKING='FALSE' \
     --workdir $WORK_DIR \

From 241c7ede9246d7c0c879feee04ea2939bc21f6ad Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Mon, 6 May 2024 15:01:19 +0200
Subject: [PATCH 066/100] rename

---
 nersc/submit_job.slr | 62 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 62 insertions(+)
 create mode 100755 nersc/submit_job.slr

diff --git a/nersc/submit_job.slr b/nersc/submit_job.slr
new file mode 100755
index 00000000..b17925ec
--- /dev/null
+++ b/nersc/submit_job.slr
@@ -0,0 +1,62 @@
+#!/bin/bash
+#SBATCH -N 1
+#SBATCH -C cpu
+#SBATCH -q shared -t 1:00:00
+# SBATCH -q debug -t 5:00  # charged for full node
+#SBATCH -J perf_ind_qkernel2
+#SBATCH --ntasks-per-node=1
+#SBATCH --cpus-per-task=6
+#SBATCH --output out1/%j.log
+#SBATCH  --licenses=scratch
+# - - - E N D    O F    SLURM    C O M M A N D S
+set -u ;  # exit  if you try to use an uninitialized variable
+
+echo 'S:start'
+
+numFeatures=3
+
+# .... container
+# salloc -q interactive -C cpu -t 4:00:00 -A nstaff
+
+#IMG=balewski/ubu22-pennylane:p5  # old
+IMG=jbowles/ubu22-pennylane:p5
+echo use image $IMG
+
+POD_PUB=//dvs_ro/cfs/cdirs/m4139/qml-benchmarks/nersc/podman/
+#POD_PUB=/dvs_ro/cfs/cdirs/nstaff/balewski/podman_common/  #old
+export PODMANHPC_ADDITIONAL_STORES=$POD_PUB
+
+#.... output sandbox 
+CODE_DIR=`pwd`
+myJob=${SLURM_JOBID}
+outPath=$SCRATCH/tmp_pennylane_jobs/$myJob
+echo outPath=$outPath
+mkdir -p $outPath
+cp -rp performance_indicators *.slr *.sh   $outPath
+
+
+#... location of input on CSF
+CFSH=/global/cfs/cdirs/m4139  # for Joseph
+# CFSH=/pscratch/sd/b/balewski/pennylane_wrk  # for Jan
+BASE_DIR=/qml-benchmarks   # here git has home
+WORK_DIR=$BASE_DIR/nersc
+
+#... the task to be executed
+CMD=" python3 -u  performance_indicators/perf_ind_kernel_jax.py --numFeatures $numFeatures --inputPath performance_indicators/linearly_separable/ "
+#CMD=" python3 -u -c \"import pennylane \"; echo done1  "
+#CMD=" python3 -u -c \"import qml_benchmarks \";  echo done2 host  "
+
+cd $outPath
+# this is hack Jan  does not like  but 'paper' dir is not part of pip-install
+# qml-benchmarks> cp -rp paper /pscratch/sd/b/balewski/tmp_pennylane_jobs
+
+#ls -l $CFSH/$BASE_DIR/paper
+#ln -s $BASE_DIR/paper ..
+
+G=1
+echo 'S:ready to run '; pwd
+srun -n $G  ./wrap_podman2.sh $IMG " $CMD "  $outPath  $CFSH $BASE_DIR $WORK_DIR
+
+sleep 1
+echo 'S:end '; date
+

From f6c48847eb95774ea4800ba1ce1ba86f24dc8be6 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Mon, 6 May 2024 15:01:48 +0200
Subject: [PATCH 067/100] rename

---
 nersc/wrap_podman.sh | 38 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 38 insertions(+)
 create mode 100755 nersc/wrap_podman.sh

diff --git a/nersc/wrap_podman.sh b/nersc/wrap_podman.sh
new file mode 100755
index 00000000..ee8ec5d0
--- /dev/null
+++ b/nersc/wrap_podman.sh
@@ -0,0 +1,38 @@
+#!/bin/bash
+echo W:myRank is $SLURM_PROCID
+IMG=$1
+CMD=$2
+outPath=$3
+CFSH=$4
+BASE_DIR=$5
+WORK_DIR=$6
+
+if [ $SLURM_PROCID -eq 0 ] ; then 
+   echo W:IMG=$IMG 
+   echo W:CMD=$CMD
+   #echo Q:fire $
+fi
+
+echo W:BASE_DIR=$BASE_DIR
+echo 'W:start podman'
+podman-hpc run -it \
+    --volume $CFSH/$BASE_DIR:/root \
+    --volume $CFSH/$BASE_DIR:$BASE_DIR \
+    --volume $CFSH/$BASE_DIR/nersc/performance_indicators/linearly_separable:/linearly_separable \
+    --volume $CFSH/$WORK_DIR:$WORK_DIR \
+    -e HDF5_USE_FILE_LOCKING='FALSE' \
+    --workdir $WORK_DIR \
+    $IMG <<EOF 
+echo P:pwd; pwd
+echo P:all
+ls -l ../..
+$CMD
+echo P:done
+exit
+EOF
+  
+echo 'W:done podman'
+
+# spare
+# 	   --volume $HOME:/home \
+    # --volume  $CFSH/qml-benchmarks:/root 

From 03cfc6b5535c1810574d8a75016b63a8d94bb7ae Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Mon, 6 May 2024 15:05:49 +0200
Subject: [PATCH 068/100] cleanup

---
 .../perf_ind_kernel_jax.py                    |   6 -
 .../perf_ind_variational_jax.py               | 174 ++++++++++--------
 nersc/submit_job.slr                          |   2 +-
 nersc/submit_job2.slr                         |  62 -------
 nersc/wrap_podman2.sh                         |  38 ----
 5 files changed, 100 insertions(+), 182 deletions(-)
 delete mode 100755 nersc/submit_job2.slr
 delete mode 100755 nersc/wrap_podman2.sh

diff --git a/nersc/performance_indicators/perf_ind_kernel_jax.py b/nersc/performance_indicators/perf_ind_kernel_jax.py
index d982f005..516e6d70 100644
--- a/nersc/performance_indicators/perf_ind_kernel_jax.py
+++ b/nersc/performance_indicators/perf_ind_kernel_jax.py
@@ -33,9 +33,6 @@ def get_parser():
 if __name__=="__main__":
     args=get_parser()
 
-    #print(os.getcwd())
-    #os.chdir('performance_indicators')
-
     #################################
     # settings for the performance indicator.
     # You only need to change this to make a different performance indicator
@@ -108,7 +105,4 @@ def get_parser():
     with open('performance_indicators/'+perf_ind_name+'/'+packages_filename, 'w') as file:
         file.write(output)
 
-    # make an empty text file to store the scontrol data
-    with open('performance_indicators/'+perf_ind_name+'/'+scontrol_filename, 'w') as file:
-        pass
     print('M:done')
diff --git a/nersc/performance_indicators/perf_ind_variational_jax.py b/nersc/performance_indicators/perf_ind_variational_jax.py
index a62b2af5..f945e674 100644
--- a/nersc/performance_indicators/perf_ind_variational_jax.py
+++ b/nersc/performance_indicators/perf_ind_variational_jax.py
@@ -10,104 +10,128 @@
 import subprocess
 from qml_benchmarks.hyperparam_search_utils import read_data
 
-print(os.getcwd())
-# os.chdir('nersc/performance_indicators')
+import argparse
 
-#################################
-# settings for the performance indicator.
-# You only need to change this to make a different performance indicator
 
-#define model
-from qml_benchmarks.models.quantum_metric_learning import QuantumMetricLearner as Model
+def get_parser():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-v", "--verbosity", type=int, choices=[0, 1, 2, 3, 4], help="increase output verbosity",
+                        default=1, dest='verb')
+    parser.add_argument("--inputPath", default='linearly_separable/', help='input data location')
+    parser.add_argument('-n', '--numFeatures', type=int, default=2, help="dataset dimension ")
 
-#implementation attributes of model
-use_jax = True
-vmap = True
-jit = True
+    args = parser.parse_args()
 
-max_steps = 100 #the number of gradient descent steps to use to estimate the step time
-perf_ind_name = 'JAX'  #a name for the performance indicator used for naming files
-n_features = 2 #dataset dimension
-n_trials = 2 #number of trials to average over
+    print('myArg-program:', parser.prog)
+    for arg in vars(args):  print('myArg:', arg, getattr(args, arg))
 
-#################################
+    # assert os.path.exists(args.outPath)
+    return args
 
-model_name = Model().__class__.__name__
 
-# get the 'worst case' hyperparameter settings for the model (those that require the most resources)
-with open('hyperparam_settings.yaml', "r") as file:
-    hp_settings = yaml.safe_load(file)
+# =================================
+# =================================
+#  M A I N
+# =================================
+# =================================
+if __name__ == "__main__":
+    args = get_parser()
 
-hyperparams = {**hp_settings[model_name], **{'use_jax':use_jax, 'vmap':vmap, 'max_steps':max_steps, 'jit': jit}}
-print(hyperparams)
+    #################################
+    # settings for the performance indicator.
+    # You only need to change this to make a different performance indicator
 
-X_train,y_train = read_data(f'linearly_separable/linearly_separable_{n_features}d_train.csv')
-X_test,y_test = read_data(f'linearly_separable/linearly_separable_{n_features}d_train.csv')
+    #define model
+    from qml_benchmarks.models.quantum_metric_learning import QuantumMetricLearner as Model
 
-first_train_steps = []
-av_consec_train_steps = []
-predict_times = []
+    #implementation attributes of model
+    use_jax = True
+    vmap = True
+    jit = True
 
-for trial in range(n_trials):
-    jax.clear_caches()
+    max_steps = 100 #the number of gradient descent steps to use to estimate the step time
+    perf_ind_name = 'JAX'  #a name for the performance indicator used for naming files
+    n_features = 2 #dataset dimension
+    n_trials = 5 #number of trials to average over
 
-    model = Model(**hyperparams)
+    #################################
 
-    #note we train for max_steps only, so we won't reach convergence.
-    model.fit(X_train, y_train)
 
-    #get step times from loss history data
-    step_times = np.array([model.loss_history_[1][i + 1] - model.loss_history_[1][i]
-                           for i in range(len(model.loss_history_[1]) - 1)])
-    step_times = np.insert(step_times, 0, [model.loss_history_[1][0]])
+    model_name = Model().__class__.__name__
 
-    #first train step
-    first_train_steps.append(step_times[0])
-    #consecutive (average) train step
-    av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
-    #prediction time
-    time0 = time.time()
-    model.predict(X_test)
-    predict_times.append(time.time() - time0)
+    # get the 'worst case' hyperparameter settings for the model (those that require the most resources)
+    with open('performance_indicators/hyperparam_settings.yaml', "r") as file:
+        hp_settings = yaml.safe_load(file)
 
+    hyperparams = {**hp_settings[model_name], **{'use_jax':use_jax, 'vmap':vmap, 'max_steps':max_steps, 'jit': jit}}
+    print(hyperparams)
 
-#calculate mean and stds
-first_train_step = np.mean(first_train_steps)
-first_train_step_std = np.std(first_train_steps)
+    X_train,y_train = read_data(f'linearly_separable/linearly_separable_{n_features}d_train.csv')
+    X_test,y_test = read_data(f'linearly_separable/linearly_separable_{n_features}d_train.csv')
 
-consec_train_step = np.mean(av_consec_train_steps)
-consec_train_step_std = np.std(av_consec_train_steps)
+    first_train_steps = []
+    av_consec_train_steps = []
+    predict_times = []
 
-predict_time = np.mean(predict_times)
-predict_time_std = np.std(predict_times)
+    for trial in range(n_trials):
+        jax.clear_caches()
 
-#write to csv
-data = [first_train_step, first_train_step_std, consec_train_step, consec_train_step_std, predict_time,
-        predict_time_std, hyperparams]
+        model = Model(**hyperparams)
 
-filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_{perf_ind_name}.csv"
-packages_filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_{perf_ind_name}_packages.txt"
-scontrol_filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_{perf_ind_name}_scontrol.txt"
+        #note we train for max_steps only, so we won't reach convergence.
+        model.fit(X_train, y_train)
 
-header = ['first_train_step', 'first_train_step_std', 'consec_train_step', 'consec_train_step_std', 'predict_time',
-          'predict_time_std', 'hyperparameters']
+        #get step times from loss history data
+        step_times = np.array([model.loss_history_[1][i + 1] - model.loss_history_[1][i]
+                               for i in range(len(model.loss_history_[1]) - 1)])
+        step_times = np.insert(step_times, 0, [model.loss_history_[1][0]])
 
-if not os.path.exists(perf_ind_name):
-    # Create the directory
-    os.mkdir(perf_ind_name)
+        #first train step
+        first_train_steps.append(step_times[0])
+        #consecutive (average) train step
+        av_consec_train_steps.append(float(jnp.mean(step_times[1:])))
+        #prediction time
+        time0 = time.time()
+        model.predict(X_test)
+        predict_times.append(time.time() - time0)
 
-#write perf indicator data
-with open(perf_ind_name+'/'+filename, mode="w", newline="") as file:
-    writer = csv.writer(file)
-    writer.writerow(header)
-    for row in [data]:
-        writer.writerow(row)
 
-# get package list and write to file
-output = subprocess.check_output(['pip', 'list']).decode('utf-8')
-with open(perf_ind_name+'/'+packages_filename, 'w') as file:
-    file.write(output)
+    #calculate mean and stds
+    first_train_step = np.mean(first_train_steps)
+    first_train_step_std = np.std(first_train_steps)
+
+    consec_train_step = np.mean(av_consec_train_steps)
+    consec_train_step_std = np.std(av_consec_train_steps)
+
+    predict_time = np.mean(predict_times)
+    predict_time_std = np.std(predict_times)
+
+    #write to csv
+    data = [first_train_step, first_train_step_std, consec_train_step, consec_train_step_std, predict_time,
+            predict_time_std, hyperparams]
+
+    filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_{perf_ind_name}.csv"
+    packages_filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_{perf_ind_name}_packages.txt"
+    scontrol_filename =  model_name+f"_linearly_separable_{n_features}d_performance_indicators_{perf_ind_name}_scontrol.txt"
+
+    header = ['first_train_step', 'first_train_step_std', 'consec_train_step', 'consec_train_step_std', 'predict_time',
+              'predict_time_std', 'hyperparameters']
+
+    if not os.path.exists(perf_ind_name):
+        # Create the directory
+        os.mkdir(perf_ind_name)
+
+    #write perf indicator data
+    with open('performance_indicators/'+perf_ind_name+'/'+filename, mode="w", newline="") as file:
+        writer = csv.writer(file)
+        writer.writerow(header)
+        for row in [data]:
+            writer.writerow(row)
+
+    # get package list and write to file
+    output = subprocess.check_output(['pip', 'list']).decode('utf-8')
+    with open('performance_indicators/'+perf_ind_name+'/'+packages_filename, 'w') as file:
+        file.write(output)
+
+    print('M:done')
 
-# make an empty text file to store the scontrol data
-with open(perf_ind_name+'/'+scontrol_filename, 'w') as file:
-    pass
diff --git a/nersc/submit_job.slr b/nersc/submit_job.slr
index b17925ec..9b7a1ad4 100755
--- a/nersc/submit_job.slr
+++ b/nersc/submit_job.slr
@@ -55,7 +55,7 @@ cd $outPath
 
 G=1
 echo 'S:ready to run '; pwd
-srun -n $G  ./wrap_podman2.sh $IMG " $CMD "  $outPath  $CFSH $BASE_DIR $WORK_DIR
+srun -n $G  ./wrap_podman.sh $IMG " $CMD "  $outPath  $CFSH $BASE_DIR $WORK_DIR
 
 sleep 1
 echo 'S:end '; date
diff --git a/nersc/submit_job2.slr b/nersc/submit_job2.slr
deleted file mode 100755
index b17925ec..00000000
--- a/nersc/submit_job2.slr
+++ /dev/null
@@ -1,62 +0,0 @@
-#!/bin/bash
-#SBATCH -N 1
-#SBATCH -C cpu
-#SBATCH -q shared -t 1:00:00
-# SBATCH -q debug -t 5:00  # charged for full node
-#SBATCH -J perf_ind_qkernel2
-#SBATCH --ntasks-per-node=1
-#SBATCH --cpus-per-task=6
-#SBATCH --output out1/%j.log
-#SBATCH  --licenses=scratch
-# - - - E N D    O F    SLURM    C O M M A N D S
-set -u ;  # exit  if you try to use an uninitialized variable
-
-echo 'S:start'
-
-numFeatures=3
-
-# .... container
-# salloc -q interactive -C cpu -t 4:00:00 -A nstaff
-
-#IMG=balewski/ubu22-pennylane:p5  # old
-IMG=jbowles/ubu22-pennylane:p5
-echo use image $IMG
-
-POD_PUB=//dvs_ro/cfs/cdirs/m4139/qml-benchmarks/nersc/podman/
-#POD_PUB=/dvs_ro/cfs/cdirs/nstaff/balewski/podman_common/  #old
-export PODMANHPC_ADDITIONAL_STORES=$POD_PUB
-
-#.... output sandbox 
-CODE_DIR=`pwd`
-myJob=${SLURM_JOBID}
-outPath=$SCRATCH/tmp_pennylane_jobs/$myJob
-echo outPath=$outPath
-mkdir -p $outPath
-cp -rp performance_indicators *.slr *.sh   $outPath
-
-
-#... location of input on CSF
-CFSH=/global/cfs/cdirs/m4139  # for Joseph
-# CFSH=/pscratch/sd/b/balewski/pennylane_wrk  # for Jan
-BASE_DIR=/qml-benchmarks   # here git has home
-WORK_DIR=$BASE_DIR/nersc
-
-#... the task to be executed
-CMD=" python3 -u  performance_indicators/perf_ind_kernel_jax.py --numFeatures $numFeatures --inputPath performance_indicators/linearly_separable/ "
-#CMD=" python3 -u -c \"import pennylane \"; echo done1  "
-#CMD=" python3 -u -c \"import qml_benchmarks \";  echo done2 host  "
-
-cd $outPath
-# this is hack Jan  does not like  but 'paper' dir is not part of pip-install
-# qml-benchmarks> cp -rp paper /pscratch/sd/b/balewski/tmp_pennylane_jobs
-
-#ls -l $CFSH/$BASE_DIR/paper
-#ln -s $BASE_DIR/paper ..
-
-G=1
-echo 'S:ready to run '; pwd
-srun -n $G  ./wrap_podman2.sh $IMG " $CMD "  $outPath  $CFSH $BASE_DIR $WORK_DIR
-
-sleep 1
-echo 'S:end '; date
-
diff --git a/nersc/wrap_podman2.sh b/nersc/wrap_podman2.sh
deleted file mode 100755
index ee8ec5d0..00000000
--- a/nersc/wrap_podman2.sh
+++ /dev/null
@@ -1,38 +0,0 @@
-#!/bin/bash
-echo W:myRank is $SLURM_PROCID
-IMG=$1
-CMD=$2
-outPath=$3
-CFSH=$4
-BASE_DIR=$5
-WORK_DIR=$6
-
-if [ $SLURM_PROCID -eq 0 ] ; then 
-   echo W:IMG=$IMG 
-   echo W:CMD=$CMD
-   #echo Q:fire $
-fi
-
-echo W:BASE_DIR=$BASE_DIR
-echo 'W:start podman'
-podman-hpc run -it \
-    --volume $CFSH/$BASE_DIR:/root \
-    --volume $CFSH/$BASE_DIR:$BASE_DIR \
-    --volume $CFSH/$BASE_DIR/nersc/performance_indicators/linearly_separable:/linearly_separable \
-    --volume $CFSH/$WORK_DIR:$WORK_DIR \
-    -e HDF5_USE_FILE_LOCKING='FALSE' \
-    --workdir $WORK_DIR \
-    $IMG <<EOF 
-echo P:pwd; pwd
-echo P:all
-ls -l ../..
-$CMD
-echo P:done
-exit
-EOF
-  
-echo 'W:done podman'
-
-# spare
-# 	   --volume $HOME:/home \
-    # --volume  $CFSH/qml-benchmarks:/root 

From e678c9608a169beee6b5200c4048b5eb8ae953e5 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Mon, 6 May 2024 15:08:39 +0200
Subject: [PATCH 069/100] cleanup

---
 nersc/performance_indicators/submit_job1.slr | 46 --------------------
 nersc/performance_indicators/wrap_podman.sh  | 23 ----------
 2 files changed, 69 deletions(-)
 delete mode 100755 nersc/performance_indicators/submit_job1.slr
 delete mode 100755 nersc/performance_indicators/wrap_podman.sh

diff --git a/nersc/performance_indicators/submit_job1.slr b/nersc/performance_indicators/submit_job1.slr
deleted file mode 100755
index 0c0aaa50..00000000
--- a/nersc/performance_indicators/submit_job1.slr
+++ /dev/null
@@ -1,46 +0,0 @@
-#!/bin/bash
-#SBATCH -N 1
-#SBATCH -C cpu
-#SBATCH -q shared -t 01:00:00
-# SBATCH -q debug -t 5:00  # charged for full node
-#SBATCH -J perf_ind_iqp15
-# SBATCH --account=jbowles
-#SBATCH --ntasks-per-node=1
-#SBATCH --cpus-per-task=6
-#SBATCH --output out1/%j.log
-#SBATCH  --licenses=scratch
-
-
-echo 'S:start'
-
-# .... container
-# salloc -q interactive -C cpu -t 4:00:00 -A nstaff
-
-IMG=balewski/ubu22-pennylane:p4
-echo use image $IMG
-POD_PUB=/dvs_ro/cfs/cdirs/nstaff/balewski/podman_common/  
-export PODMANHPC_ADDITIONAL_STORES=$POD_PUB
-
-#.... output sandbox 
-CODE_DIR=`pwd`
-myJob=${SLURM_JOBID}
-outPath=$SCRATCH/tmp_pennylane_jobs/$myJob
-echo outPath=$outPath
-mkdir -p $outPath
-cp -rp *.py wrap_podman.sh   $outPath
-
-#... location of input on CSF
-CFSH=/global/cfs/cdirs/m4139
-BASE_DIR=qml-benchmarks   # here git has home
-WORK_DIR=/$BASE_DIR/nersc
-
-#... the task to be executed
-CMD=" python -u  perf_ind_variational_jax.py "
-
-cd $outPath
-G=1
-echo 'S:ready to run '; pwd
-srun -n $G  ./wrap_podman.sh $IMG " $CMD "  $outPath  $CFSH $BASE_DIR $WORK_DIR
-
-echo 'S:end '; date
-
diff --git a/nersc/performance_indicators/wrap_podman.sh b/nersc/performance_indicators/wrap_podman.sh
deleted file mode 100755
index f601c873..00000000
--- a/nersc/performance_indicators/wrap_podman.sh
+++ /dev/null
@@ -1,23 +0,0 @@
-#!/bin/bash
-echo W:myRank is $SLURM_PROCID
-IMG=$1
-CMD=$2
-outPath=$3
-CFSH=$4
-BASE_DIR=$5
-WORK_DIR=$6
-
-if [ $SLURM_PROCID -eq 0 ] ; then 
-   echo W:IMG=$IMG 
-   echo W:CMD=$CMD
-   #echo Q:fire $
-fi
-echo 'W:start'
-time podman-hpc run -it \
-	   --volume $HOME:/home \
-	   --volume $CFSH/$BASE_DIR:/$BASE_DIR \
-	   --volume $CFSH/$WORK_DIR:$WORK_DIR \
-	   -e HDF5_USE_FILE_LOCKING='FALSE' \
-	   --workdir $WORK_DIR  \
-	   $IMG $CMD 
-echo 'W:done'

From 9cf55509faa804e5c0d743c76d62414796c57d21 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Mon, 6 May 2024 15:10:40 +0200
Subject: [PATCH 070/100] cleanup

---
 .../performance_indicators/perf_ind_variational_jax.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/nersc/performance_indicators/perf_ind_variational_jax.py b/nersc/performance_indicators/perf_ind_variational_jax.py
index f945e674..2057cac0 100644
--- a/nersc/performance_indicators/perf_ind_variational_jax.py
+++ b/nersc/performance_indicators/perf_ind_variational_jax.py
@@ -66,8 +66,14 @@ def get_parser():
     hyperparams = {**hp_settings[model_name], **{'use_jax':use_jax, 'vmap':vmap, 'max_steps':max_steps, 'jit': jit}}
     print(hyperparams)
 
-    X_train,y_train = read_data(f'linearly_separable/linearly_separable_{n_features}d_train.csv')
-    X_test,y_test = read_data(f'linearly_separable/linearly_separable_{n_features}d_train.csv')
+    assert os.path.exists(args.inputPath)
+    # inpF1=f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv'
+    inpF1 = os.path.join(args.inputPath, 'linearly_separable_%dd_train.csv' % (n_features))
+    inpF2 = inpF1.replace('train', 'test')
+    print('M:inpF1', inpF1)
+    X_train, y_train = read_data(inpF1)
+    print('M:inpF2', inpF2)
+    X_test, y_test = read_data(inpF2)
 
     first_train_steps = []
     av_consec_train_steps = []

From 49b0f8021ebfcbb0e94c5020fccfe84dfe79761c Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Mon, 6 May 2024 15:20:36 +0200
Subject: [PATCH 071/100] cleanup

---
 nersc/performance_indicators/perf_ind_variational_jax.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/nersc/performance_indicators/perf_ind_variational_jax.py b/nersc/performance_indicators/perf_ind_variational_jax.py
index 2057cac0..da2abb64 100644
--- a/nersc/performance_indicators/perf_ind_variational_jax.py
+++ b/nersc/performance_indicators/perf_ind_variational_jax.py
@@ -51,7 +51,7 @@ def get_parser():
 
     max_steps = 100 #the number of gradient descent steps to use to estimate the step time
     perf_ind_name = 'JAX'  #a name for the performance indicator used for naming files
-    n_features = 2 #dataset dimension
+    n_features = args.numFeatures #dataset dimension
     n_trials = 5 #number of trials to average over
 
     #################################

From b2f9bf0e0842e331cce3958df9e5f04a32ba38ed Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Mon, 6 May 2024 15:50:09 +0200
Subject: [PATCH 072/100] update

---
 nersc/performance_indicators/README.md | 55 +++++++++++---------------
 1 file changed, 24 insertions(+), 31 deletions(-)

diff --git a/nersc/performance_indicators/README.md b/nersc/performance_indicators/README.md
index 44a37cdc..bbecdd7a 100644
--- a/nersc/performance_indicators/README.md
+++ b/nersc/performance_indicators/README.md
@@ -1,63 +1,56 @@
 # Instructions for running performance indicators
 
 To get the data for a given model we need to run the python script `perf_ind_variational.py` (for variational models) or 
-`perf_ind_kernel.py` (for kernel models). This gets the indicators 'first train step time', 'consec. train step time',
- and 'prediction time'. 
-
-For variational models, there is an additional script `perf_ind_full_train.py` that attempts to run the model until 
-convergence in order to get the 'accuracy' and 'train steps for convergence' numbers. Since this can take a long time, it is a separate script. 
-
-The results are saved to a subdirectory specified at the
-start of the script (see e.g. JAX/). The script also saves the pip package list for
-future reference. 
+`perf_ind_kernel.py` (for kernel models). 
 
 Within the scripts there are a number of settings that can be changed at the start. This should make it easy to reuse 
 the same code when we have updated models. **Please make sure you chose a unique choice of `perf_ind_name` when gathering results for a new
 workflow** (i.e. the name of the workflow) since results will be overwritten if the same name is used twice.
 
-
 The performance indicators use the hyperparameter settings specified in `hyperparam_settings.yaml`. These are chosen to 
 be those which require the most compute and shouldn't be changed. 
 
-## determinining the number of CPUs
-To avoid wasting resources, you should first determine how many CPUs are required. This can be done by launching an 
-interative node:
+To run a performance indicator you need to edit a number of things in the file `nersc/submit_job.slr`:
+
+- Choose the number of features by editing `numFeatures=X`
 
-`salloc -q shared_interactive -C cpu -t 4:00:00 --cpus-per-task=XX`
 
-where you control the number of cpus-per-task. Start with a high number and
-monitor the CPU usage with top by SSHing to the compute node in another terminal.
+- Chose whether it is for a variational or kernel model by editing the line
 
-On the interactive node, from the /nersc directory run 
+  `CMD=" python3 -u  performance_indicators/XXX.py --numFeatures $numFeatures --inputPath performance_indicators/linearly_separable/ "`
 
-`source pm_podman.source`
+  where `XXX` is either `perf_ind_kernel.py` or `perf_ind_variational.py`
 
-to load the container. Then 
 
-`bash installs.sh`
+- Decide the maximum job time (format HH:MM:SS):
 
-To install the qml-benchmarks package (we also downgrade to a version of scipy that works with catalyst). Run the 
-performance indicator pythons script:
+  `#SBATCH -q shared -t 1:00:00`
 
-`python3 perf_ind_variational.py`
 
-Note the CPU usage in top, and repeat this process (lowering the CPUs each time) until you determine a reasonable
-number of CPUs to use, and add this to the `CPUs` row of the google sheet. 
+- Decide the number of CPUs to use (see below):
 
-## running a batch job
+   `#SBATCH --cpus-per-task=X`
 
-Once the resources have been decided, edit the file `sbatch submit_job.slr` accordingly and launch a slurm job for 
-the performance indicator by running 
+To launch a job run 
 
 `sbatch submit_job.slr`
 
-## recording memory usage
+This will add the job to the queue; when finished the results are stored in `performance_indicators/perf_ind_name`.
+Add the jobID to the google sheet for reference. 
+
+## determinining the number of CPUs
+To avoid wasting resources, you should first determine how many CPUs are required. To have an idea of 
+CPU usage, launch a job and ssh to the compute node, and run top to see the CPU usage (then kill the job). 
+Repeat the process until a reasonable number of CPUs is found (i.e. most are in use). 
+
+Add this choice to the `CPUs` row of the google sheet. 
+
+## recording memory and CPU usage
 Once the job has finished, run 
 
 `seff JOBID`
 
-where JOBID is the slurm job id. The value `Memory Utilized` is an estimate of the peak memory usage. 
-Add this to the sheet. 
+where JOBID is the slurm job id. Add the corresponding info to the google sheet. 
 
 
 

From a6b85f4d12f8f4f53f101e65ff4828343735aa7d Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Mon, 6 May 2024 15:51:39 +0200
Subject: [PATCH 073/100] update

---
 nersc/performance_indicators/README.md | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/nersc/performance_indicators/README.md b/nersc/performance_indicators/README.md
index bbecdd7a..6a11a199 100644
--- a/nersc/performance_indicators/README.md
+++ b/nersc/performance_indicators/README.md
@@ -19,7 +19,7 @@ To run a performance indicator you need to edit a number of things in the file `
 
   `CMD=" python3 -u  performance_indicators/XXX.py --numFeatures $numFeatures --inputPath performance_indicators/linearly_separable/ "`
 
-  where `XXX` is either `perf_ind_kernel.py` or `perf_ind_variational.py`
+  where `XXX` is either `perf_ind_kernel` or `perf_ind_variational`
 
 
 - Decide the maximum job time (format HH:MM:SS):
@@ -38,14 +38,14 @@ To launch a job run
 This will add the job to the queue; when finished the results are stored in `performance_indicators/perf_ind_name`.
 Add the jobID to the google sheet for reference. 
 
-## determinining the number of CPUs
+## Determinining the number of CPUs
 To avoid wasting resources, you should first determine how many CPUs are required. To have an idea of 
 CPU usage, launch a job and ssh to the compute node, and run top to see the CPU usage (then kill the job). 
 Repeat the process until a reasonable number of CPUs is found (i.e. most are in use). 
 
 Add this choice to the `CPUs` row of the google sheet. 
 
-## recording memory and CPU usage
+## Recording memory and CPU usage
 Once the job has finished, run 
 
 `seff JOBID`

From 5fe815096a2aa0f333449560971c39dfcd4b7e46 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Mon, 6 May 2024 15:52:19 +0200
Subject: [PATCH 074/100] cleanup

---
 nersc/performance_indicators/perf_ind_kernel_jax.py      | 8 ++++----
 nersc/performance_indicators/perf_ind_variational_jax.py | 9 ++++-----
 2 files changed, 8 insertions(+), 9 deletions(-)

diff --git a/nersc/performance_indicators/perf_ind_kernel_jax.py b/nersc/performance_indicators/perf_ind_kernel_jax.py
index 516e6d70..8cebe1db 100644
--- a/nersc/performance_indicators/perf_ind_kernel_jax.py
+++ b/nersc/performance_indicators/perf_ind_kernel_jax.py
@@ -33,8 +33,7 @@ def get_parser():
 if __name__=="__main__":
     args=get_parser()
 
-    #################################
-    # settings for the performance indicator.
+    ####### SETTINGS #######################
     # You only need to change this to make a different performance indicator
 
     #define the model
@@ -44,19 +43,20 @@ def get_parser():
     use_jax = True
     vmap = True
     jit = True
+    model_settings = {'use_jax': use_jax, 'vmap': vmap, 'jit': jit}
 
     perf_ind_name = 'JAX'  #a name for the performance indicator used for naming files
-    n_features = args.numFeatures  #dataset dimension
 
     #################################
 
+    n_features = args.numFeatures  # dataset dimension
     model_name = Model().__class__.__name__
 
     # get the 'worst case' hyperparameter settings for the model (those that require the most resources)
     with open('performance_indicators/hyperparam_settings.yaml', "r") as file:
         hp_settings = yaml.safe_load(file)
 
-    hyperparams = {**hp_settings[model_name], **{'use_jax':use_jax, 'vmap':vmap, 'jit': jit}}
+    hyperparams = {**hp_settings[model_name], **model_settings}
     print(hyperparams)
     assert os.path.exists(args.inputPath)
     #inpF1=f'../../paper/benchmarks/linearly_separable/linearly_separable_{n_features}d_train.csv'
diff --git a/nersc/performance_indicators/perf_ind_variational_jax.py b/nersc/performance_indicators/perf_ind_variational_jax.py
index da2abb64..11c33efc 100644
--- a/nersc/performance_indicators/perf_ind_variational_jax.py
+++ b/nersc/performance_indicators/perf_ind_variational_jax.py
@@ -37,8 +37,7 @@ def get_parser():
 if __name__ == "__main__":
     args = get_parser()
 
-    #################################
-    # settings for the performance indicator.
+    ####### SETTINGS #######################
     # You only need to change this to make a different performance indicator
 
     #define model
@@ -48,22 +47,22 @@ def get_parser():
     use_jax = True
     vmap = True
     jit = True
+    model_settings = {'use_jax': use_jax, 'vmap': vmap, 'jit': jit}
 
     max_steps = 100 #the number of gradient descent steps to use to estimate the step time
     perf_ind_name = 'JAX'  #a name for the performance indicator used for naming files
-    n_features = args.numFeatures #dataset dimension
     n_trials = 5 #number of trials to average over
 
     #################################
 
-
+    n_features = args.numFeatures  # dataset dimension
     model_name = Model().__class__.__name__
 
     # get the 'worst case' hyperparameter settings for the model (those that require the most resources)
     with open('performance_indicators/hyperparam_settings.yaml', "r") as file:
         hp_settings = yaml.safe_load(file)
 
-    hyperparams = {**hp_settings[model_name], **{'use_jax':use_jax, 'vmap':vmap, 'max_steps':max_steps, 'jit': jit}}
+    hyperparams = {**hp_settings[model_name], **model_settings}
     print(hyperparams)
 
     assert os.path.exists(args.inputPath)

From ada44af8f4cac3266790fb71e9148937fc81b36f Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Mon, 6 May 2024 15:53:21 +0200
Subject: [PATCH 075/100] rename

---
 .../{perf_ind_kernel_jax.py => perf_ind_kernel.py}                | 0
 .../{perf_ind_variational_jax.py => perf_ind_variational.py}      | 0
 2 files changed, 0 insertions(+), 0 deletions(-)
 rename nersc/performance_indicators/{perf_ind_kernel_jax.py => perf_ind_kernel.py} (100%)
 rename nersc/performance_indicators/{perf_ind_variational_jax.py => perf_ind_variational.py} (100%)

diff --git a/nersc/performance_indicators/perf_ind_kernel_jax.py b/nersc/performance_indicators/perf_ind_kernel.py
similarity index 100%
rename from nersc/performance_indicators/perf_ind_kernel_jax.py
rename to nersc/performance_indicators/perf_ind_kernel.py
diff --git a/nersc/performance_indicators/perf_ind_variational_jax.py b/nersc/performance_indicators/perf_ind_variational.py
similarity index 100%
rename from nersc/performance_indicators/perf_ind_variational_jax.py
rename to nersc/performance_indicators/perf_ind_variational.py

From d0a951e48dbd034758374d527ae757cd82abdc8b Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Mon, 6 May 2024 15:53:56 +0200
Subject: [PATCH 076/100] cleanup

---
 nersc/submit_job.slr | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/nersc/submit_job.slr b/nersc/submit_job.slr
index 9b7a1ad4..49a51610 100755
--- a/nersc/submit_job.slr
+++ b/nersc/submit_job.slr
@@ -42,7 +42,7 @@ BASE_DIR=/qml-benchmarks   # here git has home
 WORK_DIR=$BASE_DIR/nersc
 
 #... the task to be executed
-CMD=" python3 -u  performance_indicators/perf_ind_kernel_jax.py --numFeatures $numFeatures --inputPath performance_indicators/linearly_separable/ "
+CMD=" python3 -u  performance_indicators/perf_ind_kernel.py --numFeatures $numFeatures --inputPath performance_indicators/linearly_separable/ "
 #CMD=" python3 -u -c \"import pennylane \"; echo done1  "
 #CMD=" python3 -u -c \"import qml_benchmarks \";  echo done2 host  "
 

From adf4e338d811e74bcc2f5b8b4a013ea0b5d959ea Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Tue, 7 May 2024 10:43:38 +0200
Subject: [PATCH 077/100] update

---
 nersc/performance_indicators/README.md | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/nersc/performance_indicators/README.md b/nersc/performance_indicators/README.md
index 6a11a199..8c131e62 100644
--- a/nersc/performance_indicators/README.md
+++ b/nersc/performance_indicators/README.md
@@ -40,8 +40,8 @@ Add the jobID to the google sheet for reference.
 
 ## Determinining the number of CPUs
 To avoid wasting resources, you should first determine how many CPUs are required. To have an idea of 
-CPU usage, launch a job and ssh to the compute node, and run top to see the CPU usage (then kill the job). 
-Repeat the process until a reasonable number of CPUs is found (i.e. most are in use). 
+CPU usage, launch a job for a short amount of time, then kill it. Run `seff JOBID` and check the CPU usage.
+Repeat the process, decreasing the number of CPUs each time until a reasonable number of CPUs is found (i.e. most are in use). 
 
 Add this choice to the `CPUs` row of the google sheet. 
 

From 92fc797fd504f9ac72ff8cb82b21c21898f8d59b Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Wed, 29 May 2024 15:52:00 +0100
Subject: [PATCH 078/100] working with catalyst

---
 .../models/projected_quantum_kernel.py        | 73 +++++++++++--------
 1 file changed, 42 insertions(+), 31 deletions(-)

diff --git a/src/qml_benchmarks/models/projected_quantum_kernel.py b/src/qml_benchmarks/models/projected_quantum_kernel.py
index f01dca4f..4f22b1cd 100644
--- a/src/qml_benchmarks/models/projected_quantum_kernel.py
+++ b/src/qml_benchmarks/models/projected_quantum_kernel.py
@@ -175,7 +175,7 @@ def circuitZ(x):
                 return [qml.expval(qml.PauliZ(wires=i)) for i in range(self.n_qubits_)]
 
             @qjit(autograph=True)
-            def circuit_as_array(x):
+            def ccircuit_as_array(x):
                 xvals = jnp.array(circuitX(x))
                 yvals = jnp.array(circuitY(x))
                 zvals = jnp.array(circuitZ(x))
@@ -203,7 +203,7 @@ def circuit_as_array(x):
         elif "lightning" in self.dev_type and self.jit:
             # circuit_as_array = qjit(circuit_as_array)
             def batch_circuit_as_array(X):
-                return jnp.array([circuit_as_array(x) for x in X])
+                return jnp.array([ccircuit_as_array(x) for x in X])
             circuit_as_array = batch_circuit_as_array
 
         return circuit_as_array
@@ -223,23 +223,24 @@ def precompute_kernel(self, X1, X2):
         # get all of the Pauli expvals needed to constrcut the kernel
         self.circuit = self.construct_circuit()
 
-        valsX1 = np.array(self.circuit(X1))
-        valsX1 = np.reshape(valsX1, (dim1, 3, -1))
-        valsX2 = np.array(self.circuit(X2))
-        valsX2 = np.reshape(valsX2, (dim2, 3, -1))
+        if self.use_jax:
+            #not actually using JAX here but it is part of the JAX pipeline
 
-        valsX_X1 = valsX1[:, 0]
-        valsX_X2 = valsX2[:, 0]
-        valsY_X1 = valsX1[:, 1]
-        valsY_X2 = valsX2[:, 1]
-        valsZ_X1 = valsX1[:, 2]
-        valsZ_X2 = valsX2[:, 2]
+            valsX1 = np.array(self.circuit(X1))
+            valsX1 = np.reshape(valsX1, (dim1, 3, -1))
+            valsX2 = np.array(self.circuit(X2))
+            valsX2 = np.reshape(valsX2, (dim2, 3, -1))
 
-        all_vals_X1 = np.reshape(jnp.concatenate((valsX_X1, valsY_X1, valsZ_X1)), -1)
-        default_gamma = 1 / np.var(all_vals_X1) / self.n_features_
+            valsX_X1 = valsX1[:, 0]
+            valsX_X2 = valsX2[:, 0]
+            valsY_X1 = valsX1[:, 1]
+            valsY_X2 = valsX2[:, 1]
+            valsZ_X1 = valsX1[:, 2]
+            valsZ_X2 = valsX2[:, 2]
+
+            all_vals_X1 = np.reshape(np.concatenate((valsX_X1, valsY_X1, valsZ_X1)), -1)
+            default_gamma = 1 / np.var(all_vals_X1) / self.n_features_
 
-        if self.use_jax:
-            #no actually using JAX here but it is part of the JAX pipeline
             kernel_matrix = np.zeros([dim1, dim2])
             for i in range(dim1):
                 for j in range(dim2):
@@ -251,31 +252,41 @@ def precompute_kernel(self, X1, X2):
                         -default_gamma * self.gamma_factor * (sumX + sumY + sumZ)
                     )
 
-        else:
-            valsX_X1 = jnp.array(valsX_X1)
-            valsX_X2 = jnp.array(valsX_X2)
-            valsY_X1 = jnp.array(valsY_X1)
-            valsY_X2 = jnp.array(valsY_X2)
-            valsZ_X1 = jnp.array(valsZ_X1)
-            valsZ_X2 = jnp.array(valsZ_X2)
+        elif "lightning" in self.dev_type and self.jit:
+            valsX1 = jnp.array(self.circuit(X1))
+            valsX1 = jnp.reshape(valsX1, (dim1, 3, -1))
+            valsX2 = jnp.array(self.circuit(X2))
+            valsX2 = jnp.reshape(valsX2, (dim2, 3, -1))
+
+            valsX_X1 = valsX1[:, 0]
+            valsX_X2 = valsX2[:, 0]
+            valsY_X1 = valsX1[:, 1]
+            valsY_X2 = valsX2[:, 1]
+            valsZ_X1 = valsX1[:, 2]
+            valsZ_X2 = valsX2[:, 2]
+
+            all_vals_X1 = jnp.reshape(jnp.concatenate((valsX_X1, valsY_X1, valsZ_X1)), -1)
+            default_gamma = 1 / jnp.var(all_vals_X1) / self.n_features_
 
             @qjit(autograph=True)
-            def construct_kernel(valsX_X1, valsX_X2, valsY_X1, valsY_X2, valsZ_X1, valsZ_X2,
-                                 dim1, dim2, default_gamma):
+            def construct_kernel(valsX_X1, valsX_X2, valsY_X1, valsY_X2, valsZ_X1, valsZ_X2):
                 kernel_matrix = jnp.zeros([dim1, dim2])
                 for i in range(dim1):
                     for j in range(dim2):
-                        sumX = sum([(valsX_X1[i, q] - valsX_X2[j, q]) ** 2 for q in range(self.n_qubits_)])
-                        sumY = sum([(valsY_X1[i, q] - valsY_X2[j, q]) ** 2 for q in range(self.n_qubits_)])
-                        sumZ = sum([(valsZ_X1[i, q] - valsZ_X2[j, q]) ** 2 for q in range(self.n_qubits_)])
+                        kx = jnp.array([(valsX_X1[i, q] - valsX_X2[j, q]) ** 2 for q in range(self.n_qubits_)])
+                        ky = jnp.array([(valsY_X1[i, q] - valsY_X2[j, q]) ** 2 for q in range(self.n_qubits_)])
+                        kz = jnp.array([(valsZ_X1[i, q] - valsZ_X2[j, q]) ** 2 for q in range(self.n_qubits_)])
+
+                        sumX = jnp.sum(kx)
+                        sumY = jnp.sum(ky)
+                        sumZ = jnp.sum(kz)
 
-                        kernel_matrix = kernel_matrix.at[i,j].set(np.exp(
+                        kernel_matrix = kernel_matrix.at[i,j].set(jnp.exp(
                             -default_gamma * self.gamma_factor * (sumX + sumY + sumZ))
                         )
                 return kernel_matrix
 
-            kernel_matrix = construct_kernel(valsX_X1, valsX_X2, valsY_X1, valsY_X2, valsZ_X1, valsZ_X2,
-                                 dim1, dim2, default_gamma)
+            kernel_matrix = construct_kernel(valsX_X1, valsX_X2, valsY_X1, valsY_X2, valsZ_X1, valsZ_X2)
 
         return kernel_matrix
 

From 3108ca15b76e6cc37c3168dd37f0a46703b77889 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Thu, 30 May 2024 12:01:33 +0100
Subject: [PATCH 079/100] model cleanup

---
 src/qml_benchmarks/models/iqp_kernel.py       | 55 ++++++------
 src/qml_benchmarks/models/iqp_variational.py  | 15 ++--
 .../models/projected_quantum_kernel.py        | 88 ++++++++-----------
 3 files changed, 74 insertions(+), 84 deletions(-)

diff --git a/src/qml_benchmarks/models/iqp_kernel.py b/src/qml_benchmarks/models/iqp_kernel.py
index 8c8bf74f..e381bb72 100644
--- a/src/qml_benchmarks/models/iqp_kernel.py
+++ b/src/qml_benchmarks/models/iqp_kernel.py
@@ -27,17 +27,16 @@
 
 jax.config.update("jax_enable_x64", True)
 
-
 class IQPKernelClassifier(BaseEstimator, ClassifierMixin):
     def __init__(
         self,
         svm=SVC(kernel="precomputed", probability=True),
         repeats=2,
         C=1.0,
-        dev_type = 'default.qubit.jax',
-        use_jax=None,
-        vmap=None,
-        jit=None,
+        dev_type = None,
+        use_jax=False,
+        vmap=True,
+        jit=True,
         random_state=42,
         scaling=1.0,
         max_vmap=250,
@@ -83,10 +82,13 @@ def __init__(
         self.random_state = random_state
         self.rng = np.random.default_rng(random_state)
         # device-related attributes
-        self.dev_type = dev_type
-        self.use_jax = use_jax if use_jax is not None else self.dev_type == "default.qubit.jax"
-        self.vmap = vmap if vmap is not None else self.use_jax
-        self.jit = jit if jit is not None else self.use_jax
+        if dev_type is not None:
+            self.dev_type = dev_type
+        else:
+            self.dev_type = "default.qubit.jax" if use_jax else "lightning.qubit"
+        self.use_jax = use_jax
+        self.vmap = vmap
+        self.jit = jit
 
         # data-dependant attributes
         # which will be initialised by calling "fit"
@@ -96,8 +98,6 @@ def __init__(
         self.circuit = None
 
     def generate_key(self):
-        if self.use_jax:
-            return jax.random.PRNGKey(self.rng.integers(1000000))
         return self.rng.integers(1000000)
 
     def construct_circuit(self):
@@ -130,10 +130,11 @@ def circuit(x):
 
         circuit = wrapped_circuit
 
-        if self.use_jax and self.jit:
-            circuit = jax.jit(circuit)
-        elif "lightning" in self.dev_type and self.jit:
-            circuit = qjit(circuit)
+        if self.jit:
+            if self.use_jax:
+                circuit = jax.jit(circuit)
+            else:
+                circuit = qjit(circuit)
 
         self.circuit = circuit
 
@@ -153,13 +154,18 @@ def precompute_kernel(self, X1, X2):
 
         circuit = self.construct_circuit()
 
-        if self.use_jax and self.vmap:
+        if self.use_jax:
             # concatenate all pairs of vectors
             Z = np.array([np.concatenate((X1[i], X2[j])) for i in range(dim1) for j in range(dim2)])
             # if batched circuit is used
-            self.batched_circuit = chunk_vmapped_fn(
-                jax.vmap(circuit, 0), start=0, max_vmap=self.max_vmap
-            )
+            if self.vmap:
+                self.batched_circuit = chunk_vmapped_fn(
+                    jax.vmap(circuit, 0), start=0, max_vmap=self.max_vmap
+                )
+            else:
+                def batched_circuit(X):
+                    return jnp.vstack([circuit(x) for x in X])
+                self.batched_circuit = batched_circuit
             kernel_values = self.batched_circuit(Z)
             # reshape the values into the kernel matrix
             kernel_matrix = np.reshape(kernel_values, (dim1, dim2))
@@ -167,7 +173,6 @@ def precompute_kernel(self, X1, X2):
             # could also use catalyst.vmap like as above, although I think it does basically the same thing as this.
             X1 = jnp.array(X1)
             X2 = jnp.array(X2)
-            @qjit(autograph=True)
             def construct_kernel(X1,X2):
                 dim1 = len(X1)
                 dim2 = len(X2)
@@ -176,6 +181,8 @@ def construct_kernel(X1,X2):
                     for j, y in enumerate(X2):
                        kernel_matrix = kernel_matrix.at[i,j].set(circuit(jnp.concatenate((x, y))))
                 return  kernel_matrix
+            if self.jit:
+                construct_kernel = qjit(construct_kernel, autograph=True)
             kernel_matrix = construct_kernel(X1,X2)
 
         return kernel_matrix
@@ -207,13 +214,7 @@ def fit(self, X, y):
             y (np.ndarray): Labels of shape (n_samples,)
         """
 
-        if self.use_jax:
-            self.svm.random_state = int(
-                jax.random.randint(self.generate_key(), shape=(1,), minval=0, maxval=1000000)
-            )
-        else:
-            self.svm.random_state = self.generate_key()
-
+        self.svm.random_state = self.generate_key()
         self.initialize(X.shape[1], np.unique(y))
 
         self.scaler = MinMaxScaler(feature_range=(-np.pi / 2, np.pi / 2))
diff --git a/src/qml_benchmarks/models/iqp_variational.py b/src/qml_benchmarks/models/iqp_variational.py
index 015c3022..b566f70c 100644
--- a/src/qml_benchmarks/models/iqp_variational.py
+++ b/src/qml_benchmarks/models/iqp_variational.py
@@ -29,13 +29,13 @@ def __init__(
         batch_size=32,
         use_jax=False,
         jit=True,
-        vmap=False,
+        vmap=True,
         max_vmap=None,
         max_steps=10000,
         convergence_interval=200,
         random_state=42,
         scaling=1.0,
-        dev_type="default.qubit",
+        dev_type=None,
         qnode_kwargs={},
     ):
         r"""
@@ -77,7 +77,6 @@ def __init__(
         self.max_steps = max_steps
         self.convergence_interval = convergence_interval
         self.batch_size = batch_size
-        self.dev_type = dev_type
         self.qnode_kwargs = qnode_kwargs
         self.use_jax = use_jax
         self.vmap = vmap
@@ -86,10 +85,12 @@ def __init__(
         self.random_state = random_state
         self.rng = np.random.default_rng(random_state)
 
-        if max_vmap is None:
-            self.max_vmap = self.batch_size
+        if dev_type is not None:
+            self.dev_type = dev_type
         else:
-            self.max_vmap = max_vmap
+            self.dev_type = "default.qubit.jax" if use_jax else "lightning.qubit"
+
+        self.max_vmap = self.batch_size if max_vmap is None else max_vmap
 
         # data-dependant attributes
         # which will be initialised by calling "fit"
@@ -151,7 +152,6 @@ def circuit(params, x):
 
             if self.jit:
                 circuit = qjit(circuit)
-                # circuit(np.random.rand(self.n_layers, self.n_qubits_, 3), np.random.rand(self.n_qubits_))
 
             self.circuit = circuit 
 
@@ -246,7 +246,6 @@ def loss_fn(params, X, y):
                 return jnp.mean(probs)
 
             if self.jit:
-                from catalyst import qjit
                 loss_fn = qjit(loss_fn)
 
             self.params_ = train_with_catalyst(self, loss_fn, optimizer, X, y, self.generate_key)
diff --git a/src/qml_benchmarks/models/projected_quantum_kernel.py b/src/qml_benchmarks/models/projected_quantum_kernel.py
index 4f22b1cd..337d3dd1 100644
--- a/src/qml_benchmarks/models/projected_quantum_kernel.py
+++ b/src/qml_benchmarks/models/projected_quantum_kernel.py
@@ -27,7 +27,6 @@
 
 jax.config.update("jax_enable_x64", True)
 
-
 class ProjectedQuantumKernel(BaseEstimator, ClassifierMixin):
     def __init__(
         self,
@@ -39,10 +38,10 @@ def __init__(
         trotter_steps=5,
         use_jax=False,
         jit=True,
-        vmap = False,
+        vmap = True,
         max_vmap=None,
         scaling=1.0,
-        dev_type="default.qubit.jax",
+        dev_type=None,
         qnode_kwargs={},
         random_state=42,
     ):
@@ -99,16 +98,15 @@ def __init__(
         self.use_jax = use_jax
         self.vmap = vmap
         self.jit = jit
-        self.dev_type = dev_type
         self.qnode_kwargs = qnode_kwargs
         self.scaling = scaling
         self.random_state = random_state
         self.rng = np.random.default_rng(random_state)
-
-        if max_vmap is None:
-            self.max_vmap = 50
+        if dev_type is not None:
+            self.dev_type = dev_type
         else:
-            self.max_vmap = max_vmap
+            self.dev_type = "default.qubit.jax" if use_jax else "lightning.qubit"
+        self.max_vmap = 50 if max_vmap is None else max_vmap
 
         # data-dependant attributes
         # which will be initialised by calling "fit"
@@ -123,19 +121,12 @@ def generate_key(self):
         return jax.random.PRNGKey(self.rng.integers(1000000))
 
     def construct_circuit(self):
-        """
-        Constructs the circuit to get the expvals of a given qubit and Pauli operator
-        We will use JAX to parallelize over these circuits in precompute kernel.
-        Args:
-            P: a pennylane Pauli X,Y,Z operator on a given qubit
-        """
-        if self.embedding == "IQP":
 
+        if self.embedding == "IQP":
             def embedding(x):
                 qml.IQPEmbedding(x, wires=range(self.n_qubits_), n_repeats=2)
 
         elif self.embedding == "Hamiltonian":
-
             def embedding(x):
                 evol_time = self.t / self.trotter_steps * (self.n_qubits_ - 1)
                 for i in range(self.n_qubits_):
@@ -153,60 +144,60 @@ def embedding(x):
 
         dev = qml.device(self.dev_type, wires=self.n_qubits_)
 
-        if "lightning" in self.dev_type and self.jit:
+        if self.use_jax:
+            @qml.qnode(dev, **self.qnode_kwargs)
+            def circuit(x):
+                embedding(x)
+                return (
+                        [qml.expval(qml.PauliX(wires=i)) for i in range(self.n_qubits_)]
+                        + [qml.expval(qml.PauliY(wires=i)) for i in range(self.n_qubits_)]
+                        + [qml.expval(qml.PauliZ(wires=i)) for i in range(self.n_qubits_)]
+                )
+
+            def circuit_XYZ(x):
+                return jnp.array(circuit(x))
+
+        else:
             # currently only support for returning expvals of qubit-wise commuting observables,
             # so we split into three circuits
-            @qjit(autograph=True)
             @qml.qnode(dev, **self.qnode_kwargs)
             def circuitX(x):
                 embedding(x)
                 return [qml.expval(qml.PauliX(wires=i)) for i in range(self.n_qubits_)]
 
-            @qjit(autograph=True)
             @qml.qnode(dev, **self.qnode_kwargs)
             def circuitY(x):
                 embedding(x)
                 return [qml.expval(qml.PauliY(wires=i)) for i in range(self.n_qubits_)]
 
-            @qjit(autograph=True)
             @qml.qnode(dev, **self.qnode_kwargs)
             def circuitZ(x):
                 embedding(x)
                 return [qml.expval(qml.PauliZ(wires=i)) for i in range(self.n_qubits_)]
 
-            @qjit(autograph=True)
-            def ccircuit_as_array(x):
+            def circuit_XYZ(x):
                 xvals = jnp.array(circuitX(x))
                 yvals = jnp.array(circuitY(x))
                 zvals = jnp.array(circuitZ(x))
                 return jnp.concatenate((xvals, yvals, zvals))
 
+        if self.jit:
+            if self.use_jax:
+                circuit_XYZ = jax.jit(circuit_XYZ)
+            else:
+                circuit_XYZ = qjit(circuit_XYZ, autograph=True)
+
+        if self.vmap:
+            if self.use_jax:
+                batched_circuit = jax.vmap(circuit_XYZ, in_axes=(0))
+                batched_circuit = chunk_vmapped_fn(batched_circuit, 0, self.max_vmap)
+            else:
+                batched_circuit = catalyst.vmap(circuit_XYZ, in_axes=(0))
         else:
-            @qml.qnode(dev, **self.qnode_kwargs)
-            def circuit(x):
-                embedding(x)
-                return (
-                        [qml.expval(qml.PauliX(wires=i)) for i in range(self.n_qubits_)]
-                        + [qml.expval(qml.PauliY(wires=i)) for i in range(self.n_qubits_)]
-                        + [qml.expval(qml.PauliZ(wires=i)) for i in range(self.n_qubits_)]
-                )
-
-            def circuit_as_array(x):
-                return jnp.array(circuit(x))
-
+            def batched_circuit(X):
+                return jnp.vstack([circuit_XYZ(x) for x in X])
 
-        if self.use_jax and self.jit:
-            circuit_as_array = jax.jit(circuit_as_array)
-            circuit_as_array = jax.vmap(circuit_as_array, in_axes=(0))
-            circuit_as_array = chunk_vmapped_fn(circuit_as_array, 0, self.max_vmap)
-
-        elif "lightning" in self.dev_type and self.jit:
-            # circuit_as_array = qjit(circuit_as_array)
-            def batch_circuit_as_array(X):
-                return jnp.array([ccircuit_as_array(x) for x in X])
-            circuit_as_array = batch_circuit_as_array
-
-        return circuit_as_array
+        return batched_circuit
 
     def precompute_kernel(self, X1, X2):
         """
@@ -225,7 +216,6 @@ def precompute_kernel(self, X1, X2):
 
         if self.use_jax:
             #not actually using JAX here but it is part of the JAX pipeline
-
             valsX1 = np.array(self.circuit(X1))
             valsX1 = np.reshape(valsX1, (dim1, 3, -1))
             valsX2 = np.array(self.circuit(X2))
@@ -252,7 +242,7 @@ def precompute_kernel(self, X1, X2):
                         -default_gamma * self.gamma_factor * (sumX + sumY + sumZ)
                     )
 
-        elif "lightning" in self.dev_type and self.jit:
+        else:
             valsX1 = jnp.array(self.circuit(X1))
             valsX1 = jnp.reshape(valsX1, (dim1, 3, -1))
             valsX2 = jnp.array(self.circuit(X2))
@@ -268,7 +258,6 @@ def precompute_kernel(self, X1, X2):
             all_vals_X1 = jnp.reshape(jnp.concatenate((valsX_X1, valsY_X1, valsZ_X1)), -1)
             default_gamma = 1 / jnp.var(all_vals_X1) / self.n_features_
 
-            @qjit(autograph=True)
             def construct_kernel(valsX_X1, valsX_X2, valsY_X1, valsY_X2, valsZ_X1, valsZ_X2):
                 kernel_matrix = jnp.zeros([dim1, dim2])
                 for i in range(dim1):
@@ -286,6 +275,7 @@ def construct_kernel(valsX_X1, valsX_X2, valsY_X1, valsY_X2, valsZ_X1, valsZ_X2)
                         )
                 return kernel_matrix
 
+            construct_kernel = qjit(construct_kernel, autograph=True) if self.jit else construct_kernel
             kernel_matrix = construct_kernel(valsX_X1, valsX_X2, valsY_X1, valsY_X2, valsZ_X1, valsZ_X2)
 
         return kernel_matrix

From 04670ba2d1ddef74eeeb6d8d88cd80251d05f3c9 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Tue, 4 Jun 2024 18:51:44 +0200
Subject: [PATCH 080/100] catalyst support

---
 .../models/quantum_metric_learning.py         | 90 +++++++++++--------
 1 file changed, 55 insertions(+), 35 deletions(-)

diff --git a/src/qml_benchmarks/models/quantum_metric_learning.py b/src/qml_benchmarks/models/quantum_metric_learning.py
index 3fedbf3d..0ec5fca0 100644
--- a/src/qml_benchmarks/models/quantum_metric_learning.py
+++ b/src/qml_benchmarks/models/quantum_metric_learning.py
@@ -14,13 +14,15 @@
 
 import warnings
 import pennylane as qml
+import catalyst
+from catalyst import qjit
 import numpy as np
 import jax
 from jax import numpy as jnp
 import optax
 from sklearn.base import BaseEstimator, ClassifierMixin
 from sklearn.preprocessing import MinMaxScaler
-from qml_benchmarks.model_utils import chunk_vmapped_fn, train
+from qml_benchmarks.model_utils import chunk_vmapped_fn, train, train_with_catalyst
 
 jax.config.update("jax_enable_x64", True)
 
@@ -52,11 +54,11 @@ def __init__(
         batch_size=32,
         use_jax = True,
         vmap = True,
-        max_vmap=4,
+        max_vmap=None,
         jit=True,
         random_state=42,
         scaling=1.0,
-        dev_type="default.qubit.jax",
+        dev_type=None,
         qnode_kwargs={"interface": "jax-jit"},
     ):
         """
@@ -104,16 +106,17 @@ def __init__(
         self.jit = jit
         self.use_jax = use_jax
         self.vmap = vmap
-        self.dev_type = dev_type
         self.qnode_kwargs = qnode_kwargs
         self.scaling = scaling
         self.random_state = random_state
         self.rng = np.random.default_rng(random_state)
 
-        if max_vmap is None:
-            self.max_vmap = self.batch_size
+        if dev_type is not None:
+            self.dev_type = dev_type
         else:
-            self.max_vmap = max_vmap
+            self.dev_type = "default.qubit.jax" if use_jax else "lightning.qubit"
+
+        self.max_vmap = 4 if max_vmap is None else max_vmap
 
         # data-dependant attributes
         # which will be initialised by calling "fit"
@@ -121,7 +124,6 @@ def __init__(
         self.n_qubits_ = None
         self.n_features_ = None
         self.scaler = None  # data scaler will be fitted on training data
-        self.circuit = None
 
     def generate_key(self):
         return jax.random.PRNGKey(self.rng.integers(1000000))
@@ -130,29 +132,44 @@ def construct_model(self):
         dev = qml.device(self.dev_type, wires=self.n_qubits_)
         wires = range(self.n_qubits_)
 
-        @qml.qnode(dev, **self.qnode_kwargs)
-        def circuit(params, x1, x2):
-            qml.QAOAEmbedding(features=x1, weights=params["weights"], wires=wires)
-            qml.adjoint(qml.QAOAEmbedding)(features=x2, weights=params["weights"], wires=wires)
-            return qml.expval(qml.Projector(np.array([0] * self.n_qubits_), wires=wires))
+        def wrapped_circuit(params, x1, x2):
+            @qml.qnode(dev, **self.qnode_kwargs)
+            def circuit(params, x1, x2):
+                qml.QAOAEmbedding(features=x1, weights=params["weights"], wires=wires)
+                qml.adjoint(qml.QAOAEmbedding(features=x2, weights=params["weights"], wires=wires))
+                return qml.probs()
+            return circuit(params, x1, x2)[0]
 
-        self.circuit = circuit
+        circuit = wrapped_circuit
 
         if self.jit:
-            circuit = jax.jit(circuit)
-        overlaps = jax.vmap(circuit, in_axes=(None, 0, 0))
-        chunked_overlaps = chunk_vmapped_fn(overlaps, start=1, max_vmap=self.max_vmap)
+            if self.use_jax:
+                circuit = jax.jit(circuit)
+            else:
+                qjit(circuit, autograph=True)
+
+        # always vmapping for now
+        if self.use_jax:
+            batched_overlaps = jax.vmap(circuit, in_axes=(None, 0, 0))
+            chunked_overlaps = chunk_vmapped_fn(batched_overlaps, start=1, max_vmap=self.max_vmap)
+        else:
+            def batched_overlaps(params, X1, X2):
+                return jnp.array([circuit(params, elem[0], elem[1]) for elem in jnp.stack((X1, X2), axis=1)])
 
         def model(params, X1=None, X2=None):
-            res = overlaps(params, X1, X2)
-            return jnp.mean(res)
-
-        def chunked_model(params, X1=None, X2=None):
-            res = chunked_overlaps(params, X1, X2)
+            res = batched_overlaps(params, X1, X2)
             return jnp.mean(res)
 
         self.forward = model
-        self.chunked_forward = chunked_model
+
+        if self.use_jax:
+            def chunked_model(params, X1=None, X2=None):
+                res = chunked_overlaps(params, X1, X2)
+                return jnp.mean(res)
+
+            self.chunked_forward = chunked_model
+        else:
+            self.chunked_forward = self.forward
 
         return self.forward
 
@@ -214,18 +231,17 @@ def loss_fn(params, A=None, B=None):
             return 1 - d_hs
 
         if self.jit:
-            loss_fn = jax.jit(loss_fn)
+            loss_fn = jax.jit(loss_fn) if self.use_jax else qjit(loss_fn)
 
         optimizer = optax.adam
-        self.params_ = train(
-            self,
-            loss_fn,
-            optimizer,
-            A,
-            B,
-            self.generate_key,
-            convergence_interval=self.convergence_interval,
-        )
+
+        if self.use_jax:
+            self.params_ = train(self, loss_fn, optimizer, A, B, self.generate_key,
+                                 convergence_interval=self.convergence_interval)
+
+        else:
+            self.params_ = train_with_catalyst(self, loss_fn, optimizer, A, B, self.generate_key,
+                                               convergence_interval=self.convergence_interval)
 
         self.params_["examples_-1"] = A
         self.params_["examples_+1"] = B
@@ -275,8 +291,12 @@ def predict_proba(self, X):
             # create list [x, x, x, ...] to get overlaps with A_examples = [a1, a2, a3...] and B_examples
             x_tiled = jnp.tile(x, (self.n_examples_predict, 1))
 
-            pred_a = jnp.mean(self.chunked_forward(self.params_, A_examples, x_tiled))
-            pred_b = jnp.mean(self.chunked_forward(self.params_, B_examples, x_tiled))
+            if self.use_jax:
+                pred_a = jnp.mean(self.chunked_forward(self.params_, A_examples, x_tiled))
+                pred_b = jnp.mean(self.chunked_forward(self.params_, B_examples, x_tiled))
+            else:
+                pred_a = jnp.mean(self.forward(self.params_, A_examples, x_tiled))
+                pred_b = jnp.mean(self.forward(self.params_, B_examples, x_tiled))
 
             # normalise to [0,1]
             predictions.append([pred_a / (pred_a + pred_b), pred_b / (pred_a + pred_b)])

From fc72c005135ab961bb820d0fb14f34525ecc857d Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Tue, 4 Jun 2024 18:52:28 +0200
Subject: [PATCH 081/100] qjit update

---
 src/qml_benchmarks/model_utils.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/qml_benchmarks/model_utils.py b/src/qml_benchmarks/model_utils.py
index 93762bee..e751e5c6 100644
--- a/src/qml_benchmarks/model_utils.py
+++ b/src/qml_benchmarks/model_utils.py
@@ -184,7 +184,6 @@ def train_with_catalyst(model, loss_fn, optimizer, X, y, random_key_generator, c
     params = model.params_
     opt = optimizer(learning_rate=model.learning_rate)
 
-    @qjit
     def update(i, args):
         params, opt_state, X, y, loss_history, key = args
         X_batch, y_batch = get_batch(X, y, key, batch_size=model.batch_size)
@@ -197,12 +196,15 @@ def update(i, args):
         key, subkey = jax.random.split(key)
         return (params, opt_state, X, y, loss_history, subkey)
 
-    @qjit
+    update = qjit(update) if model.jit else update
+
     def optimize(params, X, y, steps, loss_history, opt_state, key):
         args = (params, opt_state, X, y, loss_history, key)
         (params, opt_state, _, _, loss_history, key) = catalyst.for_loop(0,steps,1)(update)(args)
         return params, loss_history, opt_state
 
+    optimize = qjit(optimize) if  model.jit else update
+
     def train_until_convergence(params, X, y):
         converged = False
         current_steps = 0

From 5016ec45fac81e1f93a9883ba4c916f481dbad5a Mon Sep 17 00:00:00 2001
From: josephbowles <bowles.physics@gmail.com>
Date: Wed, 5 Jun 2024 07:23:17 -0700
Subject: [PATCH 082/100] fix results dir

---
 .../performance_indicators/perf_ind_kernel.py  | 10 +++++-----
 .../perf_ind_variational.py                    | 18 ++++++++++++------
 2 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/nersc/performance_indicators/perf_ind_kernel.py b/nersc/performance_indicators/perf_ind_kernel.py
index 8cebe1db..e633af24 100644
--- a/nersc/performance_indicators/perf_ind_kernel.py
+++ b/nersc/performance_indicators/perf_ind_kernel.py
@@ -37,15 +37,15 @@ def get_parser():
     # You only need to change this to make a different performance indicator
 
     #define the model
-    from qml_benchmarks.models.iqp_kernel import IQPKernelClassifier as Model
+    from qml_benchmarks.models.projected_quantum_kernel import ProjectedQuantumKernel as Model
 
     #implementation attributes of model
-    use_jax = True
+    use_jax = False
     vmap = True
     jit = True
     model_settings = {'use_jax': use_jax, 'vmap': vmap, 'jit': jit}
 
-    perf_ind_name = 'JAX'  #a name for the performance indicator used for naming files
+    perf_ind_name = 'CAT_CPU'  #a name for the performance indicator used for naming files
 
     #################################
 
@@ -89,9 +89,9 @@ def get_parser():
 
     header = ['construct_kernel_time', 'training_time', 'predict_time', 'hyperparameters']
 
-    if not os.path.exists(perf_ind_name):
+    if not os.path.exists('performance_indicators/'+perf_ind_name):
         # Create the directory
-        os.mkdir(perf_ind_name)
+        os.mkdir('performance_indicators/'+perf_ind_name)
 
     #write perf indicator data
     with open('performance_indicators/'+perf_ind_name+'/'+filename, mode="w", newline="") as file:
diff --git a/nersc/performance_indicators/perf_ind_variational.py b/nersc/performance_indicators/perf_ind_variational.py
index 11c33efc..ba51c6b3 100644
--- a/nersc/performance_indicators/perf_ind_variational.py
+++ b/nersc/performance_indicators/perf_ind_variational.py
@@ -41,17 +41,19 @@ def get_parser():
     # You only need to change this to make a different performance indicator
 
     #define model
-    from qml_benchmarks.models.quantum_metric_learning import QuantumMetricLearner as Model
-
+    #
+    #from qml_benchmarks.models.quantum_metric_learning import QuantumMetricLearner as Model
+    from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier as Model
     #implementation attributes of model
     use_jax = True
     vmap = True
     jit = True
     model_settings = {'use_jax': use_jax, 'vmap': vmap, 'jit': jit}
 
-    max_steps = 100 #the number of gradient descent steps to use to estimate the step time
+    max_steps = 10 #the number of gradient descent steps to use to estimate the step time
     perf_ind_name = 'JAX'  #a name for the performance indicator used for naming files
-    n_trials = 5 #number of trials to average over
+    n_trials = 1 #number of trials to average over
+    n_test = 10 #number of test set points. For full test set use n_test = -1
 
     #################################
 
@@ -74,6 +76,10 @@ def get_parser():
     print('M:inpF2', inpF2)
     X_test, y_test = read_data(inpF2)
 
+    if n_test != -1:
+        X_test = X_test[:n_test]
+        y_test = y_test[:n_test]
+
     first_train_steps = []
     av_consec_train_steps = []
     predict_times = []
@@ -122,9 +128,9 @@ def get_parser():
     header = ['first_train_step', 'first_train_step_std', 'consec_train_step', 'consec_train_step_std', 'predict_time',
               'predict_time_std', 'hyperparameters']
 
-    if not os.path.exists(perf_ind_name):
+    if not os.path.exists(i'performance_indicators/'+perf_ind_name):
         # Create the directory
-        os.mkdir(perf_ind_name)
+        os.mkdir('performance_indicators/'+perf_ind_name)
 
     #write perf indicator data
     with open('performance_indicators/'+perf_ind_name+'/'+filename, mode="w", newline="") as file:

From 4b333b1bb1f93bb8639fc0461a4e6f12ad50d912 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Mon, 17 Jun 2024 12:41:23 +0200
Subject: [PATCH 083/100] profile time

---
 nersc/performance_indicators/profiling.py | 100 ++++++++++++++++++++++
 1 file changed, 100 insertions(+)
 create mode 100644 nersc/performance_indicators/profiling.py

diff --git a/nersc/performance_indicators/profiling.py b/nersc/performance_indicators/profiling.py
new file mode 100644
index 00000000..1f5d596b
--- /dev/null
+++ b/nersc/performance_indicators/profiling.py
@@ -0,0 +1,100 @@
+import qml_benchmarks
+import pennylane as qml
+import jax
+import jax.numpy as jnp
+import numpy as np
+import time
+import csv
+import os
+import pickle
+import yaml
+import subprocess
+from qml_benchmarks.hyperparam_search_utils import read_data
+
+import argparse
+
+
+def get_parser():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-v", "--verbosity", type=int, choices=[0, 1, 2, 3, 4], help="increase output verbosity",
+                        default=1, dest='verb')
+    parser.add_argument("--inputPath", default='linearly_separable/', help='input data location')
+    parser.add_argument('-n', '--numFeatures', type=int, default=2, help="dataset dimension ")
+
+    args = parser.parse_args()
+
+    print('myArg-program:', parser.prog)
+    for arg in vars(args):  print('myArg:', arg, getattr(args, arg))
+
+    # assert os.path.exists(args.outPath)
+    return args
+
+
+# =================================
+# =================================
+#  M A I N
+# =================================
+# =================================
+if __name__ == "__main__":
+    args = get_parser()
+
+    ####### SETTINGS #######################
+    # You only need to change this to make a different performance indicator
+
+    #define model
+    #
+    #from qml_benchmarks.models.quantum_metric_learning import QuantumMetricLearner as Model
+    from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier as Model
+    #implementation attributes of model
+    use_jax = True
+    vmap = True
+    jit = True
+    model_settings = {'use_jax': use_jax, 'vmap': vmap, 'jit': jit}
+
+    max_steps = 2 #the number of gradient descent steps to use to estimate the step time
+    profile_name = 'jax'  #a name for the performance indicator used for naming files
+
+    #################################
+
+    n_features = args.numFeatures
+
+    model_name = Model().__class__.__name__
+
+    # get the 'worst case' hyperparameter settings for the model (those that require the most resources)
+    with open('performance_indicators/hyperparam_settings.yaml', "r") as file:
+        hp_settings = yaml.safe_load(file)
+
+    hyperparams = {**hp_settings[model_name], **model_settings}
+    print(hyperparams)
+
+    assert os.path.exists(args.inputPath)
+
+    first_step_times = []
+    second_step_times = []
+
+    inpF1 = os.path.join(args.inputPath, 'linearly_separable_%dd_train.csv' % (n_features))
+    print('M:inpF1', inpF1)
+    X_train, y_train = read_data(inpF1)
+
+    jax.clear_caches()
+    model = Model(**hyperparams)
+    model.fit(X_train, y_train)
+
+    #get step times from loss history data
+    step_times = np.array([model.loss_history_[1][i + 1] - model.loss_history_[1][i]
+                           for i in range(len(model.loss_history_[1]) - 1)])
+    step_times = np.insert(step_times, 0, [model.loss_history_[1][0]])
+
+    file_path = f'performance_indicators/profiling/step_times_{profile_name}_{model_name}.pkl'
+    if not os.path.exists(file_path):
+        data = {}
+        with open(file_path, 'wb') as file:
+            pickle.dump(data, file)
+
+    data[n_features] = [step_times[0], step_times[1]]
+
+    with open(file_path, 'rb') as file:
+        data = pickle.load(file)
+
+    print('M:done')
+

From 1899a0c07cd9ccaf7236d43ff432994cc602d80b Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Mon, 17 Jun 2024 12:48:49 +0200
Subject: [PATCH 084/100] .

---
 nersc/performance_indicators/profiling.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/nersc/performance_indicators/profiling.py b/nersc/performance_indicators/profiling.py
index 1f5d596b..52a45372 100644
--- a/nersc/performance_indicators/profiling.py
+++ b/nersc/performance_indicators/profiling.py
@@ -91,10 +91,13 @@ def get_parser():
         with open(file_path, 'wb') as file:
             pickle.dump(data, file)
 
-    data[n_features] = [step_times[0], step_times[1]]
-
     with open(file_path, 'rb') as file:
         data = pickle.load(file)
 
+    data[n_features] = [step_times[0], step_times[1]]
+
+    with open(file_path, 'wb') as file:
+        pickle.dump(data, file)
+
     print('M:done')
 

From 3a1baa8c03ba6b22b60244b48a1b30b552166483 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Mon, 17 Jun 2024 12:52:45 +0200
Subject: [PATCH 085/100] .

---
 nersc/performance_indicators/profiling.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/nersc/performance_indicators/profiling.py b/nersc/performance_indicators/profiling.py
index 52a45372..45336536 100644
--- a/nersc/performance_indicators/profiling.py
+++ b/nersc/performance_indicators/profiling.py
@@ -85,7 +85,12 @@ def get_parser():
                            for i in range(len(model.loss_history_[1]) - 1)])
     step_times = np.insert(step_times, 0, [model.loss_history_[1][0]])
 
-    file_path = f'performance_indicators/profiling/step_times_{profile_name}_{model_name}.pkl'
+    dir_path = f'performance_indicators/profiling'
+    file_path = f'{dir_path}/step_times_{profile_name}_{model_name}.pkl'
+
+    if not os.path.exists(dir_path):
+        os.mkdir(dir_path)
+
     if not os.path.exists(file_path):
         data = {}
         with open(file_path, 'wb') as file:

From 5e26427303e1a831cb80594c1026ef19c47ab47b Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Mon, 17 Jun 2024 14:39:39 +0200
Subject: [PATCH 086/100] profile

---
 .../profiling_variational.py                  | 108 ++++++++++++++++++
 nersc/submit_profiling_jobs.slr               |  52 +++++++++
 2 files changed, 160 insertions(+)
 create mode 100644 nersc/performance_indicators/profiling_variational.py
 create mode 100755 nersc/submit_profiling_jobs.slr

diff --git a/nersc/performance_indicators/profiling_variational.py b/nersc/performance_indicators/profiling_variational.py
new file mode 100644
index 00000000..45336536
--- /dev/null
+++ b/nersc/performance_indicators/profiling_variational.py
@@ -0,0 +1,108 @@
+import qml_benchmarks
+import pennylane as qml
+import jax
+import jax.numpy as jnp
+import numpy as np
+import time
+import csv
+import os
+import pickle
+import yaml
+import subprocess
+from qml_benchmarks.hyperparam_search_utils import read_data
+
+import argparse
+
+
+def get_parser():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-v", "--verbosity", type=int, choices=[0, 1, 2, 3, 4], help="increase output verbosity",
+                        default=1, dest='verb')
+    parser.add_argument("--inputPath", default='linearly_separable/', help='input data location')
+    parser.add_argument('-n', '--numFeatures', type=int, default=2, help="dataset dimension ")
+
+    args = parser.parse_args()
+
+    print('myArg-program:', parser.prog)
+    for arg in vars(args):  print('myArg:', arg, getattr(args, arg))
+
+    # assert os.path.exists(args.outPath)
+    return args
+
+
+# =================================
+# =================================
+#  M A I N
+# =================================
+# =================================
+if __name__ == "__main__":
+    args = get_parser()
+
+    ####### SETTINGS #######################
+    # You only need to change this to make a different performance indicator
+
+    #define model
+    #
+    #from qml_benchmarks.models.quantum_metric_learning import QuantumMetricLearner as Model
+    from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier as Model
+    #implementation attributes of model
+    use_jax = True
+    vmap = True
+    jit = True
+    model_settings = {'use_jax': use_jax, 'vmap': vmap, 'jit': jit}
+
+    max_steps = 2 #the number of gradient descent steps to use to estimate the step time
+    profile_name = 'jax'  #a name for the performance indicator used for naming files
+
+    #################################
+
+    n_features = args.numFeatures
+
+    model_name = Model().__class__.__name__
+
+    # get the 'worst case' hyperparameter settings for the model (those that require the most resources)
+    with open('performance_indicators/hyperparam_settings.yaml', "r") as file:
+        hp_settings = yaml.safe_load(file)
+
+    hyperparams = {**hp_settings[model_name], **model_settings}
+    print(hyperparams)
+
+    assert os.path.exists(args.inputPath)
+
+    first_step_times = []
+    second_step_times = []
+
+    inpF1 = os.path.join(args.inputPath, 'linearly_separable_%dd_train.csv' % (n_features))
+    print('M:inpF1', inpF1)
+    X_train, y_train = read_data(inpF1)
+
+    jax.clear_caches()
+    model = Model(**hyperparams)
+    model.fit(X_train, y_train)
+
+    #get step times from loss history data
+    step_times = np.array([model.loss_history_[1][i + 1] - model.loss_history_[1][i]
+                           for i in range(len(model.loss_history_[1]) - 1)])
+    step_times = np.insert(step_times, 0, [model.loss_history_[1][0]])
+
+    dir_path = f'performance_indicators/profiling'
+    file_path = f'{dir_path}/step_times_{profile_name}_{model_name}.pkl'
+
+    if not os.path.exists(dir_path):
+        os.mkdir(dir_path)
+
+    if not os.path.exists(file_path):
+        data = {}
+        with open(file_path, 'wb') as file:
+            pickle.dump(data, file)
+
+    with open(file_path, 'rb') as file:
+        data = pickle.load(file)
+
+    data[n_features] = [step_times[0], step_times[1]]
+
+    with open(file_path, 'wb') as file:
+        pickle.dump(data, file)
+
+    print('M:done')
+
diff --git a/nersc/submit_profiling_jobs.slr b/nersc/submit_profiling_jobs.slr
new file mode 100755
index 00000000..6d1fa41a
--- /dev/null
+++ b/nersc/submit_profiling_jobs.slr
@@ -0,0 +1,52 @@
+#!/bin/bash
+#SBATCH -N 1
+#SBATCH -C cpu
+#SBATCH -q shared -t 1:00:00
+# SBATCH -q debug -t 5:00  # charged for full node
+#SBATCH -J perf_ind_qkernel2
+#SBATCH --ntasks-per-node=1
+#SBATCH --cpus-per-task=6
+#SBATCH --output out1/%j.log
+#SBATCH  --licenses=scratch
+# - - - E N D    O F    SLURM    C O M M A N D S
+set -u ;  # exit  if you try to use an uninitialized variable
+
+echo 'S:start'
+
+numFeaturesMin=2
+numFeaturesMin=16
+
+#IMG=balewski/ubu22-pennylane:p5  # old
+IMG=jbowles/ubu22-pennylane:p5
+echo use image $IMG
+
+POD_PUB=//dvs_ro/cfs/cdirs/m4139/qml-benchmarks/nersc/podman/
+#POD_PUB=/dvs_ro/cfs/cdirs/nstaff/balewski/podman_common/  #old
+export PODMANHPC_ADDITIONAL_STORES=$POD_PUB
+
+#.... output sandbox 
+CODE_DIR=`pwd`
+myJob=${SLURM_JOBID}
+outPath=$SCRATCH/tmp_pennylane_jobs/$myJob
+echo outPath=$outPath
+mkdir -p $outPath
+cp -rp performance_indicators *.slr *.sh   $outPath
+
+#... location of input on CSF
+CFSH=/global/cfs/cdirs/m4139  # for Joseph
+# CFSH=/pscratch/sd/b/balewski/pennylane_wrk  # for Jan
+BASE_DIR=/qml-benchmarks   # here git has home
+WORK_DIR=$BASE_DIR/nersc
+
+cd $outPath
+
+#... the task to be executed
+for ((numFeatures=$numFeaturesMin; numFeatures<=$numFeaturesMax; numFeatures++)); do
+    CMD=" python3 -u  performance_indicators/profiling.py --numFeatures $numFeatures --inputPath performance_indicators/linearly_separable/ "
+    G=1
+    srun -n $G  ./wrap_podman.sh $IMG " $CMD "  $outPath  $CFSH $BASE_DIR $WORK_DIR
+sleep 1
+
+
+echo 'S:end '; date
+

From 628292af0b8006a3a45715cc31d70c1e6483a1be Mon Sep 17 00:00:00 2001
From: josephbowles <bowles.physics@gmail.com>
Date: Tue, 18 Jun 2024 06:57:43 -0700
Subject: [PATCH 087/100] results

---
 .../step_times_jax_IQPVariationalClassifier.pkl   | Bin 0 -> 585 bytes
 1 file changed, 0 insertions(+), 0 deletions(-)
 create mode 100644 nersc/performance_indicators/profiling/step_times_jax_IQPVariationalClassifier.pkl

diff --git a/nersc/performance_indicators/profiling/step_times_jax_IQPVariationalClassifier.pkl b/nersc/performance_indicators/profiling/step_times_jax_IQPVariationalClassifier.pkl
new file mode 100644
index 0000000000000000000000000000000000000000..cadc017b64452c3c13a61f74da73c544d7032956
GIT binary patch
literal 585
zcmZo*nQF(x00y;FG`v}2r)czu=9T6aRO%(?7p3aumgbaXCKeSXR!-?*D^5<#Ni3Q&
zc}fo}SP76$DXA<-1&J`FSxo8dXrB@UG=;f`(PoODpP$!%AOI8I3?);NoH>B@oS0*)
z>oBDas4jygg9|Lgz;Ha<9xRmV%?fk?Tp}Zh(F8;7)FPN#He|JTb}zC*S9>6}7N(XR
zS?!1XyLRYm8@M}RYB`YAZh7nOg09x1X8}wtC$idg9<grdY8RYY4^zv9tXAvX7Y}r`
z@7|n+spUpi+f>@)kFNII!ILnxJjiM}zp;d%sXZ`z$48i2USzcmtOe0%Y85_rFhhh=
Xz4?&U2DfL#p{Y$cSjqxbTdD^D<43)?

literal 0
HcmV?d00001


From c69895b4bec36f2a7955663817aa0d4aa4080754 Mon Sep 17 00:00:00 2001
From: josephbowles <bowles.physics@gmail.com>
Date: Wed, 19 Jun 2024 01:20:53 -0700
Subject: [PATCH 088/100] results

---
 .../step_times_jax_IQPVariationalClassifier.pkl | Bin 585 -> 717 bytes
 1 file changed, 0 insertions(+), 0 deletions(-)

diff --git a/nersc/performance_indicators/profiling/step_times_jax_IQPVariationalClassifier.pkl b/nersc/performance_indicators/profiling/step_times_jax_IQPVariationalClassifier.pkl
index cadc017b64452c3c13a61f74da73c544d7032956..23fbe9cb46f9fed992784119bf6e64f5ff43a064 100644
GIT binary patch
literal 717
zcmZo*nR<wc0Ss!VXn3>4PSNNQ%`43<sMJf&FG|(REzK#(Oe`u&ten!rR-BxelUOum
z@{}G{uo57hQc_uv3KC&TvzXG^(LN;zXbN)=qs<gQKR>VkKmaDZ8A_%kIdcH*(YSk7
z$6-nvP+bN~1{YX}q2WuWJy<B!n-%B)xP-&u)yC**9b$`MYT1z0R)}U>W2lv@fvIIj
zR@<=ci#xj7hR1y{wH(N5nSP`?qpL0GodZ+LiL6%P#bGyewHp>~fT`s|Ry$`#fiJpR
zEBg~LwcN;RKmR`)g09w?;VMim53<^QUyH)f)yB`h15?Y3tX9cZI|@zhf#eV0U~2i0
z)ux&JibGSY;IoPmB9!XQkEFIC`$2drn%W0X%jIEe1(4J-FugjJf~HoZ{JSDdtss)x
UhTr##)6mpTuoN<bsx8$60Qn{8QUCw|

literal 585
zcmZo*nQF(x00y;FG`v}2r)czu=9T6aRO%(?7p3aumgbaXCKeSXR!-?*D^5<#Ni3Q&
zc}fo}SP76$DXA<-1&J`FSxo8dXrB@UG=;f`(PoODpP$!%AOI8I3?);NoH>B@oS0*)
z>oBDas4jygg9|Lgz;Ha<9xRmV%?fk?Tp}Zh(F8;7)FPN#He|JTb}zC*S9>6}7N(XR
zS?!1XyLRYm8@M}RYB`YAZh7nOg09x1X8}wtC$idg9<grdY8RYY4^zv9tXAvX7Y}r`
z@7|n+spUpi+f>@)kFNII!ILnxJjiM}zp;d%sXZ`z$48i2USzcmtOe0%Y85_rFhhh=
Xz4?&U2DfL#p{Y$cSjqxbTdD^D<43)?


From e5dbfb6c7f39997a3f4ea255a81c84b4549ccf6d Mon Sep 17 00:00:00 2001
From: josephbowles <bowles.physics@gmail.com>
Date: Wed, 19 Jun 2024 02:31:25 -0700
Subject: [PATCH 089/100] results

---
 ...times_jax_no_vmap_IQPVariationalClassifier.pkl | Bin 0 -> 144 bytes
 1 file changed, 0 insertions(+), 0 deletions(-)
 create mode 100644 nersc/performance_indicators/profiling/step_times_jax_no_vmap_IQPVariationalClassifier.pkl

diff --git a/nersc/performance_indicators/profiling/step_times_jax_no_vmap_IQPVariationalClassifier.pkl b/nersc/performance_indicators/profiling/step_times_jax_no_vmap_IQPVariationalClassifier.pkl
new file mode 100644
index 0000000000000000000000000000000000000000..6fb45d93ac5bb3ba1f15ce368e7ddeb356d7fcbc
GIT binary patch
literal 144
zcmZo*ncB(#0ku=SSz@PX^oZt_<`z`yCFd8V>gAT^lw>9r6(v?q>0v8QPRvOxnlgDx
z4=Y#+kWMM7EJy{3Fr`^c>Fj8q5;R4_o4JS4W{RJmpVxmN02AH}B~y}|IT#oUg3pA^
bbC}WwRF}b$!37dxU^p<l%pNS1TC4{EVxl;!

literal 0
HcmV?d00001


From 0328ca86e0c11c96e71423475aafb8179ff4835d Mon Sep 17 00:00:00 2001
From: josephbowles <bowles.physics@gmail.com>
Date: Wed, 19 Jun 2024 04:13:25 -0700
Subject: [PATCH 090/100] results

---
 ...mes_jax_no_vmap_IQPVariationalClassifier.pkl | Bin 144 -> 805 bytes
 1 file changed, 0 insertions(+), 0 deletions(-)

diff --git a/nersc/performance_indicators/profiling/step_times_jax_no_vmap_IQPVariationalClassifier.pkl b/nersc/performance_indicators/profiling/step_times_jax_no_vmap_IQPVariationalClassifier.pkl
index 6fb45d93ac5bb3ba1f15ce368e7ddeb356d7fcbc..5b632b22c8e573e23160efc5e21d4f987b673595 100644
GIT binary patch
literal 805
zcmZo*nJUH100y;FG`v}2r)czu=9T6aRO%(?7p3aumgbaXCKeSXR!-?*D^5<#Ni3Q&
zc}fo}SP76$DXA<-1&J`FSxo8dXrB@UG=;f`(PoODpP$!%AOI8I3?);NoH>B@JmA@{
z>oBDas4jygg9|Lgz|hrU4;D)GW(7I`E}=6m!wg+*15Y1JEgQ00j!#E0)H>))gsEjm
zRvWWI#tz-w4RJGJYB`YAPWY(fg06PP+@&zJoXBd`=K8pzt9{_J1*Vn@S#8_iA3o@6
zFQ^=aspUpi8)q^l5MAvY#cMFNJjiOpCvFZ!Q`^v6`4py>7g?>$MX?xkwe@npU~2i0
z)%JI9PDE4tU^x#jL@3ppA6c!4)w?t_wH_*~)nRG{kkmGmWe239sns~TKoh1`5Ls<9
zdqobKTECk<?l84N$ZES^?#Mz@n`3yz3#L{WNiD;2$v;)7Y8!6yS7t!fmg)fj0A2tH

delta 55
zcmZ3=Hi411fn{nd0|eAgnaG!<SP*<BWS+y6wkbhVGFURWoH;;}2WFSqgN0Iy^#Hmg
B5kmj~


From 6bc2c734cdfaa26983ced75b9cb95b77e13d2c2f Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Wed, 19 Jun 2024 13:14:43 +0200
Subject: [PATCH 091/100] no vmap train

---
 src/qml_benchmarks/model_utils.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/src/qml_benchmarks/model_utils.py b/src/qml_benchmarks/model_utils.py
index e751e5c6..59473851 100644
--- a/src/qml_benchmarks/model_utils.py
+++ b/src/qml_benchmarks/model_utils.py
@@ -73,12 +73,14 @@ def train(model, loss_fn, optimizer, X, y, random_key_generator, convergence_int
 
     # note: assumes that the loss function is a sample mean of
     # some function over the input data set
-    chunked_grad_fn = chunk_grad(grad_fn, model.max_vmap)
-    chunked_loss_fn = chunk_loss(loss_fn, model.max_vmap)
+    if model.vmap:
+        grad_fn = chunk_grad(grad_fn, model.max_vmap)
+        loss_fn = chunk_loss(loss_fn, model.max_vmap)
+
 
     def update(params, opt_state, x, y):
-        grads = chunked_grad_fn(params, x, y)
-        loss_val = chunked_loss_fn(params, x, y)
+        grads = grad_fn(params, x, y)
+        loss_val = loss_fn(params, x, y)
         updates, opt_state = opt.update(grads, opt_state)
         params = optax.apply_updates(params, updates)
         return params, opt_state, loss_val

From 362a7fe0fc28166832537a326e7330ddeb61362b Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Wed, 19 Jun 2024 14:21:28 +0200
Subject: [PATCH 092/100] results

---
 nersc/performance_indicators/profiling/iqpvar_memory_usage.txt  | 2 ++
 .../profiling/iqpvar_memory_usage_no_vmap.txt                   | 2 ++
 2 files changed, 4 insertions(+)
 create mode 100644 nersc/performance_indicators/profiling/iqpvar_memory_usage.txt
 create mode 100644 nersc/performance_indicators/profiling/iqpvar_memory_usage_no_vmap.txt

diff --git a/nersc/performance_indicators/profiling/iqpvar_memory_usage.txt b/nersc/performance_indicators/profiling/iqpvar_memory_usage.txt
new file mode 100644
index 00000000..2afb576b
--- /dev/null
+++ b/nersc/performance_indicators/profiling/iqpvar_memory_usage.txt
@@ -0,0 +1,2 @@
+#from 4 to ..
+297144, 303920, 334832, 696136, 691800, 781188, 897556, 1003588, 1095188, 1208472, 1326364, 1464716, 1646004, 33495836
\ No newline at end of file
diff --git a/nersc/performance_indicators/profiling/iqpvar_memory_usage_no_vmap.txt b/nersc/performance_indicators/profiling/iqpvar_memory_usage_no_vmap.txt
new file mode 100644
index 00000000..c76cb21f
--- /dev/null
+++ b/nersc/performance_indicators/profiling/iqpvar_memory_usage_no_vmap.txt
@@ -0,0 +1,2 @@
+#from 4 to ..
+438708, 500256, 567260, 353384, 627016, 404372, 866712, 977068, 1081744, 1189500, 1312008, 1459828, 1638916, 2730588, 4112380, 6709988
\ No newline at end of file

From bea29fa665b53ad57679e90a4eeacc7459d27dff Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Wed, 19 Jun 2024 14:22:06 +0200
Subject: [PATCH 093/100] lax batching

---
 src/qml_benchmarks/models/iqp_variational.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/src/qml_benchmarks/models/iqp_variational.py b/src/qml_benchmarks/models/iqp_variational.py
index b566f70c..d96a32a0 100644
--- a/src/qml_benchmarks/models/iqp_variational.py
+++ b/src/qml_benchmarks/models/iqp_variational.py
@@ -130,9 +130,13 @@ def circuit(params, x):
                 self.chunked_forward = chunk_vmapped_fn(self.forward, 1, self.max_vmap)
 
             else:
-                # use jax but do not batch feed the circuit
+                def apply_circuit(params, x):
+                    result = circuit(params, x)
+                    return params, result
+
                 def forward(params, X):
-                    return jnp.stack([circuit(params, x) for x in X])
+                    params, results = jax.lax.scan(apply_circuit, params, X)
+                    return results
 
                 self.forward = forward
 
@@ -153,9 +157,8 @@ def circuit(params, x):
             if self.jit:
                 circuit = qjit(circuit)
 
-            self.circuit = circuit 
+            self.circuit = circuit
 
-            # use autograd and do not batch feed the circuit
             def forward(params, X):
                 return jnp.array([circuit(params, x) for x in X])
 

From 10553dd5098688a3042ff61d620a7661488784f9 Mon Sep 17 00:00:00 2001
From: josephbowles <bowles.physics@gmail.com>
Date: Tue, 2 Jul 2024 05:21:00 -0700
Subject: [PATCH 094/100] update

---
 nersc/performance_indicators/profiling_variational.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/nersc/performance_indicators/profiling_variational.py b/nersc/performance_indicators/profiling_variational.py
index 45336536..71c2ccec 100644
--- a/nersc/performance_indicators/profiling_variational.py
+++ b/nersc/performance_indicators/profiling_variational.py
@@ -47,17 +47,20 @@ def get_parser():
     from qml_benchmarks.models.iqp_variational import IQPVariationalClassifier as Model
     #implementation attributes of model
     use_jax = True
-    vmap = True
+    vmap = False
     jit = True
     model_settings = {'use_jax': use_jax, 'vmap': vmap, 'jit': jit}
 
     max_steps = 2 #the number of gradient descent steps to use to estimate the step time
-    profile_name = 'jax'  #a name for the performance indicator used for naming files
+    profile_name = 'jax_no_vmap_lax'  #a name for the performance indicator used for naming files
+    
 
     #################################
 
     n_features = args.numFeatures
 
+    print('NUM FEATURES: ' + str(n_features))
+
     model_name = Model().__class__.__name__
 
     # get the 'worst case' hyperparameter settings for the model (those that require the most resources)
@@ -77,7 +80,7 @@ def get_parser():
     X_train, y_train = read_data(inpF1)
 
     jax.clear_caches()
-    model = Model(**hyperparams)
+    model = Model(**hyperparams, max_steps=max_steps)
     model.fit(X_train, y_train)
 
     #get step times from loss history data

From fd2efbafc8c920d9982307cebbbf4fc4dd4cc299 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Tue, 2 Jul 2024 15:26:45 +0200
Subject: [PATCH 095/100] kernel profiling

---
 .../profiling_kernel.py                       | 108 ++++++++++++++++++
 1 file changed, 108 insertions(+)
 create mode 100644 nersc/performance_indicators/profiling_kernel.py

diff --git a/nersc/performance_indicators/profiling_kernel.py b/nersc/performance_indicators/profiling_kernel.py
new file mode 100644
index 00000000..49986eb2
--- /dev/null
+++ b/nersc/performance_indicators/profiling_kernel.py
@@ -0,0 +1,108 @@
+import qml_benchmarks
+import pennylane as qml
+import jax
+import jax.numpy as jnp
+import numpy as np
+import time
+import csv
+import os
+import pickle
+import yaml
+import subprocess
+from qml_benchmarks.hyperparam_search_utils import read_data
+
+import argparse
+
+
+def get_parser():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-v", "--verbosity", type=int, choices=[0, 1, 2, 3, 4], help="increase output verbosity",
+                        default=1, dest='verb')
+    parser.add_argument("--inputPath", default='linearly_separable/', help='input data location')
+    parser.add_argument('-n', '--numFeatures', type=int, default=2, help="dataset dimension ")
+
+    args = parser.parse_args()
+
+    print('myArg-program:', parser.prog)
+    for arg in vars(args):  print('myArg:', arg, getattr(args, arg))
+
+    # assert os.path.exists(args.outPath)
+    return args
+
+
+# =================================
+# =================================
+#  M A I N
+# =================================
+# =================================
+if __name__ == "__main__":
+    args = get_parser()
+
+    ####### SETTINGS #######################
+    # You only need to change this to make a different performance indicator
+
+    #define model
+    #
+    from qml_benchmarks.models.iqp_kernel import IQPKernelClassifier as Model
+    #implementation attributes of model
+    use_jax = False
+    vmap = False
+    jit = True
+    model_settings = {'use_jax': use_jax, 'vmap': vmap, 'jit': jit}
+
+    profile_name = 'catalyst_qjit'  #a name for the performance indicator used for naming files
+
+    #################################
+
+    n_features = args.numFeatures
+
+    print('NUM FEATURES: ' + str(n_features))
+
+    model_name = Model().__class__.__name__
+
+    # get the 'worst case' hyperparameter settings for the model (those that require the most resources)
+    with open('performance_indicators/hyperparam_settings.yaml', "r") as file:
+        hp_settings = yaml.safe_load(file)
+
+    hyperparams = {**hp_settings[model_name], **model_settings}
+    print(hyperparams)
+
+    assert os.path.exists(args.inputPath)
+
+    av_circuit_times = [] #the average time after the first circuit has been compiled
+
+    inpF1 = os.path.join(args.inputPath, 'linearly_separable_%dd_train.csv' % (n_features))
+    print('M:inpF1', inpF1)
+    X_train, y_train = read_data(inpF1)
+
+    jax.clear_caches()
+    model = Model(**hyperparams)
+
+    if model_name=='ProjectedQuantumKernel':
+        first_circuit_time = model.circuit(X[0])
+        second_circuit_time = model.circuit(X[0])
+    elif model_name == 'IQPKernelClassifier':
+        first_circuit_time = model.circuit(jnp.concatenate((X[0], X[1])))
+        second_circuit_time = model.circuit(jnp.concatenate((X[0], X[1])))
+
+    dir_path = f'performance_indicators/profiling'
+    file_path = f'{dir_path}/step_times_{profile_name}_{model_name}.pkl'
+
+    if not os.path.exists(dir_path):
+        os.mkdir(dir_path)
+
+    if not os.path.exists(file_path):
+        data = {}
+        with open(file_path, 'wb') as file:
+            pickle.dump(data, file)
+
+    with open(file_path, 'rb') as file:
+        data = pickle.load(file)
+
+    data[n_features] = [first_circuit_time, second_circuit_time]
+
+    with open(file_path, 'wb') as file:
+        pickle.dump(data, file)
+
+    print('M:done')
+

From 3f01d33a121e91b34f74d5576abddb057037da9c Mon Sep 17 00:00:00 2001
From: josephbowles <bowles.physics@gmail.com>
Date: Wed, 3 Jul 2024 05:52:18 -0700
Subject: [PATCH 096/100] results

---
 ...p_times_catalyst_qjit_IQPKernelClassifier.pkl | Bin 0 -> 1187 bytes
 1 file changed, 0 insertions(+), 0 deletions(-)
 create mode 100644 nersc/performance_indicators/profiling/step_times_catalyst_qjit_IQPKernelClassifier.pkl

diff --git a/nersc/performance_indicators/profiling/step_times_catalyst_qjit_IQPKernelClassifier.pkl b/nersc/performance_indicators/profiling/step_times_catalyst_qjit_IQPKernelClassifier.pkl
new file mode 100644
index 0000000000000000000000000000000000000000..f82cbdb55190ea1a89a8098e9fd1b18c650e5879
GIT binary patch
literal 1187
zcmZo*nL2}o0Ss!VXn6C*PSNNQ%`43<sMJf&FG|(REzK#(Oe`u&ten!r6JL~?oS#=*
zQdF8;GG+3V9#*izDLw3YDG()-r+728PH|>Tn$kWc2xuImW(I2yYf4FFK`KZ+Q<}w;
z&JKtaa}T4<6hA*dum3;*CcGI+rX+PbbL84fsIC#zx1R#!X0T+igN)4Jgc~S;rX<yy
z6X<9x3VBiO%MiqB+m_`CJ9f>e!)2QS&^BZxsoq?a*j7>E``G-+OkB3<18qZAlIqP(
ziEXMXop#eM)Z?<v4rm*)l2mUVN^A>xqo=&}a2PJze1WzhD@pa{rNp+{V14cnrBS$S
zivik(tR&T&g%aD$&wpmF`h63ZZP`HEkd>r*vr=N)^tWx6X%ipfvaJee8?usAZ#GJ7
r(_}GQ8SJ<kmu;Ou+mMx{db3kvo4P;$>!Z5UaoIKlXdAMU)KWbF6LQkX

literal 0
HcmV?d00001


From 8231bd0437c46603bec4d0259f18b869a3e7aa18 Mon Sep 17 00:00:00 2001
From: josephbowles <bowles.physics@gmail.com>
Date: Wed, 3 Jul 2024 05:56:58 -0700
Subject: [PATCH 097/100] results

---
 ...imes_catalyst_qjit_ProjectedQuantumKernel.pkl | Bin 0 -> 7701 bytes
 1 file changed, 0 insertions(+), 0 deletions(-)
 create mode 100644 nersc/performance_indicators/profiling/step_times_catalyst_qjit_ProjectedQuantumKernel.pkl

diff --git a/nersc/performance_indicators/profiling/step_times_catalyst_qjit_ProjectedQuantumKernel.pkl b/nersc/performance_indicators/profiling/step_times_catalyst_qjit_ProjectedQuantumKernel.pkl
new file mode 100644
index 0000000000000000000000000000000000000000..c35b3f172a4b5505fb96f7f2f10b14f90210575e
GIT binary patch
literal 7701
zcmeHMc{EjR+doK|Qkg=bQZFgehoaO~DHLf&52fPNQJKjcPstQY4iX_65JE(#%REnq
zGnfvJdG5^+rSf{~`PTRR@wL9zde{5M`#s%zt^LQ|*S+@sP1k+zwLKSw@<<ZC4R^Av
z&f+s<*#Q2_=baQx>}*XH&fD6acP0l2nAj3ctt{>AZ5>SQO}<|YStt~?bg-~?Rxq`)
zB`R1rT(vj<qn_{2_2dw80Oz+ta=;Qxv;Us0L%2oOV7ox}S$Hbj!XfNUCUFIDn%O&B
z6aP7d{h}(_&)0{1eBnMv0Gm2lPfzdA{KEKl(jnNBFZgLBlwLfA)g@iv8cfg*UdzOn
z$71)Nx|a=QgtpV>g)I1DWc^kq<r(U^dVF%;-2<XG!ZpQe0+D+C48eIR4ad7bmrn>(
zW4RCG(d5Byym8I%Ax9V$6*nB;=rmM`^U|-~WBRIaGjAtrG<*Q0mggBugx<xd7?bq&
zige7j&T{40UXE;@gu~S<sn|+gHmuvgLa#2XwZ5(yP@jDNzvjI9kZez0pf`Zqk$C=!
z$v2380lb#yEr@0&cBJzQh~3E!_6t}^i)2YlNId@}hs5`1Y6Sn|z>>uLNaBCMBI+zz
zFxLMHPIN){1#kr^SB3b&sU7CnZ7v?EDTQZ>j5{LaIv~Wk=hm!G!&yhiF)5k*Fub%j
zl0j#J39Z}lGnBy4D3{^!b}Fis`O{njn^8P>+%!rq8PD$~<?V{4Baf9d=YVn_*4)z-
zYNO^tO2cQzs)i<13Nt^cT*ZQ*h-GWb=y{MMU0|VR(FbcL@*LwAjo{l2+#T0K`mjFA
zVjZ!v0Tv%@^}m$gi)&XMNMXOw0J%J=JAzz#FiNQSv&G|K_^$AuLI1_z5_LF#q_Mzv
zjqk3~O$#n2l*07d73EnGy`UQ$y;%5dKjyfGe=1(q57$!7s|>bQ<6=FZIkhDNz*K9A
z<uGMp@ZPrV@{KhZ_Sp52z*sY`H>-TDx`c)WeKYKBn2+m54sfOk(eOyP2`To@5UO)r
z(%|avgj>2Ale}>)@HCj8z|B>J%cGX+8y^_N9*(5m>qQNJr1JmN>_6xAi8@?A(qPqh
z4VqJ<S9%;ar@#|8#kr8D9guXTL~V;eEoAsqRNj(^$B|@D+23B2;-<@4Y$}1>aO^7A
z(|x(#Se}1O-pH^IHO}?e`=@mv_nlff=9NYi)@-(lDRf7(0I!WQiBxnL?a+BB*NPz*
zsn#1xSZHKVQn~P?9*b{nJ~0&01&Jw-j4eJif$Wsq?u4tAuv>A_;i1wTXuR*`lWpOP
z)l(B6h287W(QZKc)c_q$g8bYgji|_(Lp<?;RgDE-3yVg#r9xW1l(K<Q7BE>S8~sV0
zKbyb59t}hto*#K)?Ke;C<6<KanlczX#xseidbw9#wYME6Z+B`noo>Sm>8sq^4);J+
zxoTs)cP~hc`|i#WDTd`sF(A#c2D9g?ZicVv#U%H^zF~zfj1I1<Kh{zS+Lc-c`c)0M
zsYU%A`|lLoD^Ap@uI$0w@F9ax4;Gg4B$2fmnAow__Vz268tj~D?)C2L1V^XLk2blT
zDDsxJ|F&lW4t<THPkWU>t7PnP?S%#T&D}y{K57SCcqpXF+BSr1p1&|3uB8Cq4aYmA
z$}$ME{2I-ip`v^B=!K2*F)(HGg!w`KF(i4|aX%U^#>Zjw){3?D@IELgV0(NIsB6w1
zAK%aivN?u|d%NjBYolMx{RDGF9o`=~Z~b@Ytyys14cZFP$eAX{T{)WMw>S-^F3qBt
zUIawnR=bW#6nL4QY;EFB1yP?0t;0>hs8_@G$y=@o4~g@R%Slr(qHog0(42y1{)L2t
zFCq|++vYvp`5H8LoQusSSKz&gjk4OF*%)=DWU+V!9S^v?+3AtjfyA-VX+^&#47Btu
z;2uWwW@!d-kt`ra{uSqY=U1TUQEPgotQNQ~sO?SCD+Pl{%E>5i3W~Hx3SNz)qMe~<
zfnsVlW|g{0I4x(Q$b;BSMJX!Y-=FHTg_)0A5(xP*YH^tUa@n<`8gU>zF|&m?Ckwr|
zf0WZQsDo*Fo#ng06Gcl$?N0SI0Gr^h0ME20JW;@;>IyW0&8;>!J=NZyb>uHb{{=>g
zI($EJs?2w%sxCOSEijsTcqtvY*z(p~6>9*INe%CVVN@vit*XTIV-xtFAxF99FrY}Z
zy>Uvn5Kgu@#<1P$#!wM8O10=y%oMTO`wT~M=WJC_;CniHIh|x3?nnj>mDNgRq+~cL
z#_jR7HUmDkrMo^bZAEhx89wPq7RE%f`j*W;gRNn<W;XX|NGhyZ6!naWIrZ`RBBx8>
zqBZooTT{UtO{R^P)<b5LO~kxsH4MpdFeWnw@d%j~#M6-hGDnT17@IQSu1AP%zWf;w
z=Y6p(F25A?NeQP+n(83?xRvkL@?v<OG2<97(}ZhZYN>a}HX~`^;4D432ZtRtbB*bC
zW3FF4UtkyuCJGNf5hpdn=L|=wxuP!U7TO{$^Dz>Iu1NH&9_mHW=GcwiYiK{a&HBak
zzkz9@4)>3oE%|*%_;(Hd<idTq!Qh-|RyQf!2mQm1M_TggASq;aB19`6MIThDT}m#2
zORe*Jw^fv*&B5(PRIO^bHWgZEe7_&13=@NCkD727gKP2=X8~5-4BT?Wvl$H!Om4nw
zn+8PT(q}$9nYd2(63-{`ZZxAFb&?1yMP7L)h5I{dV4q|MWw$m14!SGx`w2dVEEUTM
zCSMD_wF)Ue97VyY&LwhZ4F}P;^^Ff>dlw{6P2Bu_unf7?IeP?}^Uz*|S5!Sb9S``j
zyY1DY;cus7Y}p%1FvWS!jogw8&#GN`dbZQyoK>P*aLLfmHq>AA{Lk0>iXVAs+xPC6
zw$L3-MwcmsIMYFP_U!2rG6kLgK5W<H6^<q1(xIB_?O<f-(lNe`34uyV%uOcsa6RSF
zN3P@yBsSX8dorUivt0hgt;iZ&AD5|Ez{A3>5UqAS77f|d_YLRIXChYaZBL82f`p=w
zJ$3r^z&}Z1yDiiYf)1e<V^V4GMG#~QqZoL_?eX*ekL}R%J0~f>l!~_*2lO5qHR6eM
z;iGo@8R%Vm^t3W}DZW^#J9^cn3g@#JZ=(mQKqqk1F)iHytS}LjmQr!YOAqR#9-0l|
zGK?mBN;kothp8fa3TkldW#QgJQY9QZHW<N>sRXWUZ?Ba1gQ#np=Pqpc8t?eaa0Hq6
z0+V)6N!Yd<Hs8>;Kb|lIWk;0+l@$v?a&=mycF6DWKK>D<;8{Hye4uC_nvaCa;^0p+
z2C=9yaVF)fa|4!*jqWscE&!D;ANtD*S~33B*+O4oDs0f14r7Y-!XCLI_Qi_^VEn>}
zW0LVW>YikGHIyGiZp}o>7Ueel=<J+$qpTN`>C(o=Sc$p3#y+<aszLSiaqCGw{MoMj
zOX2@V3KMk%ew1Xj?@6v$NHV#u%C(VD3l)~#wWnAt*sAhaJa&62<R$Sf%9<|6{0*HF
z-Dj)OL;K~OIk##sys<K3`SW`iqdYn<VHAyh9m}@_ZEeQWk?+nL71tnU1=g)Hya8I*
zJhjdU=HdExr)>tt((o~V)fLBeh2S)1Gr6R-80G>y78@pXLM@%^TGw177$|xkTH@aV
zpD&NUd#5}E`cH&hz1B1Eep0UNZTpAVn;EafAxpvN%+6#R%|Q$(l{DArFM{h^1YM%z
z$^g1pc~7mfP|@m+{wAq1B&DwZn?!XP`UMzuMNU~SbdbIS@i!9DMxRbQdY=sE90WLz
zeW74vuiCk-Z>phsoAC-|%M={2_1~!?L;=Sh4dyIo7tDvLg&M0>!)ji^a*NE@pwFGj
zxiMoD{an^e>+Ki@#=6zEgIq(HCK4?x#%h93X}mo}yv^wK9C%*xQb2TCWgphHL&F=@
zup{k7*lzZ)y-xicj^AONBuTcw*Udv#a>4DW_i}sphs!LOA1PCxR_aHIxjp-)lY22$
zbm^>*TO;0yw0u6ZiiNup$KB(Wgk#RDxn*@?<@h(52P+E3ntyho`=$I}KzX8$;E$58
z{yq6!3&|hhc8tDW+y#d$?IYiiAAzP}C_n4BVmRxtNwoR&OAr|uG~DY`k5=Vg8+WnO
zK~%gW!XVTaj@BAFJP)Mf+xJn0cIjQ{lX+Lze?<;1ic~#8mMQ^(FhK$F^{+trgt^a@
zYYh-uj|Mq4zd}h;E~nYDRA|r_6=G?#AYmp&+jQ;`l#5^UFPQ9r%U{WUrafJ7@Vdg(
zstqIHG^jxRqST3=%;$0X?x85F`MPe!qi$TYhH{{4tP0cf?ycrn{t_&niDoM;&%t*)
z?$WK^mg4ZKr-rH1_2A*ViZHmT6yw}Zl*z8lL6ZV`PHDL~%xR&O=cLx6692}whOlfX
zm2|!ycB&Z9&~kRFThoBxx8r0VPdlvWA*~nYXMn9|TUhWC8mNbE_{z9J1v9@Kuj>U*
zaK!U6A8n)s+g|MF?wTqB{*Wokm`^NL2wgDaSm=ZzU0;+rzceAapFZm<mIvJ%Pu&!e
z^~R>UbL|-$8i6=C_ezSr6wYnc*>ZWT9yGpuAjFb~&|gfbQ}RJ0DD5-g5!GnJ#G~X>
zm-Y|hrh8(9L&qv0^C_Q#x!H3(be2O<C;TCvy)_WA+$Rk?Yi}B}hqGWf!EXhZa>mcD
U`M(_TS9Aoa<=ZbvqJzRe08+^xVgLXD

literal 0
HcmV?d00001


From 670ecec6d683662fde5f9d0332f080576e5a6fb8 Mon Sep 17 00:00:00 2001
From: josephbowles <bowles.physics@gmail.com>
Date: Wed, 3 Jul 2024 06:16:09 -0700
Subject: [PATCH 098/100] results

---
 ...s_catalyst_qjit_ProjectedQuantumKernel.pkl | Bin 7701 -> 2119 bytes
 1 file changed, 0 insertions(+), 0 deletions(-)

diff --git a/nersc/performance_indicators/profiling/step_times_catalyst_qjit_ProjectedQuantumKernel.pkl b/nersc/performance_indicators/profiling/step_times_catalyst_qjit_ProjectedQuantumKernel.pkl
index c35b3f172a4b5505fb96f7f2f10b14f90210575e..d707a5a08f6e749587c14de0c098d95afe079f79 100644
GIT binary patch
delta 606
zcmbPgb6miofn}-<2Ll+?PSNmYi=Cq3ZvV32DoDiLK0Y2Slj_Y57XN(YjRRObtQsoL
zI#HmKRj6%B2G_(sgULB8(<kRLDoe6yWN>Eicr&z4$>3`Tn#u@-AVs2+vsh+nOLu0-
z)J`canG!T5)tduoa0WK@(qQ%PSRbh(smHI<8A)ZTHy7B+uMcH5fSp`Z3Ux9!E*n)q
zHf}!0t;<NzN`%sS6gzpqcD`|~0@>wmKZSV$h>_~e331EQk_NDN)-8}QP+Sl!{=37;
z0myZ?e<Ap!0mw-876OSou)c6o0CL^!e{8il02UVpiwj(qD*%Z*usxey50PL28}w}T
z5ui9wgNGA{1{%Z*7Jqa05s2$<zi^=eM4S&S{*IBO0W7}l3=2e@A1wZXIm7`hezd|4
oA}#<H|8zDJB7T)^2}E2(0wm!ew*A2ekOl`~wj~KbMrx@Z0GpDqlmGw#

literal 7701
zcmeHMc{EjR+doK|Qkg=bQZFgehoaO~DHLf&52fPNQJKjcPstQY4iX_65JE(#%REnq
zGnfvJdG5^+rSf{~`PTRR@wL9zde{5M`#s%zt^LQ|*S+@sP1k+zwLKSw@<<ZC4R^Av
z&f+s<*#Q2_=baQx>}*XH&fD6acP0l2nAj3ctt{>AZ5>SQO}<|YStt~?bg-~?Rxq`)
zB`R1rT(vj<qn_{2_2dw80Oz+ta=;Qxv;Us0L%2oOV7ox}S$Hbj!XfNUCUFIDn%O&B
z6aP7d{h}(_&)0{1eBnMv0Gm2lPfzdA{KEKl(jnNBFZgLBlwLfA)g@iv8cfg*UdzOn
z$71)Nx|a=QgtpV>g)I1DWc^kq<r(U^dVF%;-2<XG!ZpQe0+D+C48eIR4ad7bmrn>(
zW4RCG(d5Byym8I%Ax9V$6*nB;=rmM`^U|-~WBRIaGjAtrG<*Q0mggBugx<xd7?bq&
zige7j&T{40UXE;@gu~S<sn|+gHmuvgLa#2XwZ5(yP@jDNzvjI9kZez0pf`Zqk$C=!
z$v2380lb#yEr@0&cBJzQh~3E!_6t}^i)2YlNId@}hs5`1Y6Sn|z>>uLNaBCMBI+zz
zFxLMHPIN){1#kr^SB3b&sU7CnZ7v?EDTQZ>j5{LaIv~Wk=hm!G!&yhiF)5k*Fub%j
zl0j#J39Z}lGnBy4D3{^!b}Fis`O{njn^8P>+%!rq8PD$~<?V{4Baf9d=YVn_*4)z-
zYNO^tO2cQzs)i<13Nt^cT*ZQ*h-GWb=y{MMU0|VR(FbcL@*LwAjo{l2+#T0K`mjFA
zVjZ!v0Tv%@^}m$gi)&XMNMXOw0J%J=JAzz#FiNQSv&G|K_^$AuLI1_z5_LF#q_Mzv
zjqk3~O$#n2l*07d73EnGy`UQ$y;%5dKjyfGe=1(q57$!7s|>bQ<6=FZIkhDNz*K9A
z<uGMp@ZPrV@{KhZ_Sp52z*sY`H>-TDx`c)WeKYKBn2+m54sfOk(eOyP2`To@5UO)r
z(%|avgj>2Ale}>)@HCj8z|B>J%cGX+8y^_N9*(5m>qQNJr1JmN>_6xAi8@?A(qPqh
z4VqJ<S9%;ar@#|8#kr8D9guXTL~V;eEoAsqRNj(^$B|@D+23B2;-<@4Y$}1>aO^7A
z(|x(#Se}1O-pH^IHO}?e`=@mv_nlff=9NYi)@-(lDRf7(0I!WQiBxnL?a+BB*NPz*
zsn#1xSZHKVQn~P?9*b{nJ~0&01&Jw-j4eJif$Wsq?u4tAuv>A_;i1wTXuR*`lWpOP
z)l(B6h287W(QZKc)c_q$g8bYgji|_(Lp<?;RgDE-3yVg#r9xW1l(K<Q7BE>S8~sV0
zKbyb59t}hto*#K)?Ke;C<6<KanlczX#xseidbw9#wYME6Z+B`noo>Sm>8sq^4);J+
zxoTs)cP~hc`|i#WDTd`sF(A#c2D9g?ZicVv#U%H^zF~zfj1I1<Kh{zS+Lc-c`c)0M
zsYU%A`|lLoD^Ap@uI$0w@F9ax4;Gg4B$2fmnAow__Vz268tj~D?)C2L1V^XLk2blT
zDDsxJ|F&lW4t<THPkWU>t7PnP?S%#T&D}y{K57SCcqpXF+BSr1p1&|3uB8Cq4aYmA
z$}$ME{2I-ip`v^B=!K2*F)(HGg!w`KF(i4|aX%U^#>Zjw){3?D@IELgV0(NIsB6w1
zAK%aivN?u|d%NjBYolMx{RDGF9o`=~Z~b@Ytyys14cZFP$eAX{T{)WMw>S-^F3qBt
zUIawnR=bW#6nL4QY;EFB1yP?0t;0>hs8_@G$y=@o4~g@R%Slr(qHog0(42y1{)L2t
zFCq|++vYvp`5H8LoQusSSKz&gjk4OF*%)=DWU+V!9S^v?+3AtjfyA-VX+^&#47Btu
z;2uWwW@!d-kt`ra{uSqY=U1TUQEPgotQNQ~sO?SCD+Pl{%E>5i3W~Hx3SNz)qMe~<
zfnsVlW|g{0I4x(Q$b;BSMJX!Y-=FHTg_)0A5(xP*YH^tUa@n<`8gU>zF|&m?Ckwr|
zf0WZQsDo*Fo#ng06Gcl$?N0SI0Gr^h0ME20JW;@;>IyW0&8;>!J=NZyb>uHb{{=>g
zI($EJs?2w%sxCOSEijsTcqtvY*z(p~6>9*INe%CVVN@vit*XTIV-xtFAxF99FrY}Z
zy>Uvn5Kgu@#<1P$#!wM8O10=y%oMTO`wT~M=WJC_;CniHIh|x3?nnj>mDNgRq+~cL
z#_jR7HUmDkrMo^bZAEhx89wPq7RE%f`j*W;gRNn<W;XX|NGhyZ6!naWIrZ`RBBx8>
zqBZooTT{UtO{R^P)<b5LO~kxsH4MpdFeWnw@d%j~#M6-hGDnT17@IQSu1AP%zWf;w
z=Y6p(F25A?NeQP+n(83?xRvkL@?v<OG2<97(}ZhZYN>a}HX~`^;4D432ZtRtbB*bC
zW3FF4UtkyuCJGNf5hpdn=L|=wxuP!U7TO{$^Dz>Iu1NH&9_mHW=GcwiYiK{a&HBak
zzkz9@4)>3oE%|*%_;(Hd<idTq!Qh-|RyQf!2mQm1M_TggASq;aB19`6MIThDT}m#2
zORe*Jw^fv*&B5(PRIO^bHWgZEe7_&13=@NCkD727gKP2=X8~5-4BT?Wvl$H!Om4nw
zn+8PT(q}$9nYd2(63-{`ZZxAFb&?1yMP7L)h5I{dV4q|MWw$m14!SGx`w2dVEEUTM
zCSMD_wF)Ue97VyY&LwhZ4F}P;^^Ff>dlw{6P2Bu_unf7?IeP?}^Uz*|S5!Sb9S``j
zyY1DY;cus7Y}p%1FvWS!jogw8&#GN`dbZQyoK>P*aLLfmHq>AA{Lk0>iXVAs+xPC6
zw$L3-MwcmsIMYFP_U!2rG6kLgK5W<H6^<q1(xIB_?O<f-(lNe`34uyV%uOcsa6RSF
zN3P@yBsSX8dorUivt0hgt;iZ&AD5|Ez{A3>5UqAS77f|d_YLRIXChYaZBL82f`p=w
zJ$3r^z&}Z1yDiiYf)1e<V^V4GMG#~QqZoL_?eX*ekL}R%J0~f>l!~_*2lO5qHR6eM
z;iGo@8R%Vm^t3W}DZW^#J9^cn3g@#JZ=(mQKqqk1F)iHytS}LjmQr!YOAqR#9-0l|
zGK?mBN;kothp8fa3TkldW#QgJQY9QZHW<N>sRXWUZ?Ba1gQ#np=Pqpc8t?eaa0Hq6
z0+V)6N!Yd<Hs8>;Kb|lIWk;0+l@$v?a&=mycF6DWKK>D<;8{Hye4uC_nvaCa;^0p+
z2C=9yaVF)fa|4!*jqWscE&!D;ANtD*S~33B*+O4oDs0f14r7Y-!XCLI_Qi_^VEn>}
zW0LVW>YikGHIyGiZp}o>7Ueel=<J+$qpTN`>C(o=Sc$p3#y+<aszLSiaqCGw{MoMj
zOX2@V3KMk%ew1Xj?@6v$NHV#u%C(VD3l)~#wWnAt*sAhaJa&62<R$Sf%9<|6{0*HF
z-Dj)OL;K~OIk##sys<K3`SW`iqdYn<VHAyh9m}@_ZEeQWk?+nL71tnU1=g)Hya8I*
zJhjdU=HdExr)>tt((o~V)fLBeh2S)1Gr6R-80G>y78@pXLM@%^TGw177$|xkTH@aV
zpD&NUd#5}E`cH&hz1B1Eep0UNZTpAVn;EafAxpvN%+6#R%|Q$(l{DArFM{h^1YM%z
z$^g1pc~7mfP|@m+{wAq1B&DwZn?!XP`UMzuMNU~SbdbIS@i!9DMxRbQdY=sE90WLz
zeW74vuiCk-Z>phsoAC-|%M={2_1~!?L;=Sh4dyIo7tDvLg&M0>!)ji^a*NE@pwFGj
zxiMoD{an^e>+Ki@#=6zEgIq(HCK4?x#%h93X}mo}yv^wK9C%*xQb2TCWgphHL&F=@
zup{k7*lzZ)y-xicj^AONBuTcw*Udv#a>4DW_i}sphs!LOA1PCxR_aHIxjp-)lY22$
zbm^>*TO;0yw0u6ZiiNup$KB(Wgk#RDxn*@?<@h(52P+E3ntyho`=$I}KzX8$;E$58
z{yq6!3&|hhc8tDW+y#d$?IYiiAAzP}C_n4BVmRxtNwoR&OAr|uG~DY`k5=Vg8+WnO
zK~%gW!XVTaj@BAFJP)Mf+xJn0cIjQ{lX+Lze?<;1ic~#8mMQ^(FhK$F^{+trgt^a@
zYYh-uj|Mq4zd}h;E~nYDRA|r_6=G?#AYmp&+jQ;`l#5^UFPQ9r%U{WUrafJ7@Vdg(
zstqIHG^jxRqST3=%;$0X?x85F`MPe!qi$TYhH{{4tP0cf?ycrn{t_&niDoM;&%t*)
z?$WK^mg4ZKr-rH1_2A*ViZHmT6yw}Zl*z8lL6ZV`PHDL~%xR&O=cLx6692}whOlfX
zm2|!ycB&Z9&~kRFThoBxx8r0VPdlvWA*~nYXMn9|TUhWC8mNbE_{z9J1v9@Kuj>U*
zaK!U6A8n)s+g|MF?wTqB{*Wokm`^NL2wgDaSm=ZzU0;+rzceAapFZm<mIvJ%Pu&!e
z^~R>UbL|-$8i6=C_ezSr6wYnc*>ZWT9yGpuAjFb~&|gfbQ}RJ0DD5-g5!GnJ#G~X>
zm-Y|hrh8(9L&qv0^C_Q#x!H3(be2O<C;TCvy)_WA+$Rk?Yi}B}hqGWf!EXhZa>mcD
U`M(_TS9Aoa<=ZbvqJzRe08+^xVgLXD


From 9c535e57cf96cdea653d4eee2c94ecc86ae7ec53 Mon Sep 17 00:00:00 2001
From: josephbowles <bowles.physics@gmail.com>
Date: Wed, 3 Jul 2024 06:19:35 -0700
Subject: [PATCH 099/100] results

---
 ...s_catalyst_qjit_ProjectedQuantumKernel.pkl | Bin 2119 -> 1097 bytes
 1 file changed, 0 insertions(+), 0 deletions(-)

diff --git a/nersc/performance_indicators/profiling/step_times_catalyst_qjit_ProjectedQuantumKernel.pkl b/nersc/performance_indicators/profiling/step_times_catalyst_qjit_ProjectedQuantumKernel.pkl
index d707a5a08f6e749587c14de0c098d95afe079f79..b898cd4b462f7339957d9ba3a8a017ad9cd89c48 100644
GIT binary patch
delta 305
zcmX>uaFT<yfn};4%S2WU?U!+L8GyjuKKwa|W=QpBjh&+5ZvWz%K?9KMZXcWr5(bKM
zOmwJZm2R7o!98)G!Q_7C`;%K4l_l9VGPp8$y%}1kWbn67(eP#jLXaY{$-T^v>Sa1J
zWNW9CmP`qnlIqO`Hu&}2tswK<?F;uF05MX%xxwPEd=ndhTzC71x&uHOD9!^BU!w%#
zy4z1#k^o|)dUJxs->4UZxbF5jEFc;vE(jL?9sH>Q$OZb~v7{u3mFg`7l5k-99S4#4
qSsvjqc_zDx@}>VkX`o7fkT^&oFW7_+{xd-W?)Ho4MofOdo(=#GAXys#

literal 2119
zcmeH{drVVT9LHNrOA+Kz1ewg~rpu5G=#n|<L{5c4orP`+6Ea-N>uL)WSZS4~G>bgS
z>Vt7NWiVbtqRe5P<Ds(scy+=6C$t!p6-9ZDY8($yNw(d4%M$+Bk|q0Re|T<k&pDs-
zJ16J&J>UFP7N7aBX;jO?-Gr<UaW?^)N**AI0Exxa3MRCq_-u!XPw~AOOPon?jcr_1
zs;4+b9^)A$mq(@H8je_w#Uv%BDCEi*h1d*>h(M=Ar7StkGbTxndCHUt3h7e3&GLA>
z2-h&FKwM*)82i_5AuR{@r$^&?#4YK>$H~PdD>ckmMOrdOZdtM48&~J&;d_X6hKBBg
zi$tPbe-MrO326#ETJ3+rHAq|7PSc_J7agwF7muQ#!ZMcA59833aq?-q-8ec}C;chx
z^=@E^@)miP6UfNBS;2@gLg9{~tseb7@OgDcCFlMiY>Mr=<ZU?sEn^EShM*Z<nhs!U
zR}a9uC&iMA<5S?nIN;A7zm0N&{AaDJZlan(TbdQS4eU!-?cEnJ3HKPaqleqB1Fpac
z*)>*&u&8)3m5N$pofsv<V#O(vC?c~gT&W;(lCaz)RuUUh1c~i(37r17Nl9GQlINGM
zj|tg?QT{dwSK=LsBnIos)?^1u`_So3?qX5RFshC3;CXR+(YJY>T{-K%g6VTAw-3(i
z;MKT#`j(@2(e4Cxjo;;5xZ1qiBO-hZ{3Guv3beyub-dTz7~c<0frClqtyvJOIkNec
zvj$L_hJ|0c--4q3hUAw!Mj%2V*%E#FXK2rPEo7?r4m$fy<v!UXJ#w4N^gWZ%g?zad
z!Be^>)PL$oUZX4@y60w|IAvXfLn#v*8zydm_?SAYB*Fm9CM@L9NH?_n-r8n*vkqyS
z*S`@KQICuxq5TDt+sORF%f)ym17&GUxRNR=(``}}Ra7#SOGU-{pS-e(@`~REI*s=6
zfP+iPf*xI7W2#lNv}o?YV|Y$fjLJUY9fsOIbn(KuWO0@Ox#V>Ug7t+U>Y+c)b=SkL
zwbl>ZH}*mC*lc>Zv=3qnT50c`F9z5nzf!a95(<1PvZ4|1gimKSyX{mpLTP*lYi;KZ
z2uS~aTlSS<!0wyoxoSNeP0Vkx`V}B|B=8tpB17dK4b0!t8W1<THzxkaUc`>(?Wh&$
zP*_P{XlZUAIR8@Om{4ValyFrGx2_TDb(!l@?Tz64Ma6gA^#(Zgc3rxcu^GJ1(3;D6
zRdC^=-NE<#tB}*of|qquJ>+hE;w}ighUPtl_S=yPTy&<C_s6ayy5n|@N~?#E7Na4E
zqema*3}uSENB^zuo+<ht6vc!#l!HwH>7;W72xpui2bj>x?CddxisEs%NK9}CiZAB-
z@rkJbO*)LKhyf<#kett)^xHs81?X|TcMHWGDQ>TSA(`ZS)_94TTTthWf89Y~B7UbI
z*%RliqPVH|E;$vT%1W-8+fw`iovosHLyv`-+fn=>Ba}~ZU8=8{+f#g|S8C=1tZFk~
WO>(|-<DG|OJl~O39Z3u@rRSenjJYTP


From b8483819eca92d76d534c689089165edf4029905 Mon Sep 17 00:00:00 2001
From: Joseph Bowles <bowles.physics@gmail.com>
Date: Mon, 5 Aug 2024 16:59:01 +0200
Subject: [PATCH 100/100] update

---
 src/qml_benchmarks/model_utils.py            | 2 +-
 src/qml_benchmarks/models/iqp_variational.py | 3 ++-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/qml_benchmarks/model_utils.py b/src/qml_benchmarks/model_utils.py
index 59473851..d9aa1042 100644
--- a/src/qml_benchmarks/model_utils.py
+++ b/src/qml_benchmarks/model_utils.py
@@ -205,7 +205,7 @@ def optimize(params, X, y, steps, loss_history, opt_state, key):
         (params, opt_state, _, _, loss_history, key) = catalyst.for_loop(0,steps,1)(update)(args)
         return params, loss_history, opt_state
 
-    optimize = qjit(optimize) if  model.jit else update
+    optimize = qjit(optimize) if  model.jit else optimize
 
     def train_until_convergence(params, X, y):
         converged = False
diff --git a/src/qml_benchmarks/models/iqp_variational.py b/src/qml_benchmarks/models/iqp_variational.py
index d96a32a0..82bbd032 100644
--- a/src/qml_benchmarks/models/iqp_variational.py
+++ b/src/qml_benchmarks/models/iqp_variational.py
@@ -251,7 +251,8 @@ def loss_fn(params, X, y):
             if self.jit:
                 loss_fn = qjit(loss_fn)
 
-            self.params_ = train_with_catalyst(self, loss_fn, optimizer, X, y, self.generate_key)
+            self.params_ = train_with_catalyst(self, loss_fn, optimizer, X, y, self.generate_key,
+                                               convergence_interval=self.convergence_interval)
 
         return self