Skip to content

Autodiff

autodiff

circuit_to_energy_fn 

circuit_to_energy_fn(
    circuit: QASMCircuit, n_qubits: int
) -> Tuple[Callable, int]

Convert a QASMCircuit into a JAX-differentiable energy function.

QASMCircuit — from QASMParser.parse(qasm), or from the

Qiskit/PennyLane interop bridge (from_qiskit/from_pennylane).

Returns (energy_fn, n_params): energy_fn(theta, h_matrix, stato_zero=None, noise=None) -> (energy, statevector) is a pure JAX function, differentiable w.r.t. theta via jax.grad / jax.value_and_grad(energy_fn, argnums=0, has_aux=True). stato_zero defaults to |0...0> if not given. n_params is the number of parametric gates in the circuit, in the same order theta is injected — build theta as an array of that length.

noise, when given, is a registry.NoiseSpec (a JAX PyTree) applied to the statevector right after the circuit and before the energy expectation value is computed — natively inside the same traced computation as theta, not as an external step the caller has to splice in around energy_fn themselves. Because NoiseSpec carries its own jax_key as a pytree leaf, the whole thing stays jit/grad/vmap-composable with no OS-entropy fallback and no external key-management workaround:

  noise = NoiseSpec(model='depolarizing', p=0.05,
                     jax_key=jax.random.PRNGKey(0))
  energy, sv = energy_fn(theta, h_matrix, noise=noise)

This is the same engine dashboard_core.py's real VQE gradient uses internally (verified against finite differences, ~1e-11 agreement) — exposed here as public API so it's reachable without reading dashboard_core.py, and so circuits imported via from_qiskit/ from_pennylane (which are NOT differentiable on their own — see run_pennylane_circuit's docstring) have a real way to become differentiable instead of just a documented dead end.

Source code in dense_evolution/autodiff.py 
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
def circuit_to_energy_fn(
    circuit: QASMCircuit, n_qubits: int
) -> Tuple[Callable, int]:
    """
    Convert a QASMCircuit into a JAX-differentiable energy function.

    circuit : QASMCircuit — from QASMParser.parse(qasm), or from the
              Qiskit/PennyLane interop bridge (from_qiskit/from_pennylane).

    Returns (energy_fn, n_params):
      energy_fn(theta, h_matrix, stato_zero=None, noise=None) ->
      (energy, statevector) is a pure JAX function, differentiable w.r.t.
      theta via jax.grad / jax.value_and_grad(energy_fn, argnums=0,
      has_aux=True). stato_zero defaults to |0...0> if not given.
      n_params is the number of parametric gates in the circuit, in the
      same order theta is injected — build theta as an array of that
      length.

      noise, when given, is a registry.NoiseSpec (a JAX PyTree) applied
      to the statevector right after the circuit and before the energy
      expectation value is computed — natively inside the same traced
      computation as theta, not as an external step the caller has to
      splice in around energy_fn themselves. Because NoiseSpec carries
      its own jax_key as a pytree leaf, the whole thing stays
      jit/grad/vmap-composable with no OS-entropy fallback and no
      external key-management workaround:

          noise = NoiseSpec(model='depolarizing', p=0.05,
                             jax_key=jax.random.PRNGKey(0))
          energy, sv = energy_fn(theta, h_matrix, noise=noise)

    This is the same engine dashboard_core.py's real VQE gradient uses
    internally (verified against finite differences, ~1e-11 agreement) —
    exposed here as public API so it's reachable without reading
    dashboard_core.py, and so circuits imported via from_qiskit/
    from_pennylane (which are NOT differentiable on their own — see
    run_pennylane_circuit's docstring) have a real way to become
    differentiable instead of just a documented dead end.
    """
    _require_jax()
    template = _build_template(circuit, n_qubits)
    n_params = sum(1 for op in circuit.ops
                   if str(op['name']).lower().strip() in _PARAMETRIC_GATES)

    def energy_fn(theta, h_matrix, stato_zero: Optional["jnp.ndarray"] = None,
                  noise: Optional["NoiseSpec"] = None):
        if stato_zero is None:
            stato_zero = jnp.zeros(2 ** n_qubits, dtype=jnp.complex128).at[0].set(1.0)

        if n_params == 0:
            # No parametric gates -> no sentinel (-1.0) rows in template, so
            # patch_and_apply below would never take its is_param branch.
            # Skip the scan entirely rather than index into an empty theta
            # array during tracing (n_params is a static Python int, fixed
            # at circuit_to_energy_fn() call time, so this branch is
            # resolved before any tracing happens — not a jax.lax.cond).
            sv = _compile_and_run_circuit_jit(stato_zero, template)
        else:
            def patch_and_apply(carry, op):
                idx = carry
                is_param = op[3] == -1.0
                final_p = jnp.where(is_param, theta[idx], op[3])
                next_idx = jnp.where(is_param, idx + jnp.int32(1), idx)
                return next_idx, jnp.array([op[0], op[1], op[2], final_p], dtype=jnp.float64)

            _, patched_ops = jax.lax.scan(patch_and_apply, jnp.int32(0), template)
            sv = _compile_and_run_circuit_jit(stato_zero, patched_ops)

        if noise is not None:
            sv = NoiseModel.apply_to_sv(
                sv, n_qubits, model=noise.model, p=noise.p,
                jax_key=noise.jax_key, qubits=list(noise.qubits) if noise.qubits else None,
            )

        energy = jnp.real(jnp.vdot(sv, h_matrix @ sv))
        return energy, sv

    return energy_fn, n_params