tensorcircuit_backend
eigenkit.core.backends.tensorcircuit_backend
The tensorcircuit-ng backend, for tensor-network simulation in JAX.
It uses the native automatic differentiation of JAX. tensorcircuit-ng is an optional extra (pip install eigenkit[tensorcircuit]).
Its packages (tensorcircuit and JAX) are imported only when needed, so
import eigenkit does not import them. compile() raises ImportError if they
are missing.
It declares gradient_method = "adjoint": the compiled circuit has a jacobian
method, computed by reverse-mode automatic differentiation in JAX
(jax.jacrev), so the gradient code of Circuit uses it directly. Both the
NumPy variational optimisation (used by the Variational Quantum Eigensolver and
the Quantum Approximate Optimization Algorithm) and the parameter-shift fallback
of the PyTorch interface get their gradients through _gradient_jacobians,
which calls this jacobian.
The simulation contracts the tensor network exactly, through the
Circuit.expectation method of tensorcircuit. The batch axis is vectorised with
jax.vmap, and the whole forward and gradient computation is compiled with
jax.jit. This suits repeated workloads (training loops, sweeps, restarts
with an ansatz of fixed shape), where the one-time compilation is spread over
many evaluations. For a single solve, the NumPy reference simulator is faster,
because the compilation takes most of the time. The simulation is still an exact
contraction, so the gain is a constant factor: it does not make the cost
sub-exponential.
The supported gates are RX, RY, RZ, H, X, Z, CNOT (controlled-NOT)
and SWAP. To add more, extend _apply and SUPPORTED. The observables are sums
of Pauli operators. The 64-bit mode of JAX is turned on at first use, for the
complex128 statevector.
Example
from eigenkit import Circuit, Encoding, Variational, all_z
blocks = [Encoding(in_features=4), Variational(depth=2)]
# by registered name
circuit = Circuit(blocks=blocks, observables=all_z(4), provider="tensorcircuit")
# or with the config, when a device setting matters
from eigenkit.configs import TensorCircuitConfig
circuit = Circuit(
blocks=blocks,
observables=all_z(4),
provider=TensorCircuitConfig(),
)
The simulation uses tensor networks on JAX, with batching through
jax.vmap. This backend needs the tensorcircuit extra.
TensorCircuitBackend
The provider for tensorcircuit-ng, tensor-network simulation on JAX.
It is registered as "tensorcircuit" once the extra is installed. It is
internal: you select this backend with provider="tensorcircuit" or by
passing a
TensorCircuitConfig,
not by building this class.
capabilities
property
What this backend supports.
It runs the subset of gates that this backend translates to tensorcircuit
(RX, RY, RZ, H, X, Z, CNOT (controlled-NOT) and SWAP),
evaluates a batch in one call through jax.vmap, returns exact
expectation values only, and has no limit on the number of qubits. See
Capabilities.
__init__(config=None, name='tensorcircuit')
Build the backend from a device configuration.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
config
|
TensorCircuitConfig
|
The device settings. Default: |
None
|
name
|
str
|
The registry name, used in error messages. Default:
|
'tensorcircuit'
|
compile(circuit, observables)
Prepare circuit for repeated evaluation of observables.
A Circuit calls it once and reuses the result for every call, so the
expensive work belongs here.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
circuit
|
CircuitSpec
|
The circuit to prepare. |
required |
observables
|
Sequence[Observable]
|
The observables to measure. |
required |
Returns:
| Type | Description |
|---|---|
TensorCircuitCompiled
|
The prepared circuit, as a |
TensorCircuitCompiled
A circuit prepared for tensorcircuit-ng, compiled with jax.jit on first use.
It is internal:
TensorCircuitBackend.compile()
builds it and returns it, and it is never built directly. jacobian() uses
reverse-mode automatic differentiation in JAX, and the batch axis is
vectorised with jax.vmap, so the one-time compilation cost is spread over a
training loop or a sweep.
__init__(circuit, observables)
Prepare the circuit for repeated evaluation.
TensorCircuitBackend.compile() builds this object, so there is no reason
to build one directly.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
circuit
|
CircuitSpec
|
The circuit to prepare. |
required |
observables
|
Sequence[Observable]
|
The observables to measure. |
required |
jacobian(weights, inputs, *, need_input_grad=False)
Compute the jacobians with reverse-mode automatic differentiation in JAX.
The shapes are the same as for the adjoint method of the NumPy simulator:
(n_obs, n_weights) and (n_obs, n_inputs), with a leading batch axis
for 2-D inputs.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
weights
|
ndarray
|
The parameter vector (1-D). |
required |
inputs
|
ndarray
|
The input features: 1-D for one sample, or 2-D with a leading batch axis. |
required |
need_input_grad
|
bool
|
If True, also compute the jacobian with
respect to the inputs. Default: |
False
|
Returns:
| Name | Type | Description |
|---|---|---|
weight_jacobian |
ndarray
|
The jacobian with respect to the weights. |
input_jacobian |
ndarray | None
|
The jacobian with respect to the
inputs, or |
probabilities(weights, inputs)
Return the full measurement distribution over the 2**n basis states.
Qubit 0 is the most significant bit. The Quantum Approximate
Optimization Algorithm (QAOA) uses it to decode its result. It builds the
dense state, so it is only practical for a small number of qubits.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
weights
|
ndarray
|
The parameter vector (1-D). |
required |
inputs
|
ndarray
|
The input features of one sample (1-D). |
required |
Returns:
| Type | Description |
|---|---|
ndarray
|
The probabilities, of shape |
run(weights, inputs)
Evaluate the circuit and return one expectation value per observable.
weights must be 1-D. The result has shape (n_obs,) when inputs is
1-D, and (B, n_obs) when inputs is 2-D.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
weights
|
ndarray
|
The parameter vector (1-D). |
required |
inputs
|
ndarray
|
The input features: 1-D for one sample, or 2-D with a leading batch axis. |
required |
Returns:
| Type | Description |
|---|---|
ndarray
|
The expectation values, one for each observable. |
TensorCircuitConfig
dataclass
Configuration for the tensorcircuit backend.
It has no settings yet: the contraction path is the default of tensorcircuit,
and the batch axis is vectorised with jax.vmap. It exists for consistency
with the configurations of the other backends, and so that
provider=TensorCircuitConfig() works without a register call.
make_backend()
Build the backend that this configuration describes.
The provider layer calls it when a configuration is passed as provider=.
There is no reason to call it directly.
Returns:
| Type | Description |
|---|---|
'TensorCircuitBackend'
|
The |