merlin.pcvl_pytorch.slos_torchscript module

This module extends slos_torch.py with TorchScript-optimized computation graphs for photonic quantum circuit simulations. It separates the graph construction from the actual computation for improved performance.

The optimized implementation pre-builds the computation graph based on the input state configuration, which can then be reused for multiple unitary evaluations.

class merlin.pcvl_pytorch.slos_torchscript.Callable

Bases: object

class merlin.pcvl_pytorch.slos_torchscript.SLOSComputeGraph(m, n_photons, output_map_func=None, no_bunching=True, keep_keys=True, device=None, dtype=torch.float32, index_photons=None)

Bases: object

A class that builds and stores the computation graph for SLOS algorithm.

This separates the graph construction (which depends only on input state, no_bunching, and output_map_func) from the actual computation using the unitary matrix.

compute(unitary, input_state)

Compute the probability distribution using the pre-built graph.

Return type:

tuple[list[tuple[int, ...]], Tensor]

Args:

unitary (torch.Tensor): Single unitary matrix [m x m] or batch of unitaries [b x m x m]. The unitary should be provided in the complex dtype corresponding to the graph’s dtype. For example, for torch.float32, use torch.cfloat; for torch.float64, use torch.cdouble. input_state (list[int]): Input_state of length self.m with self.n_photons in the input state

Returns:
Tuple[List[Tuple[int, …]], torch.Tensor]:

  • List of tuples representing output Fock state configurations

  • Probability distribution tensor

compute_pa_inc(unitary, input_state_prev, contributions, input_state)
Return type:

tuple[list[tuple[int, ...]], Tensor]

to(dtype, device)

Moves the converter to a specific device.

Parameters:

  • dtype (dtype) – The data type to use for the tensors - one can specify either a float or complex dtype. Supported dtypes are torch.float32 or torch.complex64, torch.float64 or torch.complex128.

  • device (str | device) – The device to move the converter to.

merlin.pcvl_pytorch.slos_torchscript.build_slos_distribution_computegraph(m, n_photons, output_map_func=None, no_bunching=False, keep_keys=True, device=None, dtype=torch.float32, index_photons=None)

Build a computation graph for Strong Linear Optical Simulation (SLOS) algorithm that can be reused for multiple unitaries.

[existing docstring…]

Return type:

SLOSComputeGraph

merlin.pcvl_pytorch.slos_torchscript.compute_slos_distribution(unitary, input_state, output_map_func=None, no_bunching=False, keep_keys=True, index_photons=None)

TorchScript-optimized version of pytorch_slos_output_distribution.

This function builds the computation graph first, then uses it to compute the probabilities. For repeated calculations with the same input configuration but different unitaries, it’s more efficient to use build_slos_compute_graph() directly.

Return type:

tuple[list[tuple[int, ...]], Tensor]

Args:

unitary (torch.Tensor): Single unitary matrix [m x m] or batch of unitaries [b x m x m] input_state (List[int]): Number of photons in every mode of the circuit output_map_func (callable, optional): Function that maps output states no_bunching (bool): If True, the algorithm is optimized for no-bunching states only keep_keys (bool): If True, output state keys are returned index_photons: List of tuples (first_integer, second_integer). The first_integer is the lowest index layer a photon can take and the second_integer is the highest index

Returns:
Tuple[List[Tuple[int, …]], torch.Tensor]:

  • List of tuples representing output Fock state configurations

  • Probability distribution tensor

merlin.pcvl_pytorch.slos_torchscript.layer_compute_backward(unitary, contributions, sources, modes, p)

Compute amplitudes for a single layer using vectorized operations.

Return type:

Tensor

Args:

unitary: Batch of unitary matrices [batch_size, m, m] prev_amplitudes: Previous layer amplitudes [batch_size, prev_size] sources: Source indices for operations [num_ops] destinations: Destination indices for operations [num_ops] modes: Mode indices for operations [num_ops] p: Photon index for this layer

Returns:

Next layer amplitudes [batch_size, next_size]

merlin.pcvl_pytorch.slos_torchscript.layer_compute_vectorized(unitary, prev_amplitudes, sources, destinations, modes, p, return_contributions=False)

Compute amplitudes for a single layer using vectorized operations.

Return type:

Tensor

Args:

unitary: Batch of unitary matrices [batch_size, m, m] prev_amplitudes: Previous layer amplitudes [batch_size, prev_size] sources: Source indices for operations [num_ops] destinations: Destination indices for operations [num_ops] modes: Mode indices for operations [num_ops] p: Photon index for this layer

Returns:

Next layer amplitudes [batch_size, next_size]

merlin.pcvl_pytorch.slos_torchscript.load_slos_distribution_computegraph(path)

Load a previously saved SLOS distribution computation graph.

Args:

path: Path to the saved computation graph

Returns:

SLOSComputeGraph: Loaded computation graph ready for computations

Example:
>>> # Save a computation graph
>>> graph = build_slos_distribution_computegraph([1, 1])
>>> graph.save("hom_graph.pt")
>>>
>>> # Later, load the saved graph
>>> loaded_graph = load_slos_distribution_computegraph("hom_graph.pt")
>>>
>>> # Use the loaded graph
>>> unitary = torch.tensor([[0.7071, 0.7071], [0.7071, -0.7071]], dtype=torch.cfloat)
>>> keys, probs = loaded_graph.compute(unitary)
merlin.pcvl_pytorch.slos_torchscript.prepare_vectorized_operations(operations_list, device=None)

Convert operations list to tensors for vectorized computation.

Args:

operations_list: List of operations, each as [src_idx, dest_idx, mode_i] device: Optional device to place tensors on (defaults to CPU if None)

Returns:

Tuple of tensors: (sources, destinations, modes)