tensorcircuit.experimental

Experimental features

class tensorcircuit.experimental.DistributedContractor(nodes_fn: Callable[[Any], List[Gate]], params: Any, cotengra_options: Dict[str, Any] | None = None, devices: List[Any] | None = None, mesh: Any | None = None, tree_data: Dict[str, Any] | None = None)

Bases: object

A distributed tensor network contractor that parallelizes computations across multiple devices.

This class uses cotengra to find optimal contraction paths and distributes the computational load across multiple devices (e.g., GPUs) for efficient tensor network calculations. Particularly useful for large-scale quantum circuit simulations and variational quantum algorithms.

Example

>>> def nodes_fn(params):
...     c = tc.Circuit(4)
...     c.rx(0, theta=params[0])
...     return c.expectation_before([tc.gates.z(), [0]], reuse=False)
>>> dc = DistributedContractor(nodes_fn, params)
>>> value, grad = dc.value_and_grad(params)

Parameters:

• nodes_fn (Callable[[Tensor], List[Gate]]) – Function that takes parameters and returns a list of tensor network nodes

• params (Tensor) – Initial parameters used to determine the tensor network structure

• cotengra_options (Optional[Dict[str, Any]], optional) – Configuration options passed to the cotengra optimizer. Defaults to None

• devices (Optional[List[Any]], optional) – List of devices to use. If None, uses all available devices

• mesh (Optional[Any], optional) – Mesh object to use for distributed computation. If None, uses all available devices

__init__(nodes_fn: Callable[[Any], List[Gate]], params: Any, cotengra_options: Dict[str, Any] | None = None, devices: List[Any] | None = None, mesh: Any | None = None, tree_data: Dict[str, Any] | None = None) → None

static find_path(nodes_fn: Callable[[Any], Any], params: Any, cotengra_options: Dict[str, Any] | None = None, filepath: str | None = None) → None

classmethod from_path(filepath: str, nodes_fn: Callable[[Any], List[Gate]], devices: List[Any] | None = None, mesh: Any | None = None, params: Any | None = None) → DistributedContractor

grad(params: Any, op: Callable[[Any], Any] | None = None, output_dtype: str | None = None) → Any

value(params: Any, op: Callable[[Any], Any] | None = None, output_dtype: str | None = None) → Any

Calculates the value, compiling the pmap function for the first call.

Parameters:

• params (Tensor) – Parameters for the nodes_fn input

• op (Optional[Callable[[Tensor], Tensor]], optional) – Optional post-processing function for the output, defaults to None (corresponding to identity) op is a cache key, so dont directly pass lambda function for op

• output_dtype (Optional[str], optional) – dtype str for the output of nodes_fn, defaults to None (corresponding to dtypestr)

value_and_grad(params: Any, op: Callable[[Any], Any] | None = None, output_dtype: str | None = None) → Tuple[Any, Any]

Calculates the value and gradient, compiling the pmap function if needed for the first call.

Parameters:

• params (Tensor) – Parameters for the nodes_fn input

• op (Optional[Callable[[Tensor], Tensor]], optional) – Optional post-processing function for the output, defaults to None (corresponding to backend.real) op is a cache key, so dont directly pass lambda function for op

• output_dtype (Optional[str], optional) – dtype str for the output of nodes_fn, defaults to None (corresponding to rdtypestr)

tensorcircuit.experimental.Mesh: Any

tensorcircuit.experimental.NamedSharding: Any

tensorcircuit.experimental.P: Any

tensorcircuit.experimental.adaptive_vmap(f: Callable[[...], Any], vectorized_argnums: int | Sequence[int] = 0, static_argnums: int | Sequence[int] | None = None, chunk_size: int | None = None) → Callable[[...], Any]

Vectorized map with adaptive chunking for memory efficiency.

Parameters:

• f – Function to be vectorized

• vectorized_argnums – Arguments to be vectorized over

• static_argnums – Arguments that remain static during vectorization

• chunk_size – Size of chunks for batch processing, None means no chunking (naive vmap)

Returns:

Vectorized function

tensorcircuit.experimental.broadcast_py_object(obj: Any, shared_dir: str | None = None) → Any

Broadcast a picklable Python object from process 0 to all other processes, with fallback mechanism from gRPC to file system based approach.

This function first attempts to use gRPC-based broadcast. If that fails due to pickling issues, it falls back to a file system based approach that is more robust.

Parameters:

• obj (Any) – The Python object to broadcast. It must be picklable. This object should exist on process 0 and can be None on others.

• shared_dir (Optional[str], optional) – Directory path for shared file system broadcast fallback. If None, uses current directory. Only used in fallback mode.

Returns:

The broadcasted object, now present on all processes.

Return type:

Any

tensorcircuit.experimental.broadcast_py_object_fs(obj: Any, shared_dir: str | None = None, timeout_seconds: int = 300) → Any

Broadcast a picklable Python object from process 0 to all other processes using a shared file system approach.

This is a fallback method when gRPC-based broadcast fails. It uses UUID-based file communication to share objects between processes through a shared file system.

Parameters:

• obj (Any) – The Python object to broadcast. Must be picklable. Should exist on process 0, can be None on others.

• shared_dir (Optional[str], optional) – Directory path for shared file system communication. If None, uses current directory.

• timeout_seconds (int, optional) – Maximum time to wait for file operations before timing out. Defaults to 300 seconds.

Returns:

The broadcasted object, now present on all processes.

Return type:

Any

tensorcircuit.experimental.broadcast_py_object_jax(obj: Any) → Any

Broadcast a picklable Python object from process 0 to all other processes within jax ditribution system.

This function uses a two-step broadcast: first the size, then the data. This is necessary because broadcast_one_to_all requires the same shaped array on all hosts.

Parameters:

obj – The Python object to broadcast. It must be picklable. This object should exist on process 0 and can be None on others.

Returns:

The broadcasted object, now present on all processes.

tensorcircuit.experimental.ctg: Any

tensorcircuit.experimental.dynamics_matrix(f: Callable[[...], Any], *, kernel: str = 'dynamics', postprocess: str | None = None, mode: str = 'fwd') → Callable[[...], Any]

Compute quantum natural gradient for quantum circuit optimization.

Parameters:

• f – Function that takes parameters and returns quantum state

• kernel – Type of kernel to use (“qng” or “dynamics”), the former has the second term

• postprocess – Post-processing method (“qng” or None)

• mode – Mode of differentiation (“fwd” or “rev”)

Returns:

Function computing QNG matrix

Example:

>>> import tensorcircuit as tc
>>> def ansatz(params):
...     c = tc.Circuit(2)
...     c.rx(0, theta=params[0])
...     c.ry(1, theta=params[1])
...     return c.state()
>>> qng_fn = tc.experimental.qng(ansatz)
>>> params = tc.array_to_tensor([0.5, 0.8])
>>> qng_matrix = qng_fn(params)
>>> print(qng_matrix.shape)  # (2, 2)

tensorcircuit.experimental.dynamics_rhs(f: Callable[[...], Any], h: Any) → Callable[[...], Any]

tensorcircuit.experimental.finite_difference_differentiator(f: Callable[[...], Any], argnums: Tuple[int, ...] = (0,), shifts: Tuple[float, float] = (0.001, 0.002)) → Callable[[...], Any]

tensorcircuit.experimental.jax_func_load(filename: str) → Callable[[...], Any]

load a jitted function from file

Parameters:

filename (str) – _description_

Returns:

the loaded function

Return type:

_type_

tensorcircuit.experimental.jax_func_save(filename: str, f: Callable[[...], Any], *args: Any) → None

save a jitted jax function as a file

Parameters:

• filename (str) – _description_

• f (Callable[..., Any]) – the jitted function

• args – example function arguments for f

tensorcircuit.experimental.jax_jitted_function_load(filename: str) → Callable[[...], Any]

load a jitted function from file

Parameters:

filename (str) – _description_

Returns:

the loaded function

Return type:

_type_

tensorcircuit.experimental.jax_jitted_function_save(filename: str, f: Callable[[...], Any], *args: Any) → None

save a jitted jax function as a file

Parameters:

• filename (str) – _description_

• f (Callable[..., Any]) – the jitted function

• args – example function arguments for f

tensorcircuit.experimental.jaxlib: Any

tensorcircuit.experimental.parameter_shift_grad(f: Callable[[...], Any], argnums: int | Sequence[int] = 0, jit: bool = False, shifts: Tuple[float, float] = (1.5707963267948966, 2)) → Callable[[...], Any]

similar to grad function but using parameter shift internally instead of AD, vmap is utilized for evaluation, so the speed is still ok

Parameters:

• f (Callable[..., Tensor]) – quantum function with weights in and expectation out

• argnums (Union[int, Sequence[int]], optional) – label which args should be differentiated, defaults to 0

• jit (bool, optional) – whether jit the original function f at the beginning, defaults to False

• shifts (Tuple[float, float]) – two floats for the delta shift on the numerator and dominator, defaults to (pi/2, 2) for parameter shift

Returns:

the grad function

Return type:

Callable[…, Tensor]

tensorcircuit.experimental.parameter_shift_grad_v2(f: Callable[[...], Any], argnums: int | Sequence[int] = 0, jit: bool = False, random_argnums: Sequence[int] | None = None, shifts: Tuple[float, float] = (1.5707963267948966, 2)) → Callable[[...], Any]

similar to grad function but using parameter shift internally instead of AD, vmap is utilized for evaluation, v2 also supports random generator for finite measurememt shot, only jax backend is supported, since no vmap randomness is available in tensorflow

Parameters:

• f (Callable[..., Tensor]) – quantum function with weights in and expectation out

• argnums (Union[int, Sequence[int]], optional) – label which args should be differentiated, defaults to 0

• jit (bool, optional) – whether jit the original function f at the beginning, defaults to False

• random_argnums (Optional[Sequence[int]], optional) – label which args should be treated as random sources, including JAX PRNG keys or uniform random status tensors, defaults to None. If a JAX PRNG key is provided, it is split into independent subkeys for each shifted circuit instance evaluation. If a status tensor is provided, it should be pre-batched to match the shifted circuit instances. Specifically, if the status tensor has shape [2, size, …], the first slice [0, size, …] is used for the positive shifts and the second slice [1, size, …] is used for the negative shifts. If the status tensor has shape [size, …], it is used for both positive and negative shifts (Common Random Numbers). Here size refers to the total number of elements in the parameter tensor being differentiated (i.e., backend.sizen(args[argnum])).

• shifts (Tuple[float, float]) – two floats for the delta shift on the numerator and dominator, defaults to (pi/2, 2) for parameter shift

Returns:

the grad function

Return type:

Callable[…, Tensor]

Note

While the wrapper logic is backend-agnostic, stochastic functions (like those using sample) may only be differentiable via this wrapper on the JAX backend due to limited vmap support for such operations in other backends.

tensorcircuit.experimental.qng(f: Callable[[...], Any], kernel: str = 'qng', postprocess: str | None = 'qng', mode: str = 'fwd') → Callable[[...], Any]

Compute quantum natural gradient for quantum circuit optimization.

Parameters:

• f – Function that takes parameters and returns quantum state

• kernel – Type of kernel to use (“qng” or “dynamics”), the former has the second term

• postprocess – Post-processing method (“qng” or None)

• mode – Mode of differentiation (“fwd” or “rev”)

Returns:

Function computing QNG matrix

Example:

>>> import tensorcircuit as tc
>>> def ansatz(params):
...     c = tc.Circuit(2)
...     c.rx(0, theta=params[0])
...     c.ry(1, theta=params[1])
...     return c.state()
>>> qng_fn = tc.experimental.qng(ansatz)
>>> params = tc.array_to_tensor([0.5, 0.8])
>>> qng_matrix = qng_fn(params)
>>> print(qng_matrix.shape)  # (2, 2)

tensorcircuit.experimental.qng2(f: Callable[[...], Any], kernel: str = 'qng', postprocess: str | None = 'qng', mode: str = 'rev') → Callable[[...], Any]