From 9b0628b32966697f8379bd0abd3999a2f137e8c6 Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Wed, 5 Jan 2022 14:21:49 -0700
Subject: [PATCH 01/38] Create spsa_minimizer.py

---
 .../python/optimizers/spsa_minimizer.py       | 261 ++++++++++++++++++
 1 file changed, 261 insertions(+)
 create mode 100644 tensorflow_quantum/python/optimizers/spsa_minimizer.py

diff --git a/tensorflow_quantum/python/optimizers/spsa_minimizer.py b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
new file mode 100644
index 000000000..212284f97
--- /dev/null
+++ b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
@@ -0,0 +1,261 @@
+# Copyright 2021 The TensorFlow Quantum Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+"""The SPSA minimization algorithm"""
+import collections
+import numpy as np
+import tensorflow as tf
+
+
+def prefer_static_shape(x):
+    """Return static shape of tensor `x` if available,
+
+    else `tf.shape(x)`.
+
+    Args:
+        x: `tf.Tensor` (already converted).
+    Returns:
+        Numpy array (if static shape is obtainable), else `tf.Tensor`.
+    """
+    return prefer_static_value(tf.shape(x))
+
+
+def prefer_static_value(x):
+    """Return static value of tensor `x` if available, else `x`.
+
+    Args:
+        x: `tf.Tensor` (already converted).
+    Returns:
+        Numpy array (if static value is obtainable), else `tf.Tensor`.
+    """
+    static_x = tf.get_static_value(x)
+    if static_x is not None:
+        return static_x
+    return x
+
+
+SPSAOptimizerResults = collections.namedtuple(
+    'SPSAOptimizerResults',
+    [
+        'converged',
+        # Scalar boolean tensor indicating whether the minimum
+        # was found within tolerance.
+        'num_iterations',
+        # The number of iterations of the SPSA update.
+        'num_objective_evaluations',
+        # The total number of objective
+        # evaluations performed.
+        'position',
+        # A tensor containing the last argument value found
+        # during the search. If the search converged, then
+        # this value is the argmin of the objective function.
+        # A tensor containing the value of the objective from
+        # previous iteration
+        'objective_value_previous_iteration',
+        # Save the evaluated value of the objective function
+        # from the previous iteration
+        'objective_value',
+        # A tensor containing the value of the objective
+        # function at the `position`. If the search
+        # converged, then this is the (local) minimum of
+        # the objective function.
+        'tolerance',
+        # Define the stop criteria. Iteration will stop when the
+        # objective value difference between two iterations is
+        # smaller than tolerance
+        'a',
+        # Specifies the learning rate
+        'alpha',
+        # Specifies scaling of the learning rate
+        'c', 
+        # Specifies the size of the perturbations
+        'gamma',
+        # Specifies scaling of the size of the perturbations
+        'blocking',
+        # If true, then the optimizer will only accept updates that improve the objective function. 
+        'allowed_increase'
+        # Specifies maximum allowable increase in objective function (only applies if blocking is true).
+    ])
+
+
+def _get_initial_state(initial_position, tolerance, expectation_value_function, a, alpha, c, gamma, blocking, allowed_increase):
+    """Create SPSAOptimizerResults with initial state of search."""
+    init_args = {
+        "converged": tf.Variable(False),
+        "num_iterations": tf.Variable(0),
+        "num_objective_evaluations": tf.Variable(0),
+        "position": tf.Variable(initial_position),
+        "objective_value": tf.Variable(0.),
+        "objective_value_previous_iteration": tf.Variable(0.),
+        "tolerance": tolerance,
+        "a": tf.Variable(a),
+        "alpha": tf.Variable(alpha),
+        "c": tf.Variable(c),
+        "gamma" : tf.Variable(gamma),
+        "blocking" : tf.Variable(blocking),
+        "allowed_increase" : tf.Variable(allowed_increase)
+    }
+    return SPSAOptimizerResults(**init_args)
+
+
+def minimize(expectation_value_function,
+             initial_position,
+             tolerance=1e-7,
+             max_iterations=60,
+             alpha=0.602,
+             a=1.0,
+             c=1.0,
+             gamma=0.101,
+             blocking=True,
+             allowed_increase=0.5,
+             name=None):
+    """Applies the SPSA algorithm.
+
+    The SPSA algorithm can be used to minimize a noisy function. See:
+
+    [SPSA website](https://www.jhuapl.edu/SPSA/)
+
+    Usage:
+
+    Here is an example of optimize a function which consists summation of
+    a few sinusoids.
+
+    >>> n = 10  # Number of sinusoids
+    >>> coefficient = tf.random.uniform(shape=[n])
+    >>> min_value = -tf.math.reduce_sum(tf.abs(coefficient))
+    >>> func = lambda x:tf.math.reduce_sum(tf.sin(x) * coefficient)
+    >>> # Optimize the function with SPSA, start with random parameters
+    >>> result = tfq.optimizers.SPSA_minimize(func, np.random.random(n))
+    >>> result.converged
+    tf.Tensor(True, shape=(), dtype=bool)
+    >>> result.objective_value
+    tf.Tensor(-4.7045116, shape=(), dtype=float32)
+
+    Args:
+        expectation_value_function:  A Python callable that accepts
+            a point as a real `tf.Tensor` and returns a `tf.Tensor`s
+            of real dtype containing the value of the function.
+            The function to be minimized. The input is of shape `[n]`,
+            where `n` is the size of the trainable parameters.
+            The return value is a real `tf.Tensor` Scalar (matching shape
+            `[1]`).  This must be a linear combination of quantum
+            measurement expectation value, otherwise this algorithm cannot
+            work.
+        initial_position: Real `tf.Tensor` of shape `[n]`. The starting
+            point, or points when using batching dimensions, of the search
+            procedure. At these points the function value and the gradient
+            norm should be finite.
+        tolerance: Scalar `tf.Tensor` of real dtype. Specifies the tolerance
+            for the procedure. If the supremum norm between two iteration
+            vector is below this number, the algorithm is stopped.
+        a: Scalar `tf.Tensor` of real dtype. Specifies the learning rate
+        alpha: Scalar `tf.Tensor` of real dtype. Specifies scaling of the
+            learning rate.
+        c: Scalar `tf.Tensor` of real dtype. Specifies the size of the 
+            perturbations.
+        gamma: Scalar `tf.Tensor` of real dtype. Specifies scaling of the
+            size of the perturbations.
+        blocking: Boolean. If true, then the optimizer will only accept
+            updates that improve the objective function. 
+        allowed_increase: Scalar `tf.Tensor` of real dtype. Specifies maximum 
+            allowable increase in objective function (only applies if blocking
+            is true).
+        name: (Optional) Python `str`. The name prefixed to the ops created
+            by this function. If not supplied, the default name 'minimize'
+            is used.
+
+    Returns:
+        optimizer_results: A SPSAOptimizerResults object contains the
+            result of the optimization process.
+    """
+
+    with tf.name_scope(name or 'minimize'):
+        initial_position = tf.convert_to_tensor(initial_position,
+                                                name='initial_position',
+                                                dtype='float32')
+        dtype = initial_position.dtype.base_dtype
+        tolerance = tf.convert_to_tensor(tolerance,
+                                         dtype=dtype,
+                                         name='grad_tolerance')
+        max_iterations = tf.convert_to_tensor(max_iterations,
+                                              name='max_iterations')
+
+        a_init = tf.convert_to_tensor(a, name='initial a', dtype='float32')
+        c_init = tf.convert_to_tensor(c, name='initial c', dtype='float32')
+    
+
+        def _spsa_once(state):
+            """Caclulate single SPSA gradient estimation
+
+            Args:
+                state: A SPSAOptimizerResults object stores the
+                       current state of the minimizer.
+
+            Returns:
+                states: A list which the first element is the new state
+            """
+            delta_shift = tf.cast(2 * tf.random.uniform(shape=state.position.shape, minval=0, maxval=2, dtype=tf.int32) - 1, tf.float32)
+
+            v_m, v_p =  expectation_value_function(state.position - state.c * delta_shift),\
+                 expectation_value_function(state.position + state.c * delta_shift)
+
+            gradient_estimate = (v_p - v_m) / (2 * state.c) * delta_shift
+            update = state.a * gradient_estimate
+
+            current_obj = expectation_value_function(state.position - update)
+            #state.position.assign(tf.math.floormod(state.position - update, 2* np.pi))
+            if state.num_objective_evaluations == 0 or state.objective_value_previous_iteration < current_obj + state.allowed_increase or not state.blocking:
+                state.position.assign(state.position - update)
+
+            state.num_objective_evaluations.assign_add(2)
+            state.objective_value_previous_iteration.assign(
+                state.objective_value)
+            state.objective_value.assign(current_obj)
+
+            return [state]
+
+
+        # The `state` here is a `SPSAOptimizerResults` tuple with
+        # values for the current state of the algorithm computation.
+        def _cond(state):
+            """Continue if iterations remain and stopping condition
+            is not met."""
+            return (state.num_iterations < max_iterations) \
+                   and (not state.converged)
+
+        def _body(state):
+            """Main optimization loop."""
+            
+            new_a = a_init / ((tf.cast(state.num_iterations + 1, tf.float32) + 0.01 * tf.cast(max_iterations, tf.float32))**state.alpha)
+            new_c = c_init / (tf.cast(state.num_iterations + 1, tf.float32)**state.gamma)
+
+            state.a.assign(new_a)
+            state.c.assign(new_c)
+
+            _spsa_once(state)
+            state.num_iterations.assign_add(1)
+            state.converged.assign(
+                tf.abs(state.objective_value -
+                       state.objective_value_previous_iteration) <
+                state.tolerance)
+            return [state]
+
+        initial_state = _get_initial_state(initial_position, tolerance, expectation_value_function, a, alpha, c, gamma, blocking, allowed_increase)
+
+        initial_state.objective_value.assign(expectation_value_function(initial_state.position))
+
+        return tf.while_loop(cond=_cond,
+                             body=_body,
+                             loop_vars=[initial_state],
+                             parallel_iterations=1)[0]

From 0486c45c79a419b8a07cf9889e97d41b43f1c10a Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Wed, 5 Jan 2022 14:22:33 -0700
Subject: [PATCH 02/38] Update spsa_minimizer.py

---
 tensorflow_quantum/python/optimizers/spsa_minimizer.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow_quantum/python/optimizers/spsa_minimizer.py b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
index 212284f97..56a9b4268 100644
--- a/tensorflow_quantum/python/optimizers/spsa_minimizer.py
+++ b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
@@ -191,8 +191,8 @@ def minimize(expectation_value_function,
         max_iterations = tf.convert_to_tensor(max_iterations,
                                               name='max_iterations')
 
-        a_init = tf.convert_to_tensor(a, name='initial a', dtype='float32')
-        c_init = tf.convert_to_tensor(c, name='initial c', dtype='float32')
+        a_init = tf.convert_to_tensor(a, name='initial_a', dtype='float32')
+        c_init = tf.convert_to_tensor(c, name='initial_c', dtype='float32')
     
 
         def _spsa_once(state):

From 7b74ccbe3f463f63eb9a0789ae0280a152b5d90a Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Wed, 5 Jan 2022 14:26:58 -0700
Subject: [PATCH 03/38] Update __init__.py

---
 tensorflow_quantum/python/optimizers/__init__.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensorflow_quantum/python/optimizers/__init__.py b/tensorflow_quantum/python/optimizers/__init__.py
index 21a9d08fc..aea47fa9d 100755
--- a/tensorflow_quantum/python/optimizers/__init__.py
+++ b/tensorflow_quantum/python/optimizers/__init__.py
@@ -17,3 +17,5 @@
 # Quantum circuit specific optimizers.
 from tensorflow_quantum.python.optimizers.rotosolve_minimizer import (
     minimize as rotosolve_minimize)
+from tensorflow_quantum.python.optimizers.spsa_minimizer import (
+    minimize as spsa_minimize)

From 833b8c8f3a3c70b6dad8025354577cc78e1d1c8a Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Wed, 5 Jan 2022 14:28:28 -0700
Subject: [PATCH 04/38] Create spsa_minimizer_test.py

---
 .../python/optimizers/spsa_minimizer_test.py  | 154 ++++++++++++++++++
 1 file changed, 154 insertions(+)
 create mode 100644 tensorflow_quantum/python/optimizers/spsa_minimizer_test.py

diff --git a/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py b/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py
new file mode 100644
index 000000000..5f1f16c37
--- /dev/null
+++ b/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py
@@ -0,0 +1,154 @@
+# Copyright 2020 The TensorFlow Quantum Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+"""Test module for tfq.python.optimizers.spsa_minimizer optimizer."""
+# Remove PYTHONPATH collisions for protobuf.
+# pylint: disable=wrong-import-position
+import sys
+NEW_PATH = [x for x in sys.path if 'com_google_protobuf' not in x]
+sys.path = NEW_PATH
+# pylint: enable=wrong-import-position
+
+from operator import mul
+from functools import reduce
+import numpy as np
+import tensorflow as tf
+from absl.testing import parameterized
+import cirq
+import sympy
+from tensorflow_quantum.python.layers.high_level import pqc
+from tensorflow_quantum.python import util
+from tensorflow_quantum.python.optimizers import spsa_minimizer
+
+def loss_function_with_model_parameters(model, loss, train_x, train_y):
+    """Create a new function that assign the model parameter to the model
+    and evaluate its value.
+
+    Args:
+        model : an instance of `tf.keras.Model` or its subclasses.
+        loss : a function with signature loss_value = loss(pred_y, true_y).
+        train_x : the input part of training data.
+        train_y : the output part of training data.
+
+    Returns:
+        A function that has a signature of:
+            loss_value = f(model_parameters).
+    """
+
+    # obtain the shapes of all trainable parameters in the model
+    shapes = tf.shape_n(model.trainable_variables)
+    count = 0
+    sizes = []
+
+    # Record the shape of each parameter
+    for shape in shapes:
+        n = reduce(mul, shape)
+        sizes.append(n)
+        count += n
+
+    # Function accept the parameter and evaluate model
+    @tf.function
+    def func(params):
+        """A function that can be used by tfq.optimizer.spsa_minimize.
+
+        Args:
+           params [in]: a 1D tf.Tensor.
+
+        Returns:
+            Loss function value
+        """
+
+        # update the parameters of the model
+        start = 0
+        for i, size in enumerate(sizes):
+            model.trainable_variables[i].assign(
+                tf.reshape(params[start:start + size], shape))
+            start += size
+
+        # evaluate the loss
+        loss_value = loss(model(train_x, training=True), train_y)
+        return loss_value
+
+    return func
+
+
+class SPSAMinimizerTest(tf.test.TestCase, parameterized.TestCase):
+    """Tests for the SPSA optimization algorithm."""
+
+    def test_nonlinear_function_optimization(self):
+        """Test to optimize a non-linear function.
+        """
+        func = lambda x: x[0]**2 + x[1]**2
+
+        result = spsa_minimizer(func, tf.random.uniform(shape=[2]))
+        print(func(result.position))
+        self.assertAlmostEqual(func(result.position).numpy(), 0, delta=1e-4)
+        self.assertTrue(result.converged)
+
+
+    def test_keras_model_optimization(self):
+        """Optimizate a PQC based keras model."""
+
+        x = np.asarray([
+            [0, 0],
+            [0, 1],
+            [1, 0],
+            [1, 1],
+        ], dtype=float)
+
+        y = np.asarray([[-1], [1], [1], [-1]], dtype=np.float32)
+
+        def convert_to_circuit(input_data):
+            """Encode into quantum datapoint."""
+            values = np.ndarray.flatten(input_data)
+            qubits = cirq.GridQubit.rect(1, 2)
+            circuit = cirq.Circuit()
+            for i, value in enumerate(values):
+                if value:
+                    circuit.append(cirq.X(qubits[i]))
+            return circuit
+
+        x_circ = util.convert_to_tensor([convert_to_circuit(x) for x in x])
+
+        # Create two qubits
+        q0, q1 = cirq.GridQubit.rect(1, 2)
+
+        # Create an anzatz on these qubits.
+        a, b = sympy.symbols('a b')  # parameters for the circuit
+        circuit = cirq.Circuit(
+            cirq.rx(a).on(q0),
+            cirq.ry(b).on(q1), cirq.CNOT(control=q0, target=q1))
+
+        # Build the Keras model.
+        model = tf.keras.Sequential([
+            # The input is the data-circuit, encoded as a tf.string
+            tf.keras.layers.Input(shape=(), dtype=tf.string),
+            # The PQC layer returns the expected value of the
+            # readout gate, range [-1,1].
+            pqc.PQC(circuit, cirq.Z(q1)),
+        ])
+
+        # Initial guess of the parameter from random number
+        result = spsa_minimizer(
+            loss_function_with_model_parameters(model, tf.keras.losses.Hinge(),
+                                                x_circ, y),
+            tf.random.uniform(shape=[2]) * 2 * np.pi)
+
+        self.assertAlmostEqual(result.objective_value.numpy(), 0, delta=1e-4)
+        self.assertTrue(result.converged)
+
+
+
+if __name__ == "__main__":
+    tf.test.main()

From 522bf08b842e1f44a542dfd0b318a6e64082a3c8 Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Wed, 5 Jan 2022 17:42:13 -0700
Subject: [PATCH 05/38] Update __init__.py

---
 tensorflow_quantum/python/optimizers/__init__.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow_quantum/python/optimizers/__init__.py b/tensorflow_quantum/python/optimizers/__init__.py
index aea47fa9d..d93023205 100755
--- a/tensorflow_quantum/python/optimizers/__init__.py
+++ b/tensorflow_quantum/python/optimizers/__init__.py
@@ -17,5 +17,5 @@
 # Quantum circuit specific optimizers.
 from tensorflow_quantum.python.optimizers.rotosolve_minimizer import (
     minimize as rotosolve_minimize)
-from tensorflow_quantum.python.optimizers.spsa_minimizer import (
-    minimize as spsa_minimize)
+from tensorflow_quantum.python.optimizers.spsa_minimizer import (minimize as
+                                                                 spsa_minimize)

From 33e607f09d77c00f52127ce7741f96c09e2721ad Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Wed, 5 Jan 2022 17:53:27 -0700
Subject: [PATCH 06/38] Update spsa_minimizer.py

---
 tensorflow_quantum/python/optimizers/spsa_minimizer.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensorflow_quantum/python/optimizers/spsa_minimizer.py b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
index 56a9b4268..64fbf4c28 100644
--- a/tensorflow_quantum/python/optimizers/spsa_minimizer.py
+++ b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
@@ -102,9 +102,9 @@ def _get_initial_state(initial_position, tolerance, expectation_value_function,
         "a": tf.Variable(a),
         "alpha": tf.Variable(alpha),
         "c": tf.Variable(c),
-        "gamma" : tf.Variable(gamma),
-        "blocking" : tf.Variable(blocking),
-        "allowed_increase" : tf.Variable(allowed_increase)
+        "gamma": tf.Variable(gamma),
+        "blocking": tf.Variable(blocking),
+        "allowed_increase": tf.Variable(allowed_increase)
     }
     return SPSAOptimizerResults(**init_args)
 

From ba92437caf647b159699fbbb28450b07964f822a Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Wed, 5 Jan 2022 17:54:31 -0700
Subject: [PATCH 07/38] Update spsa_minimizer.py

---
 tensorflow_quantum/python/optimizers/spsa_minimizer.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensorflow_quantum/python/optimizers/spsa_minimizer.py b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
index 64fbf4c28..e50405cf2 100644
--- a/tensorflow_quantum/python/optimizers/spsa_minimizer.py
+++ b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
@@ -89,7 +89,8 @@ def prefer_static_value(x):
     ])
 
 
-def _get_initial_state(initial_position, tolerance, expectation_value_function, a, alpha, c, gamma, blocking, allowed_increase):
+def _get_initial_state(initial_position, tolerance, expectation_value_function,
+                       a, alpha, c, gamma, blocking, allowed_increase):
     """Create SPSAOptimizerResults with initial state of search."""
     init_args = {
         "converged": tf.Variable(False),

From 363af00ac9bb28851ca31880b668d10cbd2c80aa Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Wed, 5 Jan 2022 18:03:40 -0700
Subject: [PATCH 08/38] Update spsa_minimizer.py

---
 .../python/optimizers/spsa_minimizer.py       | 24 +++++++++++++------
 1 file changed, 17 insertions(+), 7 deletions(-)

diff --git a/tensorflow_quantum/python/optimizers/spsa_minimizer.py b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
index e50405cf2..25e0244b0 100644
--- a/tensorflow_quantum/python/optimizers/spsa_minimizer.py
+++ b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
@@ -78,12 +78,12 @@ def prefer_static_value(x):
         # Specifies the learning rate
         'alpha',
         # Specifies scaling of the learning rate
-        'c', 
+        'c',
         # Specifies the size of the perturbations
         'gamma',
         # Specifies scaling of the size of the perturbations
         'blocking',
-        # If true, then the optimizer will only accept updates that improve the objective function. 
+        # If true, then the optimizer will only accept updates that improve the objective function.
         'allowed_increase'
         # Specifies maximum allowable increase in objective function (only applies if blocking is true).
     ])
@@ -206,7 +206,11 @@ def _spsa_once(state):
             Returns:
                 states: A list which the first element is the new state
             """
-            delta_shift = tf.cast(2 * tf.random.uniform(shape=state.position.shape, minval=0, maxval=2, dtype=tf.int32) - 1, tf.float32)
+            delta_shift = tf.cast(
+                2 * tf.random.uniform(shape=state.position.shape,
+                                      minval=0,
+                                      maxval=2,
+                                      dtype=tf.int32) - 1, tf.float32)
 
             v_m, v_p =  expectation_value_function(state.position - state.c * delta_shift),\
                  expectation_value_function(state.position + state.c * delta_shift)
@@ -238,8 +242,11 @@ def _cond(state):
         def _body(state):
             """Main optimization loop."""
             
-            new_a = a_init / ((tf.cast(state.num_iterations + 1, tf.float32) + 0.01 * tf.cast(max_iterations, tf.float32))**state.alpha)
-            new_c = c_init / (tf.cast(state.num_iterations + 1, tf.float32)**state.gamma)
+            new_a = a_init / (
+                (tf.cast(state.num_iterations + 1, tf.float32) +
+                 0.01 * tf.cast(max_iterations, tf.float32))**state.alpha)
+            new_c = c_init / (tf.cast(state.num_iterations + 1, tf.float32)**
+                              state.gamma)
 
             state.a.assign(new_a)
             state.c.assign(new_c)
@@ -252,9 +259,12 @@ def _body(state):
                 state.tolerance)
             return [state]
 
-        initial_state = _get_initial_state(initial_position, tolerance, expectation_value_function, a, alpha, c, gamma, blocking, allowed_increase)
+            initial_state = _get_initial_state(initial_position, tolerance,
+                                           expectation_value_function, a, alpha,
+                                           c, gamma, blocking, allowed_increase)
 
-        initial_state.objective_value.assign(expectation_value_function(initial_state.position))
+        initial_state.objective_value.assign(
+            expectation_value_function(initial_state.position))
 
         return tf.while_loop(cond=_cond,
                              body=_body,

From ba8668a05c7d42ff15d5966d57713346d452e7d9 Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Wed, 5 Jan 2022 18:10:35 -0700
Subject: [PATCH 09/38] Update spsa_minimizer.py

---
 tensorflow_quantum/python/optimizers/spsa_minimizer.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tensorflow_quantum/python/optimizers/spsa_minimizer.py b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
index 25e0244b0..c52a6b5d8 100644
--- a/tensorflow_quantum/python/optimizers/spsa_minimizer.py
+++ b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
@@ -260,8 +260,9 @@ def _body(state):
             return [state]
 
             initial_state = _get_initial_state(initial_position, tolerance,
-                                           expectation_value_function, a, alpha,
-                                           c, gamma, blocking, allowed_increase)
+                                               expectation_value_function, a,
+                                               alpha, c, gamma, blocking,
+                                               allowed_increase)
 
         initial_state.objective_value.assign(
             expectation_value_function(initial_state.position))

From a7471e6eaf094aebd79e8f139fc4b7f48d9e5064 Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Wed, 5 Jan 2022 18:15:07 -0700
Subject: [PATCH 10/38] Update spsa_minimizer.py

---
 tensorflow_quantum/python/optimizers/spsa_minimizer.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow_quantum/python/optimizers/spsa_minimizer.py b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
index c52a6b5d8..45db50c40 100644
--- a/tensorflow_quantum/python/optimizers/spsa_minimizer.py
+++ b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
@@ -259,10 +259,10 @@ def _body(state):
                 state.tolerance)
             return [state]
 
-            initial_state = _get_initial_state(initial_position, tolerance,
-                                               expectation_value_function, a,
-                                               alpha, c, gamma, blocking,
-                                               allowed_increase)
+        initial_state = _get_initial_state(initial_position, tolerance,
+                                           expectation_value_function, a,
+                                           alpha, c, gamma, blocking,
+                                           allowed_increase)
 
         initial_state.objective_value.assign(
             expectation_value_function(initial_state.position))

From fc0e0ef3f451591d96886b4f866380a4d50b09a3 Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Wed, 5 Jan 2022 18:16:30 -0700
Subject: [PATCH 11/38] Update spsa_minimizer_test.py

---
 tensorflow_quantum/python/optimizers/spsa_minimizer_test.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py b/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py
index 5f1f16c37..e9a2eda91 100644
--- a/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py
+++ b/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py
@@ -91,7 +91,7 @@ def test_nonlinear_function_optimization(self):
         """
         func = lambda x: x[0]**2 + x[1]**2
 
-        result = spsa_minimizer(func, tf.random.uniform(shape=[2]))
+        result = spsa_minimizer.minimize(func, tf.random.uniform(shape=[2]))
         print(func(result.position))
         self.assertAlmostEqual(func(result.position).numpy(), 0, delta=1e-4)
         self.assertTrue(result.converged)
@@ -140,7 +140,7 @@ def convert_to_circuit(input_data):
         ])
 
         # Initial guess of the parameter from random number
-        result = spsa_minimizer(
+        result = spsa_minimizer.minimize(
             loss_function_with_model_parameters(model, tf.keras.losses.Hinge(),
                                                 x_circ, y),
             tf.random.uniform(shape=[2]) * 2 * np.pi)

From 01448897da9fb562314fbe9a1989f4f0508f846a Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Wed, 5 Jan 2022 18:19:04 -0700
Subject: [PATCH 12/38] Update spsa_minimizer.py

---
 tensorflow_quantum/python/optimizers/spsa_minimizer.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/tensorflow_quantum/python/optimizers/spsa_minimizer.py b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
index 45db50c40..a848fde13 100644
--- a/tensorflow_quantum/python/optimizers/spsa_minimizer.py
+++ b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
@@ -260,9 +260,8 @@ def _body(state):
             return [state]
 
         initial_state = _get_initial_state(initial_position, tolerance,
-                                           expectation_value_function, a,
-                                           alpha, c, gamma, blocking,
-                                           allowed_increase)
+                                           expectation_value_function, a, alpha,
+                                           c, gamma, blocking, allowed_increase)
 
         initial_state.objective_value.assign(
             expectation_value_function(initial_state.position))

From 728d55e9525ace803eb80904c9108988aead271d Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Wed, 5 Jan 2022 18:22:11 -0700
Subject: [PATCH 13/38] Update spsa_minimizer.py

---
 tensorflow_quantum/python/optimizers/spsa_minimizer.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/tensorflow_quantum/python/optimizers/spsa_minimizer.py b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
index a848fde13..9ae567771 100644
--- a/tensorflow_quantum/python/optimizers/spsa_minimizer.py
+++ b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
@@ -194,7 +194,6 @@ def minimize(expectation_value_function,
 
         a_init = tf.convert_to_tensor(a, name='initial_a', dtype='float32')
         c_init = tf.convert_to_tensor(c, name='initial_c', dtype='float32')
-    
 
         def _spsa_once(state):
             """Caclulate single SPSA gradient estimation
@@ -230,7 +229,6 @@ def _spsa_once(state):
 
             return [state]
 
-
         # The `state` here is a `SPSAOptimizerResults` tuple with
         # values for the current state of the algorithm computation.
         def _cond(state):
@@ -241,7 +239,6 @@ def _cond(state):
 
         def _body(state):
             """Main optimization loop."""
-            
             new_a = a_init / (
                 (tf.cast(state.num_iterations + 1, tf.float32) +
                  0.01 * tf.cast(max_iterations, tf.float32))**state.alpha)

From 5d4998105eb38935a49939c161bfdf6611778caa Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Wed, 5 Jan 2022 18:22:59 -0700
Subject: [PATCH 14/38] Update spsa_minimizer_test.py

---
 tensorflow_quantum/python/optimizers/spsa_minimizer_test.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py b/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py
index e9a2eda91..0aa08b6f4 100644
--- a/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py
+++ b/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py
@@ -31,6 +31,7 @@
 from tensorflow_quantum.python import util
 from tensorflow_quantum.python.optimizers import spsa_minimizer
 
+
 def loss_function_with_model_parameters(model, loss, train_x, train_y):
     """Create a new function that assign the model parameter to the model
     and evaluate its value.
@@ -96,7 +97,6 @@ def test_nonlinear_function_optimization(self):
         self.assertAlmostEqual(func(result.position).numpy(), 0, delta=1e-4)
         self.assertTrue(result.converged)
 
-
     def test_keras_model_optimization(self):
         """Optimizate a PQC based keras model."""
 
@@ -149,6 +149,5 @@ def convert_to_circuit(input_data):
         self.assertTrue(result.converged)
 
 
-
 if __name__ == "__main__":
     tf.test.main()

From 7daf326fb0e6e8dc0a1953bc662f6e5e7fc33739 Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Wed, 5 Jan 2022 18:23:17 -0700
Subject: [PATCH 15/38] Update spsa_minimizer.py

---
 tensorflow_quantum/python/optimizers/spsa_minimizer.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensorflow_quantum/python/optimizers/spsa_minimizer.py b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
index 9ae567771..1269198f4 100644
--- a/tensorflow_quantum/python/optimizers/spsa_minimizer.py
+++ b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
@@ -14,7 +14,6 @@
 # ==============================================================================
 """The SPSA minimization algorithm"""
 import collections
-import numpy as np
 import tensorflow as tf
 
 

From c2b9ebd538a632f5ce82ba7fe1e9db2fd78efc89 Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Wed, 5 Jan 2022 18:35:29 -0700
Subject: [PATCH 16/38] Update spsa_minimizer.py

---
 .../python/optimizers/spsa_minimizer.py       | 19 +++++++++++--------
 1 file changed, 11 insertions(+), 8 deletions(-)

diff --git a/tensorflow_quantum/python/optimizers/spsa_minimizer.py b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
index 1269198f4..9b53a6d09 100644
--- a/tensorflow_quantum/python/optimizers/spsa_minimizer.py
+++ b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
@@ -82,9 +82,11 @@ def prefer_static_value(x):
         'gamma',
         # Specifies scaling of the size of the perturbations
         'blocking',
-        # If true, then the optimizer will only accept updates that improve the objective function.
+        # If true, then the optimizer will only accept updates that improve 
+        # the objective function.
         'allowed_increase'
-        # Specifies maximum allowable increase in objective function (only applies if blocking is true).
+        # Specifies maximum allowable increase in objective function 
+        # (only applies if blocking is true).
     ])
 
 
@@ -162,13 +164,13 @@ def minimize(expectation_value_function,
         a: Scalar `tf.Tensor` of real dtype. Specifies the learning rate
         alpha: Scalar `tf.Tensor` of real dtype. Specifies scaling of the
             learning rate.
-        c: Scalar `tf.Tensor` of real dtype. Specifies the size of the 
+        c: Scalar `tf.Tensor` of real dtype. Specifies the size of the
             perturbations.
         gamma: Scalar `tf.Tensor` of real dtype. Specifies scaling of the
             size of the perturbations.
         blocking: Boolean. If true, then the optimizer will only accept
-            updates that improve the objective function. 
-        allowed_increase: Scalar `tf.Tensor` of real dtype. Specifies maximum 
+            updates that improve the objective function.
+        allowed_increase: Scalar `tf.Tensor` of real dtype. Specifies maximum
             allowable increase in objective function (only applies if blocking
             is true).
         name: (Optional) Python `str`. The name prefixed to the ops created
@@ -210,15 +212,16 @@ def _spsa_once(state):
                                       maxval=2,
                                       dtype=tf.int32) - 1, tf.float32)
 
-            v_m, v_p =  expectation_value_function(state.position - state.c * delta_shift),\
-                 expectation_value_function(state.position + state.c * delta_shift)
+            v_m =  expectation_value_function(state.position - state.c * delta_shift)
+            v_p = expectation_value_function(state.position + state.c * delta_shift)
 
             gradient_estimate = (v_p - v_m) / (2 * state.c) * delta_shift
             update = state.a * gradient_estimate
 
             current_obj = expectation_value_function(state.position - update)
             #state.position.assign(tf.math.floormod(state.position - update, 2* np.pi))
-            if state.num_objective_evaluations == 0 or state.objective_value_previous_iteration < current_obj + state.allowed_increase or not state.blocking:
+            if state.num_objective_evaluations == 0 or state.objective_value_previous_iteration \
+                < current_obj + state.allowed_increase or not state.blocking:
                 state.position.assign(state.position - update)
 
             state.num_objective_evaluations.assign_add(2)

From d145b6a32508fe0db90854a5556a5d48aaf6dd94 Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Wed, 5 Jan 2022 18:39:58 -0700
Subject: [PATCH 17/38] Update spsa_minimizer.py

---
 .../python/optimizers/spsa_minimizer.py             | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/tensorflow_quantum/python/optimizers/spsa_minimizer.py b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
index 9b53a6d09..39dbb55ea 100644
--- a/tensorflow_quantum/python/optimizers/spsa_minimizer.py
+++ b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
@@ -82,10 +82,10 @@ def prefer_static_value(x):
         'gamma',
         # Specifies scaling of the size of the perturbations
         'blocking',
-        # If true, then the optimizer will only accept updates that improve 
+        # If true, then the optimizer will only accept updates that improve
         # the objective function.
         'allowed_increase'
-        # Specifies maximum allowable increase in objective function 
+        # Specifies maximum allowable increase in objective function
         # (only applies if blocking is true).
     ])
 
@@ -212,14 +212,17 @@ def _spsa_once(state):
                                       maxval=2,
                                       dtype=tf.int32) - 1, tf.float32)
 
-            v_m =  expectation_value_function(state.position - state.c * delta_shift)
-            v_p = expectation_value_function(state.position + state.c * delta_shift)
+            v_m = expectation_value_function(state.position -
+                                             state.c * delta_shift)
+            v_p = expectation_value_function(state.position +
+                                             state.c * delta_shift)
 
             gradient_estimate = (v_p - v_m) / (2 * state.c) * delta_shift
             update = state.a * gradient_estimate
 
             current_obj = expectation_value_function(state.position - update)
-            #state.position.assign(tf.math.floormod(state.position - update, 2* np.pi))
+            # state.position.assign(tf.math.floormod(state.position - 
+            #                                        update, 2* np.pi))
             if state.num_objective_evaluations == 0 or state.objective_value_previous_iteration \
                 < current_obj + state.allowed_increase or not state.blocking:
                 state.position.assign(state.position - update)

From 54f01ea75fa501d5b78bd786c9785458ccbe5530 Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Wed, 5 Jan 2022 18:42:56 -0700
Subject: [PATCH 18/38] Update spsa_minimizer.py

---
 tensorflow_quantum/python/optimizers/spsa_minimizer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow_quantum/python/optimizers/spsa_minimizer.py b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
index 39dbb55ea..03b49739a 100644
--- a/tensorflow_quantum/python/optimizers/spsa_minimizer.py
+++ b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
@@ -221,7 +221,7 @@ def _spsa_once(state):
             update = state.a * gradient_estimate
 
             current_obj = expectation_value_function(state.position - update)
-            # state.position.assign(tf.math.floormod(state.position - 
+            # state.position.assign(tf.math.floormod(state.position -
             #                                        update, 2* np.pi))
             if state.num_objective_evaluations == 0 or state.objective_value_previous_iteration \
                 < current_obj + state.allowed_increase or not state.blocking:

From 5a2dd0f8b93f6fdc59625e8dd912a4d5d722a383 Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Wed, 5 Jan 2022 18:43:28 -0700
Subject: [PATCH 19/38] Update spsa_minimizer.py

---
 tensorflow_quantum/python/optimizers/spsa_minimizer.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensorflow_quantum/python/optimizers/spsa_minimizer.py b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
index 03b49739a..347d77c51 100644
--- a/tensorflow_quantum/python/optimizers/spsa_minimizer.py
+++ b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
@@ -223,7 +223,8 @@ def _spsa_once(state):
             current_obj = expectation_value_function(state.position - update)
             # state.position.assign(tf.math.floormod(state.position -
             #                                        update, 2* np.pi))
-            if state.num_objective_evaluations == 0 or state.objective_value_previous_iteration \
+            if state.num_objective_evaluations == 0 or \
+                state.objective_value_previous_iteration \
                 < current_obj + state.allowed_increase or not state.blocking:
                 state.position.assign(state.position - update)
 

From ad542c08651887e49cf268d9745486bc4d3a2219 Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Wed, 5 Jan 2022 18:53:06 -0700
Subject: [PATCH 20/38] Update import_test.py

---
 scripts/import_test.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/scripts/import_test.py b/scripts/import_test.py
index 48a4d4b8c..cff5917f4 100644
--- a/scripts/import_test.py
+++ b/scripts/import_test.py
@@ -81,6 +81,7 @@ def test_imports():
 
     #Optimizers
     _ = tfq.optimizers.rotosolve_minimize
+    _ = tfq.optimizers.spsa_minimize
 
 
 if __name__ == "__main__":

From 35558bb55192ac5b5cd08f6918091483a000f54a Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Wed, 5 Jan 2022 19:10:59 -0700
Subject: [PATCH 21/38] Update BUILD

---
 tensorflow_quantum/python/optimizers/BUILD | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/tensorflow_quantum/python/optimizers/BUILD b/tensorflow_quantum/python/optimizers/BUILD
index 582041dcd..e632eb168 100755
--- a/tensorflow_quantum/python/optimizers/BUILD
+++ b/tensorflow_quantum/python/optimizers/BUILD
@@ -11,6 +11,7 @@ py_library(
     srcs_version = "PY3",
     deps = [
         ":rotosolve_minimizer",
+        ":spsa_minimizer",
     ],
 )
 
@@ -19,6 +20,11 @@ py_library(
     srcs = ["rotosolve_minimizer.py"],
 )
 
+py_library(
+    name = "spsa_minimizer",
+    srcs = ["spsa_minimizer.py"],
+)
+
 py_test(
     name = "rotosolve_minimizer_test",
     srcs = ["rotosolve_minimizer_test.py"],
@@ -29,3 +35,14 @@ py_test(
         "//tensorflow_quantum/python/layers/high_level:pqc",
     ],
 )
+
+py_test(
+    name = "spsa_minimizer_test",
+    srcs = ["spsa_minimizer_test.py"],
+    python_version = "PY3",
+    deps = [
+        ":spsa_minimizer",
+        "//tensorflow_quantum/core/ops:tfq_ps_util_ops_py",
+        "//tensorflow_quantum/python/layers/high_level:pqc",
+    ],
+)

From 47074771bcdf2dbb85e3f43a8a88d9e86ae859dd Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Wed, 5 Jan 2022 19:12:47 -0700
Subject: [PATCH 22/38] Update BUILD

---
 release/BUILD | 1 +
 1 file changed, 1 insertion(+)

diff --git a/release/BUILD b/release/BUILD
index deeae5a42..7212cb1d3 100644
--- a/release/BUILD
+++ b/release/BUILD
@@ -65,5 +65,6 @@ sh_binary(
         "//tensorflow_quantum/python:quantum_context",
         "//tensorflow_quantum/python:util",
         "//tensorflow_quantum/python/optimizers:rotosolve_minimizer",
+        "//tensorflow_quantum/python/optimizers:spsa_minimizer",
     ],
 )

From 64e5c9afc05392385d1eedca762bac6703a43da6 Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Wed, 5 Jan 2022 19:13:22 -0700
Subject: [PATCH 23/38] Update build_docs.py

---
 scripts/build_docs.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/scripts/build_docs.py b/scripts/build_docs.py
index 81b4b46ea..dc47df224 100644
--- a/scripts/build_docs.py
+++ b/scripts/build_docs.py
@@ -67,7 +67,7 @@ def main(unused_argv):
                 "parameter_shift_util", "adjoint"
             ],
             "tfq.datasets": ["cluster_state"],
-            "tfq.optimizers": ["rotosolve_minimizer"],
+            "tfq.optimizers": ["rotosolve_minimizer", "spsa_minimizer"],
             "tfq.util": [
                 "from_tensor", "convert_to_tensor", "exp_identity",
                 "check_commutability", "kwargs_cartesian_product",

From 09679ce67526112d011ea93be516642b19d26bef Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Thu, 13 Jan 2022 13:52:08 -0700
Subject: [PATCH 24/38] Update spsa_minimizer.py

---
 .../python/optimizers/spsa_minimizer.py       | 78 +++++++++----------
 1 file changed, 35 insertions(+), 43 deletions(-)

diff --git a/tensorflow_quantum/python/optimizers/spsa_minimizer.py b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
index 347d77c51..ded342cbb 100644
--- a/tensorflow_quantum/python/optimizers/spsa_minimizer.py
+++ b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
@@ -73,11 +73,11 @@ def prefer_static_value(x):
         # Define the stop criteria. Iteration will stop when the
         # objective value difference between two iterations is
         # smaller than tolerance
-        'a',
+        'lr',
         # Specifies the learning rate
         'alpha',
         # Specifies scaling of the learning rate
-        'c',
+        'perturb',
         # Specifies the size of the perturbations
         'gamma',
         # Specifies scaling of the size of the perturbations
@@ -91,7 +91,7 @@ def prefer_static_value(x):
 
 
 def _get_initial_state(initial_position, tolerance, expectation_value_function,
-                       a, alpha, c, gamma, blocking, allowed_increase):
+                       lr, alpha, perturb, gamma, blocking, allowed_increase):
     """Create SPSAOptimizerResults with initial state of search."""
     init_args = {
         "converged": tf.Variable(False),
@@ -101,9 +101,9 @@ def _get_initial_state(initial_position, tolerance, expectation_value_function,
         "objective_value": tf.Variable(0.),
         "objective_value_previous_iteration": tf.Variable(0.),
         "tolerance": tolerance,
-        "a": tf.Variable(a),
+        "lr": tf.Variable(lr),
         "alpha": tf.Variable(alpha),
-        "c": tf.Variable(c),
+        "perturb": tf.Variable(perturb),
         "gamma": tf.Variable(gamma),
         "blocking": tf.Variable(blocking),
         "allowed_increase": tf.Variable(allowed_increase)
@@ -113,13 +113,13 @@ def _get_initial_state(initial_position, tolerance, expectation_value_function,
 
 def minimize(expectation_value_function,
              initial_position,
-             tolerance=1e-7,
-             max_iterations=60,
+             tolerance=1e-5,
+             max_iterations=200,
              alpha=0.602,
-             a=1.0,
-             c=1.0,
+             lr=1.0,
+             perturb=1.0,
              gamma=0.101,
-             blocking=True,
+             blocking=False,
              allowed_increase=0.5,
              name=None):
     """Applies the SPSA algorithm.
@@ -130,30 +130,25 @@ def minimize(expectation_value_function,
 
     Usage:
 
-    Here is an example of optimize a function which consists summation of
-    a few sinusoids.
+    Here is an example of optimize a function which consists the 
+    summation of a few quadratics.
 
-    >>> n = 10  # Number of sinusoids
-    >>> coefficient = tf.random.uniform(shape=[n])
-    >>> min_value = -tf.math.reduce_sum(tf.abs(coefficient))
-    >>> func = lambda x:tf.math.reduce_sum(tf.sin(x) * coefficient)
+    >>> n = 5  # Number of quadractics
+    >>> coefficient = tf.random.uniform(minval=0, maxval=1, shape=[n])
+    >>> min_value = 0
+    >>> func = func = lambda x : tf.math.reduce_sum(np.power(x, 2) * coefficient)
     >>> # Optimize the function with SPSA, start with random parameters
-    >>> result = tfq.optimizers.SPSA_minimize(func, np.random.random(n))
+    >>> result = tfq.optimizers.spsa_minimize(func, np.random.random(n))
     >>> result.converged
     tf.Tensor(True, shape=(), dtype=bool)
     >>> result.objective_value
-    tf.Tensor(-4.7045116, shape=(), dtype=float32)
+    tf.Tensor(0.0013349084, shape=(), dtype=float32)
 
     Args:
-        expectation_value_function:  A Python callable that accepts
-            a point as a real `tf.Tensor` and returns a `tf.Tensor`s
-            of real dtype containing the value of the function.
-            The function to be minimized. The input is of shape `[n]`,
-            where `n` is the size of the trainable parameters.
-            The return value is a real `tf.Tensor` Scalar (matching shape
-            `[1]`).  This must be a linear combination of quantum
-            measurement expectation value, otherwise this algorithm cannot
-            work.
+        expectation_value_function:  Python callable that accepts a real
+            valued tf.Tensor with shape [n] where n is the number of function
+            parameters. The return value is a real `tf.Tensor` Scalar
+            (matching shape `[1]`).
         initial_position: Real `tf.Tensor` of shape `[n]`. The starting
             point, or points when using batching dimensions, of the search
             procedure. At these points the function value and the gradient
@@ -193,8 +188,8 @@ def minimize(expectation_value_function,
         max_iterations = tf.convert_to_tensor(max_iterations,
                                               name='max_iterations')
 
-        a_init = tf.convert_to_tensor(a, name='initial_a', dtype='float32')
-        c_init = tf.convert_to_tensor(c, name='initial_c', dtype='float32')
+        lr_init = tf.convert_to_tensor(lr, name='initial_a', dtype='float32')
+        perturb_init = tf.convert_to_tensor(perturb, name='initial_c', dtype='float32')
 
         def _spsa_once(state):
             """Caclulate single SPSA gradient estimation
@@ -211,18 +206,15 @@ def _spsa_once(state):
                                       minval=0,
                                       maxval=2,
                                       dtype=tf.int32) - 1, tf.float32)
-
             v_m = expectation_value_function(state.position -
-                                             state.c * delta_shift)
+                                             state.perturb * delta_shift)
             v_p = expectation_value_function(state.position +
-                                             state.c * delta_shift)
+                                             state.perturb * delta_shift)
 
-            gradient_estimate = (v_p - v_m) / (2 * state.c) * delta_shift
-            update = state.a * gradient_estimate
+            gradient_estimate = (v_p - v_m) / (2 * state.perturb) * delta_shift
+            update = state.lr * gradient_estimate
 
             current_obj = expectation_value_function(state.position - update)
-            # state.position.assign(tf.math.floormod(state.position -
-            #                                        update, 2* np.pi))
             if state.num_objective_evaluations == 0 or \
                 state.objective_value_previous_iteration \
                 < current_obj + state.allowed_increase or not state.blocking:
@@ -245,14 +237,14 @@ def _cond(state):
 
         def _body(state):
             """Main optimization loop."""
-            new_a = a_init / (
+            new_lr = lr_init / (
                 (tf.cast(state.num_iterations + 1, tf.float32) +
                  0.01 * tf.cast(max_iterations, tf.float32))**state.alpha)
-            new_c = c_init / (tf.cast(state.num_iterations + 1, tf.float32)**
+            new_perturb = perturb_init / (tf.cast(state.num_iterations + 1, tf.float32)**
                               state.gamma)
 
-            state.a.assign(new_a)
-            state.c.assign(new_c)
+            state.lr.assign(new_lr)
+            state.perturb.assign(new_perturb)
 
             _spsa_once(state)
             state.num_iterations.assign_add(1)
@@ -263,11 +255,11 @@ def _body(state):
             return [state]
 
         initial_state = _get_initial_state(initial_position, tolerance,
-                                           expectation_value_function, a, alpha,
-                                           c, gamma, blocking, allowed_increase)
+                                           expectation_value_function, lr, alpha,
+                                           perturb, gamma, blocking, allowed_increase)
 
         initial_state.objective_value.assign(
-            expectation_value_function(initial_state.position))
+            tf.cast(expectation_value_function(initial_state.position), tf.float32))
 
         return tf.while_loop(cond=_cond,
                              body=_body,

From 3fdc0c3b062d4706a83e0c9f60117a19e285d39d Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Thu, 13 Jan 2022 13:53:47 -0700
Subject: [PATCH 25/38] Update spsa_minimizer_test.py

---
 .../python/optimizers/spsa_minimizer_test.py  | 84 +++++++++++++++++--
 1 file changed, 78 insertions(+), 6 deletions(-)

diff --git a/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py b/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py
index 0aa08b6f4..fc26dbcdc 100644
--- a/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py
+++ b/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py
@@ -31,7 +31,6 @@
 from tensorflow_quantum.python import util
 from tensorflow_quantum.python.optimizers import spsa_minimizer
 
-
 def loss_function_with_model_parameters(model, loss, train_x, train_y):
     """Create a new function that assign the model parameter to the model
     and evaluate its value.
@@ -79,6 +78,8 @@ def func(params):
 
         # evaluate the loss
         loss_value = loss(model(train_x, training=True), train_y)
+        if loss_value.shape != ():
+            loss_value = tf.cast(tf.math.reduce_mean(loss_value), tf.float32)
         return loss_value
 
     return func
@@ -92,11 +93,81 @@ def test_nonlinear_function_optimization(self):
         """
         func = lambda x: x[0]**2 + x[1]**2
 
-        result = spsa_minimizer.minimize(func, tf.random.uniform(shape=[2]))
-        print(func(result.position))
+        result = spsa_minimizer(func, tf.random.uniform(shape=[2]))
         self.assertAlmostEqual(func(result.position).numpy(), 0, delta=1e-4)
         self.assertTrue(result.converged)
 
+    def test_quadratic_function_optimization(self):
+        """Test to optimize a sum of quadratic function.
+        """
+        n = 2
+        coefficient = tf.random.uniform(minval=0, maxval=1, shape=[n])
+        func = lambda x : tf.math.reduce_sum(np.power(x, 2) * coefficient)
+
+        result = spsa_minimizer(func, tf.random.uniform(shape=[n]))
+        self.assertAlmostEqual(func(result.position).numpy(), 0, delta=2e-4)
+        self.assertTrue(result.converged)
+
+    def test_noisy_sin_function_optimization(self):
+        """Test noisy ssinusoidal function
+        """
+        n = 10
+        func = lambda x : tf.math.reduce_sum(tf.math.sin(x) + tf.random.uniform(minval=-0.1, maxval=0.1, shape=[n]))
+
+        result = spsa_minimizer(func, tf.random.uniform(shape=[n]))
+        self.assertLessEqual(func(result.position).numpy(), -n + 0.1 * n)
+
+    def test_failure_optimization(self):
+        """Test a function that is completely random and cannot be minimized
+        """
+        n = 100
+        func = lambda x : np.random.uniform(-10, 10, 1)[0]
+        it = 50
+
+        result = spsa_minimizer(func, tf.random.uniform(shape=[n]), max_iterations=it)
+        self.assertFalse(result.converged)
+        self.assertEqual(result.num_iterations, it)
+
+    def test_3_qubit_circuit(self):
+        """Test quantum circuit optimization, adapted from Qiskit SPSA testing
+        https://github.com/Qiskit/qiskit-terra/blob/main/test/python/algorithms/optimizers/test_spsa.py#L37
+        """
+        
+        qubits = [cirq.GridQubit(0, i) for i in range(3)]
+        params = sympy.symbols("q0:9")
+        circuit = cirq.Circuit()
+        for i in qubits:
+            circuit += cirq.ry(np.pi/4).on(i)
+        circuit += cirq.ry(params[0]).on(qubits[0])
+        circuit += cirq.ry(params[1]).on(qubits[1])
+        circuit += cirq.rz(params[2]).on(qubits[2])
+
+        circuit += cirq.CZ(qubits[0], qubits[1])
+        circuit += cirq.CZ(qubits[1], qubits[2])
+        circuit += cirq.rz(params[3]).on(qubits[0])
+        circuit += cirq.rz(params[4]).on(qubits[1])
+        circuit += cirq.rx(params[5]).on(qubits[2])
+
+        # Build the Keras model.
+        model = tf.keras.Sequential([
+            # The input is the data-circuit, encoded as a tf.string
+            tf.keras.layers.Input(shape=(), dtype=tf.string),
+            # The PQC layer returns the expected value of the
+            # readout gate, range [-1,1].
+            pqc.PQC(circuit, cirq.Z(qubits[0]) * cirq.Z(qubits[1]), repetitions=1024),
+        ])
+
+        initial_point = np.array(
+            [0.82311034, 0.02611798, 0.21077064, 0.61842177, 0.09828447, 0.62013131]
+        )
+
+        result = spsa_minimizer(loss_function_with_model_parameters(model, lambda x, y : x[0][0], 
+            util.convert_to_tensor([cirq.Circuit()]), None), initial_point, max_iterations=100)
+
+        self.assertTrue(result.converged)
+        self.assertLess(result.objective_value.numpy(), -0.95)
+
+
     def test_keras_model_optimization(self):
         """Optimizate a PQC based keras model."""
 
@@ -136,18 +207,19 @@ def convert_to_circuit(input_data):
             tf.keras.layers.Input(shape=(), dtype=tf.string),
             # The PQC layer returns the expected value of the
             # readout gate, range [-1,1].
-            pqc.PQC(circuit, cirq.Z(q1)),
+            pqc.PQC(circuit, 2 * cirq.Z(q1)),
         ])
 
         # Initial guess of the parameter from random number
-        result = spsa_minimizer.minimize(
+        result = spsa_minimizer(
             loss_function_with_model_parameters(model, tf.keras.losses.Hinge(),
                                                 x_circ, y),
             tf.random.uniform(shape=[2]) * 2 * np.pi)
 
-        self.assertAlmostEqual(result.objective_value.numpy(), 0, delta=1e-4)
+        self.assertAlmostEqual(result.objective_value.numpy(), 0)
         self.assertTrue(result.converged)
 
 
+
 if __name__ == "__main__":
     tf.test.main()

From 65b6b3bb40c2c1de807df7189e250950d273cf69 Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Thu, 13 Jan 2022 14:00:33 -0700
Subject: [PATCH 26/38] Update spsa_minimizer_test.py

---
 .../python/optimizers/spsa_minimizer_test.py  | 40 ++++++++++++-------
 1 file changed, 25 insertions(+), 15 deletions(-)

diff --git a/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py b/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py
index fc26dbcdc..7427970ba 100644
--- a/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py
+++ b/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py
@@ -93,7 +93,7 @@ def test_nonlinear_function_optimization(self):
         """
         func = lambda x: x[0]**2 + x[1]**2
 
-        result = spsa_minimizer(func, tf.random.uniform(shape=[2]))
+        result = spsa_minimizer.minimize(func, tf.random.uniform(shape=[2]))
         self.assertAlmostEqual(func(result.position).numpy(), 0, delta=1e-4)
         self.assertTrue(result.converged)
 
@@ -102,9 +102,9 @@ def test_quadratic_function_optimization(self):
         """
         n = 2
         coefficient = tf.random.uniform(minval=0, maxval=1, shape=[n])
-        func = lambda x : tf.math.reduce_sum(np.power(x, 2) * coefficient)
+        func = lambda x: tf.math.reduce_sum(np.power(x, 2) * coefficient)
 
-        result = spsa_minimizer(func, tf.random.uniform(shape=[n]))
+        result = spsa_minimizer.minimize(func, tf.random.uniform(shape=[n]))
         self.assertAlmostEqual(func(result.position).numpy(), 0, delta=2e-4)
         self.assertTrue(result.converged)
 
@@ -112,19 +112,23 @@ def test_noisy_sin_function_optimization(self):
         """Test noisy ssinusoidal function
         """
         n = 10
-        func = lambda x : tf.math.reduce_sum(tf.math.sin(x) + tf.random.uniform(minval=-0.1, maxval=0.1, shape=[n]))
+        func = lambda x: tf.math.reduce_sum(
+            tf.math.sin(x) + tf.random.uniform(
+                minval=-0.1, maxval=0.1, shape=[n]))
 
-        result = spsa_minimizer(func, tf.random.uniform(shape=[n]))
+        result = spsa_minimizer.minimize(func, tf.random.uniform(shape=[n]))
         self.assertLessEqual(func(result.position).numpy(), -n + 0.1 * n)
 
     def test_failure_optimization(self):
         """Test a function that is completely random and cannot be minimized
         """
         n = 100
-        func = lambda x : np.random.uniform(-10, 10, 1)[0]
+        func = lambda x: np.random.uniform(-10, 10, 1)[0]
         it = 50
 
-        result = spsa_minimizer(func, tf.random.uniform(shape=[n]), max_iterations=it)
+        result = spsa_minimizer(func,
+                                tf.random.uniform(shape=[n]),
+                                max_iterations=it)
         self.assertFalse(result.converged)
         self.assertEqual(result.num_iterations, it)
 
@@ -137,7 +141,7 @@ def test_3_qubit_circuit(self):
         params = sympy.symbols("q0:9")
         circuit = cirq.Circuit()
         for i in qubits:
-            circuit += cirq.ry(np.pi/4).on(i)
+            circuit += cirq.ry(np.pi / 4).on(i)
         circuit += cirq.ry(params[0]).on(qubits[0])
         circuit += cirq.ry(params[1]).on(qubits[1])
         circuit += cirq.rz(params[2]).on(qubits[2])
@@ -154,15 +158,21 @@ def test_3_qubit_circuit(self):
             tf.keras.layers.Input(shape=(), dtype=tf.string),
             # The PQC layer returns the expected value of the
             # readout gate, range [-1,1].
-            pqc.PQC(circuit, cirq.Z(qubits[0]) * cirq.Z(qubits[1]), repetitions=1024),
+            pqc.PQC(circuit,
+                    cirq.Z(qubits[0]) * cirq.Z(qubits[1]),
+                    repetitions=1024),
         ])
 
-        initial_point = np.array(
-            [0.82311034, 0.02611798, 0.21077064, 0.61842177, 0.09828447, 0.62013131]
-        )
+        initial_point = np.array([
+            0.82311034, 0.02611798, 0.21077064, 0.61842177, 0.09828447,
+            0.62013131
+        ])
 
-        result = spsa_minimizer(loss_function_with_model_parameters(model, lambda x, y : x[0][0], 
-            util.convert_to_tensor([cirq.Circuit()]), None), initial_point, max_iterations=100)
+        result = spsa_minimizer(loss_function_with_model_parameters(
+            model, lambda x, y: x[0][0],
+            util.convert_to_tensor([cirq.Circuit()]), None),
+                                initial_point,
+                                max_iterations=100)
 
         self.assertTrue(result.converged)
         self.assertLess(result.objective_value.numpy(), -0.95)
@@ -211,7 +221,7 @@ def convert_to_circuit(input_data):
         ])
 
         # Initial guess of the parameter from random number
-        result = spsa_minimizer(
+        result = spsa_minimizer.minimize(
             loss_function_with_model_parameters(model, tf.keras.losses.Hinge(),
                                                 x_circ, y),
             tf.random.uniform(shape=[2]) * 2 * np.pi)

From 4a3edbf2f61950cad6cb20cdabe0e17854ac25be Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Thu, 13 Jan 2022 14:00:47 -0700
Subject: [PATCH 27/38] Update spsa_minimizer.py

---
 .../python/optimizers/spsa_minimizer.py       | 21 +++++++++++++------
 1 file changed, 15 insertions(+), 6 deletions(-)

diff --git a/tensorflow_quantum/python/optimizers/spsa_minimizer.py b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
index ded342cbb..dfbdacb25 100644
--- a/tensorflow_quantum/python/optimizers/spsa_minimizer.py
+++ b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
@@ -121,6 +121,7 @@ def minimize(expectation_value_function,
              gamma=0.101,
              blocking=False,
              allowed_increase=0.5,
+             seed=None,
              name=None):
     """Applies the SPSA algorithm.
 
@@ -178,6 +179,9 @@ def minimize(expectation_value_function,
     """
 
     with tf.name_scope(name or 'minimize'):
+        if seed is not None:
+            tf.random.set_seed(seed)
+
         initial_position = tf.convert_to_tensor(initial_position,
                                                 name='initial_position',
                                                 dtype='float32')
@@ -189,7 +193,9 @@ def minimize(expectation_value_function,
                                               name='max_iterations')
 
         lr_init = tf.convert_to_tensor(lr, name='initial_a', dtype='float32')
-        perturb_init = tf.convert_to_tensor(perturb, name='initial_c', dtype='float32')
+        perturb_init = tf.convert_to_tensor(perturb,
+                                            name='initial_c',
+                                            dtype='float32')
 
         def _spsa_once(state):
             """Caclulate single SPSA gradient estimation
@@ -240,8 +246,9 @@ def _body(state):
             new_lr = lr_init / (
                 (tf.cast(state.num_iterations + 1, tf.float32) +
                  0.01 * tf.cast(max_iterations, tf.float32))**state.alpha)
-            new_perturb = perturb_init / (tf.cast(state.num_iterations + 1, tf.float32)**
-                              state.gamma)
+            new_perturb = perturb_init / (tf.cast(state.num_iterations + 1,
+                                                  tf.float32)**state.gamma)
+ 
 
             state.lr.assign(new_lr)
             state.perturb.assign(new_perturb)
@@ -255,11 +262,13 @@ def _body(state):
             return [state]
 
         initial_state = _get_initial_state(initial_position, tolerance,
-                                           expectation_value_function, lr, alpha,
-                                           perturb, gamma, blocking, allowed_increase)
+                                           expectation_value_function, lr,
+                                           alpha, perturb, gamma, blocking,
+                                           allowed_increase)
 
         initial_state.objective_value.assign(
-            tf.cast(expectation_value_function(initial_state.position), tf.float32))
+            tf.cast(expectation_value_function(initial_state.position),
+                    tf.float32))
 
         return tf.while_loop(cond=_cond,
                              body=_body,

From 412adf4499d2d88f7a081f53907dbfe2dec8ca98 Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Thu, 13 Jan 2022 14:05:03 -0700
Subject: [PATCH 28/38] Update spsa_minimizer.py

---
 tensorflow_quantum/python/optimizers/spsa_minimizer.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensorflow_quantum/python/optimizers/spsa_minimizer.py b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
index dfbdacb25..a8d77b210 100644
--- a/tensorflow_quantum/python/optimizers/spsa_minimizer.py
+++ b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
@@ -248,7 +248,6 @@ def _body(state):
                  0.01 * tf.cast(max_iterations, tf.float32))**state.alpha)
             new_perturb = perturb_init / (tf.cast(state.num_iterations + 1,
                                                   tf.float32)**state.gamma)
- 
 
             state.lr.assign(new_lr)
             state.perturb.assign(new_perturb)

From 654f398d86d756d1b7f9da5397d11ab0860c7957 Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Thu, 13 Jan 2022 14:45:52 -0700
Subject: [PATCH 29/38] Update spsa_minimizer_test.py

---
 .../python/optimizers/spsa_minimizer_test.py  | 36 ++++++++++++++++---
 1 file changed, 31 insertions(+), 5 deletions(-)

diff --git a/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py b/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py
index 7427970ba..ed23ccbb0 100644
--- a/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py
+++ b/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py
@@ -31,6 +31,7 @@
 from tensorflow_quantum.python import util
 from tensorflow_quantum.python.optimizers import spsa_minimizer
 
+
 def loss_function_with_model_parameters(model, loss, train_x, train_y):
     """Create a new function that assign the model parameter to the model
     and evaluate its value.
@@ -126,17 +127,44 @@ def test_failure_optimization(self):
         func = lambda x: np.random.uniform(-10, 10, 1)[0]
         it = 50
 
-        result = spsa_minimizer(func,
+        result = spsa_minimizer.minimize(func,
                                 tf.random.uniform(shape=[n]),
                                 max_iterations=it)
         self.assertFalse(result.converged)
         self.assertEqual(result.num_iterations, it)
 
+    def test_blocking(self):
+        """Test the blocking functionality.
+        """
+        n = 10
+        it = 50
+
+        init = 1
+        self.incr = 0
+        def func(params):
+            self.incr += init
+            return self.incr
+
+        result = spsa_minimizer.minimize(func, tf.random.uniform(shape=[n]), blocking=True, allowed_increase=0.5, max_iterations=it)
+        self.assertFalse(result.converged)
+        self.assertEqual(result.num_iterations, it)
+        self.assertEqual(result.objective_value, init * 4) # function executd 3 (in step) + 1 (initial evaluation) times
+
+        init = 1/6 * 0.49
+        self.incr = 0
+        def func(params):
+            self.incr += init
+            return self.incr
+
+        result = spsa_minimizer.minimize(func, tf.random.uniform(shape=[n]), blocking=True, allowed_increase=0.5, max_iterations=it)
+        self.assertFalse(result.converged)
+        self.assertEqual(result.num_iterations, it)
+        self.assertEqual(result.objective_value, init * 3 * it + init)
+
     def test_3_qubit_circuit(self):
         """Test quantum circuit optimization, adapted from Qiskit SPSA testing
         https://github.com/Qiskit/qiskit-terra/blob/main/test/python/algorithms/optimizers/test_spsa.py#L37
         """
-        
         qubits = [cirq.GridQubit(0, i) for i in range(3)]
         params = sympy.symbols("q0:9")
         circuit = cirq.Circuit()
@@ -168,7 +196,7 @@ def test_3_qubit_circuit(self):
             0.62013131
         ])
 
-        result = spsa_minimizer(loss_function_with_model_parameters(
+        result = spsa_minimizer.minimize(loss_function_with_model_parameters(
             model, lambda x, y: x[0][0],
             util.convert_to_tensor([cirq.Circuit()]), None),
                                 initial_point,
@@ -177,7 +205,6 @@ def test_3_qubit_circuit(self):
         self.assertTrue(result.converged)
         self.assertLess(result.objective_value.numpy(), -0.95)
 
-
     def test_keras_model_optimization(self):
         """Optimizate a PQC based keras model."""
 
@@ -230,6 +257,5 @@ def convert_to_circuit(input_data):
         self.assertTrue(result.converged)
 
 
-
 if __name__ == "__main__":
     tf.test.main()

From 14e9e04851bad30c462272fcdbcaa565f0ea4b4f Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Thu, 13 Jan 2022 14:46:06 -0700
Subject: [PATCH 30/38] Update spsa_minimizer.py

---
 .../python/optimizers/spsa_minimizer.py         | 17 ++++++++---------
 1 file changed, 8 insertions(+), 9 deletions(-)

diff --git a/tensorflow_quantum/python/optimizers/spsa_minimizer.py b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
index a8d77b210..3eb57d533 100644
--- a/tensorflow_quantum/python/optimizers/spsa_minimizer.py
+++ b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
@@ -99,7 +99,7 @@ def _get_initial_state(initial_position, tolerance, expectation_value_function,
         "num_objective_evaluations": tf.Variable(0),
         "position": tf.Variable(initial_position),
         "objective_value": tf.Variable(0.),
-        "objective_value_previous_iteration": tf.Variable(0.),
+        "objective_value_previous_iteration": tf.Variable(np.inf),
         "tolerance": tolerance,
         "lr": tf.Variable(lr),
         "alpha": tf.Variable(alpha),
@@ -220,16 +220,15 @@ def _spsa_once(state):
             gradient_estimate = (v_p - v_m) / (2 * state.perturb) * delta_shift
             update = state.lr * gradient_estimate
 
+            state.num_objective_evaluations.assign_add(2)
+
             current_obj = expectation_value_function(state.position - update)
-            if state.num_objective_evaluations == 0 or \
-                state.objective_value_previous_iteration \
-                < current_obj + state.allowed_increase or not state.blocking:
+            if state.objective_value_previous_iteration + state.allowed_increase \
+                >= current_obj or not state.blocking:
                 state.position.assign(state.position - update)
-
-            state.num_objective_evaluations.assign_add(2)
-            state.objective_value_previous_iteration.assign(
-                state.objective_value)
-            state.objective_value.assign(current_obj)
+                state.objective_value_previous_iteration.assign(
+                    state.objective_value)
+                state.objective_value.assign(current_obj)
 
             return [state]
 

From 00d982fc87ae7f124b60709b0b784f8d9a00aed8 Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Thu, 13 Jan 2022 14:55:57 -0700
Subject: [PATCH 31/38] Update spsa_minimizer_test.py

---
 .../python/optimizers/spsa_minimizer_test.py  | 31 +++++++++++++------
 1 file changed, 21 insertions(+), 10 deletions(-)

diff --git a/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py b/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py
index ed23ccbb0..86b1f3289 100644
--- a/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py
+++ b/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py
@@ -128,8 +128,8 @@ def test_failure_optimization(self):
         it = 50
 
         result = spsa_minimizer.minimize(func,
-                                tf.random.uniform(shape=[n]),
-                                max_iterations=it)
+                                         tf.random.uniform(shape=[n]),
+                                         max_iterations=it)
         self.assertFalse(result.converged)
         self.assertEqual(result.num_iterations, it)
 
@@ -141,22 +141,33 @@ def test_blocking(self):
 
         init = 1
         self.incr = 0
-        def func(params):
+        def block_func1(params):
             self.incr += init
             return self.incr
 
-        result = spsa_minimizer.minimize(func, tf.random.uniform(shape=[n]), blocking=True, allowed_increase=0.5, max_iterations=it)
+        result = spsa_minimizer.minimize(block_func1,
+                                         tf.random.uniform(shape=[n]),
+                                         blocking=True,
+                                         allowed_increase=0.5,
+                                         max_iterations=it)
         self.assertFalse(result.converged)
         self.assertEqual(result.num_iterations, it)
-        self.assertEqual(result.objective_value, init * 4) # function executd 3 (in step) + 1 (initial evaluation) times
+        self.assertEqual(
+            result.objective_value, init *
+            4)  # function executd 3 (in step) +
+            # 1 (initial evaluation) times
 
-        init = 1/6 * 0.49
+        init = 1 / 6 * 0.49
         self.incr = 0
-        def func(params):
+        def block_func2(params):
             self.incr += init
             return self.incr
 
-        result = spsa_minimizer.minimize(func, tf.random.uniform(shape=[n]), blocking=True, allowed_increase=0.5, max_iterations=it)
+        result = spsa_minimizer.minimize(block_func2,
+                                         tf.random.uniform(shape=[n]),
+                                         blocking=True,
+                                         allowed_increase=0.5,
+                                         max_iterations=it)
         self.assertFalse(result.converged)
         self.assertEqual(result.num_iterations, it)
         self.assertEqual(result.objective_value, init * 3 * it + init)
@@ -199,8 +210,8 @@ def test_3_qubit_circuit(self):
         result = spsa_minimizer.minimize(loss_function_with_model_parameters(
             model, lambda x, y: x[0][0],
             util.convert_to_tensor([cirq.Circuit()]), None),
-                                initial_point,
-                                max_iterations=100)
+                                         initial_point,
+                                         max_iterations=100)
 
         self.assertTrue(result.converged)
         self.assertLess(result.objective_value.numpy(), -0.95)

From e71d8e2b0945183229fb7882319685445fa6506f Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Thu, 13 Jan 2022 14:56:10 -0700
Subject: [PATCH 32/38] Update spsa_minimizer.py

---
 tensorflow_quantum/python/optimizers/spsa_minimizer.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/tensorflow_quantum/python/optimizers/spsa_minimizer.py b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
index 3eb57d533..7cd4def60 100644
--- a/tensorflow_quantum/python/optimizers/spsa_minimizer.py
+++ b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
@@ -15,6 +15,7 @@
 """The SPSA minimization algorithm"""
 import collections
 import tensorflow as tf
+import numpy as np
 
 
 def prefer_static_shape(x):
@@ -131,13 +132,14 @@ def minimize(expectation_value_function,
 
     Usage:
 
-    Here is an example of optimize a function which consists the 
+    Here is an example of optimize a function which consists the
     summation of a few quadratics.
 
     >>> n = 5  # Number of quadractics
     >>> coefficient = tf.random.uniform(minval=0, maxval=1, shape=[n])
     >>> min_value = 0
-    >>> func = func = lambda x : tf.math.reduce_sum(np.power(x, 2) * coefficient)
+    >>> func = func = lambda x : tf.math.reduce_sum(np.power(x, 2) * \
+            coefficient)
     >>> # Optimize the function with SPSA, start with random parameters
     >>> result = tfq.optimizers.spsa_minimize(func, np.random.random(n))
     >>> result.converged
@@ -223,8 +225,8 @@ def _spsa_once(state):
             state.num_objective_evaluations.assign_add(2)
 
             current_obj = expectation_value_function(state.position - update)
-            if state.objective_value_previous_iteration + state.allowed_increase \
-                >= current_obj or not state.blocking:
+            if state.objective_value_previous_iteration + \
+                state.allowed_increase >= current_obj or not state.blocking:
                 state.position.assign(state.position - update)
                 state.objective_value_previous_iteration.assign(
                     state.objective_value)

From 73f1e4f1f7ceddd5d7b51d54977cd43ed772eb35 Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Thu, 13 Jan 2022 15:04:34 -0700
Subject: [PATCH 33/38] Update rotosolve_minimizer_test.py

---
 .../optimizers/rotosolve_minimizer_test.py    | 155 ++++++++++++++----
 1 file changed, 126 insertions(+), 29 deletions(-)

diff --git a/tensorflow_quantum/python/optimizers/rotosolve_minimizer_test.py b/tensorflow_quantum/python/optimizers/rotosolve_minimizer_test.py
index 7ddc54e28..22619f9e0 100755
--- a/tensorflow_quantum/python/optimizers/rotosolve_minimizer_test.py
+++ b/tensorflow_quantum/python/optimizers/rotosolve_minimizer_test.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Test module for tfq.python.optimizers.rotosolve_minimizer optimizer."""
+"""Test module for tfq.python.optimizers.spsa_minimizer optimizer."""
 # Remove PYTHONPATH collisions for protobuf.
 # pylint: disable=wrong-import-position
 import sys
@@ -29,7 +29,7 @@
 import sympy
 from tensorflow_quantum.python.layers.high_level import pqc
 from tensorflow_quantum.python import util
-from tensorflow_quantum.python.optimizers import rotosolve_minimizer
+from tensorflow_quantum.python.optimizers import spsa_minimizer
 
 
 def loss_function_with_model_parameters(model, loss, train_x, train_y):
@@ -61,7 +61,7 @@ def loss_function_with_model_parameters(model, loss, train_x, train_y):
     # Function accept the parameter and evaluate model
     @tf.function
     def func(params):
-        """A function that can be used by tfq.optimizer.rotosolve_minimize.
+        """A function that can be used by tfq.optimizer.spsa_minimize.
 
         Args:
            params [in]: a 1D tf.Tensor.
@@ -79,44 +79,141 @@ def func(params):
 
         # evaluate the loss
         loss_value = loss(model(train_x, training=True), train_y)
+        if loss_value.shape != ():
+            loss_value = tf.cast(tf.math.reduce_mean(loss_value), tf.float32)
         return loss_value
 
     return func
 
 
-class RotosolveMinimizerTest(tf.test.TestCase, parameterized.TestCase):
-    """Tests for the rotosolve optimization algorithm."""
+class SPSAMinimizerTest(tf.test.TestCase, parameterized.TestCase):
+    """Tests for the SPSA optimization algorithm."""
 
-    def test_function_optimization(self):
-        """Optimize a simple sinusoid function."""
-
-        n = 10  # Number of parameters to be optimized
-        coefficient = tf.random.uniform(shape=[n])
-        min_value = -tf.math.reduce_sum(tf.abs(coefficient))
+    def test_nonlinear_function_optimization(self):
+        """Test to optimize a non-linear function.
+        """
+        func = lambda x: x[0]**2 + x[1]**2
 
-        func = lambda x: tf.math.reduce_sum(tf.sin(x) * coefficient)
+        result = spsa_minimizer.minimize(func, tf.random.uniform(shape=[2]))
+        self.assertAlmostEqual(func(result.position).numpy(), 0, delta=1e-4)
+        self.assertTrue(result.converged)
 
-        result = rotosolve_minimizer.minimize(func, np.random.random(n))
+    def test_quadratic_function_optimization(self):
+        """Test to optimize a sum of quadratic function.
+        """
+        n = 2
+        coefficient = tf.random.uniform(minval=0, maxval=1, shape=[n])
+        func = lambda x: tf.math.reduce_sum(np.power(x, 2) * coefficient)
 
-        self.assertAlmostEqual(func(result.position), min_value)
-        self.assertAlmostEqual(result.objective_value, min_value)
+        result = spsa_minimizer.minimize(func, tf.random.uniform(shape=[n]))
+        self.assertAlmostEqual(func(result.position).numpy(), 0, delta=2e-4)
         self.assertTrue(result.converged)
-        self.assertLess(result.num_iterations,
-                        50)  # 50 is the default max iteration
 
-    def test_nonlinear_function_optimization(self):
-        """Test to optimize a non-linear function.
-        A non-linear function cannot be optimized by rotosolve,
-        therefore the optimization must never converge.
+    def test_noisy_sin_function_optimization(self):
+        """Test noisy ssinusoidal function
         """
-        func = lambda x: x[0]**2 + x[1]**2
+        n = 10
+        func = lambda x: tf.math.reduce_sum(
+            tf.math.sin(x) + tf.random.uniform(
+                minval=-0.1, maxval=0.1, shape=[n]))
+
+        result = spsa_minimizer.minimize(func, tf.random.uniform(shape=[n]))
+        self.assertLessEqual(func(result.position).numpy(), -n + 0.1 * n)
+
+    def test_failure_optimization(self):
+        """Test a function that is completely random and cannot be minimized
+        """
+        n = 100
+        func = lambda x: np.random.uniform(-10, 10, 1)[0]
+        it = 50
 
-        result = rotosolve_minimizer.minimize(func,
-                                              tf.random.uniform(shape=[2]))
+        result = spsa_minimizer.minimize(func,
+                                         tf.random.uniform(shape=[n]),
+                                         max_iterations=it)
+        self.assertFalse(result.converged)
+        self.assertEqual(result.num_iterations, it)
 
+    def test_blocking(self):
+        """Test the blocking functionality.
+        """
+        n = 10
+        it = 50
+
+        init = 1
+        self.incr = 0
+        def block_func1(params):
+            self.incr += init
+            return self.incr
+
+        result = spsa_minimizer.minimize(block_func1,
+                                         tf.random.uniform(shape=[n]),
+                                         blocking=True,
+                                         allowed_increase=0.5,
+                                         max_iterations=it)
+        self.assertFalse(result.converged)
+        self.assertEqual(result.num_iterations, it)
+        self.assertEqual(result.objective_value,
+                         init * 4)  # function executd 3 (in step) +
+        # 1 (initial evaluation) times
+
+        init = 1 / 6 * 0.49
+        self.incr = 0
+        def block_func2(params):
+            self.incr += init
+            return self.incr
+
+        result = spsa_minimizer.minimize(block_func2,
+                                         tf.random.uniform(shape=[n]),
+                                         blocking=True,
+                                         allowed_increase=0.5,
+                                         max_iterations=it)
         self.assertFalse(result.converged)
-        self.assertEqual(result.num_iterations,
-                         50)  # 50 is the default max iteration
+        self.assertEqual(result.num_iterations, it)
+        self.assertEqual(result.objective_value, init * 3 * it + init)
+
+    def test_3_qubit_circuit(self):
+        """Test quantum circuit optimization, adapted from Qiskit SPSA testing
+        https://github.com/Qiskit/qiskit-terra/blob/main/test/python/algorithms/optimizers/test_spsa.py#L37
+        """
+        qubits = [cirq.GridQubit(0, i) for i in range(3)]
+        params = sympy.symbols("q0:9")
+        circuit = cirq.Circuit()
+        for i in qubits:
+            circuit += cirq.ry(np.pi / 4).on(i)
+        circuit += cirq.ry(params[0]).on(qubits[0])
+        circuit += cirq.ry(params[1]).on(qubits[1])
+        circuit += cirq.rz(params[2]).on(qubits[2])
+
+        circuit += cirq.CZ(qubits[0], qubits[1])
+        circuit += cirq.CZ(qubits[1], qubits[2])
+        circuit += cirq.rz(params[3]).on(qubits[0])
+        circuit += cirq.rz(params[4]).on(qubits[1])
+        circuit += cirq.rx(params[5]).on(qubits[2])
+
+        # Build the Keras model.
+        model = tf.keras.Sequential([
+            # The input is the data-circuit, encoded as a tf.string
+            tf.keras.layers.Input(shape=(), dtype=tf.string),
+            # The PQC layer returns the expected value of the
+            # readout gate, range [-1,1].
+            pqc.PQC(circuit,
+                    cirq.Z(qubits[0]) * cirq.Z(qubits[1]),
+                    repetitions=1024),
+        ])
+
+        initial_point = np.array([
+            0.82311034, 0.02611798, 0.21077064, 0.61842177, 0.09828447,
+            0.62013131
+        ])
+
+        result = spsa_minimizer.minimize(loss_function_with_model_parameters(
+            model, lambda x, y: x[0][0],
+            util.convert_to_tensor([cirq.Circuit()]), None),
+                                         initial_point,
+                                         max_iterations=100)
+
+        self.assertTrue(result.converged)
+        self.assertLess(result.objective_value.numpy(), -0.95)
 
     def test_keras_model_optimization(self):
         """Optimizate a PQC based keras model."""
@@ -157,16 +254,16 @@ def convert_to_circuit(input_data):
             tf.keras.layers.Input(shape=(), dtype=tf.string),
             # The PQC layer returns the expected value of the
             # readout gate, range [-1,1].
-            pqc.PQC(circuit, cirq.Z(q1)),
+            pqc.PQC(circuit, 2 * cirq.Z(q1)),
         ])
 
         # Initial guess of the parameter from random number
-        result = rotosolve_minimizer.minimize(
+        result = spsa_minimizer.minimize(
             loss_function_with_model_parameters(model, tf.keras.losses.Hinge(),
                                                 x_circ, y),
             tf.random.uniform(shape=[2]) * 2 * np.pi)
 
-        self.assertAlmostEqual(result.objective_value, 0)
+        self.assertAlmostEqual(result.objective_value.numpy(), 0)
         self.assertTrue(result.converged)
 
 

From 2663722771156b42570d94ed95b59088bd3f2c50 Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Thu, 13 Jan 2022 15:10:05 -0700
Subject: [PATCH 34/38] Update spsa_minimizer_test.py

---
 .../python/optimizers/spsa_minimizer_test.py               | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py b/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py
index 86b1f3289..9e98663b3 100644
--- a/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py
+++ b/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py
@@ -152,10 +152,9 @@ def block_func1(params):
                                          max_iterations=it)
         self.assertFalse(result.converged)
         self.assertEqual(result.num_iterations, it)
-        self.assertEqual(
-            result.objective_value, init *
-            4)  # function executd 3 (in step) +
-            # 1 (initial evaluation) times
+        self.assertEqual(result.objective_value,
+                         init * 4)  # function executd 3 (in step) +
+        # 1 (initial evaluation) times
 
         init = 1 / 6 * 0.49
         self.incr = 0

From d47a968d9acc9b75a36f212fab84752e660a046c Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Thu, 13 Jan 2022 15:15:59 -0700
Subject: [PATCH 35/38] Update spsa_minimizer_test.py

---
 tensorflow_quantum/python/optimizers/spsa_minimizer_test.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py b/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py
index 9e98663b3..2aef27eb3 100644
--- a/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py
+++ b/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py
@@ -141,6 +141,7 @@ def test_blocking(self):
 
         init = 1
         self.incr = 0
+
         def block_func1(params):
             self.incr += init
             return self.incr
@@ -158,6 +159,7 @@ def block_func1(params):
 
         init = 1 / 6 * 0.49
         self.incr = 0
+        
         def block_func2(params):
             self.incr += init
             return self.incr

From 490dd7a25a57bf15b21d92d863a978bc3bc8713a Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Thu, 13 Jan 2022 15:16:26 -0700
Subject: [PATCH 36/38] Update rotosolve_minimizer_test.py

---
 .../optimizers/rotosolve_minimizer_test.py    | 159 ++++--------------
 1 file changed, 29 insertions(+), 130 deletions(-)

diff --git a/tensorflow_quantum/python/optimizers/rotosolve_minimizer_test.py b/tensorflow_quantum/python/optimizers/rotosolve_minimizer_test.py
index 22619f9e0..1d3a2965f 100755
--- a/tensorflow_quantum/python/optimizers/rotosolve_minimizer_test.py
+++ b/tensorflow_quantum/python/optimizers/rotosolve_minimizer_test.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Test module for tfq.python.optimizers.spsa_minimizer optimizer."""
+"""Test module for tfq.python.optimizers.rotosolve_minimizer optimizer."""
 # Remove PYTHONPATH collisions for protobuf.
 # pylint: disable=wrong-import-position
 import sys
@@ -29,19 +29,17 @@
 import sympy
 from tensorflow_quantum.python.layers.high_level import pqc
 from tensorflow_quantum.python import util
-from tensorflow_quantum.python.optimizers import spsa_minimizer
+from tensorflow_quantum.python.optimizers import rotosolve_minimizer
 
 
 def loss_function_with_model_parameters(model, loss, train_x, train_y):
     """Create a new function that assign the model parameter to the model
     and evaluate its value.
-
     Args:
         model : an instance of `tf.keras.Model` or its subclasses.
         loss : a function with signature loss_value = loss(pred_y, true_y).
         train_x : the input part of training data.
         train_y : the output part of training data.
-
     Returns:
         A function that has a signature of:
             loss_value = f(model_parameters).
@@ -61,11 +59,9 @@ def loss_function_with_model_parameters(model, loss, train_x, train_y):
     # Function accept the parameter and evaluate model
     @tf.function
     def func(params):
-        """A function that can be used by tfq.optimizer.spsa_minimize.
-
+        """A function that can be used by tfq.optimizer.rotosolve_minimize.
         Args:
            params [in]: a 1D tf.Tensor.
-
         Returns:
             Loss function value
         """
@@ -79,141 +75,44 @@ def func(params):
 
         # evaluate the loss
         loss_value = loss(model(train_x, training=True), train_y)
-        if loss_value.shape != ():
-            loss_value = tf.cast(tf.math.reduce_mean(loss_value), tf.float32)
         return loss_value
 
     return func
 
 
-class SPSAMinimizerTest(tf.test.TestCase, parameterized.TestCase):
-    """Tests for the SPSA optimization algorithm."""
-
-    def test_nonlinear_function_optimization(self):
-        """Test to optimize a non-linear function.
-        """
-        func = lambda x: x[0]**2 + x[1]**2
-
-        result = spsa_minimizer.minimize(func, tf.random.uniform(shape=[2]))
-        self.assertAlmostEqual(func(result.position).numpy(), 0, delta=1e-4)
-        self.assertTrue(result.converged)
-
-    def test_quadratic_function_optimization(self):
-        """Test to optimize a sum of quadratic function.
-        """
-        n = 2
-        coefficient = tf.random.uniform(minval=0, maxval=1, shape=[n])
-        func = lambda x: tf.math.reduce_sum(np.power(x, 2) * coefficient)
+class RotosolveMinimizerTest(tf.test.TestCase, parameterized.TestCase):
+    """Tests for the rotosolve optimization algorithm."""
 
-        result = spsa_minimizer.minimize(func, tf.random.uniform(shape=[n]))
-        self.assertAlmostEqual(func(result.position).numpy(), 0, delta=2e-4)
-        self.assertTrue(result.converged)
+    def test_function_optimization(self):
+        """Optimize a simple sinusoid function."""
 
-    def test_noisy_sin_function_optimization(self):
-        """Test noisy ssinusoidal function
-        """
-        n = 10
-        func = lambda x: tf.math.reduce_sum(
-            tf.math.sin(x) + tf.random.uniform(
-                minval=-0.1, maxval=0.1, shape=[n]))
+        n = 10  # Number of parameters to be optimized
+        coefficient = tf.random.uniform(shape=[n])
+        min_value = -tf.math.reduce_sum(tf.abs(coefficient))
 
-        result = spsa_minimizer.minimize(func, tf.random.uniform(shape=[n]))
-        self.assertLessEqual(func(result.position).numpy(), -n + 0.1 * n)
+        func = lambda x: tf.math.reduce_sum(tf.sin(x) * coefficient)
 
-    def test_failure_optimization(self):
-        """Test a function that is completely random and cannot be minimized
-        """
-        n = 100
-        func = lambda x: np.random.uniform(-10, 10, 1)[0]
-        it = 50
+        result = rotosolve_minimizer.minimize(func, np.random.random(n))
 
-        result = spsa_minimizer.minimize(func,
-                                         tf.random.uniform(shape=[n]),
-                                         max_iterations=it)
-        self.assertFalse(result.converged)
-        self.assertEqual(result.num_iterations, it)
+        self.assertAlmostEqual(func(result.position), min_value)
+        self.assertAlmostEqual(result.objective_value, min_value)
+        self.assertTrue(result.converged)
+        self.assertLess(result.num_iterations,
+                        50)  # 50 is the default max iteration
 
-    def test_blocking(self):
-        """Test the blocking functionality.
+    def test_nonlinear_function_optimization(self):
+        """Test to optimize a non-linear function.
+        A non-linear function cannot be optimized by rotosolve,
+        therefore the optimization must never converge.
         """
-        n = 10
-        it = 50
-
-        init = 1
-        self.incr = 0
-        def block_func1(params):
-            self.incr += init
-            return self.incr
-
-        result = spsa_minimizer.minimize(block_func1,
-                                         tf.random.uniform(shape=[n]),
-                                         blocking=True,
-                                         allowed_increase=0.5,
-                                         max_iterations=it)
-        self.assertFalse(result.converged)
-        self.assertEqual(result.num_iterations, it)
-        self.assertEqual(result.objective_value,
-                         init * 4)  # function executd 3 (in step) +
-        # 1 (initial evaluation) times
+        func = lambda x: x[0]**2 + x[1]**2
 
-        init = 1 / 6 * 0.49
-        self.incr = 0
-        def block_func2(params):
-            self.incr += init
-            return self.incr
+        result = rotosolve_minimizer.minimize(func,
+                                              tf.random.uniform(shape=[2]))
 
-        result = spsa_minimizer.minimize(block_func2,
-                                         tf.random.uniform(shape=[n]),
-                                         blocking=True,
-                                         allowed_increase=0.5,
-                                         max_iterations=it)
         self.assertFalse(result.converged)
-        self.assertEqual(result.num_iterations, it)
-        self.assertEqual(result.objective_value, init * 3 * it + init)
-
-    def test_3_qubit_circuit(self):
-        """Test quantum circuit optimization, adapted from Qiskit SPSA testing
-        https://github.com/Qiskit/qiskit-terra/blob/main/test/python/algorithms/optimizers/test_spsa.py#L37
-        """
-        qubits = [cirq.GridQubit(0, i) for i in range(3)]
-        params = sympy.symbols("q0:9")
-        circuit = cirq.Circuit()
-        for i in qubits:
-            circuit += cirq.ry(np.pi / 4).on(i)
-        circuit += cirq.ry(params[0]).on(qubits[0])
-        circuit += cirq.ry(params[1]).on(qubits[1])
-        circuit += cirq.rz(params[2]).on(qubits[2])
-
-        circuit += cirq.CZ(qubits[0], qubits[1])
-        circuit += cirq.CZ(qubits[1], qubits[2])
-        circuit += cirq.rz(params[3]).on(qubits[0])
-        circuit += cirq.rz(params[4]).on(qubits[1])
-        circuit += cirq.rx(params[5]).on(qubits[2])
-
-        # Build the Keras model.
-        model = tf.keras.Sequential([
-            # The input is the data-circuit, encoded as a tf.string
-            tf.keras.layers.Input(shape=(), dtype=tf.string),
-            # The PQC layer returns the expected value of the
-            # readout gate, range [-1,1].
-            pqc.PQC(circuit,
-                    cirq.Z(qubits[0]) * cirq.Z(qubits[1]),
-                    repetitions=1024),
-        ])
-
-        initial_point = np.array([
-            0.82311034, 0.02611798, 0.21077064, 0.61842177, 0.09828447,
-            0.62013131
-        ])
-
-        result = spsa_minimizer.minimize(loss_function_with_model_parameters(
-            model, lambda x, y: x[0][0],
-            util.convert_to_tensor([cirq.Circuit()]), None),
-                                         initial_point,
-                                         max_iterations=100)
-
-        self.assertTrue(result.converged)
-        self.assertLess(result.objective_value.numpy(), -0.95)
+        self.assertEqual(result.num_iterations,
+                         50)  # 50 is the default max iteration
 
     def test_keras_model_optimization(self):
         """Optimizate a PQC based keras model."""
@@ -254,16 +153,16 @@ def convert_to_circuit(input_data):
             tf.keras.layers.Input(shape=(), dtype=tf.string),
             # The PQC layer returns the expected value of the
             # readout gate, range [-1,1].
-            pqc.PQC(circuit, 2 * cirq.Z(q1)),
+            pqc.PQC(circuit, cirq.Z(q1)),
         ])
 
         # Initial guess of the parameter from random number
-        result = spsa_minimizer.minimize(
+        result = rotosolve_minimizer.minimize(
             loss_function_with_model_parameters(model, tf.keras.losses.Hinge(),
                                                 x_circ, y),
             tf.random.uniform(shape=[2]) * 2 * np.pi)
 
-        self.assertAlmostEqual(result.objective_value.numpy(), 0)
+        self.assertAlmostEqual(result.objective_value, 0)
         self.assertTrue(result.converged)
 
 

From de57af9330ac7f5bf8298d202081c29fb2a956b8 Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Thu, 13 Jan 2022 15:19:41 -0700
Subject: [PATCH 37/38] Update spsa_minimizer_test.py

---
 tensorflow_quantum/python/optimizers/spsa_minimizer_test.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py b/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py
index 2aef27eb3..1de86d86c 100644
--- a/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py
+++ b/tensorflow_quantum/python/optimizers/spsa_minimizer_test.py
@@ -159,7 +159,7 @@ def block_func1(params):
 
         init = 1 / 6 * 0.49
         self.incr = 0
-        
+
         def block_func2(params):
             self.incr += init
             return self.incr

From 9bd74f0844aa84595d90d7f0adec54e0bd6b98f6 Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Tue, 18 Jan 2022 13:54:42 -0700
Subject: [PATCH 38/38] Update spsa_minimizer.py

---
 .../python/optimizers/spsa_minimizer.py              | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/tensorflow_quantum/python/optimizers/spsa_minimizer.py b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
index 7cd4def60..24f3a5428 100644
--- a/tensorflow_quantum/python/optimizers/spsa_minimizer.py
+++ b/tensorflow_quantum/python/optimizers/spsa_minimizer.py
@@ -159,10 +159,10 @@ def minimize(expectation_value_function,
         tolerance: Scalar `tf.Tensor` of real dtype. Specifies the tolerance
             for the procedure. If the supremum norm between two iteration
             vector is below this number, the algorithm is stopped.
-        a: Scalar `tf.Tensor` of real dtype. Specifies the learning rate
+        lr: Scalar `tf.Tensor` of real dtype. Specifies the learning rate
         alpha: Scalar `tf.Tensor` of real dtype. Specifies scaling of the
             learning rate.
-        c: Scalar `tf.Tensor` of real dtype. Specifies the size of the
+        perturb: Scalar `tf.Tensor` of real dtype. Specifies the size of the
             perturbations.
         gamma: Scalar `tf.Tensor` of real dtype. Specifies scaling of the
             size of the perturbations.
@@ -171,6 +171,8 @@ def minimize(expectation_value_function,
         allowed_increase: Scalar `tf.Tensor` of real dtype. Specifies maximum
             allowable increase in objective function (only applies if blocking
             is true).
+        seed: (Optional) Python integer. Used to create a random seed for the
+            perturbations.
         name: (Optional) Python `str`. The name prefixed to the ops created
             by this function. If not supplied, the default name 'minimize'
             is used.
@@ -182,7 +184,9 @@ def minimize(expectation_value_function,
 
     with tf.name_scope(name or 'minimize'):
         if seed is not None:
-            tf.random.set_seed(seed)
+            generator = tf.random.Generator.from_seed(seed)
+        else:
+            generator = tf.random
 
         initial_position = tf.convert_to_tensor(initial_position,
                                                 name='initial_position',
@@ -210,7 +214,7 @@ def _spsa_once(state):
                 states: A list which the first element is the new state
             """
             delta_shift = tf.cast(
-                2 * tf.random.uniform(shape=state.position.shape,
+                2 * generator.uniform(shape=state.position.shape,
                                       minval=0,
                                       maxval=2,
                                       dtype=tf.int32) - 1, tf.float32)