[Bug] Repeated PSR (#46)

* add uuid to parametric and analog

* make parameterize test + add extra op in differentiation

* _param uuid in initialization

* use suffix for gpsr

* fix analog _gpsr suffix

* fix _gpsr parametric

* adding gate-wise psr for no shots

* simplify computation no shots

* add finite shots fwd notimplementederror for repeated case

* change syntax for defjvp by using a shift argument in gates

* change test for repeated params

* tried gate ind but not working

* repeated param case working

* add test with jit

* more docstr

* fix shape jax random

* use tuple of int for values

* add chex

* using chex

* add spectral gap via gates

* adding spectral gap from gates

* also assert hessian

* rm hessian parts

* decorate hessian fcts with variants

* rm forward_mode

* rm also strenum

* hessian too unstable

* jax.hessian

* rm hessian - still unstable

* fixing shift add for hessian

* add second derivatives in test

* change naming align eigenvectors

* replace chex with if for n_shots

* refactor

* remove n_shots None and small syntax changed

* reput checkify

* adding python 3.13

---------

Co-authored-by: Charles MOUSSA <[email protected]>

Author: chMoussa

.github/workflows/run-tests-and-mypy.yml

@@ -22,7 +22,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.9", "3.10", "3.11", "3.12"]
+        python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
     steps:
         - name: Checkout main code and submodules
           uses: actions/checkout@v4

horqrux/analog.py

@@ -20,9 +20,15 @@ class _HamiltonianEvolution(Primitive):
     generator_name: str
     target: QubitSupport
     control: QubitSupport
+    shift: float = 0.0
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
 
     def _unitary(self, values: dict[str, Array] = dict()) -> Array:
-        return expm(values["hamiltonian"] * (-1j * values["time_evolution"]))
+        # note: GPSR trick when the same param_name is used in many operations
+        time_val = values["time_evolution"] + self.shift
+        return expm(values["hamiltonian"] * (-1j * time_val))
 
 
 def HamiltonianEvolution(

horqrux/api.py

@@ -1,12 +1,11 @@
 from __future__ import annotations
 
 from collections import Counter
-from typing import Any, Optional
+from typing import Any
 
 import jax
 import jax.numpy as jnp
 from jax import Array
-from jax.experimental import checkify
 
 from horqrux.composite import Observable, OpSequence
 from horqrux.differentiation.ad import ad_expectation
@@ -15,7 +14,6 @@
 from horqrux.utils import (
     DensityMatrix,
     DiffMode,
-    ForwardMode,
     State,
     num_qubits,
     probabilities,
@@ -96,8 +94,7 @@ def expectation(
     observables: list[Observable],
     values: dict[str, float],
     diff_mode: DiffMode = DiffMode.AD,
-    forward_mode: ForwardMode = ForwardMode.EXACT,
-    n_shots: Optional[int] = None,
+    n_shots: int = 0,
     key: Any = jax.random.PRNGKey(0),
 ) -> Array:
     """Run 'state' through a sequence of 'gates' given parameters 'values'
@@ -109,8 +106,7 @@ def expectation(
         observables (list[Observable]): List of observables.
         values (dict[str, float]): Parameter values.
         diff_mode (DiffMode, optional): Differentiation mode. Defaults to DiffMode.AD.
-        forward_mode (ForwardMode, optional): Type of forward method. Defaults to ForwardMode.EXACT.
-        n_shots (Optional[int], optional): Number of shots. Defaults to None.
+        n_shots (int): Number of shots. Defaults to 0 for no shots.
         key (Any, optional): Random key. Defaults to jax.random.PRNGKey(0).
 
     Returns:
@@ -123,25 +119,21 @@ def expectation(
             raise TypeError("Adjoint does not support density matrices.")
         return adjoint_expectation(state, circuit, observables, values)
     elif diff_mode == DiffMode.GPSR:
-        if forward_mode == ForwardMode.SHOTS:
-            checkify.check(
-                type(n_shots) is int and n_shots > 0,
-                "Number of shots must be an integer for finite shots.",
-            )
-            # Type checking is disabled because mypy doesn't parse checkify.check.
-            # type: ignore
-            return finite_shots_fwd(
+        if n_shots < 0:
+            raise ValueError("The number of shots should be positive.")
+        if n_shots == 0:
+            return no_shots_fwd(
                 state=state,
                 gates=circuit.operations,
                 observables=observables,
                 values=values,
-                n_shots=n_shots,
-                key=key,
             )
         else:
-            return no_shots_fwd(
+            return finite_shots_fwd(
                 state=state,
                 gates=circuit.operations,
                 observables=observables,
                 values=values,
+                n_shots=n_shots,
+                key=key,
             )