add pennylane tensor as return_format option (#24)

Author: phoeenniixx

pyqit/core/trainer/trainer.py

@@ -167,8 +167,11 @@ def predict(
             Used for inference only; weights are not touched.
         datamodule : DataModule
             Test split if present, else validation, else train.
-        return_format : {"auto", "numpy", "torch"}, default "auto"
-            ``"auto"`` follows whatever the model emits.
+        return_format : {"auto", "numpy", "torch", "pennylane"}, default "auto"
+            ``"auto"`` follows whatever the model emits. ``"pennylane"``
+            returns a ``pennylane.numpy.tensor`` rather than a bare
+            ``ndarray``, for callers composing the result into further
+            pnp-based code (a custom cost function, ``utils.diagnostic``).
 
         Returns
         -------
@@ -292,14 +295,30 @@ def _collect(all_preds: list, return_format: str):
             return torch.cat(
                 [p if _is_torch(p) else torch.as_tensor(p) for p in all_preds], dim=0
             )
+        elif target in ("numpy", "pennylane"):
+            # Both share this merge: a batch can be a torch tensor, a pnp
+            # tensor or a mid-trace ArrayBox depending on backend, and
+            # "pennylane" only needs one extra wrap over the merged result
+            # rather than repeating that per-batch normalisation itself.
+            merged = np.concatenate(
+                [
+                    p.detach().cpu().numpy() if _is_torch(p) else np.asarray(p)
+                    for p in all_preds
+                ],
+                axis=0,
+            )
 
-        return np.concatenate(
-            [
-                p.detach().cpu().numpy() if _is_torch(p) else np.asarray(p)
-                for p in all_preds
-            ],
-            axis=0,
-        )
+            if target == "pennylane":
+                import pennylane.numpy as pnp
+
+                return pnp.array(merged, requires_grad=False)
+
+            return merged
+        else:
+            raise ValueError(
+                f"Unknown return_format {return_format!r}; expected one of "
+                "'auto', 'numpy', 'torch', 'pennylane'."
+            )
 
     def __repr__(self) -> str:
         return (

pyqit/tests/test_trainer.py

@@ -339,6 +339,30 @@ def test_predict_torch_format_requires_torch(monkeypatch):
         Trainer(verbose=0).predict(_model(), _dm(), return_format="torch")
 
 
+@pytest.mark.parametrize("return_format", ["auto", "numpy", "torch", "pennylane"])
+def test_predict_return_format(return_format):
+    """Each explicit format is requestable regardless of the active backend.
+
+    ``"auto"`` and ``"numpy"`` both collapse pennylane's native
+    ``pnp.tensor`` predictions down to a bare ``ndarray`` via ``np.asarray``;
+    ``"torch"`` and ``"pennylane"`` are opt-ins that instead preserve (or
+    convert into) an autograd-carrying tensor type.
+    """
+    import pennylane.numpy as pnp
+
+    expected_type = {
+        "auto": np.ndarray,
+        "numpy": np.ndarray,
+        "torch": pytest.importorskip("torch").Tensor,
+        "pennylane": pnp.tensor,
+    }[return_format]
+
+    pyqit.set_backend("pennylane")
+    preds = Trainer(verbose=0).predict(_model(), _dm(), return_format=return_format)
+
+    assert type(preds) is expected_type
+
+
 def _load(path):
     if path.endswith(".ckpt"):
         import torch