Apply pending pepsy updates

Author: rezah

.github/workflows/ci.yml

@@ -17,8 +17,12 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           pip install -e .[dev]
+      - name: Validate source syntax
+        run: python -m compileall -q -f src/pepsy
       - name: Run tests
-        run: pytest -q
+        run: |
+          NUMBA_CACHE_DIR=/tmp PYTHONPYCACHEPREFIX=/tmp \
+          pytest -q
 
   docs:
     runs-on: ubuntu-latest

README.md

@@ -10,10 +10,9 @@ Current package version: `0.1.1` (from `pyproject.toml` / `pepsy.__version__`).
 - `src/pepsy/`: installable library code
   - `boundary_states.py`: boundary state initialization (`BdyMPS`)
   - `boundary_sweeps.py`: sweep/contraction runner (`CompBdy`)
-  - `boundary_metrics.py`: input preparation + contraction (`prepare_boundary_inputs`, `ContractBoundary`)
-  - `optimize_sweep.py`, `optimize_global.py`, `gate.py`, `gradient_solver.py`, `debug.py`
-  - `fit.py`, `core.py`, `linalg_registrations.py`
-- `example/`: example notebooks
+  - `boundary_metrics.py`: input preparation + contraction (`build_bra_ket`, `contract_boundary`, `BoundaryContractResult`)
+  - `optimize_sweep.py`, `optimize_global.py`, `optimize_energy.py`, `gate.py`, `gradient_solver.py`
+  - `fit.py`, `core.py`, `_backend_utils.py`, `_backend_linalg.py`
 - `docs/`: Sphinx documentation source
 - `tests/`: package tests
 
@@ -25,23 +24,32 @@ pip install -e .
 # Optional backends:
 # pip install -e .[torch]
 # pip install -e .[solvers]
+# jax backend (manual, platform-specific wheels):
+# pip install jax jaxlib
 # Optional plotting helpers:
 # pip install -e .[viz]
 ```
 
 ## Quick Usage
 ```python
 import pepsy
-from pepsy import BdyMPS, CompBdy, ContractBoundary, prepare_boundary_inputs
+import quimb.tensor as qtn
 
-print(pepsy.__version__)
+ket = qtn.PEPS.rand(Lx=3, Ly=3, bond_dim=2, seed=1, dtype="complex128")
+ket_tagged, norm = pepsy.build_bra_ket(ket=ket)
+
+bdy = pepsy.BdyMPS(tn_flat=ket_tagged, tn_double=norm, chi=32, single_layer=False)
+res = pepsy.contract_boundary(norm=norm, bdy=bdy, direction="y", n_iter=2)
+
+print(pepsy.__version__, res.cost)
 ```
 
 ## Documentation
 Build docs locally:
 
 ```bash
 pip install -e .[docs]
+NUMBA_CACHE_DIR=/tmp PYTHONPYCACHEPREFIX=/tmp \
 sphinx-build -W -b html docs docs/_build/html
 ```
 
@@ -52,10 +60,6 @@ Main docs sections:
 - `howto/`
 - `api/`
 
-Guided notebook example:
-
-- `example/norm.ipynb`
-
 ## Notes
 - `.gitattributes` marks notebooks as binary to avoid noisy diffs.
 - `.gitignore` excludes checkpoints, caches, `cash/`, and `nohup.out`.

docs/examples.md

@@ -1,10 +1,10 @@
 # Examples
 
-## Primary notebook
+## Primary walkthrough
 
-- `example/norm.ipynb`
+- [Tutorial: Contract a PEPS Norm](tutorials/contract_norm.md)
 
-This is the recommended first notebook for interactive usage.
+This is the recommended first end-to-end walkthrough for interactive usage.
 
 It demonstrates:
 
@@ -13,9 +13,10 @@ It demonstrates:
 - running `contract_boundary(...)`
 - inspecting `BoundaryContractResult.cost` and `.fidel`
 
-## Additional notebooks
+## Additional walkthroughs
 
-- `example/peps_norm_.ipynb`: historical exploratory notebook
-- `example/peps_boundary_states.ipynb`: boundary-state focused experiments
+- [Tutorial: Fidelity Diagnostics](tutorials/fidelity_diagnostics.md)
+- [How-To: Choose Parameters](howto/choose_parameters.md)
+- [How-To: Tune Sweep Solvers](howto/solver_tuning.md)
 
 For a cleaner, docs-first narrative, start from [tutorials](tutorials/index.md).

docs/howto/troubleshooting.md

@@ -43,3 +43,14 @@ Cause:
 Fix:
 
 - Set `track_boundary_fidelity=True`.
+
+## `RuntimeError: cannot cache function ... quimb/core.py`
+
+Cause:
+
+- Some environments do not allow Numba cache writes in default locations.
+
+Fix:
+
+- Run commands with explicit cache env vars, for example:
+  `NUMBA_CACHE_DIR=/tmp PYTHONPYCACHEPREFIX=/tmp pytest -q`

src/pepsy/__init__.py

@@ -53,7 +53,6 @@
         set_default_grad_backend,
         tns_align,
     )
-    from .linalg_registrations import reg_complex_svd_jax, reg_complex_svd_torch
     from .fit import FIT
     from .gate import (
         gate,
@@ -122,8 +121,6 @@
     "set_default_grad_backend",
     "get_default_grad_backend",
     "register_torch_linalg",
-    "reg_complex_svd_torch",
-    "reg_complex_svd_jax",
     "reset_default_backends",
     "SweepOptimizer",
     "EnergyOptimizer",
@@ -469,17 +466,6 @@ def __getattr__(name):
             "reset_default_backends": reset_default_backends,
         }[name]
 
-    if name in ("reg_complex_svd_torch", "reg_complex_svd_jax"):
-        from .linalg_registrations import (  # pylint: disable=import-outside-toplevel
-            reg_complex_svd_jax,
-            reg_complex_svd_torch,
-        )
-
-        return {
-            "reg_complex_svd_torch": reg_complex_svd_torch,
-            "reg_complex_svd_jax": reg_complex_svd_jax,
-        }[name]
-
     if name == "CompBdy":
         from .boundary_sweeps import CompBdy  # pylint: disable=import-outside-toplevel
 

src/pepsy/_backend_utils.py

@@ -175,6 +175,56 @@ def _to_cupy(
     return _to_cupy
 
 
+def _build_to_jax(sample_data, dtype_name, *, cast_complex_to_real=False):
+    import jax  # pylint: disable=import-outside-toplevel
+    import jax.numpy as jnp  # pylint: disable=import-outside-toplevel
+
+    dtype_map = {
+        "complex128": jnp.complex128,
+        "complex64": jnp.complex64,
+        "float64": jnp.float64,
+        "float32": jnp.float32,
+        "float16": jnp.float16,
+        "int64": jnp.int64,
+        "int32": jnp.int32,
+    }
+    if dtype_name not in dtype_map:
+        raise ValueError(f"Unsupported dtype '{dtype_name}' for jax backend.")
+
+    dtype = getattr(sample_data, "dtype", None) or dtype_map[dtype_name]
+    device = getattr(sample_data, "device", None)
+
+    def _to_jax(
+        x,
+        dtype=dtype,
+        device=device,
+        cast_complex_to_real=cast_complex_to_real,
+    ):
+        # Torch tensors need explicit host conversion before jnp.asarray.
+        try:
+            import torch  # pylint: disable=import-outside-toplevel
+        except ImportError:  # pragma: no cover - optional dependency
+            torch = None
+        if torch is not None and isinstance(x, torch.Tensor):
+            x = x.detach().cpu().numpy()
+
+        arr = jnp.asarray(x)
+
+        if cast_complex_to_real and jnp.issubdtype(dtype, jnp.floating) and jnp.iscomplexobj(arr):
+            arr = arr.real
+
+        target_dtype = dtype
+        if (not cast_complex_to_real) and jnp.issubdtype(target_dtype, jnp.floating) and jnp.iscomplexobj(arr):
+            target_dtype = jnp.result_type(target_dtype, jnp.complex64)
+
+        out = jnp.asarray(arr, dtype=target_dtype)
+        if device is not None:
+            out = jax.device_put(out, device)
+        return out
+
+    return _to_jax
+
+
 def dispatch_backend_converter(
     *,
     backend,
@@ -204,6 +254,12 @@ def dispatch_backend_converter(
             dtype_name,
             cast_complex_to_real=cast_complex_to_real,
         )
+    if backend == "jax":
+        return _build_to_jax(
+            sample_data,
+            dtype_name,
+            cast_complex_to_real=cast_complex_to_real,
+        )
 
     raise ValueError(f"Unsupported backend: {backend}")
 

src/pepsy/boundary_states.py

@@ -151,6 +151,15 @@ def _build_to_cupy(sample_data, dtype_name):
             cast_complex_to_real=True,
         )
 
+    @staticmethod
+    def _build_to_jax(sample_data, dtype_name):
+        return dispatch_backend_converter(
+            backend="jax",
+            dtype_name=dtype_name,
+            sample_data=sample_data,
+            cast_complex_to_real=True,
+        )
+
     def _dispatch_backend_converter(self, backend, dtype_name, sample_data):
         """Return a conversion callable for the detected backend."""
         if sample_data is None:
@@ -164,6 +173,8 @@ def _dispatch_backend_converter(self, backend, dtype_name, sample_data):
             return self._build_to_torch(sample_data, dtype_name)
         if backend == "cupy":
             return self._build_to_cupy(sample_data, dtype_name)
+        if backend == "jax":
+            return self._build_to_jax(sample_data, dtype_name)
 
         raise ValueError(f"Unsupported backend: {backend}")
 

src/pepsy/fit.py

@@ -584,7 +584,7 @@ def run_gate(
                     raise TypeError("Unexpected effective tensor type during run_gate.")
 
                 norm_f = (f.H & f).contract(all) ** 0.5
-                self.local_norm_trace.append(complex(norm_f).real)
+                self.local_norm_trace.append(ar.do("real", norm_f))
 
                 # Update tensor data
                 psi[site].modify(data=f.data)

src/pepsy/gate.py

@@ -1105,13 +1105,10 @@ def gate(tn, gates, where=None, **kwargs):
     tn : qtn.TensorNetwork
         Input tensor network.
     gates : array_like | tuple | sequence
-        Gate payload in one of these forms:
-        1) single gate with explicit ``where`` argument:
-           ``gate(tn, G, where=...)``
-        2) canonical bundled stream:
-           ``gate(tn, ((G1, where1), (G2, where2), ...))``
-        The single bundled alias ``(gate, where)`` is intentionally rejected
-        to avoid ambiguity with a plain rank-2 gate tensor.
+        Gate payload. Use ``gate(tn, G, where=...)`` for a single gate, or
+        ``gate(tn, ((G1, where1), (G2, where2), ...))`` for bundled gates.
+        The single bundled alias ``(gate, where)`` is rejected to avoid
+        ambiguity with a plain rank-2 gate tensor.
     where : object, optional
         Target location for single-gate form.
         - 1D: ``1``, ``(1,)``, ``(1, 2)``
@@ -1623,13 +1620,12 @@ def build_pepo_from_gates(
     Parameters
     ----------
     gates : array_like | sequence | tuple
-        Gate payload in one of these forms:
-        1) single gate with explicit ``where`` argument:
-           ``build_pepo_from_gates(G, where=...)``
-        2) canonical bundled stream:
-           ``build_pepo_from_gates(((G1, where1), (G2, where2), ...))``
-        3) legacy parallel form with ``wheres``:
-           ``build_pepo_from_gates([G1, G2], [where1, where2])``
+        Gate payload. Accepted forms are:
+        ``build_pepo_from_gates(G, where=...)`` for a single gate,
+        ``build_pepo_from_gates(((G1, where1), (G2, where2), ...))`` for the
+        canonical bundled stream, and
+        ``build_pepo_from_gates([G1, G2], [where1, where2])`` for the legacy
+        parallel ``wheres`` form.
     wheres : sequence[tuple] | None, optional
         Legacy parallel where stream aligned with ``gates``.
     where : object, optional
@@ -1735,13 +1731,12 @@ def build_mpo_from_gates(
     Parameters
     ----------
     gates : array_like | sequence | tuple
-        Gate payload in one of these forms:
-        1) single gate with explicit ``where`` argument:
-           ``build_mpo_from_gates(G, where=...)``
-        2) canonical bundled stream:
-           ``build_mpo_from_gates(((G1, where1), (G2, where2), ...))``
-        3) legacy parallel form with ``wheres``:
-           ``build_mpo_from_gates([G1, G2], [where1, where2])``
+        Gate payload. Accepted forms are:
+        ``build_mpo_from_gates(G, where=...)`` for a single gate,
+        ``build_mpo_from_gates(((G1, where1), (G2, where2), ...))`` for the
+        canonical bundled stream, and
+        ``build_mpo_from_gates([G1, G2], [where1, where2])`` for the legacy
+        parallel ``wheres`` form.
     wheres : sequence[tuple[int, ...]] | None, optional
         Legacy parallel where stream aligned with ``gates``.
     where : object, optional

src/pepsy/gradient_solver.py

@@ -90,32 +90,9 @@ def _require_optax():
 class GradSolverResult:
     """Structured optimization result returned by :class:`GradientOptimizer`.
 
-    Attributes
-    ----------
-    params : dict[str, Any]
-        Optimised parameters on their original devices, detached from
-        autograd. Concrete element type is backend-dependent: a
-        ``torch.Tensor`` for the torch/scipy/nlopt solvers shipped here,
-        a ``jax.Array`` / ``jnp.ndarray`` for a JAX-based solver, or any
-        other array-like a future backend produces.
-    history : list[float]
-        Per-step loss trace. Length matches ``n_steps``.
-    solver : str
-        Normalised solver name actually used (e.g. ``"scipy"``).
-    n_steps : int
-        Number of entries recorded in ``history`` (iterations for scipy,
-        function evaluations for nlopt/torch).
-    best_loss : float
-        Lowest loss value observed during the run.
-    final_loss : float
-        Loss after the final restored state (== ``best_loss`` when
-        ``restore_best=True``).
-    convergence_reason : str
-        ``"maxiter"``, ``"patience"``, ``"bad_max"``, ``"empty_params"``,
-        ``"nan_x"``, or ``"nlopt_error:<ExcType>"``.
-    n_evals : int
-        Number of objective evaluations (``-1`` if the backend does not
-        report it separately).
+    This dataclass stores backend-agnostic outputs (parameters, loss history,
+    convergence metadata, and evaluation counts). Field docs are derived from
+    the dataclass members listed below.
     """
 
     # Backend-agnostic: holds torch.Tensor, jax.Array, np.ndarray, etc.