pyqrack-complex128 1.78.3__py3-none-macosx_14_0_arm64.whl → 1.82.0__py3-none-macosx_14_0_arm64.whl

pyqrack/qrack_neuron_torch_layer.py

@@ -1,4 +1,4 @@
-# (C) Daniel Strano and the Qrack contributors 2017-2025. All rights reserved.
+# (C) Daniel Strano and the Qrack contributors 2017-2026. All rights reserved.
 #
 # Initial draft by Elara (OpenAI custom GPT)
 # Refined and architecturally clarified by Dan Strano
@@ -6,8 +6,10 @@
 # Use of this source code is governed by an MIT-style license that can be
 # found in the LICENSE file or at https://opensource.org/licenses/MIT.
 
+import ctypes
 import itertools
 import math
+import random
 import sys
 
 _IS_TORCH_AVAILABLE = True
@@ -21,102 +23,90 @@ except ImportError:
 from .pauli import Pauli
 from .qrack_neuron import QrackNeuron
 from .qrack_simulator import QrackSimulator
+from .qrack_system import Qrack
 from .neuron_activation_fn import NeuronActivationFn
 
 
-# Should be safe for 16-bit
-angle_eps = math.pi * (2 ** -8)
+# Parameter-shift rule
+param_shift_eps = math.pi / 2
+# Neuron angle initialization
+init_phi = math.asin(0.5)
+# Systemic floating-point type
+fp_type = ctypes.c_float if Qrack.fppow <= 5 else ctypes.c_double
 
 
-if not _IS_TORCH_AVAILABLE:
-    class TorchContextMock(object):
-        def __init__(self):
-            pass
-
-        def save_for_backward(self, *args):
-            self.saved_tensors = args
-
 class QrackNeuronTorchFunction(Function if _IS_TORCH_AVAILABLE else object):
     """Static forward/backward/apply functions for QrackNeuronTorch"""
 
-    if not _IS_TORCH_AVAILABLE:
-        @staticmethod
-        def apply(x, neuron_wrapper):
-            return forward(TorchContextMock(), x, neuron_wrapper)
-
     @staticmethod
-    def forward(ctx, x, neuron_wrapper):
-        ctx.neuron_wrapper = neuron_wrapper
+    def forward(ctx, x, neuron):
+        ctx.neuron = neuron
+        ctx.simulator = neuron.simulator
         ctx.save_for_backward(x)
-        neuron = neuron_wrapper.neuron
 
-        angles = (x.detach().cpu().numpy() if x.requires_grad else x.numpy()) if _IS_TORCH_AVAILABLE else x
-        neuron.set_angles(angles)
+        # Baseline probability BEFORE applying this neuron's unitary
+        pre_prob = neuron.simulator.prob(neuron.target)
+
+        angles = x.detach().cpu().numpy() if x.requires_grad else x.numpy()
+        neuron.angles = angles.ctypes.data_as(ctypes.POINTER(fp_type))
         neuron.predict(True, False)
+
+        # Probability AFTER applying this neuron's unitary
         post_prob = neuron.simulator.prob(neuron.target)
-        if _IS_TORCH_AVAILABLE:
-            post_prob = torch.tensor([post_prob], dtype=torch.float32, device=x.device)
+        ctx.post_prob = post_prob
+
+        delta = math.asin(post_prob) - math.asin(pre_prob)
+        ctx.delta = delta
 
-        return post_prob
+        # Return shape: (1,)
+        return x.new_tensor([delta])
 
     @staticmethod
-    def _backward(x, neuron_wrapper):
-        neuron = neuron_wrapper.neuron
-        angles = (x.detach().cpu().numpy() if x.requires_grad else x.numpy()) if _IS_TORCH_AVAILABLE else x
+    def backward(ctx, grad_output):
+        (x,) = ctx.saved_tensors
+        neuron = ctx.neuron
+        neuron.set_simulator(ctx.simulator)
+        post_prob = ctx.post_prob
+
+        angles = x.detach().cpu().numpy() if x.requires_grad else x.numpy()
 
-        # Uncompute
-        neuron.set_angles(angles)
+        # Restore simulator to state BEFORE this neuron's unitary
+        neuron.angles = angles.ctypes.data_as(ctypes.POINTER(fp_type))
         neuron.unpredict()
         pre_sim = neuron.simulator
-        pre_prob = pre_sim.prob(neuron.target)
 
-        param_count = 1 << len(neuron.controls)
-        delta = [0.0] * param_count
-        for param in range(param_count):
-            angle = angles[param]
+        grad_x = torch.zeros_like(x)
 
-            # x + angle_eps
-            angles[param] = angle + angle_eps
+        for i in range(x.shape[0]):
+            angle = angles[i]
+
+            # θ + π/2
+            angles[i] = angle + param_shift_eps
             neuron.set_angles(angles)
             neuron.simulator = pre_sim.clone()
             neuron.predict(True, False)
             p_plus = neuron.simulator.prob(neuron.target)
 
-            # x - angle_eps
-            angles[param] = angle - angle_eps
+            # θ − π/2
+            angles[i] = angle - param_shift_eps
             neuron.set_angles(angles)
             neuron.simulator = pre_sim.clone()
             neuron.predict(True, False)
             p_minus = neuron.simulator.prob(neuron.target)
 
-            # Central difference
-            delta[param] = (p_plus - p_minus) / (2 * angle_eps)
+            # Parameter-shift gradient
+            grad_x[i] = 0.5 * (p_plus - p_minus)
 
-            angles[param] = angle
+            angles[i] = angle
 
-        neuron.simulator = pre_sim
+        # Restore simulator
+        neuron.set_simulator(pre_sim)
 
-        if _IS_TORCH_AVAILABLE:
-            delta = torch.tensor(delta, dtype=torch.float32, device=x.device)
+        # Apply chain rule and upstream gradient
+        grad_x *= grad_output[0] / math.sqrt(max(1.0 - post_prob * post_prob, 1e-6))
 
-        return delta
+        return grad_x, None
 
-    @staticmethod
-    def backward(ctx, grad_output):
-        (x,) = ctx.saved_tensors
-        neuron_wrapper = ctx.neuron_wrapper
-        delta = _backward(x, neuron_wrapper, grad_output)
-        if _IS_TORCH_AVAILABLE:
-            # grad_output: (O,)
-            # delta:       (O, I)
-            grad_input = torch.matmul(grad_output, delta)  # result: (I,)
-        else:
-            grad_input = [
-                sum(o * d for o, d in zip(grad_output, col))
-                for col in zip(*delta)
-            ]
-
-        return grad_input, None
 
 class QrackNeuronTorch(nn.Module if _IS_TORCH_AVAILABLE else object):
     """Torch wrapper for QrackNeuron
@@ -125,19 +115,20 @@ class QrackNeuronTorch(nn.Module if _IS_TORCH_AVAILABLE else object):
         neuron(QrackNeuron): QrackNeuron backing this torch wrapper
     """
 
-    def __init__(self, neuron):
+    def __init__(self, neuron, x):
         super().__init__()
         self.neuron = neuron
+        self.weights = nn.Parameter(x)
 
-    def forward(self, x):
-        return QrackNeuronTorchFunction.apply(x, self.neuron)
+    def forward(self):
+        return QrackNeuronTorchFunction.apply(self.weights, self.neuron)
 
 
 class QrackNeuronTorchLayer(nn.Module if _IS_TORCH_AVAILABLE else object):
     """Torch layer wrapper for QrackNeuron (with maximally expressive set of neurons between inputs and outputs)
 
     Attributes:
-        simulator (QrackSimulator): Prototype simulator that batching copies to use with QrackNeuron instances
+        simulator (QrackSimulator): Prototype simulator that batching copies to use with QrackNeuron instances. (You may customize or overwrite the initialization or reference, before calling forward(x).)
         simulators (list[QrackSimulator]): In-flight copies of prototype simulator corresponding to batch count
         input_indices (list[int], read-only): simulator qubit indices used as QrackNeuron inputs
         output_indices (list[int], read-only): simulator qubit indices used as QrackNeuron outputs
@@ -145,7 +136,7 @@ class QrackNeuronTorchLayer(nn.Module if _IS_TORCH_AVAILABLE else object):
         neurons (ModuleList[QrackNeuronTorch]): QrackNeuronTorch wrappers (for PyQrack QrackNeurons) in this layer, corresponding to weights
         weights (ParameterList): List of tensors corresponding one-to-one with weights of list of neurons
         apply_fn (Callable[Tensor, QrackNeuronTorch]): Corresponds to QrackNeuronTorchFunction.apply(x, neuron_wrapper) (or override with a custom implementation)
-        backward_fn (Callable[Tensor, Tensor]): Corresponds to QrackNeuronTorchFunction._backward(x, neuron_wrapper) (or override with a custom implementation)
+        post_init_fn (Callable[QrackSimulator]): Function that is applied after forward(x) state initialization, before inference. (As the function depends on nothing but the simulator, it's differentiable.)
     """
 
     def __init__(
@@ -157,6 +148,8 @@ class QrackNeuronTorchLayer(nn.Module if _IS_TORCH_AVAILABLE else object):
         highest_combo_count=2,
         activation=int(NeuronActivationFn.Generalized_Logistic),
         parameters=None,
+        post_init_fn=lambda simulator: None,
+        **kwargs
     ):
         """
         Initialize a QrackNeuron layer for PyTorch with a power set of neurons connecting inputs to outputs.
@@ -166,195 +159,99 @@ class QrackNeuronTorchLayer(nn.Module if _IS_TORCH_AVAILABLE else object):
             sim (QrackSimulator): Simulator into which predictor features are loaded
             input_qubits (int): Count of inputs (1 per qubit)
             output_qubits (int): Count of outputs (1 per qubit)
-            hidden_qubits (int): Count of "hidden" inputs (1 per qubit, always initialized to |+>, suggested to be same a highest_combo_count)
-            lowest_combo_count (int): Lowest combination count of input qubits iterated (0 is bias)
-            highest_combo_count (int): Highest combination count of input qubits iterated
-            activation (int): Integer corresponding to choice of activation function from NeuronActivationFn
+            hidden_qubits (int): (Optional) Count of "hidden" inputs (1 per qubit, always initialized to |+>, suggested to be same a highest_combo_count)
+            lowest_combo_count (int): (Optional) Lowest combination count of input qubits iterated (0 is bias)
+            highest_combo_count (int): (Optional) Highest combination count of input qubits iterated
+            activation (int): (Optional) Integer corresponding to choice of activation function from NeuronActivationFn
             parameters (list[float]): (Optional) Flat list of initial neuron parameters, corresponding to little-endian basis states of input + hidden qubits, repeated for ascending combo count, repeated for each output index
+            post_init_fn (Callable[QrackSimulator]): (Optional) Function that is applied after forward(x) state initialization, before inference. (As the function depends on nothing but the simulator, it's differentiable.)
         """
         super(QrackNeuronTorchLayer, self).__init__()
         if hidden_qubits is None:
             hidden_qubits = highest_combo_count
-        self.simulator = QrackSimulator(input_qubits + hidden_qubits + output_qubits)
+        self.simulator = QrackSimulator(input_qubits + hidden_qubits + output_qubits, **kwargs)
         self.simulators = []
         self.input_indices = list(range(input_qubits))
         self.hidden_indices = list(range(input_qubits, input_qubits + hidden_qubits))
-        self.output_indices = list(range(input_qubits + hidden_qubits, input_qubits + hidden_qubits + output_qubits))
+        self.output_indices = list(
+            range(input_qubits + hidden_qubits, input_qubits + hidden_qubits + output_qubits)
+        )
         self.activation = NeuronActivationFn(activation)
+        self.dtype = torch.float if Qrack.fppow <= 5 else torch.double
         self.apply_fn = QrackNeuronTorchFunction.apply
-        self.backward_fn = QrackNeuronTorchFunction._backward
+        self.post_init_fn = post_init_fn
 
         # Create neurons from all input combinations, projecting to coherent output qubits
-        neurons = [
-            QrackNeuronTorch(
-                QrackNeuron(self.simulator, input_subset, output_id, activation)
-            )
-            for output_id in self.output_indices
-            for k in range(lowest_combo_count, highest_combo_count + 1)
-            for input_subset in itertools.combinations(self.input_indices + self.hidden_indices, k)
-        ]
-        self.neurons = nn.ModuleList(neurons) if _IS_TORCH_AVAILABLE else neurons
-
-        # Set Qrack's internal parameters:
-        if parameters:
-            param_count = 0
-            self.weights = nn.ParameterList() if _IS_TORCH_AVAILABLE else []
-            for neuron_wrapper in self.neurons:
-                neuron = neuron_wrapper.neuron
-                p_count = 1 << len(neuron.controls)
-                neuron.set_angles(parameters[param_count : (param_count + p_count)])
-                self.weights.append(
-                    nn.Parameter(torch.tensor(parameters[param_count : (param_count + p_count)]))
-                    if _IS_TORCH_AVAILABLE else parameters[param_count : (param_count + p_count)]
-                )
-                param_count += p_count
-        else:
-            self.weights = nn.ParameterList() if _IS_TORCH_AVAILABLE else []
-            for neuron_wrapper in self.neurons:
-                neuron = neuron_wrapper.neuron
-                p_count = 1 << len(neuron.controls)
-                self.weights.append(nn.Parameter(torch.zeros(p_count)) if _IS_TORCH_AVAILABLE else ([0.0] * p_count))
+        neurons = []
+        param_count = 0
+        for output_id in self.output_indices:
+            for k in range(lowest_combo_count, highest_combo_count + 1):
+                for input_subset in itertools.combinations(self.input_indices, k):
+                    p_count = 1 << len(input_subset)
+                    angles = (
+                        (
+                            torch.tensor(
+                                parameters[param_count : (param_count + p_count)], dtype=self.dtype
+                            )
+                            if parameters
+                            else torch.zeros(p_count, dtype=self.dtype)
+                        )
+                    )
+                    neurons.append(
+                        QrackNeuronTorch(
+                            QrackNeuron(self.simulator, input_subset, output_id, activation), angles
+                        )
+                    )
+                    param_count += p_count
+        self.neurons = nn.ModuleList(neurons)
+
+        # Prepare the state before feed-forward:
+
+        # Prepare hidden predictors
+        for hidden_id in self.hidden_indices:
+            self.simulator.h(hidden_id)
+        # Prepare a maximally uncertain output state.
+        for output_id in self.output_indices:
+            self.simulator.h(output_id)
 
     def forward(self, x):
-        return QrackNeuronTorchLayerFunction.apply(x, self)
+        B = x.shape[0]
+        x = x.view(B, -1)
 
+        self.simulators.clear()
 
-class QrackNeuronTorchLayerFunction(Function if _IS_TORCH_AVAILABLE else object):
-    """Static forward/backward/apply functions for QrackNeuronTorch"""
+        # Group neurons by output target once
+        by_out = {out: [] for out in self.output_indices}
+        for neuron_wrapper in self.neurons:
+            by_out[neuron_wrapper.neuron.target].append(neuron_wrapper)
 
-    @staticmethod
-    def forward(ctx, x, neuron_layer):
-        # Save for backward
-        ctx.save_for_backward(x)
-        ctx.neuron_layer = neuron_layer
-
-        input_indices = neuron_layer.input_indices
-        hidden_indices = neuron_layer.hidden_indices
-        output_indices = neuron_layer.output_indices
-        simulators = neuron_layer.simulators
-        weights = neuron_layer.weights
-
-        if _IS_TORCH_AVAILABLE:
-            B = x.shape[0]
-            x = x.view(B, -1)
-        else:
-            B = len(x)
-
-        simulators.clear()
-        if _IS_TORCH_AVAILABLE:
-            for b in range(B):
-                simulator = neuron_layer.simulator.clone()
-                simulators.append(simulator)
-                for q, input_id in enumerate(input_indices):
-                    simulator.r(Pauli.PauliY, math.pi * x[b, q].item(), q)
-        else:
-            for b in range(B):
-                simulator = neuron_layer.simulator.clone()
-                simulators.append(simulator)
-                for q, input_id in enumerate(input_indices):
-                    simulator.r(Pauli.PauliY, math.pi * x[b][q], q)
-
-        y = [([0.0] * len(output_indices)) for _ in range(B)]
+        batch_rows = []
         for b in range(B):
-            simulator = simulators[b]
-            # Prepare a maximally uncertain output state.
-            for output_id in output_indices:
-                simulator.h(output_id)
-            # Prepare hidden predictors
-            for h in hidden_indices:
-                simulator.h(h)
+            simulator = self.simulator.clone()
+            self.simulators.append(simulator)
 
-            # Set Qrack's internal parameters:
-            for idx, neuron_wrapper in enumerate(neuron_layer.neurons):
-                neuron_wrapper.neuron.simulator = simulator
-                neuron_layer.apply_fn(weights[idx], neuron_wrapper)
+            # Apply feed-forward
+            for q, input_id in enumerate(self.input_indices):
+                simulator.r(Pauli.PauliY, math.pi * x[b, q].item(), input_id)
 
-            for q, output_id in enumerate(output_indices):
-                y[b][q] = simulator.prob(output_id)
+            # Differentiable post-initialization:
+            self.post_init_fn(simulator)
 
-        if _IS_TORCH_AVAILABLE:
-            y = torch.tensor(y, dtype=torch.float32, device=x.device)
+            row = []
+            for out in self.output_indices:
+                phi = torch.tensor(init_phi, device=x.device, dtype=x.dtype)
 
-        return y
+                for neuron_wrapper in by_out[out]:
+                    neuron_wrapper.neuron.set_simulator(simulator)
+                    phi += self.apply_fn(
+                        neuron_wrapper.weights,
+                        neuron_wrapper.neuron
+                    ).squeeze()
 
-    @staticmethod
-    def backward(ctx, grad_output):
-        (x,) = ctx.saved_tensors
-        neuron_layer = ctx.neuron_layer
-
-        input_indices = neuron_layer.input_indices
-        hidden_indices = neuron_layer.hidden_indices
-        output_indices = neuron_layer.output_indices
-        simulators = neuron_layer.simulators
-        neurons = neuron_layer.neurons
-        backward_fn = neuron_layer.backward_fn
-
-        input_count = len(input_indices)
-        output_count = len(output_indices)
-
-        if _IS_TORCH_AVAILABLE:
-            B = x.shape[0]
-            x = x.view(B, -1)
-        else:
-            B = len(x)
-
-        # Uncompute prediction
-        if _IS_TORCH_AVAILABLE:
-            delta = torch.zeros((B, output_count, input_count), dtype=torch.float32, device=x.device)
-            for b in range(B):
-                simulator = simulators[b]
-                for neuron_wrapper in neurons:
-                    neuron = neuron_wrapper.neuron
-                    neuron.simulator = simulator
-                    angles = torch.tensor(neuron.get_angles(), dtype=torch.float32, device=x.device, requires_grad=True)
-                    o = output_indices.index(neuron.target)
-                    neuron_grad = backward_fn(angles, neuron_wrapper)
-                    for idx, c in enumerate(neuron.controls):
-                        if c not in input_indices:
-                            continue
-                        i = input_indices.index(c)
-                        delta[b, o, i] += neuron_grad[idx]
-        else:
-            delta = [[[0.0] * input_count for _ in range(output_count)] for _ in range(B)]
-            for b in range(B):
-                simulator = simulators[b]
-                for neuron_wrapper in neurons:
-                    neuron = neuron_wrapper.neuron
-                    neuron.simulator = simulator
-                    angles = neuron.get_angles()
-                    o = output_indices.index(neuron.target)
-                    neuron_grad = backward_fn(angles, neuron_wrapper)
-                    for idx, c in enumerate(neuron.controls):
-                        if c not in input_indices:
-                            continue
-                        i = input_indices.index(c)
-                        delta[b][o][i] += neuron_grad[idx]
-
-        # Uncompute output state prep
-        for simulator in simulators:
-            for output_id in output_indices:
-                simulator.h(output_id)
-            for h in hidden_indices:
-                simulator.h(output_id)
-
-        if _IS_TORCH_AVAILABLE:
-            for b in range(B):
-                simulator = simulators[b]
-                for q, input_id in enumerate(input_indices):
-                    simulator.r(Pauli.PauliY, -math.pi * x[b, q].item(), q)
-        else:
-            for b in range(B):
-                simulator = simulators[b]
-                for q, input_id in enumerate(input_indices):
-                    simulator.r(Pauli.PauliY, -math.pi * x[b][q].item(), q)
-
-        if _IS_TORCH_AVAILABLE:
-            grad_input = torch.matmul(grad_output.view(B, 1, -1), delta).view_as(x)
-        else:
-            grad_input = [[0.0] * output_count for _ in range(B)]
-            for b in range(B):
-                for o in range(output_indices):
-                    for i in range(input_indices):
-                        grad_input[b][o] += grad_output[b][o] * delta[b][o][i]
-
-        return grad_input, None
+                # Convert angle back to probability
+                p = torch.clamp(torch.sin(phi), min=0.0)
+                row.append(p)
+
+            batch_rows.append(torch.stack(row))
+
+        return torch.stack(batch_rows)