{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Distributed Circuit Simulation and TensorNetwork Contraction\n", "\n", "## Overview\n", "\n", "Simulating large quantum circuits or computing expectation values for complex Hamiltonians often involves contracting a massive tensor network. The computational cost (both time and memory) of this contraction can be a significant bottleneck, especially for systems with many qubits.\n", "\n", "TensorCircuit provides an experimental feature, `DistributedContractor`, designed to tackle this challenge. It leverages multiple devices (e.g., GPUs) to parallelize the tensor network contraction. The core idea is:\n", "\n", "1. **Pathfinding with `cotengra`**: It first uses the powerful `cotengra` library to find an optimal or near-optimal contraction path for the given tensor network. This path often involves \"slicing\" the network, which breaks the single large contraction into many smaller, independent contractions.\n", "2. **Distributed Computation**: It then distributes these smaller contraction tasks across all available devices. Each device computes a subset of the slices in parallel.\n", "3. **Aggregation**: Finally, the results from all devices are aggregated to produce the final value (e.g., an expectation value or a state amplitude).\n", "\n", "This approach allows us to tackle much larger problems than would be possible on a single device, significantly reducing the wall-clock time for expensive computations.\n", "\n", "In this tutorial, we will demonstrate how to use `DistributedContractor` for two common tasks:\n", "- Calculating the amplitude of a specific bitstring for a large quantum state.\n", "- Running a Variational Quantum Eigensolver (VQE) optimization for a transverse-field Ising model." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Setup\n", "\n", "First, let's configure JAX to use multiple (virtual) devices and import the necessary libraries.\n" ] }, { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "JAX is using 4 devices.\n" ] } ], "source": [ "import os\n", "\n", "# Set this for multiple virtual devices\n", "NUM_DEVICES = 4\n", "os.environ[\"XLA_FLAGS\"] = f\"--xla_force_host_platform_device_count={NUM_DEVICES}\"\n", "\n", "import time\n", "import jax\n", "from jax import numpy as jnp\n", "import numpy as np\n", "import optax\n", "import tensorcircuit as tc\n", "from tensorcircuit.experimental import DistributedContractor\n", "\n", "K = tc.set_backend(\"jax\")\n", "tc.set_dtype(\"complex64\")\n", "\n", "# Verify that JAX sees the configured number of devices\n", "print(f\"JAX is using {jax.local_device_count()} devices.\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## `DistributedContractor`: Mechanism\n", "\n", "Before diving into examples, let's understand the core components and the inner workings of the `DistributedContractor`.\n", "\n", "### The `nodes_fn` Template\n", "\n", "The central requirement for using `DistributedContractor` is to provide a function, which we conventionally call `nodes_fn`. This function serves as a template that defines the structure of the tensor network.\n", "\n", "- **Input**: `nodes_fn` must accept a single argument, typically a dictionary or a JAX PyTree of parameters (`params`). These parameters are what you would typically vary in your computation (e.g., the variational parameters of a circuit).\n", "- **Output**: It must return the list of tensors nodes (`tc.Gate` or `tn.Node` objects, which contain tensors) that constitute the tensor network *before the final contraction*. `tensorcircuit-ng` provides convenient methods like `.expectation_before()` and `.amplitude_before()` for this purpose.\n", "\n", "The `DistributedContractor` calls this `nodes_fn` once during its initialization (`__init__`) with a set of template parameters. It does this to understand the network's connectivity and size, which is necessary for the `cotengra` pathfinder. The *actual values* in the tensors from this initial call are discarded; only the *structure* is used.\n", "\n", "### The Internal Mechanism\n", "\n", "Here's a step-by-step breakdown of what happens inside `DistributedContractor`:\n", "\n", "1. **Initialization (`__init__`)**:\n", " - You provide the `nodes_fn` and a set of `params`.\n", " - `DistributedContractor` calls `nodes_fn(params)` to get the tensor network structure.\n", " - It passes this structure to `cotengra`'s `ReusableHyperOptimizer`.\n", " - **Pathfinding**: `cotengra` then performs an exhaustive search for an efficient contraction path. A key part of this is **slicing**. If the largest intermediate tensor in the best path exceeds a memory limit (which you can control via `cotengra_options`), `cotengra` will \"slice\" one or more of the largest tensor edges. Slicing means the contraction is repeated for each possible value of the sliced indices, and the results are summed up. This trades a massive increase in computation for a drastic reduction in memory.\n", " - The final output of this step is a `ContractionTree`, a plan that details the sequence of pairwise contractions and the slicing strategy.\n", " - **Task Distribution**: The contractor then divides the list of slices evenly among the available JAX devices.\n", "\n", "2. **Execution (`.value()` or `.value_and_grad()`)**:\n", " - You call the method with a *new* set of `params`.\n", " - The contractor uses JAX's `pmap` to send the contraction plan and the new `params` to all devices.\n", " - **Parallel Execution**: Each device, in parallel:\n", " - Calls your `nodes_fn` with the new `params` to generate the tensors with their *current numerical values*.\n", " - Iterates through its assigned subset of slices.\n", " - For each slice, it performs the small, memory-efficient contraction as prescribed by the `ContractionTree`.\n", " - It sums up the results of all its assigned slices.\n", " - **Aggregation**: The final results from each device are postprocessed via `op` function which can provided in `value()` and `value_and_grad()` methods and summed up on the host to produce the total value. If `.value_and_grad()` was called, the gradients are also aggregated in the same way.\n", "\n", "This design is powerful because the most expensive step—pathfinding—is done only once. All subsequent calls with different parameters reuse the same optimized path, leading to very fast execution, especially in iterative algorithms like VQE." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Example 1: Calculating State Amplitudes\n", "\n", "One fundamental task is to compute the amplitude of a specific basis state, $\\langle s | \\psi \\rangle$, where $|s\\rangle$ is a bitstring like $|0110\\dots\\rangle$ and $|\\psi\\rangle$ is the state produced by a quantum circuit.\n", "\n", "### Defining the `nodes_fn`\n" ] }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [], "source": [ "N_QUBITS_AMP = 14\n", "DEPTH_AMP = 7\n", "\n", "\n", "def circuit_ansatz(n, d, params):\n", " \"\"\"A standard hardware-efficient ansatz.\"\"\"\n", " c = tc.Circuit(n)\n", " c.h(range(n))\n", " for i in range(d):\n", " for j in range(0, n - 1):\n", " c.rzz(j, j + 1, theta=params[j, i, 0])\n", " for j in range(n):\n", " c.rx(j, theta=params[j, i, 1])\n", " for j in range(n):\n", " c.ry(j, theta=params[j, i, 2])\n", " return c\n", "\n", "\n", "def get_nodes_fn_amp(n, d):\n", " \"\"\"\n", " This function returns another function that will be the input for DistributedContractor.\n", " The inner function takes a dictionary of parameters and returns the tensor for a single amplitude.\n", " \"\"\"\n", "\n", " def nodes_fn(params):\n", " psi = circuit_ansatz(n, d, params[\"circuit\"])\n", " # `amplitude_before` gives us the tensor network before final contraction\n", " return psi.amplitude_before(params[\"amplitude\"])\n", "\n", " return nodes_fn\n", "\n", "\n", "def get_binary_representation(i: int, N: int) -> jax.Array:\n", " \"\"\"Helper function to convert an integer to its binary representation.\"\"\"\n", " shifts = jnp.arange(N - 1, -1, -1)\n", " return ((i >> shifts) & 1).astype(jnp.int32)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Initializing the `DistributedContractor`" ] }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Initializing DistributedContractor for amplitude calculation...\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "F=4.88 C=6.02 S=9.00 P=11.21: 100%|██████████| 64/64 [00:08<00:00, 7.70it/s] " ] }, { "name": "stdout", "output_type": "stream", "text": [ "\n", "--- Contraction Path Info ---\n", "Path found with 1 slices.\n", "Arithmetic Intensity (higher is better): 4.94\n", "flops (TFlops): 1.7143975128419697e-08\n", "write (GB): 0.00011376291513442993\n", "size (GB): 3.814697265625e-06\n", "-----------------------------\n", "\n", "Distributing across 4 devices. Each device will sequentially process up to 1 slices.\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\n" ] } ], "source": [ "nodes_fn_amp = get_nodes_fn_amp(N_QUBITS_AMP, DEPTH_AMP)\n", "\n", "# We need some initial parameters to define the network structure\n", "key = jax.random.PRNGKey(42)\n", "params_circuit_amp = (\n", " jax.random.normal(key, shape=[N_QUBITS_AMP, DEPTH_AMP, 3], dtype=tc.rdtypestr) * 0.1\n", ")\n", "initial_params_amp = {\n", " \"circuit\": params_circuit_amp,\n", " \"amplitude\": get_binary_representation(0, N_QUBITS_AMP),\n", "}\n", "\n", "print(\"Initializing DistributedContractor for amplitude calculation...\")\n", "# cotengra_options allow fine-tuning of the pathfinding process.\n", "# `target_size` in `slicing_reconf_opts` controls the memory size of each slice.\n", "DC_amp = DistributedContractor(\n", " nodes_fn=nodes_fn_amp,\n", " params=initial_params_amp,\n", " cotengra_options={\n", " \"slicing_reconf_opts\": {\"target_size\": 2**14},\n", " \"max_repeats\": 64,\n", " \"progbar\": True,\n", " \"minimize\": \"write\", # Optimizes for memory write operations\n", " \"parallel\": 4,\n", " },\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Calculating Multiple Amplitudes\n" ] }, { "cell_type": "code", "execution_count": 5, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Starting amplitude loop...\n", "Bitstring: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] | Amp (DC): 0.00353913+0.00129915j | Time: 0.0011 s\n", "Bitstring: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1] | Amp (DC): 0.00360568+0.00364827j | Time: 0.0007 s\n", "Bitstring: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0] | Amp (DC): 0.00032678+0.00287017j | Time: 0.0007 s\n", "Bitstring: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1] | Amp (DC): 0.00228170+0.00377162j | Time: 0.0006 s\n", "Bitstring: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0] | Amp (DC): 0.00194864+0.00337264j | Time: 0.0006 s\n", "Bitstring: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1] | Amp (DC): 0.00013618+0.00505736j | Time: 0.0005 s\n", "Bitstring: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0] | Amp (DC): 0.00230806+0.00277415j | Time: 0.0006 s\n", "Bitstring: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1] | Amp (DC): 0.00485482+0.00290096j | Time: 0.0005 s\n", "Bitstring: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0] | Amp (DC): 0.00410894+0.00097546j | Time: 0.0005 s\n", "Bitstring: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1] | Amp (DC): 0.00449419+0.00358756j | Time: 0.0006 s\n" ] } ], "source": [ "n_amp = 10\n", "print(\"Starting amplitude loop...\")\n", "for i in range(n_amp):\n", " bs_vector = get_binary_representation(i, N_QUBITS_AMP)\n", " params = {\"circuit\": params_circuit_amp, \"amplitude\": bs_vector}\n", "\n", " t0 = time.time()\n", " amp = DC_amp.value(params)\n", " t1 = time.time()\n", "\n", " print(\n", " f\"Bitstring: {bs_vector.tolist()} | \"\n", " f\"Amp (DC): {amp:.8f} | \"\n", " f\"Time: {t1 - t0:.4f} s\"\n", " )" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Example 2: Variational Quantum Eigensolver (VQE)\n", "\n", "`DistributedContractor` is especially powerful for variational algorithms, where we need to repeatedly compute expectation values and their gradients.\n", "\n", "### Defining the `nodes_fn` for Expectation Value" ] }, { "cell_type": "code", "execution_count": 6, "metadata": {}, "outputs": [], "source": [ "N_QUBITS_VQE = 10\n", "DEPTH_VQE = 4\n", "\n", "\n", "def get_tfi_mpo(n):\n", " \"\"\"Gets the MPO for the 1D Transverse-Field Ising model Hamiltonian.\"\"\"\n", " import tensornetwork as tn\n", "\n", " Jx = np.ones(n - 1)\n", " Bz = -1.0 * np.ones(n)\n", " tn_mpo = tn.matrixproductstates.mpo.FiniteTFI(Jx, Bz, dtype=np.complex64)\n", " return tc.quantum.tn2qop(tn_mpo)\n", "\n", "\n", "def get_nodes_fn_vqe(n, d, mpo):\n", " \"\"\"\n", " The nodes_fn for VQE expectation value.\n", " It returns the list of tensors for .\n", " \"\"\"\n", "\n", " def nodes_fn(params):\n", " psi = circuit_ansatz(n, d, params).get_quvector()\n", " expression = psi.adjoint() @ mpo @ psi\n", " return expression.nodes\n", "\n", " return nodes_fn" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### VQE Optimization Loop\n" ] }, { "cell_type": "code", "execution_count": 8, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "\n", "Initializing DistributedContractor for VQE...\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "F=5.61 C=6.81 S=7.58 P=11.10 $=16.00: 100%|██████████| 16/16 [00:08<00:00, 1.98it/s]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "\n", "--- Contraction Path Info ---\n", "Path found with 16 slices.\n", "Arithmetic Intensity (higher is better): 4.30\n", "flops (TFlops): 9.178620530292392e-08\n", "write (GB): 4.3742358684539795e-05\n", "size (GB): 1.430511474609375e-06\n", "-----------------------------\n", "\n", "Distributing across 4 devices. Each device will sequentially process up to 4 slices.\n", "\n", "Starting VQE optimization loop...\n", "Step 010 | Loss: -3.29106593 | Time: 0.0022 s\n", "Step 020 | Loss: -8.78426552 | Time: 0.0021 s\n", "Step 030 | Loss: -10.85906601 | Time: 0.0022 s\n", "Step 040 | Loss: -11.47075844 | Time: 0.0020 s\n", "Step 050 | Loss: -11.72393227 | Time: 0.0022 s\n", "Step 060 | Loss: -11.91652107 | Time: 0.0021 s\n", "Step 070 | Loss: -12.04074574 | Time: 0.0020 s\n", "Step 080 | Loss: -12.11665630 | Time: 0.0024 s\n", "Step 090 | Loss: -12.15851784 | Time: 0.0021 s\n", "Step 100 | Loss: -12.17830086 | Time: 0.0022 s\n", "\n", "Optimization finished.\n" ] } ], "source": [ "tfi_mpo = get_tfi_mpo(N_QUBITS_VQE)\n", "nodes_fn_vqe = get_nodes_fn_vqe(N_QUBITS_VQE, DEPTH_VQE, tfi_mpo)\n", "\n", "# Initial parameters for VQE\n", "key = jax.random.PRNGKey(42)\n", "params_vqe = (\n", " jax.random.normal(key, shape=[N_QUBITS_VQE, DEPTH_VQE, 3], dtype=tc.rdtypestr) * 0.1\n", ")\n", "\n", "print(\"\\nInitializing DistributedContractor for VQE...\")\n", "DC_vqe = DistributedContractor(\n", " nodes_fn=nodes_fn_vqe,\n", " params=params_vqe,\n", " cotengra_options={\n", " \"slicing_reconf_opts\": {\n", " \"target_size\": 2**8\n", " }, # Smaller target size for VQE network\n", " \"max_repeats\": 16,\n", " \"progbar\": True,\n", " \"minimize\": \"write\",\n", " \"parallel\": 4,\n", " },\n", ")\n", "\n", "# Setup Optax optimizer\n", "optimizer = optax.adam(2e-2)\n", "opt_state = optimizer.init(params_vqe)\n", "\n", "\n", "@K.jit\n", "def opt_update(params, opt_state, grads):\n", " updates, new_opt_state = optimizer.update(grads, opt_state, params)\n", " new_params = optax.apply_updates(params, updates)\n", " return new_params, new_opt_state\n", "\n", "\n", "n_steps_vqe = 100\n", "print(\"\\nStarting VQE optimization loop...\")\n", "for i in range(n_steps_vqe):\n", " t0 = time.time()\n", " loss, grads = DC_vqe.value_and_grad(params_vqe)\n", "\n", " params_vqe, opt_state = opt_update(params_vqe, opt_state, grads)\n", " t1 = time.time()\n", "\n", " if (i + 1) % 10 == 0:\n", " print(f\"Step {i+1:03d} | \" f\"Loss: {loss:.8f} | \" f\"Time: {t1 - t0:.4f} s\")\n", "\n", "print(\"\\nOptimization finished.\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Conclusion\n", "\n", "The `DistributedContractor` provides a powerful and streamlined interface for scaling up tensor network contractions to multiple devices. By abstracting away the complexities of pathfinding, slicing, and parallel execution, it allows researchers to focus on the physics of their problem while leveraging the full computational power of their hardware. This is particularly advantageous for simulating large quantum circuits and accelerating the convergence of variational quantum algorithms." ] } ], "metadata": { "kernelspec": { "display_name": "2502", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.12.0" } }, "nbformat": 4, "nbformat_minor": 2 }