Upload 9 files

__init__.py

@@ -0,0 +1,6 @@
+from .dima import DIMA
+from .dmap import DMAP
+from .gplm import GPLM
+from .ddpm import DDPM
+
+__all__ = ["DIMA", "DMAP", "GPLM", "DDPM"]

ann.py

@@ -0,0 +1,300 @@
+# src/dima/ann.py
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Tuple
+
+import numpy as np
+
+from .utils import as_contig_f32, sqdist_ab
+
+
+# -------------------------
+# Public API
+# -------------------------
+ANNBackend = str  # "auto" | "faiss" | "pynndescent" | "sklearn" | "brute"
+
+
+class ANNBase:
+    """Minimal ANN interface used by DMAP/GPLM."""
+    def build(self, X: np.ndarray) -> "ANNBase":
+        raise NotImplementedError
+
+    def search(self, Q: np.ndarray, k: int) -> Tuple[np.ndarray, np.ndarray]:
+        """
+        Returns:
+          idx: (B,k) int64
+          D2 : (B,k) float32  (squared Euclidean distances)
+        """
+        raise NotImplementedError
+
+
+def make_ann(
+    backend: ANNBackend = "auto",
+    ann_params: Optional[Dict[str, Any]] = None,
+    n_jobs: int = -1,
+) -> Tuple[ANNBase, str]:
+    """
+    Create an ANN implementation.
+    backend:
+      - "auto": prefers faiss, then pynndescent, then sklearn, else brute
+      - "faiss": FAISS (if installed)
+      - "pynndescent": NNDescent (if installed)
+      - "sklearn": sklearn NearestNeighbors (if installed)
+      - "brute": exact brute force
+    ann_params:
+      - for faiss:
+          index: "flat" | "hnsw" | "ivf_flat"
+          hnsw_M: int (default 32)
+          ef_search: int (default 64)
+          ef_construction: int (default 200)
+          ivf_nlist: int (default 1024)
+          ivf_nprobe: int (default 16)
+          use_float16: bool (default False; GPU only typically)
+      - for pynndescent:
+          n_trees: int
+          n_iters: int
+          metric: str (default "euclidean")
+      - for sklearn:
+          algorithm: str (default "auto")
+          leaf_size: int (default 40)
+          metric: str (default "euclidean")
+    """
+    ann_params = {} if ann_params is None else dict(ann_params)
+    b = (backend or "auto").lower()
+
+    if b == "auto":
+        for cand in ("faiss", "pynndescent", "sklearn", "brute"):
+            ann, used = make_ann(cand, ann_params=ann_params, n_jobs=n_jobs)
+            if used != "brute" or cand == "brute":
+                return ann, used
+        return BruteANN(), "brute"
+
+    if b == "faiss":
+        try:
+            return FaissANN(ann_params=ann_params), "faiss"
+        except Exception as e:
+            raise ImportError(
+                "FAISS backend requested but faiss is not available or failed to initialize. "
+                "Install with: pip install dima[faiss]"
+            ) from e
+
+    if b == "pynndescent":
+        try:
+            return PyNNDescentANN(ann_params=ann_params), "pynndescent"
+        except Exception as e:
+            raise ImportError(
+                "pynndescent backend requested but pynndescent is not available. "
+                "Install with: pip install pynndescent"
+            ) from e
+
+    if b == "sklearn":
+        try:
+            return SklearnANN(n_jobs=n_jobs, ann_params=ann_params), "sklearn"
+        except Exception as e:
+            raise ImportError(
+                "sklearn backend requested but scikit-learn is not available. "
+                "Install with: pip install scikit-learn"
+            ) from e
+
+    if b == "brute":
+        return BruteANN(), "brute"
+
+    raise ValueError(f"Unknown ANN backend: {backend!r}")
+
+
+# -------------------------
+# Brute-force (no deps)
+# -------------------------
+class BruteANN(ANNBase):
+    def __init__(self):
+        self.X = None
+
+    def build(self, X: np.ndarray) -> "BruteANN":
+        self.X = as_contig_f32(X)
+        return self
+
+    def search(self, Q: np.ndarray, k: int) -> Tuple[np.ndarray, np.ndarray]:
+        if self.X is None:
+            raise RuntimeError("BruteANN.search called before build().")
+        X = self.X
+        Q = as_contig_f32(Q)
+        k = int(k)
+        if k <= 0:
+            raise ValueError("k must be >= 1")
+        if k > X.shape[0]:
+            k = X.shape[0]
+
+        D2 = sqdist_ab(Q, X)  # (B,N)
+        idx = np.argpartition(D2, kth=k - 1, axis=1)[:, :k]
+        rows = np.arange(Q.shape[0])[:, None]
+        d2 = D2[rows, idx]
+
+        # sort within k
+        ordk = np.argsort(d2, axis=1)
+        idx = idx[rows, ordk].astype(np.int64)
+        d2 = d2[rows, ordk].astype(np.float32)
+        return idx, d2
+
+
+# -------------------------
+# FAISS
+# -------------------------
+class FaissANN(ANNBase):
+    def __init__(self, ann_params: Optional[Dict[str, Any]] = None):
+        self.ann_params = {} if ann_params is None else dict(ann_params)
+        self.index = None
+        self.X = None  # keep reference for possible rebuild
+
+        # delayed import
+        import faiss  # type: ignore
+        self.faiss = faiss
+
+    def _build_index(self, d: int):
+        p = self.ann_params
+        faiss = self.faiss
+
+        index_kind = str(p.get("index", "flat")).lower()
+
+        if index_kind == "flat":
+            index = faiss.IndexFlatL2(d)
+
+        elif index_kind == "hnsw":
+            M = int(p.get("hnsw_M", 32))
+            index = faiss.IndexHNSWFlat(d, M)
+            # optional tuning
+            ef_search = int(p.get("ef_search", 64))
+            ef_constr = int(p.get("ef_construction", 200))
+            index.hnsw.efSearch = ef_search
+            index.hnsw.efConstruction = ef_constr
+
+        elif index_kind == "ivf_flat":
+            nlist = int(p.get("ivf_nlist", 1024))
+            quantizer = faiss.IndexFlatL2(d)
+            index = faiss.IndexIVFFlat(quantizer, d, nlist, faiss.METRIC_L2)
+            nprobe = int(p.get("ivf_nprobe", 16))
+            index.nprobe = nprobe
+
+        else:
+            raise ValueError(f"Unknown faiss index kind: {index_kind!r}")
+
+        return index
+
+    def build(self, X: np.ndarray) -> "FaissANN":
+        X = as_contig_f32(X)
+        self.X = X
+        faiss = self.faiss
+        d = int(X.shape[1])
+
+        index = self._build_index(d)
+
+        # IVF needs training
+        if hasattr(index, "is_trained") and not index.is_trained:
+            index.train(X)
+
+        index.add(X)
+        self.index = index
+        return self
+
+    def search(self, Q: np.ndarray, k: int) -> Tuple[np.ndarray, np.ndarray]:
+        if self.index is None:
+            raise RuntimeError("FaissANN.search called before build().")
+        Q = as_contig_f32(Q)
+        k = int(k)
+        if k <= 0:
+            raise ValueError("k must be >= 1")
+
+        # FAISS returns (distances, indices); for L2 these are squared distances
+        D2, I = self.index.search(Q, k)
+        return I.astype(np.int64), D2.astype(np.float32)
+
+
+# -------------------------
+# PyNNDescent
+# -------------------------
+class PyNNDescentANN(ANNBase):
+    def __init__(self, ann_params: Optional[Dict[str, Any]] = None):
+        self.ann_params = {} if ann_params is None else dict(ann_params)
+        self.index = None
+        self.X = None
+
+        from pynndescent import NNDescent  # type: ignore
+        self.NNDescent = NNDescent
+
+    def build(self, X: np.ndarray) -> "PyNNDescentANN":
+        X = as_contig_f32(X)
+        self.X = X
+        p = self.ann_params
+
+        metric = p.get("metric", "euclidean")
+        n_trees = p.get("n_trees", None)
+        n_iters = p.get("n_iters", None)
+
+        kwargs: Dict[str, Any] = {"metric": metric}
+        if n_trees is not None:
+            kwargs["n_trees"] = int(n_trees)
+        if n_iters is not None:
+            kwargs["n_iters"] = int(n_iters)
+
+        self.index = self.NNDescent(X, **kwargs)
+        return self
+
+    def search(self, Q: np.ndarray, k: int) -> Tuple[np.ndarray, np.ndarray]:
+        if self.index is None:
+            raise RuntimeError("PyNNDescentANN.search called before build().")
+        Q = as_contig_f32(Q)
+        k = int(k)
+        if k <= 0:
+            raise ValueError("k must be >= 1")
+
+        # NNDescent returns (indices, distances) with euclidean distances (not squared)
+        I, d = self.index.query(Q, k=k)
+        D2 = (d.astype(np.float32) ** 2)
+        return I.astype(np.int64), D2
+
+
+# -------------------------
+# scikit-learn NearestNeighbors
+# -------------------------
+class SklearnANN(ANNBase):
+    def __init__(self, n_jobs: int = -1, ann_params: Optional[Dict[str, Any]] = None):
+        self.ann_params = {} if ann_params is None else dict(ann_params)
+        self.n_jobs = int(n_jobs)
+        self.nn = None
+        self.X = None
+
+        from sklearn.neighbors import NearestNeighbors  # type: ignore
+        self.NearestNeighbors = NearestNeighbors
+
+    def build(self, X: np.ndarray) -> "SklearnANN":
+        X = as_contig_f32(X)
+        self.X = X
+        p = self.ann_params
+
+        algorithm = p.get("algorithm", "auto")
+        leaf_size = int(p.get("leaf_size", 40))
+        metric = p.get("metric", "euclidean")
+
+        self.nn = self.NearestNeighbors(
+            n_neighbors=1,  # set later in search
+            algorithm=algorithm,
+            leaf_size=leaf_size,
+            metric=metric,
+            n_jobs=self.n_jobs,
+        )
+        self.nn.fit(X)
+        return self
+
+    def search(self, Q: np.ndarray, k: int) -> Tuple[np.ndarray, np.ndarray]:
+        if self.nn is None:
+            raise RuntimeError("SklearnANN.search called before build().")
+        Q = as_contig_f32(Q)
+        k = int(k)
+        if k <= 0:
+            raise ValueError("k must be >= 1")
+
+        self.nn.set_params(n_neighbors=k)
+        d, I = self.nn.kneighbors(Q, return_distance=True)
+        # sklearn distances are euclidean; square them
+        D2 = (d.astype(np.float32) ** 2)
+        return I.astype(np.int64), D2

ddim.py

@@ -0,0 +1,463 @@
+# --- src/dima/ddpm.py
+# +++ src/dima/ddpm.py
+# @@ -368,4 +368,204 @@
+#          return self.reverse_from_T(noise)
+# src/dima/ddpm.py
+from __future__ import annotations
+
+from typing import Any, Optional
+
+import jax
+import jax.numpy as jnp
+from jax import random
+
+from flax import linen as nn
+from flax.training import train_state
+from flax import struct
+import optax
+from .ann import ANNBackend, make_ann
+from .ddpm import DDPM
+# +from .ddpm import DDPM, DDIM
+from .dmap import DMAP
+from .gplm import GPLM
+
+__all__ = ["DDPM", "EpsMLP", "cosine_schedule", "sinusoidal_embedding"]
+
+
+# ------------------------------
+# DDIM (deterministic sampler; eta=0.0)
+# ------------------------------
+
+class DDIM(DDPM):
+    """
+    DDIM sampler for D-dimensional latents.
+
+    Notes
+    -----
+    * Training is identical to DDPM (epsilon prediction).
+    * Sampling / refinement uses the DDIM update:
+        x_{t_prev} = sqrt(a_bar_prev) * x0_hat
+                   + sqrt(1 - a_bar_prev - sigma^2) * eps_pred
+                   + sigma * z,
+      where sigma = eta * sqrt( (1-a_bar_prev)/(1-a_bar_t) * (1 - a_bar_t/a_bar_prev) ).
+
+    With eta=0.0 the reverse process is deterministic conditioned on x_t.
+
+    Added (HF I/O):
+      - state_dict / save_local / load_local / from_state
+      - upload_to_huggingface / download_from_huggingface
+    """
+
+    def __init__(
+        self,
+        Z_iX: jnp.ndarray,
+        *,
+        eta: float = 0.0,
+        steps: Optional[int] = None,
+        **ddpm_kwargs: Any,
+    ):
+        super().__init__(Z_iX, **ddpm_kwargs)
+        self.eta = float(eta)
+        self.steps = None if steps is None else int(steps)
+
+    @staticmethod
+    def _make_ddim_step(params_ema, apply_fn, alpha_bar_s, eps: float, eta: float):
+        eta = float(eta)
+
+        @jax.jit
+        def step(carry, t_pair):
+            key, x = carry  # x: (B, D)
+            t, t_prev = t_pair  # scalar int32
+
+            # gather ā_t and ā_prev (supports skipping)
+            a_bar_t = jnp.clip(alpha_bar_s[t], eps, 1.0)
+            a_bar_prev = jnp.clip(alpha_bar_s[t_prev], eps, 1.0)
+
+            B = x.shape[0]
+            t_batch = jnp.full((B, 1), t, dtype=jnp.int32)
+
+            eps_pred = apply_fn({"params": params_ema}, x, t_batch)  # (B, D)
+
+            # x0 estimate
+            sqrt_one_minus = jnp.sqrt(jnp.clip(1.0 - a_bar_t, eps, 1.0))
+            x0_hat = (x - sqrt_one_minus * eps_pred) / jnp.sqrt(a_bar_t)
+
+            # sigma_t (controls stochasticity). eta=0 => sigma=0 (deterministic)
+            frac = (1.0 - a_bar_prev) / jnp.clip(1.0 - a_bar_t, eps, 1.0)
+            inside = 1.0 - a_bar_t / jnp.clip(a_bar_prev, eps, 1.0)
+            sigma = eta * jnp.sqrt(jnp.clip(frac * inside, 0.0, 1.0))
+
+            # direction term
+            dir_coeff = jnp.sqrt(jnp.clip(1.0 - a_bar_prev - sigma**2, 0.0, 1.0))
+            x_prev = jnp.sqrt(a_bar_prev) * x0_hat + dir_coeff * eps_pred
+
+            # optional noise injection (eta>0)
+            if eta > 0.0:
+                key, k = random.split(key)
+                z = random.normal(k, x.shape)
+                x_prev = x_prev + sigma * z
+
+            return (key, x_prev), x_prev
+
+        return step
+
+    def _t_schedule(self, t_start: int, *, steps: Optional[int] = None) -> jnp.ndarray:
+        """
+        Descending integer timesteps from t_start -> 0.
+
+        If steps is provided (or self.steps), we use a reduced DDIM schedule
+        by rounding a linear grid and making it unique.
+        """
+        t_start = int(t_start)
+        if t_start <= 0:
+            return jnp.array([0], dtype=jnp.int32)
+
+        S = self.steps if steps is None else int(steps)
+        if S is None or S >= (t_start + 1):
+            return jnp.arange(t_start, -1, -1, dtype=jnp.int32)
+
+        grid = jnp.linspace(float(t_start), 0.0, int(S), dtype=jnp.float32)
+        ts = jnp.unique(jnp.round(grid).astype(jnp.int32))
+        ts = ts[::-1]  # descending
+
+        # guarantee endpoints
+        if ts[0] != t_start:
+            ts = jnp.concatenate([jnp.array([t_start], dtype=jnp.int32), ts])
+        if ts[-1] != 0:
+            ts = jnp.concatenate([ts, jnp.array([0], dtype=jnp.int32)])
+        return ts
+
+    def refine_latents(
+        self,
+        z0: jnp.ndarray,
+        t_start: int = 10,
+        key: Optional[jax.Array] = None,
+        add_noise: bool = True,
+        *,
+        eta: Optional[float] = None,
+        steps: Optional[int] = None,
+    ) -> jnp.ndarray:
+        """
+        Refine latents by:
+          (optional) forward-noise z0 to step t_start
+          DDIM reverse from t_start -> 0 using EMA params.
+        """
+        z0 = jnp.asarray(z0, dtype=jnp.float32)
+        if z0.ndim != 2 or z0.shape[1] != self.D:
+            raise ValueError(f"z0 must have shape (B,{self.D}).")
+        if not (0 <= int(t_start) < self.T):
+            raise ValueError(f"t_start must be in [0, {self.T-1}]")
+        t_start = int(t_start)
+
+        # RNG
+        if key is None:
+            self.key, key = random.split(self.key)
+        else:
+            # advance internal RNG too
+            self.key, _ = random.split(key)
+
+        # optional forward noise
+        if add_noise and t_start > 0:
+            key, k_eps = random.split(key)
+            eps_noise = random.normal(k_eps, z0.shape)
+            a_bar_t = jnp.clip(self.alpha_bar_s[t_start], self.eps, 1.0) # Corrected attribute name
+            z_t = jnp.sqrt(a_bar_t) * z0 + jnp.sqrt(1.0 - a_bar_t) * eps_noise
+        else:
+            z_t = z0
+
+        # schedule
+        ts = self._t_schedule(t_start, steps=steps)
+        if ts.shape[0] == 1:
+            return z_t
+
+        t_pairs = jnp.stack([ts[:-1], ts[1:]], axis=1)  # (K,2)
+
+        step = self._make_ddim_step(
+            self.state.ema_params,
+            self.state.apply_fn,
+            self.alpha_bar_s, # Corrected attribute name
+            self.eps,
+            self.eta if eta is None else float(eta),
+        )
+
+        (final_key, _), trace = jax.lax.scan(step, (key, z_t), xs=t_pairs)
+
+        self.key = final_key
+        return trace[-1]
+
+    def __call__(
+        self,
+        z0: jnp.ndarray,
+        t_start: int = 10,
+        key: Optional[jax.Array] = None,
+        add_noise: bool = True,
+        *,
+        eta: Optional[float] = None,
+        steps: Optional[int] = None,
+    ) -> jnp.ndarray:
+        return self.refine_latents(
+            z0, t_start=t_start, key=key, add_noise=add_noise, eta=eta, steps=steps
+        )
+
+    def reverse_from_T(self, x_T: jnp.ndarray, *, eta: Optional[float] = None, steps: Optional[int] = None) -> jnp.ndarray:
+        x_T = jnp.asarray(x_T, dtype=jnp.float32)
+        if x_T.ndim != 2 or x_T.shape[1] != self.D:
+            raise ValueError(f"x_T must have shape (B,{self.D}).")
+
+        self.key, k0 = random.split(self.key)
+
+        ts = self._t_schedule(self.T - 1, steps=steps)
+        t_pairs = jnp.stack([ts[:-1], ts[1:]], axis=1)
+
+        step = self._make_ddim_step(
+            self.state.ema_params,
+            self.state.apply_fn,
+            self.alpha_bar_s, # Corrected attribute name
+            self.eps,
+            self.eta if eta is None else float(eta),
+        )
+
+        (_, _), trace = jax.lax.scan(step, (k0, x_T), xs=t_pairs)
+        return trace[-1]
+
+    def sample(self, N: int = 10_000, *, eta: Optional[float] = None, steps: Optional[int] = None) -> jnp.ndarray:
+        self.key, k = random.split(self.key)
+        noise = random.normal(k, (int(N), self.D))
+        return self.reverse_from_T(noise, eta=eta, steps=steps)
+
+    # ------------------------------------------------------------------
+    # Serialization + Hugging Face Hub I/O (integrated)
+    # ------------------------------------------------------------------
+
+    def state_dict(self):
+        """
+        Msgpack-safe checkpoint dict (includes DDIM-specific eta/steps).
+        Uses flax.serialization.to_state_dict for TrainStateEMA to avoid tuple issues (opt_state).
+        """
+        from flax import serialization as flax_ser  # type: ignore
+        import numpy as np
+
+        return dict(
+            # DDPM core hyperparameters
+            T=int(self.T),
+            D=int(self.D),
+            hidden_dim=int(getattr(self, "hidden_dim", 0) or self.hidden_dim),
+            t_embed_dim=int(getattr(self, "t_embed_dim", 0) or self.t_embed_dim),
+            learning_rate=float(self.learning_rate),
+            ema_decay=float(self.ema_decay),
+            beta_max=float(self.beta_max),
+            batch_size=None if (self.batch_size is None) else int(self.batch_size),
+            eps=float(self.eps),
+
+            # DDIM specifics
+            eta=float(self.eta),
+            steps=None if (self.steps is None) else int(self.steps),
+
+            # RNG key (portable)
+            key=np.asarray(self.key),
+
+            # TrainStateEMA (params, ema_params, opt_state, step)
+            train_state=flax_ser.to_state_dict(self.state),
+        )
+
+    def save_local(self, weights_file: str = "ddim.msgpack", config_file: Optional[str] = "ddim_config.json") -> None:
+        from flax import serialization as flax_ser  # type: ignore
+        import json as _json
+
+        ckpt = self.state_dict()
+        blob = flax_ser.msgpack_serialize(ckpt)
+
+        with open(weights_file, "wb") as f:
+            f.write(blob)
+
+        if config_file is not None:
+            cfg = dict(
+                T=int(ckpt["T"]),
+                D=int(ckpt["D"]),
+                hidden_dim=int(ckpt["hidden_dim"]),
+                t_embed_dim=int(ckpt["t_embed_dim"]),
+                learning_rate=float(ckpt["learning_rate"]),
+                ema_decay=float(ckpt["ema_decay"]),
+                beta_max=float(ckpt["beta_max"]),
+                batch_size=ckpt["batch_size"],
+                eps=float(ckpt["eps"]),
+                eta=float(ckpt["eta"]),
+                steps=ckpt["steps"],
+            )
+            with open(config_file, "w") as f:
+                _json.dump(cfg, f, indent=2)
+
+    @classmethod
+    def from_state(cls, ckpt, *, key: Optional[jax.Array] = None) -> "DDIM":
+        """
+        Rehydrate a DDIM instance without retraining:
+          - rebuilds a skeleton DDIM (which builds schedules/model/optimizer)
+          - restores TrainStateEMA via flax.serialization.from_state_dict
+        """
+        from flax import serialization as flax_ser  # type: ignore
+        import numpy as np
+
+        T = int(ckpt["T"])
+        D = int(ckpt["D"])
+        hidden_dim = int(ckpt["hidden_dim"])
+        t_embed_dim = int(ckpt["t_embed_dim"])
+        learning_rate = float(ckpt["learning_rate"])
+        ema_decay = float(ckpt["ema_decay"])
+        beta_max = float(ckpt["beta_max"])
+        batch_size = ckpt.get("batch_size", None)
+        eps = float(ckpt.get("eps", 1e-5))
+
+        eta = float(ckpt.get("eta", 0.0))
+        steps = ckpt.get("steps", None)
+        steps = None if (steps is None) else int(steps)
+
+        if key is None:
+            if "key" in ckpt:
+                key = jnp.asarray(np.asarray(ckpt["key"]))
+            else:
+                key = random.PRNGKey(0)
+
+        # skeleton (no training)
+        dummy = jnp.zeros((1, D), dtype=jnp.float32)
+        obj = cls(
+            dummy,
+            T=T,
+            hidden_dim=hidden_dim,
+            t_embed_dim=t_embed_dim,
+            learning_rate=learning_rate,
+            n_iter=0,
+            ema_decay=ema_decay,
+            beta_max=beta_max,
+            batch_size=batch_size,
+            key=key,
+            verbose_every=0,
+            eps=eps,
+            eta=eta,
+            steps=steps,
+        )
+
+        obj.state = flax_ser.from_state_dict(obj.state, ckpt["train_state"])
+
+        # restore RNG key and DDIM attrs
+        if "key" in ckpt:
+            obj.key = jnp.asarray(np.asarray(ckpt["key"]))
+        obj.eta = eta
+        obj.steps = steps
+
+        return obj
+
+    @classmethod
+    def load_local(cls, weights_file: str = "ddim.msgpack", *, key: Optional[jax.Array] = None) -> "DDIM":
+        from flax import serialization as flax_ser  # type: ignore
+
+        with open(weights_file, "rb") as f:
+            ckpt = flax_ser.msgpack_restore(f.read())
+        return cls.from_state(ckpt, key=key)
+
+    def upload_to_huggingface(
+        self,
+        repo_id: str,
+        *,
+        weights_file: str = "ddim.msgpack",
+        config_file: str = "ddim_config.json",
+        token: Optional[str] = None,
+        repo_type: str = "model",
+        revision: Optional[str] = None,
+    ) -> None:
+        """
+        Upload DDIM checkpoint to HF Hub (msgpack + JSON).
+        Mirrors the DMAP/DDPM pattern.
+        """
+        try:
+            from huggingface_hub import HfApi, HfFolder, upload_file  # type: ignore
+        except Exception as e:
+            raise RuntimeError("huggingface_hub not installed. Install it (or `pip install dima[hf]`).") from e
+
+        self.save_local(weights_file=weights_file, config_file=config_file)
+
+        if token is None:
+            token = HfFolder.get_token()
+        if token is None:
+            raise RuntimeError("No HF token found. Run `huggingface-cli login`, or pass `token=...`.")
+
+        import os as _os
+
+        api = HfApi()
+        api.create_repo(repo_id=repo_id, repo_type=repo_type, exist_ok=True, token=token)
+
+        wf = _os.path.basename(weights_file)
+        cf = _os.path.basename(config_file)
+
+        upload_file(
+            path_or_fileobj=weights_file,
+            path_in_repo=wf,
+            repo_id=repo_id,
+            token=token,
+            repo_type=repo_type,
+            revision=revision,
+        )
+        upload_file(
+            path_or_fileobj=config_file,
+            path_in_repo=cf,
+            repo_id=repo_id,
+            token=token,
+            repo_type=repo_type,
+            revision=revision,
+        )
+
+    @classmethod
+    def download_from_huggingface(
+        cls,
+        repo_id: str,
+        *,
+        weights_file: str = "ddim.msgpack",
+        token: Optional[str] = None,
+        repo_type: str = "model",
+        revision: Optional[str] = None,
+        key: Optional[jax.Array] = None,
+    ) -> "DDIM":
+        """
+        Download DDIM checkpoint from HF Hub and return a rehydrated DDIM instance.
+        """
+        try:
+            from huggingface_hub import hf_hub_download  # type: ignore
+        except Exception as e:
+            raise RuntimeError("huggingface_hub not installed. Install it (or `pip install dima[hf]`).") from e
+
+        path = hf_hub_download(
+            repo_id=repo_id,
+            filename=weights_file,
+            token=token,
+            repo_type=repo_type,
+            revision=revision,
+        )
+        return cls.load_local(path, key=key)
+
+    # optional alias for naming symmetry
+    download_to_huggingface = download_from_huggingface
+
+
+__all__ = ["DDPM", "DDIM", "EpsMLP", "cosine_schedule", "sinusoidal_embedding"]
+
+# --- src/dima/dima.py
+# +++ src/dima/dima.py
+# @@ -15,7 +15,7 @@
+#  from flax import serialization as flax_ser
+ 
+
+ 
+# @@ -374,7 +374,13 @@
+#          ddpm_init["eps"] = float(ddpm_eps)
+ 
+#          with jax.default_device(self.ddpm_device):
+# -            self.dm = DDPM(Z_ix_j, **ddpm_init)
+# +            sampler = str(ddpm_init.pop("sampler", "ddpm")).lower()
+# +            if sampler == "ddim":
+# +                eta = float(ddpm_init.pop("eta", 0.0))
+# +                steps = ddpm_init.pop("steps", None)
+# +                self.dm = DDIM(Z_ix_j, eta=eta, steps=steps, **ddpm_init)
+# +            else:
+# +                self.dm = DDPM(Z_ix_j, **ddpm_init)
+ 
+#          self.training_time = time.time() - t0
+ 

ddpm.py

@@ -0,0 +1,591 @@
+# src/dima/ddpm.py
+from __future__ import annotations
+
+from typing import Any, Optional
+
+import jax
+import jax.numpy as jnp
+from jax import random
+
+from flax import linen as nn
+from flax.training import train_state
+from flax import struct
+import optax
+
+
+# ------------------------------
+# Schedules & embeddings
+# ------------------------------
+def cosine_schedule(T: int, s: float = 0.008):
+    """
+    Nichol & Dhariwal cosine schedule.
+    Returns:
+        alpha:     (T,)
+        beta:      (T,)
+        alpha_bar: (T,)
+    """
+    steps = jnp.arange(T + 1, dtype=jnp.float32)
+    f = jnp.cos(((steps / T + s) / (1.0 + s)) * jnp.pi / 2.0) ** 2
+    alpha_bar_all = f / f[0]
+    alpha_bar = alpha_bar_all[1:]  # (T,)
+    alpha = alpha_bar / jnp.concatenate([jnp.array([1.0], dtype=jnp.float32), alpha_bar[:-1]])
+    beta = 1.0 - alpha
+    return alpha, beta, alpha_bar
+
+
+def sinusoidal_embedding(t_idx: jnp.ndarray, dim: int) -> jnp.ndarray:
+    """
+    t_idx: (B,1) int32
+    returns: (B,dim)
+    """
+    if t_idx.ndim != 2 or t_idx.shape[1] != 1:
+        raise ValueError("t_idx must have shape (B,1)")
+    t = t_idx.astype(jnp.float32)
+    half = dim // 2
+    denom = float(max(half - 1, 1))
+    freqs = jnp.exp(-jnp.log(10_000.0) * jnp.arange(half, dtype=jnp.float32) / denom)
+    args = t * freqs
+    emb = jnp.concatenate([jnp.sin(args), jnp.cos(args)], axis=-1)
+    if dim % 2 == 1:
+        emb = jnp.pad(emb, ((0, 0), (0, 1)))
+    return emb
+
+
+# ------------------------------
+# Model: epsilon predictor
+# ------------------------------
+class EpsMLP(nn.Module):
+    """Simple MLP epsilon-predictor for DDPM in R^D."""
+    hidden: int
+    t_dim: int
+    data_dim: int
+
+    @nn.compact
+    def __call__(self, x: jnp.ndarray, t_idx: jnp.ndarray) -> jnp.ndarray:
+        # x: (B,D), t_idx: (B,1)
+        t_emb = sinusoidal_embedding(t_idx, self.t_dim)
+        t_h = nn.Dense(self.hidden)(t_emb)
+        t_h = nn.gelu(t_h)
+
+        h = nn.Dense(self.hidden)(x)
+        h = nn.gelu(h + t_h)
+
+        t_h2 = nn.Dense(self.hidden)(t_h)
+        h = nn.Dense(self.hidden)(h)
+        h = nn.gelu(h + t_h2)
+
+        out = nn.Dense(self.data_dim)(h)
+        return out
+
+
+# ------------------------------
+# TrainState with EMA
+# ------------------------------
+@struct.dataclass
+class TrainStateEMA(train_state.TrainState):
+    """Flax TrainState extended with EMA params."""
+    ema_params: Any = struct.field(pytree_node=True)
+
+    def apply_gradients(self, *, grads, ema_decay: float):
+        updates, new_opt_state = self.tx.update(grads, self.opt_state, self.params)
+        new_params = optax.apply_updates(self.params, updates)
+        new_ema = optax.incremental_update(new_params, self.ema_params, step_size=1.0 - ema_decay)
+        return self.replace(
+            step=self.step + 1,
+            params=new_params,
+            opt_state=new_opt_state,
+            ema_params=new_ema,
+        )
+
+
+# ------------------------------
+# DDPM
+# ------------------------------
+class DDPM:
+    """
+    DDPM for D-dimensional latents.
+
+    API expected by your DIMA wrapper:
+      - DDPM(Z_train, ...): trains in __init__ (n_iter can be 0 to skip)
+      - refine_latents(z0, t_start, key, add_noise) -> z_refined
+      - __call__(...) delegates to refine_latents
+      - sample(N) -> latent samples
+
+    Added:
+      - state_dict / save_local / load_local / from_state
+      - upload_to_huggingface / download_from_huggingface
+    """
+
+    def __init__(
+        self,
+        Z_iX: jnp.ndarray,
+        *,
+        T: int = 100,
+        hidden_dim: int = 128,
+        t_embed_dim: int = 64,
+        learning_rate: float = 1e-3,
+        n_iter: int = 20_000,
+        ema_decay: float = 0.999,
+        beta_max: float = 0.02,
+        batch_size: Optional[int] = None,
+        key: jax.Array = random.PRNGKey(0),
+        verbose_every: int = 0,
+        eps: float = 1e-5,
+    ):
+        Z_iX = jnp.asarray(Z_iX, dtype=jnp.float32)
+        if Z_iX.ndim != 2:
+            raise ValueError("Z_iX must be 2D (N,D).")
+
+        self.D = int(Z_iX.shape[1])
+        self.T = int(T)
+
+        self.key = key
+        self.ema_decay = float(ema_decay)
+        self.batch_size = batch_size
+        self.verbose_every = int(verbose_every)
+        self.eps = float(eps)
+        self.beta_max = float(beta_max)
+
+        # store scalars for serialization/rebuild
+        self.learning_rate = float(learning_rate)
+        self.hidden_dim = int(hidden_dim)
+        self.t_embed_dim = int(t_embed_dim)
+
+        # schedule (cosine, clipped by beta_max)
+        alpha, beta, alpha_bar = cosine_schedule(self.T)
+        beta = jnp.minimum(beta, self.beta_max)
+        alpha = 1.0 - beta
+        alpha_bar = jnp.cumprod(alpha)
+
+        self.alpha_s = alpha.astype(jnp.float32)
+        self.beta_s = beta.astype(jnp.float32)
+        self.alpha_bar_s = alpha_bar.astype(jnp.float32)
+
+        # model + optimizer + EMA
+        self.model = EpsMLP(hidden=self.hidden_dim, t_dim=self.t_embed_dim, data_dim=self.D)
+
+        params = self.model.init(
+            self.key,
+            jnp.zeros((1, self.D), dtype=jnp.float32),
+            jnp.zeros((1, 1), dtype=jnp.int32),
+        )["params"]
+
+        tx = optax.adam(self.learning_rate)
+        self.state = TrainStateEMA.create(apply_fn=self.model.apply, params=params, tx=tx, ema_params=params)
+
+        # train
+        if int(n_iter) > 0:
+            self._train(Z_iX, int(n_iter))
+
+    # ---------- training ----------
+
+    @staticmethod
+    def _loss(params, apply_fn, x_t, t_idx, eps_true):
+        eps_pred = apply_fn({"params": params}, x_t, t_idx)
+        return jnp.mean((eps_pred - eps_true) ** 2)
+
+    @staticmethod
+    @jax.jit
+    def _train_step(
+        state: TrainStateEMA,
+        x0_batch: jnp.ndarray,
+        key: jax.Array,
+        alpha_bar_s: jnp.ndarray,
+        ema_decay: float,
+        eps: float,
+    ):
+        B = x0_batch.shape[0]
+        key, k_eps, k_t = random.split(key, 3);
+
+        eps_noise = random.normal(k_eps, shape=x0_batch.shape)  # (B,D)
+        t_idx = random.randint(k_t, shape=(B, 1), minval=0, maxval=alpha_bar_s.shape[0])
+
+        a_bar_t = jnp.take(alpha_bar_s, t_idx.squeeze(-1))[:, None]
+        a_bar_t = jnp.clip(a_bar_t, eps, 1.0)
+
+        x_t = jnp.sqrt(a_bar_t) * x0_batch + jnp.sqrt(1.0 - a_bar_t) * eps_noise
+
+        def loss_fn(p):
+            return DDPM._loss(p, state.apply_fn, x_t, t_idx, eps_noise)
+
+        loss, grads = jax.value_and_grad(loss_fn)(state.params)
+        new_state = state.apply_gradients(grads=grads, ema_decay=ema_decay)
+        return new_state, loss, key
+
+    def _train(self, Z_iX: jnp.ndarray, n_iter: int):
+        N = int(Z_iX.shape[0])
+        bs = N if (self.batch_size is None) else min(int(self.batch_size), N)
+
+        for it in range(n_iter):
+            if bs >= N:
+                batch = Z_iX
+            else:
+                self.key, k_perm = random.split(self.key)
+                idx = random.permutation(k_perm, N)[:bs]
+                batch = Z_iX[idx]
+
+            self.state, loss, self.key = self._train_step(
+                self.state,
+                batch,
+                self.key,
+                self.alpha_bar_s,
+                self.ema_decay,
+                self.eps,
+            )
+
+            if self.verbose_every and (it % self.verbose_every == 0 or it == n_iter - 1):
+                print(f"iter {it:6d} loss {float(loss):.6f}", end="\r")
+
+        if self.verbose_every:
+            print("\ntraining complete.")
+
+    # ---------- diffusion utilities ----------
+
+    @staticmethod
+    def _posterior_variance(alpha_s, beta_s, alpha_bar_s, t):
+        a_bar_t = alpha_bar_s[t]
+        a_bar_prev = jnp.where(t > 0, alpha_bar_s[t - 1], jnp.array(1.0, dtype=alpha_bar_s.dtype))
+        return ((1.0 - a_bar_prev) / (1.0 - a_bar_t)) * beta_s[t]
+
+    @staticmethod
+    def _make_sampler_step(params_ema, apply_fn, alpha_s, beta_s, alpha_bar_s, eps: float):
+        @jax.jit
+        def step(carry, _):
+            key, t, x = carry  # x: (B,D)
+            key, k = random.split(key)
+
+            alpha_t = jnp.clip(alpha_s[t], eps, 1.0)
+            a_bar_t = jnp.clip(alpha_bar_s[t], eps, 1.0)
+
+            sqrt_alpha = jnp.sqrt(alpha_t)
+            sqrt_one_minus_a_bar = jnp.sqrt(jnp.clip(1.0 - a_bar_t, eps, 1.0))
+
+            B = x.shape[0]
+            t_batch = jnp.full((B, 1), t, dtype=jnp.int32)
+
+            eps_pred = apply_fn({"params": params_ema}, x, t_batch)  # (B,D)
+
+            # predict x0
+            x0_hat = (x - sqrt_one_minus_a_bar * eps_pred) / jnp.sqrt(a_bar_t)
+
+            a_bar_prev = jnp.where(t > 0, alpha_bar_s[t - 1], jnp.array(1.0, dtype=alpha_bar_s.dtype))
+            denom = jnp.clip(1.0 - a_bar_t, eps, 1.0)
+
+            coef1 = jnp.sqrt(jnp.clip(a_bar_prev, eps, 1.0)) * beta_s[t] / denom
+            coef2 = sqrt_alpha * (1.0 - a_bar_prev) / denom
+            mean = coef1 * x0_hat + coef2 * x
+
+            beta_tilde = DDPM._posterior_variance(alpha_s, beta_s, alpha_bar_s, t)
+            sigma = jnp.sqrt(jnp.clip(beta_tilde, 0.0, 1.0))
+
+            z = random.normal(k, x.shape)
+            z = jnp.where(t == 0, 0.0, z);
+
+            x_prev = mean + sigma * z
+            return (key, t - 1, x_prev), x_prev
+
+        return step
+
+    # ---------- API ----------
+
+    def refine_latents(
+        self,
+        z0: jnp.ndarray,
+        t_start: int = 10,
+        key: Optional[jax.Array] = None,
+        add_noise: bool = True,
+    ) -> jnp.ndarray:
+        """
+        Refine latents by:
+          (optional) forward-noise z0 to step t_start
+          reverse-diffuse from t_start -> 0 using EMA params.
+        """
+        z0 = jnp.asarray(z0, dtype=jnp.float32)
+        if z0.ndim != 2 or z0.shape[1] != self.D:
+            raise ValueError(f"z0 must have shape (B,{self.D}).")
+
+        if not (0 <= int(t_start) < self.T):
+            raise ValueError(f"t_start must be in [0, {self.T-1}]")
+        t_start = int(t_start)
+
+        if key is None:
+            self.key, key = random.split(self.key)
+        else:
+            # advance internal RNG too
+            self.key, _ = random.split(key)
+
+        # forward diffuse to t_start
+        key, k_eps = random.split(key)
+        eps_noise = random.normal(k_eps, z0.shape)
+
+        a_bar_t = jnp.clip(self.alpha_bar_s[t_start], self.eps, 1.0)
+
+        if add_noise:
+            z_t = jnp.sqrt(a_bar_t) * z0 + jnp.sqrt(1.0 - a_bar_t) * eps_noise
+        else:
+            z_t = z0
+
+        step = self._make_sampler_step(
+            self.state.ema_params,
+            self.state.apply_fn,
+            self.alpha_s,
+            self.beta_s,
+            self.alpha_bar_s,
+            self.eps,
+        )
+
+        (final_key, _, _), trace = jax.lax.scan(
+            step,
+            (key, t_start, z_t),
+            xs=None,
+            length=t_start + 1,
+        )
+
+        self.key = final_key
+        return trace[-1]
+
+    def __call__(
+        self,
+        z0: jnp.ndarray,
+        t_start: int = 10,
+        key: Optional[jax.Array] = None,
+        add_noise: bool = True,
+    ) -> jnp.ndarray:
+        return self.refine_latents(z0, t_start=t_start, key=key, add_noise=add_noise)
+
+    def reverse_from_T(self, x_T: jnp.ndarray) -> jnp.ndarray:
+        x_T = jnp.asarray(x_T, dtype=jnp.float32)
+        if x_T.ndim != 2 or x_T.shape[1] != self.D:
+            raise ValueError(f"x_T must have shape (B,{self.D}).")
+
+        step = self._make_sampler_step(
+            self.state.ema_params,
+            self.state.apply_fn,
+            self.alpha_s,
+            self.beta_s,
+            self.alpha_bar_s,
+            self.eps,
+        )
+
+        self.key, k0 = random.split(self.key)
+
+        (_, _, _), trace = jax.lax.scan(
+            step,
+            (k0, self.T - 1, x_T),
+            xs=None,
+            length=self.T,
+        )
+
+        return trace[-1]
+
+    def sample(self, N: int = 10_000) -> jnp.ndarray:
+        self.key, k = random.split(self.key)
+        noise = random.normal(k, (int(N), self.D))
+        return self.reverse_from_T(noise)
+
+    # -------------------------
+    # Serialization helpers
+    # -------------------------
+
+    def state_dict(self):
+        """
+        Msgpack-safe checkpoint dict.
+        Uses flax.serialization.to_state_dict to avoid tuple/namedtuple issues in opt_state.
+        """
+        from flax import serialization as flax_ser  # type: ignore
+        import numpy as np
+
+        return dict(
+            # hyperparameters needed to rebuild the object
+            T=int(self.T),
+            D=int(self.D),
+            hidden_dim=int(self.hidden_dim),
+            t_embed_dim=int(self.t_embed_dim),
+            learning_rate=float(self.learning_rate),
+            ema_decay=float(self.ema_decay),
+            beta_max=float(self.beta_max),
+            batch_size=None if (self.batch_size is None) else int(self.batch_size),
+            eps=float(self.eps),
+
+            # RNG key (store as numpy for portability)
+            key=np.asarray(self.key),
+
+            # TrainStateEMA (params, ema_params, opt_state, step) in a serializable form
+            train_state=flax_ser.to_state_dict(self.state),
+        )
+
+
+    def save_local(self, weights_file: str = "ddpm.msgpack", config_file: Optional[str] = "ddpm_config.json") -> None:
+        from flax import serialization as flax_ser  # type: ignore
+        import json as _json
+
+        ckpt = self.state_dict()
+        blob = flax_ser.msgpack_serialize(ckpt)
+
+        with open(weights_file, "wb") as f:
+            f.write(blob)
+
+        if config_file is not None:
+            cfg = dict(
+                T=int(ckpt["T"]),
+                D=int(ckpt["D"]),
+                hidden_dim=int(ckpt["hidden_dim"]),
+                t_embed_dim=int(ckpt["t_embed_dim"]),
+                learning_rate=float(ckpt["learning_rate"]),
+                ema_decay=float(ckpt["ema_decay"]),
+                beta_max=float(ckpt["beta_max"]),
+                batch_size=ckpt["batch_size"],
+                eps=float(ckpt["eps"]),
+            )
+            with open(config_file, "w") as f:
+                _json.dump(cfg, f, indent=2)
+
+
+    @classmethod
+    
+    def from_state(cls, ckpt, *, key: Optional[jax.Array] = None) -> "DDPM":
+        """
+        Rehydrate a DDPM instance without retraining, restoring full TrainStateEMA
+        (including opt_state) via flax.serialization.from_state_dict.
+        """
+        from flax import serialization as flax_ser  # type: ignore
+        import numpy as np
+
+        T = int(ckpt["T"])
+        D = int(ckpt["D"])
+        hidden_dim = int(ckpt["hidden_dim"])
+        t_embed_dim = int(ckpt["t_embed_dim"])
+        learning_rate = float(ckpt["learning_rate"])
+        ema_decay = float(ckpt["ema_decay"])
+        beta_max = float(ckpt["beta_max"])
+        batch_size = ckpt.get("batch_size", None)
+        eps = float(ckpt.get("eps", 1e-5))
+
+        if key is None:
+            # prefer explicit passed key; else use checkpoint key; else default
+            if "key" in ckpt:
+                key = jnp.asarray(np.asarray(ckpt["key"]))
+            else:
+                key = random.PRNGKey(0)
+
+        # Build a skeleton instance (no training)
+        dummy = jnp.zeros((1, D), dtype=jnp.float32)
+        obj = cls(
+            dummy,
+            T=T,
+            hidden_dim=hidden_dim,
+            t_embed_dim=t_embed_dim,
+            learning_rate=learning_rate,
+            n_iter=0,
+            ema_decay=ema_decay,
+            beta_max=beta_max,
+            batch_size=batch_size,
+            key=key,
+            verbose_every=0,
+            eps=eps,
+        )
+
+        # Restore TrainStateEMA (including opt_state tuples) safely
+        obj.state = flax_ser.from_state_dict(obj.state, ckpt["train_state"])
+
+        # Restore RNG key if present (keeps reproducibility)
+        if "key" in ckpt:
+            obj.key = jnp.asarray(np.asarray(ckpt["key"]))
+
+        return obj
+
+
+    @classmethod
+    def load_local(cls, weights_file: str = "ddpm.msgpack", *, key: Optional[jax.Array] = None) -> "DDPM":
+        from flax import serialization as flax_ser  # type: ignore
+    
+        with open(weights_file, "rb") as f:
+            ckpt = flax_ser.msgpack_restore(f.read())
+    
+        return cls.from_state(ckpt, key=key)
+
+
+    def upload_to_huggingface(
+        self,
+        repo_id: str,
+        *,
+        weights_file: str = "ddpm.msgpack",
+        config_file: str = "ddpm_config.json",
+        token: Optional[str] = None,
+        repo_type: str = "model",
+        revision: Optional[str] = None,
+    ) -> None:
+        """
+        Upload DDPM weights to the Hugging Face Hub.
+        Saves locally first (msgpack + JSON), then creates repo (if needed) and uploads files.
+        """
+        try:
+            from huggingface_hub import HfApi, HfFolder, upload_file  # type: ignore
+        except Exception as e:
+            raise RuntimeError("huggingface_hub not installed. Install it (or `pip install dima[hf]`).") from e
+
+        self.save_local(weights_file=weights_file, config_file=config_file)
+
+        if token is None:
+            token = HfFolder.get_token()
+        if token is None:
+            raise RuntimeError("No HF token found. Run `huggingface-cli login`, or pass `token=...`.")
+
+        import os as _os
+
+        api = HfApi()
+        api.create_repo(repo_id=repo_id, repo_type=repo_type, exist_ok=True, token=token)
+
+        wf = _os.path.basename(weights_file)
+        cf = _os.path.basename(config_file)
+
+        upload_file(
+            path_or_fileobj=weights_file,
+            path_in_repo=wf,
+            repo_id=repo_id,
+            token=token,
+            repo_type=repo_type,
+            revision=revision,
+        )
+        upload_file(
+            path_or_fileobj=config_file,
+            path_in_repo=cf,
+            repo_id=repo_id,
+            token=token,
+            repo_type=repo_type,
+            revision=revision,
+        )
+
+    @classmethod
+    def download_from_huggingface(
+        cls,
+        repo_id: str,
+        *,
+        weights_file: str = "ddpm.msgpack",
+        token: Optional[str] = None,
+        repo_type: str = "model",
+        revision: Optional[str] = None,
+        key: Optional[jax.Array] = None,
+    ) -> "DDPM":
+        """
+        Download DDPM weights from the Hugging Face Hub and return a rehydrated DDPM instance.
+        """
+        try:
+            from huggingface_hub import hf_hub_download  # type: ignore
+        except Exception as e:
+            raise RuntimeError("huggingface_hub not installed. Install it (or `pip install dima[hf]`).") from e
+
+        path = hf_hub_download(
+            repo_id=repo_id,
+            filename=weights_file,
+            token=token,
+            repo_type=repo_type,
+            revision=revision,
+        )
+        return cls.load_local(path, key=key)
+
+    # Optional backward-compatible alias (if you ever used the other name)
+    download_to_huggingface = download_from_huggingface
+
+
+__all__ = ["DDPM", "EpsMLP", "cosine_schedule", "sinusoidal_embedding"]

ddpmx.py

@@ -0,0 +1,745 @@
+# how to use 
+# ddpm = DDPM(Z_train, T=100, n_iter=20_000, key=random.PRNGKey(0))
+# 
+# Fast unconditional samples (try 15–30 first)
+# z = ddpm.sample_dpmpp(N=4096, num_steps=20)
+#
+# Fast “refine latents” starting from an intermediate noise level
+#z_ref = ddpm.refine_latents_dpmpp(z0, t_start=30, num_steps=20, add_noise=True)
+
+
+
+# src/dima/ddpmx.py
+from __future__ import annotations
+
+from typing import Any, Optional, Dict
+import os
+import json
+
+import numpy as np
+import jax
+import jax.numpy as jnp
+from jax import random
+
+from flax import linen as nn
+from flax.training import train_state
+from flax import struct, serialization as flax_ser
+import optax
+
+
+# ---------------------------------------------------------------------
+# Helpers (as in ddpmx.py)
+# ---------------------------------------------------------------------
+def _sigma_to_alpha_sigma_t(sigma: jnp.ndarray) -> tuple[jnp.ndarray, jnp.ndarray]:
+    """
+    EDM-style sigma parameterization:
+      alpha_t = 1 / sqrt(1 + sigma^2)
+      sigma_t = sigma * alpha_t
+    so that x = alpha_t * x0 + sigma_t * eps
+    """
+    alpha_t = 1.0 / jnp.sqrt(1.0 + sigma**2)
+    sigma_t = sigma * alpha_t
+    return alpha_t, sigma_t
+
+
+def _make_lu_sigma_schedule(sigma_start: float, sigma_end: float, num_steps: int) -> np.ndarray:
+    """
+    "Lu" schedule uniform in lambda = -log(sigma).
+    """
+    sigma_start = float(max(sigma_start, 1e-12))
+    sigma_end = float(max(sigma_end, 1e-12))
+    lam_start = -np.log(sigma_start)
+    lam_end = -np.log(sigma_end)
+    lambdas = np.linspace(lam_start, lam_end, int(num_steps), dtype=np.float32)
+    sigmas = np.exp(-lambdas).astype(np.float32)
+    return sigmas
+
+
+def cosine_schedule(T: int, s: float = 0.008):
+    """
+    Nichol & Dhariwal cosine schedule.
+    Returns alpha, beta, alpha_bar with shape (T,).
+    """
+    steps = jnp.arange(T + 1, dtype=jnp.float32)
+    f = jnp.cos(((steps / T + s) / (1.0 + s)) * jnp.pi / 2.0) ** 2
+    alpha_bar_all = f / f[0]
+    alpha_bar = alpha_bar_all[1:]  # (T,)
+    alpha = alpha_bar / jnp.concatenate([jnp.array([1.0], dtype=jnp.float32), alpha_bar[:-1]])
+    beta = 1.0 - alpha
+    return alpha, beta, alpha_bar
+
+
+def sinusoidal_embedding(t_idx: jnp.ndarray, dim: int) -> jnp.ndarray:
+    """
+    t_idx: (B,1) int32 or float32
+    returns: (B,dim)
+    """
+    if t_idx.ndim != 2 or t_idx.shape[1] != 1:
+        raise ValueError("t_idx must have shape (B,1)")
+    t = t_idx.astype(jnp.float32)
+    half = dim // 2
+    denom = float(max(half - 1, 1))
+    freqs = jnp.exp(-jnp.log(10_000.0) * jnp.arange(half, dtype=jnp.float32) / denom)
+    args = t * freqs
+    emb = jnp.concatenate([jnp.sin(args), jnp.cos(args)], axis=-1)
+    if dim % 2 == 1:
+        emb = jnp.pad(emb, ((0, 0), (0, 1)))
+    return emb
+
+
+class EpsMLP(nn.Module):
+    """Simple MLP epsilon-predictor for DDPM in R^D."""
+    hidden: int
+    t_dim: int
+    data_dim: int
+
+    @nn.compact
+    def __call__(self, x: jnp.ndarray, t_idx: jnp.ndarray) -> jnp.ndarray:
+        t_emb = sinusoidal_embedding(t_idx, self.t_dim)
+        t_h = nn.Dense(self.hidden)(t_emb)
+        t_h = nn.gelu(t_h)
+
+        h = nn.Dense(self.hidden)(x)
+        h = nn.gelu(h + t_h)
+
+        t_h2 = nn.Dense(self.hidden)(t_h)
+        h = nn.Dense(self.hidden)(h)
+        h = nn.gelu(h + t_h2)
+
+        out = nn.Dense(self.data_dim)(h)
+        return out
+
+
+@struct.dataclass
+class TrainStateEMA(train_state.TrainState):
+    """Flax TrainState extended with EMA params."""
+    ema_params: Any = struct.field(pytree_node=True)
+
+    def apply_gradients(self, *, grads, ema_decay: float):
+        updates, new_opt_state = self.tx.update(grads, self.opt_state, self.params)
+        new_params = optax.apply_updates(self.params, updates)
+        new_ema = optax.incremental_update(new_params, self.ema_params, step_size=1.0 - ema_decay)
+        return self.replace(
+            step=self.step + 1,
+            params=new_params,
+            opt_state=new_opt_state,
+            ema_params=new_ema,
+        )
+
+
+# ---------------------------------------------------------------------
+# DDPMX with HF upload/download integrated
+# ---------------------------------------------------------------------
+class DDPM:
+    """
+    DDPM (+ fast DPM-Solver++(2M) sampler utilities) for D-dimensional latents.
+
+    Added persistence utilities:
+      - save_local / load_local
+      - upload_to_huggingface / download_from_huggingface
+
+    Serialization is done via flax.serialization.to_state_dict / from_state_dict
+    to avoid msgpack failures with non-serializable Python objects (e.g., tuples).
+    """
+
+    def __init__(
+        self,
+        Z_iX: jnp.ndarray,
+        *,
+        T: int = 100,
+        hidden_dim: int = 128,
+        t_embed_dim: int = 64,
+        learning_rate: float = 1e-3,
+        n_iter: int = 20_000,
+        ema_decay: float = 0.999,
+        beta_max: float = 0.02,
+        batch_size: Optional[int] = None,
+        key: jax.Array = random.PRNGKey(0),
+        verbose_every: int = 0,
+        eps: float = 1e-5,
+    ):
+        Z_iX = jnp.asarray(Z_iX, dtype=jnp.float32)
+        if Z_iX.ndim != 2:
+            raise ValueError("Z_iX must be 2D (N,D).")
+
+        self.D = int(Z_iX.shape[1])
+        self.T = int(T)
+
+        # store config for checkpointing
+        self.hidden_dim = int(hidden_dim)
+        self.t_embed_dim = int(t_embed_dim)
+        self.learning_rate = float(learning_rate)
+        self.ema_decay = float(ema_decay)
+        self.beta_max = float(beta_max)
+        self.batch_size = batch_size
+        self.verbose_every = int(verbose_every)
+        self.eps = float(eps)
+
+        self.key = key
+
+        # VP schedule (cosine + clip)
+        alpha, beta, alpha_bar = cosine_schedule(self.T)
+        beta = jnp.minimum(beta, self.beta_max)
+        alpha = 1.0 - beta
+        alpha_bar = jnp.cumprod(alpha)
+
+        self.alpha_s = alpha.astype(jnp.float32)
+        self.beta_s = beta.astype(jnp.float32)
+        self.alpha_bar_s = alpha_bar.astype(jnp.float32)
+
+        # Precompute EDM-style "sigma_in" for DPM++ schedule interpolation
+        # sigma_in = sqrt((1 - a_bar) / a_bar)
+        self.sigma_in_train = jnp.sqrt(
+            jnp.clip(
+                (1.0 - self.alpha_bar_s) / jnp.clip(self.alpha_bar_s, self.eps, 1.0),
+                self.eps,
+                1e12,
+            )
+        ).astype(jnp.float32)
+
+        # model + optimizer + EMA state
+        self.model = EpsMLP(hidden=self.hidden_dim, t_dim=self.t_embed_dim, data_dim=self.D)
+        params = self.model.init(
+            self.key,
+            jnp.zeros((1, self.D), dtype=jnp.float32),
+            jnp.zeros((1, 1), dtype=jnp.int32),
+        )["params"]
+
+        tx = optax.adam(self.learning_rate)
+        self.state = TrainStateEMA.create(
+            apply_fn=self.model.apply,
+            params=params,
+            tx=tx,
+            ema_params=params,
+        )
+
+        if int(n_iter) > 0:
+            self._train(Z_iX, int(n_iter))
+
+    # -------------------------
+    # Training
+    # -------------------------
+    @staticmethod
+    def _loss(params, apply_fn, x_t, t_idx, eps_true):
+        eps_pred = apply_fn({"params": params}, x_t, t_idx)
+        return jnp.mean((eps_pred - eps_true) ** 2)
+
+    @staticmethod
+    @jax.jit
+    def _train_step(
+        state: TrainStateEMA,
+        x0_batch: jnp.ndarray,
+        key: jax.Array,
+        alpha_bar_s: jnp.ndarray,
+        ema_decay: float,
+        eps: float,
+    ):
+        B = x0_batch.shape[0]
+        key, k_eps, k_t = random.split(key, 3)
+        eps_noise = random.normal(k_eps, shape=x0_batch.shape)
+        t_idx = random.randint(k_t, shape=(B, 1), minval=0, maxval=alpha_bar_s.shape[0])
+
+        a_bar_t = jnp.take(alpha_bar_s, t_idx.squeeze(-1))[:, None]
+        a_bar_t = jnp.clip(a_bar_t, eps, 1.0)
+        x_t = jnp.sqrt(a_bar_t) * x0_batch + jnp.sqrt(1.0 - a_bar_t) * eps_noise
+
+        def loss_fn(p):
+            return DDPM._loss(p, state.apply_fn, x_t, t_idx, eps_noise)
+
+        loss, grads = jax.value_and_grad(loss_fn)(state.params)
+        new_state = state.apply_gradients(grads=grads, ema_decay=ema_decay)
+        return new_state, loss, key
+
+    def _train(self, Z_iX: jnp.ndarray, n_iter: int):
+        N = int(Z_iX.shape[0])
+        bs = N if (self.batch_size is None) else min(int(self.batch_size), N)
+
+        for it in range(n_iter):
+            if bs >= N:
+                batch = Z_iX
+            else:
+                self.key, k_perm = random.split(self.key)
+                idx = random.permutation(k_perm, N)[:bs]
+                batch = Z_iX[idx]
+
+            self.state, loss, self.key = self._train_step(
+                self.state,
+                batch,
+                self.key,
+                self.alpha_bar_s,
+                self.ema_decay,
+                self.eps,
+            )
+
+            if self.verbose_every and (it % self.verbose_every == 0 or it == n_iter - 1):
+                print(f"iter {it:6d} loss {float(loss):.6f}", end="\r")
+
+        if self.verbose_every:
+            print("\ntraining complete.")
+
+    # -------------------------
+    # Standard DDPM refine/sample
+    # -------------------------
+    @staticmethod
+    def _posterior_variance(alpha_s, beta_s, alpha_bar_s, t):
+        a_bar_t = alpha_bar_s[t]
+        a_bar_prev = jnp.where(t > 0, alpha_bar_s[t - 1], jnp.array(1.0, dtype=alpha_bar_s.dtype))
+        return ((1.0 - a_bar_prev) / (1.0 - a_bar_t)) * beta_s[t]
+
+    @staticmethod
+    def _make_sampler_step(params_ema, apply_fn, alpha_s, beta_s, alpha_bar_s, eps: float):
+        @jax.jit
+        def step(carry, _):
+            key, t, x = carry
+            key, k = random.split(key)
+
+            alpha_t = jnp.clip(alpha_s[t], eps, 1.0)
+            a_bar_t = jnp.clip(alpha_bar_s[t], eps, 1.0)
+
+            sqrt_alpha = jnp.sqrt(alpha_t)
+            sqrt_one_minus = jnp.sqrt(jnp.clip(1.0 - a_bar_t, eps, 1.0))
+
+            B = x.shape[0]
+            t_batch = jnp.full((B, 1), t, dtype=jnp.int32)
+
+            eps_pred = apply_fn({"params": params_ema}, x, t_batch)
+            x0_hat = (x - sqrt_one_minus * eps_pred) / jnp.sqrt(a_bar_t)
+
+            a_bar_prev = jnp.where(t > 0, alpha_bar_s[t - 1], jnp.array(1.0, dtype=alpha_bar_s.dtype))
+            denom = jnp.clip(1.0 - a_bar_t, eps, 1.0)
+
+            coef1 = jnp.sqrt(jnp.clip(a_bar_prev, eps, 1.0)) * beta_s[t] / denom
+            coef2 = sqrt_alpha * (1.0 - a_bar_prev) / denom
+            mean = coef1 * x0_hat + coef2 * x
+
+            beta_tilde = DDPM._posterior_variance(alpha_s, beta_s, alpha_bar_s, t)
+            sigma = jnp.sqrt(jnp.clip(beta_tilde, 0.0, 1.0))
+
+            z = random.normal(k, x.shape)
+            z = jnp.where(t == 0, 0.0, z)
+            x_prev = mean + sigma * z
+
+            return (key, t - 1, x_prev), x_prev
+
+        return step
+
+    def refine_latents(
+        self,
+        z0: jnp.ndarray,
+        t_start: int = 10,
+        key: Optional[jax.Array] = None,
+        add_noise: bool = True,
+    ) -> jnp.ndarray:
+        z0 = jnp.asarray(z0, dtype=jnp.float32)
+        if z0.ndim != 2 or z0.shape[1] != self.D:
+            raise ValueError(f"z0 must have shape (B,{self.D}).")
+        if not (0 <= int(t_start) < self.T):
+            raise ValueError(f"t_start must be in [0, {self.T-1}]")
+        t_start = int(t_start)
+
+        if key is None:
+            self.key, key = random.split(self.key)
+        else:
+            self.key, _ = random.split(key)
+
+        key, k_eps = random.split(key)
+        eps_noise = random.normal(k_eps, z0.shape)
+        a_bar_t = jnp.clip(self.alpha_bar_s[t_start], self.eps, 1.0)
+
+        if add_noise:
+            z_t = jnp.sqrt(a_bar_t) * z0 + jnp.sqrt(1.0 - a_bar_t) * eps_noise
+        else:
+            z_t = z0
+
+        step = self._make_sampler_step(
+            self.state.ema_params,
+            self.state.apply_fn,
+            self.alpha_s,
+            self.beta_s,
+            self.alpha_bar_s,
+            self.eps,
+        )
+
+        (final_key, _, _), trace = jax.lax.scan(
+            step,
+            (key, t_start, z_t),
+            xs=None,
+            length=t_start + 1,
+        )
+        self.key = final_key
+        return trace[-1]
+
+    def __call__(
+        self,
+        z0: jnp.ndarray,
+        t_start: int = 10,
+        key: Optional[jax.Array] = None,
+        add_noise: bool = True,
+    ) -> jnp.ndarray:
+        return self.refine_latents(z0, t_start=t_start, key=key, add_noise=add_noise)
+
+    def reverse_from_T(self, x_T: jnp.ndarray) -> jnp.ndarray:
+        x_T = jnp.asarray(x_T, dtype=jnp.float32)
+        if x_T.ndim != 2 or x_T.shape[1] != self.D:
+            raise ValueError(f"x_T must have shape (B,{self.D}).")
+
+        step = self._make_sampler_step(
+            self.state.ema_params,
+            self.state.apply_fn,
+            self.alpha_s,
+            self.beta_s,
+            self.alpha_bar_s,
+            self.eps,
+        )
+
+        self.key, k0 = random.split(self.key)
+        (_, _, _), trace = jax.lax.scan(
+            step,
+            (k0, self.T - 1, x_T),
+            xs=None,
+            length=self.T,
+        )
+        return trace[-1]
+
+    def sample(self, N: int = 10_000) -> jnp.ndarray:
+        self.key, k = random.split(self.key)
+        noise = random.normal(k, (int(N), self.D)).astype(jnp.float32)
+        return self.reverse_from_T(noise)
+
+    # -------------------------
+    # DPM-Solver++(2M) schedule + sampler
+    # -------------------------
+    def _make_dpmpp_schedule(self, *, num_steps: int, t_start: int) -> tuple[jnp.ndarray, jnp.ndarray]:
+        """
+        Returns:
+          sigmas_in: (K+1,) float32 decreasing, last one is 0
+          t_cont:    (K,)   float32 continuous "time" indices for model calls
+        """
+        t_start = int(t_start)
+        if not (0 <= t_start < self.T):
+            raise ValueError(f"t_start must be in [0, {self.T-1}]")
+        if int(num_steps) < 1:
+            raise ValueError("num_steps must be >= 1")
+
+        sigma_start = float(self.sigma_in_train[t_start])
+        sigma_end = float(self.sigma_in_train[0])
+
+        sigmas_k = _make_lu_sigma_schedule(sigma_start, sigma_end, int(num_steps))
+        sigmas = np.concatenate([sigmas_k, np.array([0.0], np.float32)], axis=0)
+
+        sigma_train = np.array(self.sigma_in_train).astype(np.float32)  # (T,)
+        log_sig_train = np.log(np.maximum(sigma_train, 1e-12))
+        t_train = np.arange(self.T, dtype=np.float32)
+
+        log_sig = np.log(np.maximum(sigmas[:-1], 1e-12))
+        t_cont = np.interp(log_sig, log_sig_train, t_train).astype(np.float32)
+
+        return jnp.array(sigmas, dtype=jnp.float32), jnp.array(t_cont, dtype=jnp.float32)
+
+    @staticmethod
+    @jax.jit
+    def _dpmpp_2m_midpoint_sample(
+        params_ema: Any,
+        apply_fn: Any,
+        x_start: jnp.ndarray,   # (B,D)
+        sigmas_in: jnp.ndarray, # (K+1,)
+        t_cont: jnp.ndarray,    # (K,)
+        eps: float,
+    ) -> jnp.ndarray:
+        """
+        DPM-Solver++ (2M, midpoint) sampler.
+        """
+        sigma_s = sigmas_in[:-1]   # (K,)
+        sigma_t = sigmas_in[1:]    # (K,)
+
+        alpha_s, sigma_s_t = _sigma_to_alpha_sigma_t(sigma_s)
+        alpha_t, sigma_t_t = _sigma_to_alpha_sigma_t(sigma_t)
+
+        lambda_s = jnp.log(alpha_s) - jnp.log(sigma_s_t)
+        lambda_t = jnp.log(alpha_t) - jnp.log(sigma_t_t)
+
+        K = t_cont.shape[0]
+        is_first = jnp.arange(K) == 0
+        is_last = jnp.arange(K) == (K - 1)
+
+        def step(carry, inp):
+            x, m_prev, lam_prev = carry
+            (a_s, s_s, a_t, s_t, lam_s_i, lam_t_i, t_i, first_i, last_i) = inp
+
+            B = x.shape[0]
+            t_batch = jnp.full((B, 1), t_i, dtype=jnp.float32)
+
+            eps_pred = apply_fn({"params": params_ema}, x, t_batch)
+
+            a_s_b = jnp.clip(a_s, eps, 1.0)
+            x0 = (x - s_s * eps_pred) / a_s_b
+
+            h = lam_t_i - lam_s_i
+            exp_neg_h = jnp.exp(-h)
+
+            # 1st-order
+            x_first = (s_t / s_s) * x - (a_t * (exp_neg_h - 1.0)) * x0
+
+            def do_second(_):
+                h0 = lam_s_i - lam_prev
+                r0 = h0 / jnp.clip(h, 1e-12)
+                D1 = (x0 - m_prev) / jnp.clip(r0, 1e-12)
+                x_second = (s_t / s_s) * x - (a_t * (exp_neg_h - 1.0)) * (x0 + 0.5 * D1)
+                return x_second
+
+            x_next = jax.lax.cond(first_i | last_i, lambda _: x_first, do_second, operand=None)
+            return (x_next, x0, lam_s_i), x_next
+
+        xs = (
+            alpha_s, sigma_s_t,
+            alpha_t, sigma_t_t,
+            lambda_s, lambda_t,
+            t_cont, is_first, is_last
+        )
+
+        x0_init = jnp.zeros_like(x_start)
+        lam_init = jnp.array(0.0, dtype=jnp.float32)
+
+        (x_final, _, _), _ = jax.lax.scan(step, (x_start, x0_init, lam_init), xs)
+        return x_final
+
+    def refine_latents_dpmpp(
+        self,
+        z0: jnp.ndarray,
+        *,
+        t_start: int = 10,
+        num_steps: int = 20,
+        key: Optional[jax.Array] = None,
+        add_noise: bool = True,
+    ) -> jnp.ndarray:
+        z0 = jnp.asarray(z0, dtype=jnp.float32)
+        if z0.ndim != 2 or z0.shape[1] != self.D:
+            raise ValueError(f"z0 must have shape (B,{self.D}).")
+        if not (0 <= int(t_start) < self.T):
+            raise ValueError(f"t_start must be in [0, {self.T-1}]")
+
+        if key is None:
+            self.key, key = random.split(self.key)
+        else:
+            self.key, _ = random.split(key)
+
+        # forward-noise to t_start
+        key, k_eps = random.split(key)
+        eps_noise = random.normal(k_eps, z0.shape)
+        a_bar = jnp.clip(self.alpha_bar_s[int(t_start)], self.eps, 1.0)
+
+        if add_noise:
+            x_start = jnp.sqrt(a_bar) * z0 + jnp.sqrt(1.0 - a_bar) * eps_noise
+        else:
+            x_start = z0
+
+        sigmas_in, t_cont = self._make_dpmpp_schedule(num_steps=int(num_steps), t_start=int(t_start))
+        x_final = self._dpmpp_2m_midpoint_sample(
+            self.state.ema_params,
+            self.state.apply_fn,
+            x_start,
+            sigmas_in,
+            t_cont,
+            self.eps,
+        )
+        return x_final
+
+    def reverse_from_T_dpmpp(self, x_T: jnp.ndarray, *, num_steps: int = 20) -> jnp.ndarray:
+        x_T = jnp.asarray(x_T, dtype=jnp.float32)
+        if x_T.ndim != 2 or x_T.shape[1] != self.D:
+            raise ValueError(f"x_T must have shape (B,{self.D}).")
+
+        sigmas_in, t_cont = self._make_dpmpp_schedule(num_steps=int(num_steps), t_start=self.T - 1)
+        return self._dpmpp_2m_midpoint_sample(
+            self.state.ema_params,
+            self.state.apply_fn,
+            x_T,
+            sigmas_in,
+            t_cont,
+            self.eps,
+        )
+
+    def sample_dpmpp(self, N: int = 10_000, *, num_steps: int = 20) -> jnp.ndarray:
+        self.key, k = random.split(self.key)
+        x_T = random.normal(k, (int(N), self.D)).astype(jnp.float32)
+        return self.reverse_from_T_dpmpp(x_T, num_steps=int(num_steps))
+
+    # -----------------------------------------------------------------
+    # Persistence: state_dict / save_local / load_local
+    # -----------------------------------------------------------------
+    def _config_dict(self) -> Dict[str, Any]:
+        return {
+            "class_name": "DDPMX",
+            "T": int(self.T),
+            "D": int(self.D),
+            "hidden_dim": int(self.hidden_dim),
+            "t_embed_dim": int(self.t_embed_dim),
+            "learning_rate": float(self.learning_rate),
+            "ema_decay": float(self.ema_decay),
+            "beta_max": float(self.beta_max),
+            "batch_size": None if self.batch_size is None else int(self.batch_size),
+            "eps": float(self.eps),
+            "key": np.array(self.key).tolist(),
+        }
+
+    def save_local(self, weights_file: str = "ddpmx_weights.msgpack", config_file: str = "ddpmx_config.json") -> None:
+        """
+        Saves:
+          - config_file: JSON with hyperparams + PRNG key
+          - weights_file: msgpack with flax state_dict of TrainStateEMA
+        """
+        cfg = self._config_dict()
+        with open(config_file, "w", encoding="utf-8") as f:
+            json.dump(cfg, f, indent=2, ensure_ascii=False)
+
+        # Robust serialization (avoid msgpack tuple errors)
+        state_sd = flax_ser.to_state_dict(self.state)
+        blob = flax_ser.msgpack_serialize(state_sd)
+        with open(weights_file, "wb") as f:
+            f.write(blob)
+
+    @classmethod
+    def load_local(
+        cls,
+        weights_file: str,
+        config_file: str,
+        *,
+        Z_iX: Optional[jnp.ndarray] = None,
+    ) -> "DDPM":
+        """
+        Reconstructs a DDPMX instance from local files.
+
+        Z_iX is only used to provide shape (N,D) for initialization; training is skipped.
+        If Z_iX is None, a dummy array of shape (1,D) is created.
+        """
+        with open(config_file, "r", encoding="utf-8") as f:
+            cfg = json.load(f)
+
+        D = int(cfg["D"])
+        if Z_iX is None:
+            Z_iX = jnp.zeros((1, D), dtype=jnp.float32)
+
+        # Build a fresh instance with the same architecture, skip training
+        obj = cls(
+            Z_iX,
+            T=int(cfg["T"]),
+            hidden_dim=int(cfg["hidden_dim"]),
+            t_embed_dim=int(cfg["t_embed_dim"]),
+            learning_rate=float(cfg["learning_rate"]),
+            n_iter=0,
+            ema_decay=float(cfg["ema_decay"]),
+            beta_max=float(cfg["beta_max"]),
+            batch_size=cfg["batch_size"],
+            key=random.PRNGKey(0),
+            verbose_every=0,
+            eps=float(cfg["eps"]),
+        )
+
+        with open(weights_file, "rb") as f:
+            state_sd = flax_ser.msgpack_restore(f.read())
+
+        obj.state = flax_ser.from_state_dict(obj.state, state_sd)
+
+        key_list = cfg.get("key", None)
+        if key_list is not None:
+            obj.key = jnp.array(key_list, dtype=jnp.uint32)
+
+        return obj
+
+    # -----------------------------------------------------------------
+    # Hugging Face Hub: upload / download
+    # -----------------------------------------------------------------
+    def upload_to_huggingface(
+        self,
+        repo_id: str,
+        *,
+        token: Optional[str] = None,
+        weights_file: str = "ddpmx_weights.msgpack",
+        config_file: str = "ddpmx_config.json",
+        repo_type: str = "model",
+        revision: Optional[str] = None,
+    ) -> Dict[str, str]:
+        """
+        Saves locally and uploads (weights_file, config_file) to Hugging Face Hub.
+        """
+        try:
+            from huggingface_hub import create_repo, upload_file
+        except Exception as e:
+            raise RuntimeError(
+                "huggingface_hub not installed. Install it (e.g., `pip install huggingface_hub`)."
+            ) from e
+
+        self.save_local(weights_file=weights_file, config_file=config_file)
+
+        create_repo(repo_id, token=token, repo_type=repo_type, exist_ok=True)
+
+        w_name = os.path.basename(weights_file)
+        c_name = os.path.basename(config_file)
+
+        upload_file(
+            path_or_fileobj=weights_file,
+            path_in_repo=w_name,
+            repo_id=repo_id,
+            repo_type=repo_type,
+            token=token,
+            revision=revision,
+        )
+        upload_file(
+            path_or_fileobj=config_file,
+            path_in_repo=c_name,
+            repo_id=repo_id,
+            repo_type=repo_type,
+            token=token,
+            revision=revision,
+        )
+
+        return {"repo_id": repo_id, "weights": w_name, "config": c_name}
+
+    @classmethod
+    def download_from_huggingface(
+        cls,
+        repo_id: str,
+        *,
+        token: Optional[str] = None,
+        weights_file: str = "ddpmx_weights.msgpack",
+        config_file: str = "ddpmx_config.json",
+        repo_type: str = "model",
+        revision: Optional[str] = None,
+        cache_dir: Optional[str] = None,
+        Z_iX: Optional[jnp.ndarray] = None,
+    ) -> "DDPM":
+        """
+        Downloads (weights_file, config_file) from Hugging Face Hub and reconstructs the class.
+        """
+        try:
+            from huggingface_hub import hf_hub_download
+        except Exception as e:
+            raise RuntimeError(
+                "huggingface_hub not installed. Install it (e.g., `pip install huggingface_hub`)."
+            ) from e
+
+        w_name = os.path.basename(weights_file)
+        c_name = os.path.basename(config_file)
+
+        w_path = hf_hub_download(
+            repo_id=repo_id,
+            filename=w_name,
+            repo_type=repo_type,
+            token=token,
+            revision=revision,
+            cache_dir=cache_dir,
+        )
+        c_path = hf_hub_download(
+            repo_id=repo_id,
+            filename=c_name,
+            repo_type=repo_type,
+            token=token,
+            revision=revision,
+            cache_dir=cache_dir,
+        )
+
+        return cls.load_local(w_path, c_path, Z_iX=Z_iX)
+
+
+
+__all__ = ["DDPM", "EpsMLP", "cosine_schedule", "sinusoidal_embedding"]

dima.py

@@ -0,0 +1,905 @@
+# src/dima/dima.py
+from __future__ import annotations
+
+import json
+import time
+from dataclasses import asdict, dataclass
+from typing import Any, Dict, Optional, Tuple, Union
+
+import numpy as np
+
+import jax
+import jax.numpy as jnp
+from jax import random
+
+from flax import serialization as flax_ser
+
+from .ann import ANNBackend, make_ann
+from .ddpm import DDPM
+from .dmap import DMAP
+from .gplm import GPLM
+
+
+# ----------------------------
+# Optional: Hugging Face Hub
+# ----------------------------
+try:
+    from huggingface_hub import HfApi, HfFolder, upload_file, hf_hub_download  # type: ignore
+    _HAS_HF = True
+except Exception:
+    _HAS_HF = False
+
+
+_UNSET = object()
+
+
+def _select_device(prefer: str = "auto"):
+    """
+    Safe JAX device selection.
+    prefer: "auto" | "gpu" | "cpu"
+    """
+    prefer = (prefer or "auto").lower()
+    devs = jax.devices()
+    gpu = [d for d in devs if d.platform == "gpu"]
+    cpu = [d for d in devs if d.platform == "cpu"]
+
+    if prefer in ("auto", "gpu"):
+        return gpu[0] if gpu else (cpu[0] if cpu else devs[0])
+    if prefer == "cpu":
+        return cpu[0] if cpu else devs[0]
+    return gpu[0] if gpu else (cpu[0] if cpu else devs[0])
+
+
+def _np_dtype_str(x) -> str:
+    try:
+        return str(np.dtype(x))
+    except Exception:
+        return "float32"
+
+
+# ----------------------------
+# Frozen inference-only models
+# ----------------------------
+class FrozenDMAP:
+    """
+    Inference-only Nyström DMAP embedder built from saved DMAP state.
+    Uses kNN in ambient space against reference points.
+    """
+
+    def __init__(
+        self,
+        state: Dict[str, Any],
+        *,
+        ann_backend: ANNBackend = "auto",
+        ann_params: Optional[Dict[str, Any]] = None,
+        n_jobs: int = -1,
+    ):
+        self.k = int(state["k"])
+        self.beta = float(state["beta"])
+        self.β = self.beta
+        self.alpha = float(state["alpha"])
+        self.α = self.alpha
+        self.eps = float(state["eps"])
+        self.ε = self.eps
+        self.dtype = np.dtype(state.get("dtype", "float32"))
+
+        self.R_iX = np.ascontiguousarray(np.asarray(state["R_iX"]).astype(self.dtype, copy=False))
+        self.qalpha_i = np.asarray(state["qalpha_i"], dtype=np.float64)
+        self.qα_i = self.qalpha_i  # alias
+        self.R_over_lam_ix = np.asarray(state["R_over_lam_ix"], dtype=np.float64)  # (Nref,d)
+        self.R_over_λ_ix = self.R_over_lam_ix  # alias
+
+        self.ann, self.ann_backend = make_ann(ann_backend, ann_params=ann_params, n_jobs=n_jobs)
+        self.ann.build(self.R_iX)
+
+    def __call__(self, R_aX: Union[np.ndarray, list], *, batch_size: Optional[int] = None) -> np.ndarray:
+        R_aX = np.asarray(R_aX)
+        single = (R_aX.ndim == 1)
+        if single:
+            R_aX = R_aX[None, :]
+
+        R_aX = np.ascontiguousarray(R_aX.astype(self.dtype, copy=False))
+
+        if batch_size is None:
+            Z = self._embed(R_aX)
+        else:
+            out = []
+            bs = int(batch_size)
+            for s in range(0, R_aX.shape[0], bs):
+                out.append(self._embed(R_aX[s : s + bs]))
+            Z = np.vstack(out)
+
+        return Z[0] if single else Z
+
+    def _embed(self, R_aX: np.ndarray) -> np.ndarray:
+        j_aK, D2_aK = self.ann.search(R_aX, self.k)  # (a,k)
+        K_ai = np.exp(-self.beta * (D2_aK.astype(np.float64) / self.eps))  # (a,k)
+
+        q_a = np.maximum(K_ai.sum(axis=1), 1e-30)
+        qalpha_a = np.maximum(np.power(q_a, self.alpha), 1e-30)
+
+        qalpha_i = np.maximum(self.qalpha_i[j_aK], 1e-30)
+        Kalpha_ai = K_ai / (qalpha_a[:, None] * qalpha_i)
+
+        d_a = np.maximum(Kalpha_ai.sum(axis=1), 1e-30)
+        P_ai = Kalpha_ai / d_a[:, None]
+
+        R_over = self.R_over_lam_ix[j_aK, :]              # (a,k,d)
+        Z_ax = (P_ai[:, :, None] * R_over).sum(axis=1)    # (a,d)
+        return Z_ax
+
+
+def _restore_gplm_as_object(
+    state: Dict[str, Any],
+    *,
+    ann_backend: ANNBackend = "auto",
+    ann_params: Optional[Dict[str, Any]] = None,
+    n_jobs: int = -1,
+) -> GPLM:
+    """
+    Rehydrate a GPLM instance from saved state WITHOUT retraining.
+    This is deliberately done as a true GPLM instance so you also get GPLM.flow(...)
+    (assuming your GPLM class implements .flow()).
+    """
+    obj = GPLM.__new__(GPLM)  # type: ignore
+
+    obj.beta = float(state["beta"])
+    obj.β = obj.beta
+    obj.eps = float(state["eps"])
+    obj.ε = obj.eps
+    obj.pred_k = None if state.get("pred_k", None) is None else int(state["pred_k"])
+    obj.pred_κ = obj.pred_k
+    obj.dtype = np.dtype(state.get("dtype", "float32"))
+
+    obj.mean_X = np.asarray(state["mean_X"], dtype=np.float64)
+    obj.M_mX = np.asarray(state["M_mX"], dtype=np.float64)
+
+    obj.lat_mean_x = np.asarray(state["lat_mean_x"], dtype=np.float64)
+    obj.lat_std_x = np.asarray(state["lat_std_x"], dtype=np.float64)
+
+    obj.Z_mx_w = np.ascontiguousarray(np.asarray(state["Z_mx_w"]).astype(np.float64, copy=False))
+    obj.m = int(obj.Z_mx_w.shape[0])
+
+    obj.ann_Z, _ = make_ann(ann_backend, ann_params=ann_params, n_jobs=n_jobs)
+    obj.ann_Z.build(obj.Z_mx_w.astype(obj.dtype, copy=False))
+
+    return obj
+
+
+# ----------------------------
+# Config
+# ----------------------------
+@dataclass
+class DIMAConfig:
+    d: int = 32
+    beta: float = 1.0
+    ddpm_device: str = "auto"   # "auto" | "cpu" | "gpu"
+    version: str = "0.2.0"
+
+
+# ----------------------------
+# Main wrapper
+# ----------------------------
+class DIMA:
+    """
+    DIMA: DMAP encoder + (latent DDPM) + GPLM decoder.
+    Public “user-facing” convention in this wrapper:
+      - Raw DMAP coordinates are the *public latent* (np.ndarray): R_ax (a,d)
+      - Normalized latents are the DDPM coordinates (jax/np):      Z_ax (a,d)
+    Minimal user API (what you asked for):
+        dima = DIMA(R_iX, d=20, beta=2.0)
+        R_ax = dima(R_aX)        # encode ambient -> raw latents
+        Q_aX = dima(R_ax)        # decode raw latents -> ambient
+    You can still pass full dict overrides for any submodule:
+        dima = DIMA(..., dmap_kwargs={...}, gplm_kwargs={...}, ddpm_kwargs={...})
+    and you can also tweak the “headline” DDPM knobs directly in __init__ (below).
+    """
+
+    def __init__(
+        self,
+        R_iX: np.ndarray,
+        *,
+        # main knobs
+        d: int = 32,
+        beta: float = 1.0,
+        # allow per-module override (if None -> uses global beta)
+        dmap_beta: Optional[float] = None,
+        gplm_beta: Optional[float] = None,
+        # DMAP headline knobs (everything else via dmap_kwargs)
+        dmap_alpha: float = 0.0,
+        dmap_t: float = 1.0,
+        dmap_k: Optional[int] = None,
+        # GPLM headline knobs (everything else via gplm_kwargs)
+        gplm_m: int = 1024,
+        gplm_pred_k: Optional[int] = None,
+        # DDPM headline knobs (the ones worth surfacing)
+        ddpm_T: int = 200,
+        ddpm_hidden_dim: int = 128,
+        ddpm_t_embed_dim: int = 64,
+        ddpm_learning_rate: float = 3e-4,
+        ddpm_n_iter: int = 200_000,
+        ddpm_ema_decay: float = 0.999,
+        ddpm_beta_max: float = 0.02,
+        ddpm_batch_size: int = 256,
+        ddpm_verbose_every: int = 0,
+        ddpm_eps: float = 1e-5,
+        # runtime
+        ddpm_device: str = "auto",
+        key: Optional[jax.Array] = None,
+        # ann
+        ann_backend: ANNBackend = "auto",
+        ann_params: Optional[Dict[str, Any]] = None,
+        n_jobs: int = -1,
+        # “escape hatches”
+        dmap_kwargs: Optional[Dict[str, Any]] = None,
+        gplm_kwargs: Optional[Dict[str, Any]] = None,
+        ddpm_kwargs: Optional[Dict[str, Any]] = None,
+        # unicode aliases (β)
+        **kwargs: Any,
+    ):
+        # ---- unicode aliases ----
+        if "β" in kwargs:
+            beta = float(kwargs.pop("β"))
+        if "β_dmap" in kwargs:
+            dmap_beta = float(kwargs.pop("β_dmap"))
+        if "β_gplm" in kwargs:
+            gplm_beta = float(kwargs.pop("β_gplm"))
+        if kwargs:
+            raise TypeError(f"Unexpected kwargs: {sorted(kwargs.keys())}")
+
+        self.config = DIMAConfig(d=int(d), beta=float(beta), ddpm_device=str(ddpm_device))
+
+        # devices + rng
+        self.ddpm_device = _select_device(ddpm_device)
+        self.cpu_device = _select_device("cpu")
+        self.rng = random.PRNGKey(0) if key is None else key
+
+        # training data
+        self.R_iX = np.asarray(R_iX)
+        if self.R_iX.ndim != 2:
+            raise ValueError("R_iX must be 2D (N,D).")
+        self.N, self.D = self.R_iX.shape
+        self.d = int(d)
+
+        # betas
+        self.beta = float(beta)
+        self.β = self.beta
+
+        self.dmap_beta = float(self.beta if dmap_beta is None else dmap_beta)
+        self.gplm_beta = float(self.beta if gplm_beta is None else gplm_beta)
+
+        t0 = time.time()
+
+        # -------------------------
+        # 1) Train DMAP (CPU)
+        # -------------------------
+        dmap_init = dict(
+            d=self.d,
+            beta=self.dmap_beta,
+            alpha=float(dmap_alpha),
+            t=float(dmap_t),
+            k=dmap_k,
+            ann_backend=ann_backend,
+            ann_params=ann_params,
+            n_jobs=n_jobs,
+        )
+        if dmap_kwargs:
+            dmap_init.update(dict(dmap_kwargs))
+        # enforce headline knobs
+        dmap_init["d"] = self.d
+        dmap_init["beta"] = self.dmap_beta
+        dmap_init["alpha"] = float(dmap_alpha)
+        dmap_init["t"] = float(dmap_t)
+        dmap_init["k"] = dmap_k
+
+        self.enc = DMAP(self.R_iX, **dmap_init)
+
+        # raw DMAP coordinates for *all* training points (Nyström OOS on training set)
+        R_ix = np.asarray(self.enc(self.R_iX), dtype=np.float64)  # (N,d)
+
+        # -------------------------
+        # 2) Latent normalization for DDPM
+        # -------------------------
+        self.lat_mean_np = R_ix.mean(axis=0)
+        self.lat_std_np = np.maximum(R_ix.std(axis=0), 1e-12)
+
+        self.lat_mean_j = jax.device_put(jnp.asarray(self.lat_mean_np, dtype=jnp.float32), self.ddpm_device)
+        self.lat_std_j = jax.device_put(jnp.asarray(self.lat_std_np, dtype=jnp.float32), self.ddpm_device)
+
+        Z_ix = (R_ix - self.lat_mean_np) / self.lat_std_np  # (N,d)
+
+        # -------------------------
+        # 3) Train GPLM (CPU)
+        # -------------------------
+        gplm_init = dict(
+            beta=self.gplm_beta,
+            m=int(gplm_m),
+            pred_k=gplm_pred_k,
+            ann_backend=ann_backend,
+            ann_params=ann_params,
+            n_jobs=n_jobs,
+        )
+        if gplm_kwargs:
+            gplm_init.update(dict(gplm_kwargs))
+        # enforce headline knobs
+        gplm_init["beta"] = self.gplm_beta
+        gplm_init["m"] = int(gplm_m)
+        gplm_init["pred_k"] = gplm_pred_k
+
+        self.dec = GPLM(R_ix.astype(np.float32, copy=False), self.R_iX, **gplm_init)
+
+        # -------------------------
+        # 4) Train DDPM on normalized latents (DDPM device)
+        # -------------------------
+        Z_ix_j = jax.device_put(jnp.asarray(Z_ix, dtype=jnp.float32), self.ddpm_device)
+
+        ddpm_init = dict(
+            T=int(ddpm_T),
+            hidden_dim=int(ddpm_hidden_dim),
+            t_embed_dim=int(ddpm_t_embed_dim),
+            learning_rate=float(ddpm_learning_rate),
+            n_iter=int(ddpm_n_iter),
+            ema_decay=float(ddpm_ema_decay),
+            beta_max=float(ddpm_beta_max),
+            batch_size=int(ddpm_batch_size),
+            key=self.rng,
+            verbose_every=int(ddpm_verbose_every),
+            eps=float(ddpm_eps),
+        )
+        if ddpm_kwargs:
+            ddpm_init.update(dict(ddpm_kwargs))
+        # enforce headline knobs
+        ddpm_init["T"] = int(ddpm_T)
+        ddpm_init["hidden_dim"] = int(ddpm_hidden_dim)
+        ddpm_init["t_embed_dim"] = int(ddpm_t_embed_dim)
+        ddpm_init["learning_rate"] = float(ddpm_learning_rate)
+        ddpm_init["n_iter"] = int(ddpm_n_iter)
+        ddpm_init["ema_decay"] = float(ddpm_ema_decay)
+        ddpm_init["beta_max"] = float(ddpm_beta_max)
+        ddpm_init["batch_size"] = int(ddpm_batch_size)
+        ddpm_init["verbose_every"] = int(ddpm_verbose_every)
+        ddpm_init["eps"] = float(ddpm_eps)
+
+        with jax.default_device(self.ddpm_device):
+            self.dm = DDPM(Z_ix_j, **ddpm_init)
+
+        self.training_time = time.time() - t0
+
+    # -------------------------
+    # Latent conversions
+    # -------------------------
+    def normalize(self, R_ax: Union[np.ndarray, jnp.ndarray]) -> jnp.ndarray:
+        """raw latents (R) -> normalized latents (Z) on ddpm_device."""
+        R = np.asarray(R_ax, dtype=np.float64)
+        if R.ndim == 1:
+            R = R[None, :]
+        Z = (R - self.lat_mean_np) / self.lat_std_np
+        Zj = jnp.asarray(Z, dtype=jnp.float32)
+        return jax.device_put(Zj, self.ddpm_device)
+
+    def unnormalize(self, Z_ax: Union[np.ndarray, jnp.ndarray]) -> np.ndarray:
+        """normalized latents (Z) -> raw latents (R) on CPU (np)."""
+        if isinstance(Z_ax, jax.Array):
+            Z_np = np.asarray(jax.device_get(Z_ax))
+        else:
+            Z_np = np.asarray(Z_ax)
+        if Z_np.ndim == 1:
+            Z_np = Z_np[None, :]
+        R = Z_np * self.lat_std_np + self.lat_mean_np
+        return np.asarray(R)
+
+    # -------------------------
+    # Encode / Decode (public: raw latents)
+    # -------------------------
+    def encode(self, R_aX: Union[np.ndarray, jnp.ndarray], *, normalize: bool = False) -> Union[np.ndarray, jnp.ndarray]:
+        """
+        ambient -> raw DMAP latents (np) by default.
+        If normalize=True, returns normalized latents (jnp) on ddpm_device.
+        """
+        X = np.asarray(R_aX)
+        R_raw = np.asarray(self.enc(X))  # CPU, (a,d)
+        if not normalize:
+            return R_raw
+        return self.normalize(R_raw)
+
+    def decode(
+        self,
+        R_ax: Union[np.ndarray, jnp.ndarray],
+        *,
+        refine: bool = False,
+        t_start: int = 10,
+        add_noise: bool = True,
+        key: Optional[jax.Array] = None,
+        batch_size: Optional[int] = None,
+    ) -> np.ndarray:
+        """
+        raw latent -> (optional DDPM refine in normalized coords) -> raw latent -> ambient.
+        Returns ambient np.ndarray on CPU.
+        """
+        R_raw = np.asarray(R_ax, dtype=np.float64)
+        single = (R_raw.ndim == 1)
+        if single:
+            R_raw = R_raw[None, :]
+
+        if refine:
+            Z = self.normalize(R_raw)  # on device
+            Z = self.dm.refine_latents(Z, t_start=int(t_start), key=key, add_noise=bool(add_noise))
+            R_raw = self.unnormalize(Z)  # back to CPU raw
+
+        X_hat = self.dec(R_raw.astype(np.float32, copy=False), batch_size=batch_size)
+        X_hat = np.asarray(X_hat)
+        return X_hat[0] if single else X_hat
+
+    def reconstruct(
+        self,
+        R_aX: Union[np.ndarray, jnp.ndarray],
+        *,
+        refine: bool = False,
+        t_start: int = 10,
+        add_noise: bool = True,
+        key: Optional[jax.Array] = None,
+        batch_size: Optional[int] = None,
+    ) -> np.ndarray:
+        """decode(encode(X))."""
+        R_raw = self.encode(R_aX, normalize=False)
+        return self.decode(R_raw, refine=refine, t_start=t_start, add_noise=add_noise, key=key, batch_size=batch_size)
+
+    def sample(
+        self,
+        n: int,
+        *,
+        decode: bool = True,
+        batch_size: Optional[int] = None,
+    ) -> Union[np.ndarray, np.ndarray]:
+        """
+        Unconditional samples from latent DDPM.
+        If decode=True: returns ambient samples (np) on CPU.
+        If decode=False: returns raw latents (np) on CPU.
+        """
+        with jax.default_device(self.ddpm_device):
+            Z = self.dm.sample(int(n))
+        R = self.unnormalize(Z)  # raw (np)
+        if not decode:
+            return R
+        return self.dec(R.astype(np.float32, copy=False), batch_size=batch_size)
+
+    # -------------------------
+    # Geodesic-ish flow wrapper (delegates to GPLM.flow)
+    # -------------------------
+    def flow(
+        self,
+        R_ax: Union[np.ndarray, jnp.ndarray],
+        v_ax: Union[np.ndarray, jnp.ndarray],
+        *,
+        dt: float = 0.05,
+        reg: float = 1e-8,
+        keep_speed: bool = True,
+        # optional DDPM projection step (in normalized coords)
+        refine: bool = False,
+        t_start: int = 10,
+        add_noise: bool = True,
+        key: Optional[jax.Array] = None,
+        # decode return
+        decode: bool = False,
+        batch_size: Optional[int] = None,
+    ) -> Union[Tuple[np.ndarray, np.ndarray], Tuple[np.ndarray, np.ndarray, np.ndarray]]:
+        """
+        One step of latent flow in *raw* coordinates.
+        Requires: your GPLM class implements:
+            R_next, v_next = gplm.flow(R, v, dt=..., reg=..., keep_speed=...)
+        If refine=True, we project the *position* through DDPM in normalized coords after the step.
+        (Velocity after projection is left unchanged—projection isn’t a deterministic diffeo.)
+        Returns:
+          if decode=False:
+            (R_next, v_next)        both np arrays
+          if decode=True:
+            (X_next, R_next, v_next)
+        """
+        R = np.asarray(R_ax, dtype=np.float64)
+        v = np.asarray(v_ax, dtype=np.float64)
+        single = (R.ndim == 1)
+        if single:
+            R = R[None, :]
+            v = v[None, :]
+
+        if not hasattr(self.dec, "flow"):
+            raise AttributeError(
+                "Decoder does not have .flow(). Make sure you updated GPLM to include flow()."
+            )
+
+        Rn, vn = self.dec.flow(R, v, dt=float(dt), reg=float(reg), keep_speed=bool(keep_speed))
+
+        if refine:
+            Z = self.normalize(Rn)  # device
+            Z = self.dm.refine_latents(Z, t_start=int(t_start), key=key, add_noise=bool(add_noise))
+            Rn = self.unnormalize(Z)
+
+        if not decode:
+            if single:
+                return np.asarray(Rn[0]), np.asarray(vn[0])
+            return np.asarray(Rn), np.asarray(vn)
+
+        Xn = self.dec(Rn.astype(np.float32, copy=False), batch_size=batch_size)
+        Xn = np.asarray(Xn)
+        if single:
+            return Xn[0], np.asarray(Rn[0]), np.asarray(vn[0])
+        return Xn, np.asarray(Rn), np.asarray(vn)
+
+    # -------------------------
+    # Convenience __call__
+    # -------------------------
+    def __call__(
+        self,
+        A: Union[np.ndarray, jnp.ndarray],
+        *,
+        refine: bool = False,
+        t_start: int = 10,
+        add_noise: bool = True,
+        key: Optional[jax.Array] = None,
+        batch_size: Optional[int] = None,
+        normalize_latent: bool = False,
+    ) -> Union[np.ndarray, jnp.ndarray]:
+        """
+        Dispatch by last dimension:
+          - if A is (a,D): encode -> raw latents (np) by default
+          - if A is (a,d): decode -> ambient (np)
+        Options:
+          - normalize_latent=True only affects encoding (returns Z on device)
+          - refine/t_start/add_noise/key only affect decoding
+        """
+        A_np = np.asarray(A)
+        if A_np.ndim == 1:
+            A_np = A_np[None, :]
+
+        if A_np.shape[1] == self.D:
+            return self.encode(A_np, normalize=bool(normalize_latent))
+
+        if A_np.shape[1] == self.d:
+            return self.decode(
+                A_np,
+                refine=bool(refine),
+                t_start=int(t_start),
+                add_noise=bool(add_noise),
+                key=key,
+                batch_size=batch_size,
+            )
+
+        raise ValueError(f"Input has last-dim {A_np.shape[1]}, expected D={self.D} or d={self.d}.")
+
+    # -------------------------
+    # Save / Load
+    # -------------------------
+    def _pack_encoder(self) -> Dict[str, Any]:
+        enc = self.enc
+    
+        # qalpha_i (ascii/unicode)
+        qalpha_i = np.asarray(getattr(enc, "qalpha_i", getattr(enc, "qα_i")))
+    
+        # 1) Try to read R_over_* from whichever name exists
+        R_over = getattr(enc, "R_over_lam_ix", None)
+        if R_over is None:
+            R_over = getattr(enc, "R_over_λ_ix", None)
+    
+        # 2) If missing, compute it from R_ix and λ_x (most robust)
+        if R_over is None:
+            R_ix = getattr(enc, "R_ix", None)
+    
+            # λ_x is unicode in upstream; add fallbacks for safety
+            lam = getattr(enc, "λ_x", None)
+            if lam is None:
+                lam = getattr(enc, "lam_x", None)
+            if lam is None:
+                lam = getattr(enc, "lambda_x", None)
+    
+            if R_ix is None or lam is None:
+                raise AttributeError(
+                    "DMAP encoder is missing R_over_{λ,lam}_ix and also lacks (R_ix, λ_x) "
+                    "to reconstruct it. Please ensure your DMAP computes diffusion coords."
+                )
+    
+            lam = np.asarray(lam, dtype=np.float64)
+            lam = np.maximum(lam, 1e-30)  # avoid division by 0
+            R_ix = np.asarray(R_ix, dtype=np.float64)
+    
+            R_over = (R_ix / lam[None, :]).astype(np.float64, copy=False)
+    
+        return dict(
+            R_iX=np.asarray(enc.R_iX),
+            qalpha_i=qalpha_i,
+            # Always store with ASCII key expected by FrozenDMAP loader
+            R_over_lam_ix=np.asarray(R_over, dtype=np.float64),
+            k=int(enc.k),
+            beta=float(getattr(enc, "beta", getattr(enc, "β"))),
+            alpha=float(getattr(enc, "alpha", getattr(enc, "α"))),
+            eps=float(getattr(enc, "eps", getattr(enc, "ε"))),
+            dtype=_np_dtype_str(getattr(enc, "dtype", np.float32)),
+        )
+    
+    
+    def _pack_decoder(self) -> Dict[str, Any]:
+        return dict(
+            Z_mx_w=np.asarray(getattr(self.dec, "Z_mx_w", None)),
+            M_mX=np.asarray(self.dec.M_mX),
+            mean_X=np.asarray(getattr(self.dec, "mean_X", np.zeros((self.D,), dtype=np.float64))),
+            lat_mean_x=np.asarray(getattr(self.dec, "lat_mean_x", np.zeros((self.d,), dtype=np.float64))),
+            lat_std_x=np.asarray(getattr(self.dec, "lat_std_x", np.ones((self.d,), dtype=np.float64))),
+            beta=float(getattr(self.dec, "beta", getattr(self.dec, "β"))),
+            eps=float(getattr(self.dec, "eps", getattr(self.dec, "ε"))),
+            pred_k=getattr(self.dec, "pred_k", getattr(self.dec, "pred_κ", None)),
+            dtype=_np_dtype_str(getattr(self.dec, "dtype", np.float32)),
+        )
+
+    def state_dict(self) -> Dict[str, Any]:
+        dd = dict(
+            T=int(self.dm.T),
+            D=int(self.dm.D),
+            hidden_dim=int(getattr(self.dm.model, "hidden", 128)),
+            t_embed_dim=int(getattr(self.dm.model, "t_dim", 64)),
+            ema_decay=float(getattr(self.dm, "ema_decay", 0.999)),
+            beta_max=float(getattr(self.dm, "beta_max", 0.02)),
+            eps=float(getattr(self.dm, "eps", 1e-5)),
+            params=self.dm.state.params,
+            ema_params=self.dm.state.ema_params,
+        )
+
+        state = dict(
+            meta=dict(
+                N=int(self.N),
+                D=int(self.D),
+                d=int(self.d),
+                training_time=float(getattr(self, "training_time", 0.0)),
+            ),
+            config=asdict(self.config),
+            latent_norm=dict(
+                mean=np.asarray(self.lat_mean_np, dtype=np.float64),
+                std=np.asarray(self.lat_std_np, dtype=np.float64),
+            ),
+            encoder=self._pack_encoder(),
+            decoder=self._pack_decoder(),
+            ddpm=dd,
+        )
+        return state
+
+    def save_local(self, weights_file: str = "dima.msgpack", config_file: str = "config.json") -> None:
+        """
+        Save full DIMA state to a msgpack file + a readable JSON config.
+
+        This version is robust to accidental tuples inside the state tree
+        (msgpack cannot serialize tuples by default).
+        """
+        def _sanitize(x):
+            # Convert tuples -> lists recursively (msgpack-safe)
+            if isinstance(x, tuple):
+                return [_sanitize(v) for v in x]
+            if isinstance(x, list):
+                return [_sanitize(v) for v in x]
+            if isinstance(x, dict):
+                return {k: _sanitize(v) for k, v in x.items()}
+            return x
+
+        state = _sanitize(self.state_dict())
+        blob = flax_ser.msgpack_serialize(state)
+
+        with open(weights_file, "wb") as f:
+            f.write(blob)
+
+        with open(config_file, "w") as f:
+            json.dump(state["config"], f, indent=2)
+
+        return None
+
+
+    @classmethod
+    def load_local(
+        cls,
+        weights_file: str = "dima.msgpack",
+        *,
+        ddpm_device: str = "auto",
+        ann_backend: ANNBackend = "auto",
+        ann_params: Optional[Dict[str, Any]] = None,
+        n_jobs: int = -1,
+        key: Optional[jax.Array] = None,
+    ) -> "DIMA":
+        with open(weights_file, "rb") as f:
+            state = flax_ser.msgpack_restore(f.read())
+
+        obj = cls.__new__(cls)  # bypass __init__
+
+        obj.config = DIMAConfig(**state["config"])
+        obj.ddpm_device = _select_device(ddpm_device)
+        obj.cpu_device = _select_device("cpu")
+        obj.training_time = float(state["meta"].get("training_time", 0.0))
+
+        obj.N = int(state["meta"]["N"])
+        obj.D = int(state["meta"]["D"])
+        obj.d = int(state["meta"]["d"])
+
+        # RNG
+        obj.rng = random.PRNGKey(0) if key is None else key
+
+        # beta (for display / convenience)
+        obj.beta = float(obj.config.beta)
+        obj.β = obj.beta
+
+        # latent norm
+        obj.lat_mean_np = np.asarray(state["latent_norm"]["mean"], dtype=np.float64)
+        obj.lat_std_np = np.asarray(state["latent_norm"]["std"], dtype=np.float64)
+
+        obj.lat_mean_j = jax.device_put(jnp.asarray(obj.lat_mean_np, dtype=jnp.float32), obj.ddpm_device)
+        obj.lat_std_j = jax.device_put(jnp.asarray(obj.lat_std_np, dtype=jnp.float32), obj.ddpm_device)
+
+        # frozen encoder + rehydrated decoder-as-GPLM (so flow works if GPLM.flow exists)
+        obj.enc = FrozenDMAP(state["encoder"], ann_backend=ann_backend, ann_params=ann_params, n_jobs=n_jobs)
+        obj.dec = _restore_gplm_as_object(state["decoder"], ann_backend=ann_backend, ann_params=ann_params, n_jobs=n_jobs)
+
+        # rebuild DDPM skeleton with dummy data, then load params
+        dd = state["ddpm"]
+        T = int(dd["T"])
+        D = int(dd["D"])
+        hidden_dim = int(dd["hidden_dim"])
+        t_embed_dim = int(dd["t_embed_dim"])
+        ema_decay = float(dd.get("ema_decay", 0.999))
+        beta_max = float(dd.get("beta_max", 0.02))
+        eps = float(dd.get("eps", 1e-5))
+
+        dummy = jnp.zeros((1, D), dtype=jnp.float32)
+        with jax.default_device(obj.ddpm_device):
+            obj.dm = DDPM(
+                dummy,
+                T=T,
+                hidden_dim=hidden_dim,
+                t_embed_dim=t_embed_dim,
+                learning_rate=1e-3,
+                n_iter=0,               # skip training on load
+                ema_decay=ema_decay,
+                beta_max=beta_max,
+                batch_size=1,
+                key=obj.rng,
+                verbose_every=0,
+                eps=eps,
+            )
+            obj.dm.state = obj.dm.state.replace(params=dd["params"], ema_params=dd["ema_params"])
+
+        obj.R_iX = None  # training data not stored by default
+        return obj
+
+    # -------------------------
+    # (Optional) HF helpers
+    # -------------------------
+
+    def upload_to_huggingface(
+        self,
+        repo_id: str,
+        *,
+        weights_file: str = "dima.msgpack",
+        config_file: str = "config.json",
+        token: Optional[str] = None,
+        repo_type: str = "model",
+        revision: Optional[str] = None,
+    ) -> None:
+        """
+        Serialize locally (weights + config) and upload them to Hugging Face Hub.
+
+        Parameters
+        ----------
+        repo_id : str
+            e.g. "username/my-dima-model"
+        weights_file : str
+            Filename used both locally and in the HF repo.
+        config_file : str
+            Human-readable JSON config filename (also uploaded).
+        token : Optional[str]
+            HF token. If None, uses HfFolder.get_token().
+        repo_type : str
+            Usually "model".
+        revision : Optional[str]
+            Optional target branch/revision (if your hub client supports it).
+        """
+        if not _HAS_HF:
+            raise RuntimeError("huggingface_hub not installed. Install it (or `pip install dima[hf]`).")
+
+        # 1) Save artifacts locally
+        self.save_local(weights_file=weights_file, config_file=config_file)
+
+        # 2) Resolve token
+        if token is None:
+            token = HfFolder.get_token()
+        if token is None:
+            raise RuntimeError("No HF token found. Provide `token=...` or run `huggingface-cli login`.")
+
+        # 3) Create repo if needed
+        api = HfApi()
+        api.create_repo(repo_id=repo_id, repo_type=repo_type, exist_ok=True, token=token)
+
+        # 4) Upload files
+        common_kwargs = dict(repo_id=repo_id, repo_type=repo_type, token=token)
+        if revision is not None:
+            common_kwargs["revision"] = revision
+
+        upload_file(
+            path_or_fileobj=weights_file,
+            path_in_repo=weights_file,
+            **common_kwargs,
+        )
+        upload_file(
+            path_or_fileobj=config_file,
+            path_in_repo=config_file,
+            **common_kwargs,
+        )
+        return None
+
+
+    @classmethod
+    def download_from_huggingface(
+        cls,
+        repo_id: str,
+        *,
+        weights_file: str = "dima.msgpack",
+        ddpm_device: str = "auto",
+        ann_backend: "ANNBackend" = "auto",
+        ann_params: Optional[Dict[str, Any]] = None,
+        n_jobs: int = -1,
+        key: Optional["jax.Array"] = None,
+        token: Optional[str] = None,
+        repo_type: str = "model",
+        revision: Optional[str] = None,
+    ) -> "DIMA":
+        """
+        Download weights from HF Hub and rehydrate a DIMA object via load_local.
+
+        Returns
+        -------
+        DIMA
+            A ready-to-use (inference) DIMA instance.
+        """
+        if not _HAS_HF:
+            raise RuntimeError("huggingface_hub not installed. Install it (or `pip install dima[hf]`).")
+
+        if token is None:
+            token = HfFolder.get_token()
+
+        dl_kwargs = dict(repo_id=repo_id, filename=weights_file, repo_type=repo_type)
+        if token is not None:
+            dl_kwargs["token"] = token
+        if revision is not None:
+            dl_kwargs["revision"] = revision
+
+        path = hf_hub_download(**dl_kwargs)
+
+        return cls.load_local(
+            path,
+            ddpm_device=ddpm_device,
+            ann_backend=ann_backend,
+            ann_params=ann_params,
+            n_jobs=n_jobs,
+            key=key,
+        )
+
+
+    # Backward-compatible aliases (the upstream file uses save_hf/load_hf)
+    def save_hf(self, repo_id: str, weights_file: str = "dima.msgpack", config_file: str = "config.json") -> None:
+        return self.upload_to_huggingface(repo_id, weights_file=weights_file, config_file=config_file)
+
+    @classmethod
+    def load_hf(
+        cls,
+        repo_id: str,
+        *,
+        weights_file: str = "dima.msgpack",
+        ddpm_device: str = "auto",
+        ann_backend: "ANNBackend" = "auto",
+        ann_params: Optional[Dict[str, Any]] = None,
+        n_jobs: int = -1,
+        key: Optional["jax.Array"] = None,
+    ) -> "DIMA":
+        return cls.download_from_huggingface(
+            repo_id,
+            weights_file=weights_file,
+            ddpm_device=ddpm_device,
+            ann_backend=ann_backend,
+            ann_params=ann_params,
+            n_jobs=n_jobs,
+            key=key,
+        )
+
+__all__ = ["DIMA", "DIMAConfig"]

dmap.py

@@ -0,0 +1,645 @@
+# src/dima/dmap.py
+from __future__ import annotations
+import json
+from typing import Any, Dict, Optional, Tuple, Union
+
+import numpy as np
+import scipy.sparse as sp
+from scipy.sparse.linalg import eigsh, LinearOperator, lobpcg
+from huggingface_hub import HfApi, HfFolder, upload_file, hf_hub_download
+
+from .ann import ANNBackend, make_ann
+from .utils import median_eps_from_knn_d2
+
+
+def k_ideal(d: int, N: int) -> int:
+    """
+    Heuristic for kNN graph size in diffusion maps.
+    Stable default: grows slowly with N and linearly with d.
+    """
+    d = int(max(1, d))
+    N = int(max(2, N))
+    k = int(np.ceil(2.0 * d * np.log2(N)))
+    return int(min(max(8, k), N - 1))
+
+
+def _sqdist_ab(A: np.ndarray, B: np.ndarray) -> np.ndarray:
+    """
+    Squared Euclidean distances between rows:
+      A: (a,d), B: (b,d) -> D2: (a,b)
+    """
+    A = np.asarray(A, dtype=np.float64)
+    B = np.asarray(B, dtype=np.float64)
+    A2 = np.sum(A * A, axis=1, keepdims=True)
+    B2 = np.sum(B * B, axis=1, keepdims=True).T
+    G = A @ B.T
+    return np.maximum(A2 + B2 - 2.0 * G, 0.0)
+
+
+class DMAP:
+    """
+    Diffusion Maps encoder with Nyström out-of-sample extension.
+
+    Notation (arrays named by indices):
+      R_iX: reference ambient data
+      K_ij: kernel on graph edges (sparse CSR)
+      q_i = Σ_j K_ij
+      qα_i = (q_i)^α
+      Kα_ij = K_ij / (qα_i qα_j)
+      d_i = Σ_j Kα_ij
+      A_ij = Kα_ij / sqrt(d_i d_j) (symmetric)
+      eigsh(A) -> λ_x, u_ix
+      ψ_ix = u_ix / sqrt(d_i)
+      R_ix = (λ_x)^t ψ_ix
+
+    Nyström OOS for novel ambient R_aX:
+      K_ai = exp(-β * D2_ai / ε)
+      q_a = Σ_i K_ai, qα_a = (q_a)^α
+      Kα_ai = K_ai / (qα_a qα_i)
+      d_a = Σ_i Kα_ai
+      P_ai = Kα_ai / d_a
+      R_ax = Σ_i P_ai * (R_ix / λ_x)
+
+    This implementation uses a kNN graph, sparse eigensolver, and
+    provides optional dense refinement via streaming LOBPCG.
+    """
+
+    def __init__(
+        self,
+        R_iX: np.ndarray,
+        *,
+        d: int = 32,
+        beta: float = 1.0,
+        alpha: float = 0.0,
+        t: float = 1.0,
+        k: Optional[int] = None,
+        eps: Optional[float] = None,
+        eps_use_kth: bool = True,
+        eps_mul: float = 1.0,
+        drop_trivial: bool = True,
+        seed: int = 0,
+        sym: str = "max",
+        dtype: Any = np.float32,
+        ann_backend: ANNBackend = "auto",
+        ann_params: Optional[Dict[str, Any]] = None,
+        n_jobs: int = -1,
+        # dense refinement (O(N^2) compute, streaming memory)
+        refine_dense: bool = False,
+        stream_block: int = 4096,
+        lobpcg_maxiter: int = 3,
+        lobpcg_tol: float = 1e-6,
+        use_symmetry: bool = True,
+        # allow unicode kwargs (β, α, ε, ε_mul, ε_use_kth, ...)
+        **kwargs: Any,
+    ):
+        # ---- map unicode kwargs -> ascii ----
+        if "β" in kwargs:
+            beta = kwargs.pop("β")
+        if "α" in kwargs:
+            alpha = kwargs.pop("α")
+        if "ε" in kwargs:
+            eps = kwargs.pop("ε")
+        if "ε_use_kth" in kwargs:
+            eps_use_kth = kwargs.pop("ε_use_kth")
+        if "ε_mul" in kwargs:
+            eps_mul = kwargs.pop("ε_mul")
+        if "sym" in kwargs:
+            sym = kwargs.pop("sym")
+        if kwargs:
+            raise TypeError(f"Unexpected kwargs: {sorted(kwargs.keys())}")
+
+        self.d = int(d)
+        self.k = int(k_ideal(self.d, int(np.asarray(R_iX).shape[0])) if k is None else int(k))
+
+        self.beta = float(beta)
+        self.alpha = float(alpha)
+        self.t = float(t)
+        self.drop_trivial = bool(drop_trivial)
+        self.seed = int(seed)
+        self.sym = str(sym)
+        self.dtype = dtype
+
+        # unicode aliases (so older code + packers can find them)
+        self.β = self.beta
+        self.α = self.alpha
+
+        self.refine_dense = bool(refine_dense)
+        self.stream_block = int(stream_block)
+        self.lobpcg_maxiter = int(lobpcg_maxiter)
+        self.lobpcg_tol = float(lobpcg_tol)
+        self.use_symmetry = bool(use_symmetry)
+
+        rng = np.random.default_rng(self.seed)
+
+        # reference data
+        R_iX = np.asarray(R_iX)
+        if R_iX.ndim != 2:
+            raise ValueError(f"R_iX must be 2D array, got shape {R_iX.shape}")
+        self.R_iX = np.ascontiguousarray(R_iX.astype(self.dtype, copy=False))
+
+        Nref, D = self.R_iX.shape
+        self.Nref = int(Nref)
+        self.D = int(D)
+
+        if not (2 <= self.k < Nref):
+            raise ValueError(f"Invalid k={self.k} for Nref={Nref}")
+
+        # ANN index (ambient)
+        self.ann, self.ann_backend = make_ann(ann_backend, ann_params=ann_params, n_jobs=n_jobs)
+        self.ann.build(self.R_iX)
+
+        # kNN of reference points
+        j_iK, D2_iK = self.ann.search(self.R_iX, self.k)  # (Nref,k)
+
+        # eps selection
+        if eps is None:
+            self.eps = float(median_eps_from_knn_d2(D2_iK, use_kth=bool(eps_use_kth)) * float(eps_mul))
+        else:
+            self.eps = float(eps)
+        self.ε = self.eps  # unicode alias
+
+        # -------------------------
+        # 1) Build sparse graph and warm-start eigensolver
+        # -------------------------
+        K_iK = np.exp(-self.beta * (D2_iK / self.eps)).astype(np.float64, copy=False)
+
+        indptr = (np.arange(Nref + 1, dtype=np.int64) * self.k)
+        indices = j_iK.reshape(-1).astype(np.int64, copy=False)
+        data = K_iK.reshape(-1)
+
+        K_ij = sp.csr_matrix((data, indices, indptr), shape=(Nref, Nref), dtype=np.float64)
+
+        # symmetrize
+        if self.sym == "max":
+            K_ij = K_ij.maximum(K_ij.T)
+        elif self.sym == "mean":
+            K_ij = (K_ij + K_ij.T) * 0.5
+        else:
+            raise ValueError(f"Unknown sym={self.sym!r}")
+
+        # degrees q_i and qalpha_i (warm)
+        q_i_warm = np.asarray(K_ij.sum(axis=1)).ravel()
+        q_i_warm = np.maximum(q_i_warm, 1e-30)
+        qalpha_i_warm = np.maximum(np.power(q_i_warm, self.alpha), 1e-30)
+
+        Qinv = sp.diags(1.0 / qalpha_i_warm, format="csr")
+        Kalpha_ij = Qinv @ K_ij @ Qinv
+
+        d_i_warm = np.asarray(Kalpha_ij.sum(axis=1)).ravel()
+        d_i_warm = np.maximum(d_i_warm, 1e-30)
+
+        Dinv_sqrt = sp.diags(1.0 / np.sqrt(d_i_warm), format="csr")
+        A_ij = Dinv_sqrt @ Kalpha_ij @ Dinv_sqrt
+
+        nev = self.d + (1 if self.drop_trivial else 0)
+        v0 = rng.normal(size=Nref).astype(np.float64)
+        lam0, u0 = eigsh(A_ij, k=nev, which="LA", v0=v0)
+
+        ord0 = np.argsort(lam0)[::-1]
+        lam0 = lam0[ord0]
+        u0 = u0[:, ord0]
+
+        # LOBPCG warm-start block (orthonormalize)
+        X0, _ = np.linalg.qr(u0.astype(np.float64, copy=False))
+
+        # -------------------------
+        # 2) Optional refinement: streaming dense LOBPCG on dense PSD operator
+        # -------------------------
+        if self.refine_dense:
+            try:
+                # Rebuild dense operator using streaming blocks without allocating full N^2 matrix
+                X = self.R_iX.astype(np.float64, copy=False)
+                X2 = np.sum(X * X, axis=1, keepdims=True)
+
+                # compute q_i, qalpha_i, d_i for dense kernel operator:
+                # q_i = Σ_j K_ij,  K_ij = exp(-β*||xi-xj||^2 / eps)
+                # Kα_ij = K_ij/(qα_i qα_j)
+                # d_i = Σ_j Kα_ij
+
+                # Step A: q_i
+                ones = np.ones((Nref, 1), dtype=np.float64)
+                q_i = self._K_matmat_dense(X, X2, ones).ravel()
+                q_i = np.maximum(q_i, 1e-30)
+                qalpha_i = np.maximum(np.power(q_i, self.alpha), 1e-30)
+
+                # Step B: d_i
+                inv_qalpha = 1.0 / qalpha_i
+                V = (ones * inv_qalpha[:, None])  # (N,1)  actually N x 1
+                tmp = self._K_matmat_dense(X, X2, V).ravel()  # Σ_j K_ij * inv_qalpha_j
+                d_i = inv_qalpha * tmp  # Σ_j K_ij/(qα_i qα_j)
+                d_i = np.maximum(d_i, 1e-30)
+
+                # Build symmetric operator A(v) = D^{-1/2} Q^{-1} K Q^{-1} D^{-1/2} v
+                inv_sqrt_d = 1.0 / np.sqrt(d_i)
+                inv_qalpha = 1.0 / qalpha_i
+
+                def matvec(v: np.ndarray) -> np.ndarray:
+                    v = v.astype(np.float64, copy=False).reshape(-1, 1)  # (N,1)
+                    w = v * inv_sqrt_d[:, None]
+                    w = w * inv_qalpha[:, None]
+                    y = self._K_matmat_dense(X, X2, w)
+                    y = y * inv_qalpha[:, None]
+                    y = y * inv_sqrt_d[:, None]
+                    return y.ravel()
+
+                Aop = LinearOperator((Nref, Nref), matvec=matvec, dtype=np.float64)
+
+                # LOBPCG refine using warm-start X0
+                lam, u = lobpcg(Aop, X0, largest=True, maxiter=self.lobpcg_maxiter, tol=self.lobpcg_tol)
+
+                ord1 = np.argsort(lam)[::-1]
+                lam = lam[ord1]
+                u = u[:, ord1]
+
+                self.q_i = q_i.astype(np.float64, copy=False)
+                self.qalpha_i = qalpha_i.astype(np.float64, copy=False)
+                self.d_i = d_i.astype(np.float64, copy=False)
+            except Exception:
+                # fallback to warm start if refinement fails
+                lam, u = lam0, u0
+                self.q_i = q_i_warm.astype(np.float64, copy=False)
+                self.qalpha_i = qalpha_i_warm.astype(np.float64, copy=False)
+                self.d_i = d_i_warm.astype(np.float64, copy=False)
+        else:
+            lam, u = lam0, u0
+            self.q_i = q_i_warm.astype(np.float64, copy=False)
+            self.qalpha_i = qalpha_i_warm.astype(np.float64, copy=False)
+            self.d_i = d_i_warm.astype(np.float64, copy=False)
+
+        # provide unicode aliases for packers / older code
+        self.qα_i = self.qalpha_i
+        self.λ_x = lam.astype(np.float64, copy=False)
+
+        # psi and drop trivial
+        psi = u / np.sqrt(self.d_i)[:, None]
+
+        if self.drop_trivial:
+            lam = lam[1:]
+            psi = psi[:, 1:]
+            u = u[:, 1:]
+
+        # diffusion coords
+        R_ix = psi * (lam ** self.t)[None, :]
+
+        # store
+        self.λ_x = lam.astype(np.float64, copy=False)      # (d,)
+        self.u_ix = u.astype(np.float64, copy=False)      # (Nref,d)
+        self.ψ_ix = psi.astype(np.float64, copy=False)    # (Nref,d)
+        self.R_ix = R_ix.astype(np.float64, copy=False)   # (Nref,d)
+        self.π_i = (self.d_i / self.d_i.sum()).astype(np.float64, copy=False)
+
+        # for Nyström: R_ix / λ_x
+        self.R_over_λ_ix = (self.R_ix / self.λ_x[None, :]).astype(np.float64, copy=False)
+
+    # --------- Dense kernel streaming utilities ----------
+
+    def _rbf_block(self, Xi: np.ndarray, Xj: np.ndarray, X2i: np.ndarray, X2j: np.ndarray) -> np.ndarray:
+        # ||xi-xj||^2 = xi^2 + xj^2 - 2 xi·xj
+        G = Xi @ Xj.T
+        D2 = np.maximum(X2i + X2j.T - 2.0 * G, 0.0)
+        return np.exp(-self.beta * (D2 / self.eps))
+
+    def _K_matmat_dense(self, X: np.ndarray, X2: np.ndarray, V: np.ndarray) -> np.ndarray:
+        """
+        Compute (dense) K @ V without forming K explicitly, using streaming blocks.
+        X: (N,D), X2: (N,1), V: (N,m) -> out: (N,m)
+        """
+        X = np.asarray(X, dtype=np.float64)
+        X2 = np.asarray(X2, dtype=np.float64)
+        V = np.asarray(V, dtype=np.float64)
+        N = X.shape[0]
+        bs = self.stream_block
+
+        out = np.zeros((N, V.shape[1]), dtype=np.float64)
+
+        if not self.use_symmetry:
+            for i0 in range(0, N, bs):
+                i1 = min(N, i0 + bs)
+                Xi = X[i0:i1]
+                X2i = X2[i0:i1]
+                acc = np.zeros((i1 - i0, V.shape[1]), dtype=np.float64)
+                for j0 in range(0, N, bs):
+                    j1 = min(N, j0 + bs)
+                    Xj = X[j0:j1]
+                    X2j = X2[j0:j1]
+                    Kij = self._rbf_block(Xi, Xj, X2i, X2j)
+                    acc += Kij @ V[j0:j1]
+                out[i0:i1] = acc
+            return out
+
+        # symmetric tiling
+        for i0 in range(0, N, bs):
+            i1 = min(N, i0 + bs)
+            Xi = X[i0:i1]
+            X2i = X2[i0:i1]
+            Vi = V[i0:i1]
+
+            # diagonal tile
+            Kii = self._rbf_block(Xi, Xi, X2i, X2i)
+            out[i0:i1] += Kii @ Vi
+
+            for j0 in range(i1, N, bs):
+                j1 = min(N, j0 + bs)
+                Xj = X[j0:j1]
+                X2j = X2[j0:j1]
+                Vj = V[j0:j1]
+
+                Kij = self._rbf_block(Xi, Xj, X2i, X2j)
+                out[i0:i1] += Kij @ Vj
+                out[j0:j1] += Kij.T @ Vi
+
+        return out
+
+    # --------- Nyström embedding ----------
+
+    def __call__(self, R_aX: Union[np.ndarray, list], *, batch_size: Optional[int] = None) -> np.ndarray:
+        R_aX = np.asarray(R_aX)
+        single = (R_aX.ndim == 1)
+        if single:
+            R_aX = R_aX[None, :]
+
+        R_aX = np.ascontiguousarray(R_aX.astype(self.dtype, copy=False))
+
+        if batch_size is None:
+            R_ax = self._embed(R_aX)
+        else:
+            bs = int(batch_size)
+            out = []
+            for s in range(0, R_aX.shape[0], bs):
+                out.append(self._embed(R_aX[s:s + bs]))
+            R_ax = np.vstack(out)
+
+        return R_ax[0] if single else R_ax
+
+    def _embed(self, R_aX: np.ndarray) -> np.ndarray:
+        # kNN for novel points
+        j_aK, D2_aK = self.ann.search(R_aX, self.k)  # (a,k)
+
+        K_ai = np.exp(-self.beta * (D2_aK.astype(np.float64) / self.eps))  # (a,k)
+
+        q_a = np.maximum(K_ai.sum(axis=1), 1e-30)
+        qalpha_a = np.maximum(np.power(q_a, self.alpha), 1e-30)
+
+        qalpha_i = np.maximum(self.qalpha_i[j_aK], 1e-30)
+        Kalpha_ai = K_ai / (qalpha_a[:, None] * qalpha_i)
+
+        d_a = np.maximum(Kalpha_ai.sum(axis=1), 1e-30)
+        P_ai = Kalpha_ai / d_a[:, None]
+
+        R_over = self.R_over_λ_ix[j_aK, :]  # (a,k,d)
+        R_ax = (P_ai[:, :, None] * R_over).sum(axis=1)
+        return R_ax
+
+    # -------------------------
+    # Serialization helpers
+    # -------------------------
+
+    @staticmethod
+    def _np_dtype_str(x) -> str:
+        """Best-effort numpy dtype string for serialization."""
+        try:
+            return str(np.dtype(x))
+        except Exception:
+            return "float32"
+
+    def state_dict(self) -> Dict[str, Any]:
+        """
+        Minimal, inference-sufficient state for Nyström out-of-sample embedding.
+
+        The resulting dict is intentionally compatible with `FrozenDMAP` in `dima.py`
+        (keys: R_iX, qalpha_i, R_over_lam_ix, k, beta, alpha, eps, dtype).
+        """
+        R_over = getattr(self, "R_over_λ_ix", None)
+        if R_over is None:
+            # Backward/alternate name safety
+            R_over = getattr(self, "R_over_lam_ix", None)
+
+        if R_over is None:
+            raise AttributeError("DMAP object has no Nyström matrix `R_over_λ_ix` (did training finish?).")
+
+        state: Dict[str, Any] = dict(
+            k=int(self.k),
+            beta=float(self.beta),
+            alpha=float(self.alpha),
+            eps=float(self.eps),
+            dtype=self._np_dtype_str(getattr(self, "dtype", "float32")),
+            # arrays
+            R_iX=np.asarray(self.R_iX),
+            qalpha_i=np.asarray(self.qalpha_i, dtype=np.float64),
+            R_over_lam_ix=np.asarray(R_over, dtype=np.float64),
+            # small meta (optional)
+            meta=dict(
+                Nref=int(np.asarray(self.R_iX).shape[0]),
+                D=int(np.asarray(self.R_iX).shape[1]),
+                d=int(np.asarray(R_over).shape[1]),
+                t=float(getattr(self, "t", 1.0)),
+                drop_trivial=bool(getattr(self, "drop_trivial", False)),
+                sym=str(getattr(self, "sym", "max")),
+            ),
+        )
+        return state
+
+    def save_local(self, weights_file: str = "dmap.msgpack", config_file: Optional[str] = "dmap_config.json") -> None:
+        """
+        Save DMAP weights (and optionally a lightweight JSON config) locally.
+
+        - `weights_file`: binary msgpack with arrays (via flax.serialization)
+        - `config_file`: JSON with scalar metadata only (no large arrays)
+        """
+        # local import keeps `dima` usable without flax unless you call this
+        from flax import serialization as flax_ser  # type: ignore
+        import json as _json
+
+        state = self.state_dict()
+        blob = flax_ser.msgpack_serialize(state)
+
+        with open(weights_file, "wb") as f:
+            f.write(blob)
+
+        if config_file is not None:
+            cfg = dict(
+                k=int(state["k"]),
+                beta=float(state["beta"]),
+                alpha=float(state["alpha"]),
+                eps=float(state["eps"]),
+                dtype=str(state.get("dtype", "float32")),
+                **(state.get("meta", {}) or {}),
+            )
+            with open(config_file, "w") as f:
+                _json.dump(cfg, f, indent=2)
+
+        return None
+
+    @classmethod
+    def from_state(
+        cls,
+        state: Dict[str, Any],
+        *,
+        ann_backend: ANNBackend = "auto",
+        ann_params: Optional[Dict[str, Any]] = None,
+        n_jobs: int = -1,
+    ) -> "DMAP":
+        """
+        Rehydrate a DMAP object from `state_dict()` output WITHOUT retraining.
+
+        This bypasses `__init__` and rebuilds only what is needed for out-of-sample Nyström embedding:
+        reference points, normalization factors, Nyström matrix, and ANN index.
+        """
+        # allow passing a full DIMA state dict
+        if "encoder" in state and isinstance(state["encoder"], dict):
+            state = state["encoder"]  # type: ignore[assignment]
+
+        # tolerate unicode key variants
+        k = int(state["k"])
+        beta = float(state.get("beta", state.get("β")))
+        alpha = float(state.get("alpha", state.get("α")))
+        eps = float(state.get("eps", state.get("ε")))
+
+        dtype = np.dtype(state.get("dtype", "float32"))
+        R_iX = np.ascontiguousarray(np.asarray(state["R_iX"]).astype(dtype, copy=False))
+        qalpha_i = np.asarray(state.get("qalpha_i", state.get("qα_i")), dtype=np.float64)
+
+        R_over = state.get("R_over_lam_ix", state.get("R_over_λ_ix"))
+        if R_over is None:
+            raise KeyError("State must contain 'R_over_lam_ix' (or 'R_over_λ_ix').")
+        R_over = np.asarray(R_over, dtype=np.float64)
+
+        obj = cls.__new__(cls)  # bypass __init__
+
+        # required for embedding
+        obj.k = k
+        obj.beta = beta
+        obj.alpha = alpha
+        obj.eps = eps
+        obj.dtype = dtype
+        obj.R_iX = R_iX
+        obj.qalpha_i = qalpha_i
+        obj.qα_i = obj.qalpha_i  # alias
+        obj.R_over_λ_ix = R_over
+        obj.R_over_lam_ix = obj.R_over_λ_ix  # alias for external code
+        obj.β = obj.beta
+        obj.α = obj.alpha
+        obj.ε = obj.eps
+
+        # optional metadata for introspection
+        meta = state.get("meta", {}) if isinstance(state.get("meta", {}), dict) else {}
+        obj.d = int(meta.get("d", R_over.shape[1]))
+        obj.t = float(meta.get("t", 1.0))
+        obj.drop_trivial = bool(meta.get("drop_trivial", False))
+        obj.sym = str(meta.get("sym", "max"))
+
+        # rebuild ANN index
+        obj.ann, obj.ann_backend = make_ann(ann_backend, ann_params=ann_params, n_jobs=n_jobs)
+        obj.ann.build(obj.R_iX)
+
+        return obj
+
+    @classmethod
+    def load_local(
+        cls,
+        weights_file: str = "dmap.msgpack",
+        *,
+        ann_backend: ANNBackend = "auto",
+        ann_params: Optional[Dict[str, Any]] = None,
+        n_jobs: int = -1,
+    ) -> "DMAP":
+        """Load DMAP weights saved by `save_local()`."""
+        from flax import serialization as flax_ser  # type: ignore
+
+        with open(weights_file, "rb") as f:
+            state = flax_ser.msgpack_restore(f.read())
+
+        return cls.from_state(state, ann_backend=ann_backend, ann_params=ann_params, n_jobs=n_jobs)
+
+    def upload_to_huggingface(
+        self,
+        repo_id: str,
+        *,
+        weights_file: str = "dmap.msgpack",
+        config_file: str = "dmap_config.json",
+        token: Optional[str] = None,
+        repo_type: str = "model",
+        revision: Optional[str] = None,
+    ) -> None:
+        """
+        Upload DMAP weights to the Hugging Face Hub.
+
+        This follows the same pattern as `DIMA.save_hf`:
+        - saves locally first (msgpack + JSON)
+        - creates the repo if needed
+        - uploads both files
+        """
+        try:
+            from huggingface_hub import HfApi, HfFolder, upload_file  # type: ignore
+        except Exception as e:
+            raise RuntimeError("huggingface_hub not installed. Install extras: `pip install dima[hf]`.") from e
+
+        self.save_local(weights_file=weights_file, config_file=config_file)
+
+        if token is None:
+            token = HfFolder.get_token()
+        if token is None:
+            raise RuntimeError("No HF token found. Run `huggingface-cli login`, or pass `token=...`.")
+
+        import os as _os
+
+        api = HfApi()
+        api.create_repo(repo_id=repo_id, repo_type=repo_type, exist_ok=True, token=token)
+
+        wf = _os.path.basename(weights_file)
+        cf = _os.path.basename(config_file)
+
+        upload_file(
+            path_or_fileobj=weights_file,
+            path_in_repo=wf,
+            repo_id=repo_id,
+            token=token,
+            repo_type=repo_type,
+            revision=revision,
+        )
+        upload_file(
+            path_or_fileobj=config_file,
+            path_in_repo=cf,
+            repo_id=repo_id,
+            token=token,
+            repo_type=repo_type,
+            revision=revision,
+        )
+
+        return None
+
+    @classmethod
+    def download_from_huggingface(
+        cls,
+        repo_id: str,
+        *,
+        weights_file: str = "dmap.msgpack",
+        ann_backend: ANNBackend = "auto",
+        ann_params: Optional[Dict[str, Any]] = None,
+        n_jobs: int = -1,
+        token: Optional[str] = None,
+        repo_type: str = "model",
+        revision: Optional[str] = None,
+    ) -> "DMAP":
+        """
+        Download DMAP weights from Hugging Face Hub and rehydrate a DMAP instance.
+
+        Note: despite the name, this *downloads from* the Hub.
+        """
+        try:
+            from huggingface_hub import hf_hub_download  # type: ignore
+        except Exception as e:
+            raise RuntimeError("huggingface_hub not installed. Install extras: `pip install dima[hf]`.") from e
+
+        path = hf_hub_download(
+            repo_id=repo_id,
+            filename=weights_file,
+            token=token,
+            repo_type=repo_type,
+            revision=revision,
+        )
+        return cls.load_local(path, ann_backend=ann_backend, ann_params=ann_params, n_jobs=n_jobs)
+
+
+
+__all__ = ["DMAP", "k_ideal"]

gplm.py

@@ -0,0 +1,712 @@
+# src/dima/gplm.py
+from __future__ import annotations
+
+from typing import Any, Dict, Literal, Optional, Tuple, Union
+
+import numpy as np
+import scipy.linalg as la
+import json
+import os
+import tempfile
+
+from .ann import ANNBackend, make_ann
+from .utils import fps_indices, median_eps_from_knn_d2, sqdist_ab
+
+
+InducingMode = Literal["random_subset", "fps", "kmeans_medoids", "given"]
+
+
+
+
+def _kmeans2_safe(Z: np.ndarray, m: int, seed: int = 0) -> np.ndarray:
+    """
+    KMeans centers with a safe fallback.
+    Uses scipy.cluster.vq.kmeans2 if available; otherwise samples points.
+    """
+    Z = np.asarray(Z)
+    m = int(min(max(1, m), Z.shape[0]))
+    try:
+        from scipy.cluster.vq import kmeans2  # type: ignore
+
+        C, _ = kmeans2(Z.astype(np.float64, copy=False), m, minit="points", seed=seed)
+        return C.astype(Z.dtype, copy=False)
+    except Exception:
+        rng = np.random.default_rng(seed)
+        idx = rng.choice(Z.shape[0], size=m, replace=False)
+        return Z[idx]
+
+
+def _to_builtin(x: Any) -> Any:
+    """
+    Convierte tipos numpy / tuplas / dicts anidados a tipos nativos serializables.
+    - Para msgpack (flax.serialization) las ndarrays se dejan como ndarrays.
+    - Para JSON, este helper se usa solo sobre objetos escalares/dicts (sin ndarrays).
+    """
+    if x is None:
+        return None
+    if isinstance(x, (bool, int, float, str)):
+        return x
+    if isinstance(x, (np.integer, np.floating)):
+        return x.item()
+    if isinstance(x, (tuple, list)):
+        return [_to_builtin(v) for v in x]
+    if isinstance(x, dict):
+        return {str(k): _to_builtin(v) for k, v in x.items()}
+    # fallback conservador
+    return str(x)
+
+
+class GPLM:
+    """
+    Inducing-point / Nyström GP (kernel ridge) decoder on latents.
+
+    Entrenamiento:
+      - Construye puntos inductores Z_mx_w (en latente blanqueado o no),
+      - Estima eps (si no se provee) a partir de distancias kNN,
+      - Resuelve M_mX por Cholesky (Nyström KRR/GP mean).
+
+    Además:
+      - decode() vía __call__.
+      - flow() integra un paso tipo generalized-leapfrog (geodesic flow) sobre el pullback manifold.
+      - NUEVO: save_local/load_local + upload_to_huggingface/download_from_huggingface.
+    """
+
+    def __init__(
+        self,
+        R_ix: np.ndarray,
+        R_iX: np.ndarray,
+        *,
+        # ASCII
+        beta: float = 1.0,
+        # eps estimation
+        eps: Optional[float] = None,
+        k_eps: int = 256,
+        eps_use_kth: bool = True,
+        eps_mul: float = 1.0,
+        # regularization
+        sigma2: float = 1e-5,
+        jitter: float = 1e-8,
+        # inducing
+        m: int = 1024,
+        inducing: InducingMode = "kmeans_medoids",
+        Z_mx: Optional[np.ndarray] = None,
+        seed: int = 0,
+        # preprocess
+        center_X: bool = True,
+        whiten_latent: bool = False,
+        dtype: Any = np.float32,
+        # compute/memory
+        fit_block: int = 8192,
+        # inference
+        pred_k: Optional[int] = None,
+        ann_backend: ANNBackend = "auto",
+        ann_params: Optional[Dict[str, Any]] = None,
+        n_jobs: int = -1,
+        # unicode aliases
+        **kwargs: Any,
+    ):
+        # ---- map unicode kwargs -> ascii
+        if "β" in kwargs:
+            beta = kwargs.pop("β")
+        if "ε" in kwargs:
+            eps = kwargs.pop("ε")
+        if "κ_eps" in kwargs:
+            k_eps = kwargs.pop("κ_eps")
+        if "ε_use_kth" in kwargs:
+            eps_use_kth = kwargs.pop("ε_use_kth")
+        if "ε_mul" in kwargs:
+            eps_mul = kwargs.pop("ε_mul")
+        if "σ2" in kwargs:
+            sigma2 = kwargs.pop("σ2")
+        if "pred_κ" in kwargs:
+            pred_k = kwargs.pop("pred_κ")
+        if kwargs:
+            raise TypeError(f"Unexpected kwargs: {sorted(kwargs.keys())}")
+
+        # ---- store params (and keep enough meta to reconstruct on load)
+        self.beta = float(beta)
+        self.β = self.beta
+
+        self.sigma2 = float(sigma2)
+        self.σ2 = self.sigma2
+
+        self.jitter = float(jitter)
+        self.seed = int(seed)
+
+        self.dtype = dtype
+        self.dtype_str = str(np.dtype(dtype))
+
+        self.fit_block = int(fit_block)
+        self.center_X = bool(center_X)
+        self.whiten_latent = bool(whiten_latent)
+        self.inducing = inducing
+
+        self.ann_backend = ann_backend
+        self.ann_params = ann_params if ann_params is None else dict(ann_params)
+        self.n_jobs = int(n_jobs)
+
+        # ---- validate / cast
+        R_ix = np.ascontiguousarray(np.asarray(R_ix).astype(self.dtype, copy=False))
+        R_iX = np.ascontiguousarray(np.asarray(R_iX).astype(self.dtype, copy=False))
+        if R_ix.ndim != 2 or R_iX.ndim != 2 or R_ix.shape[0] != R_iX.shape[0]:
+            raise ValueError("R_ix must be (N,d) and R_iX must be (N,D) with same N.")
+
+        self.R_ix = R_ix
+        self.R_iX = R_iX
+        self.N, self.d_lat = R_ix.shape
+        _, self.D = R_iX.shape
+
+        # ---- center output
+        if self.center_X:
+            self.mean_X = R_iX.mean(axis=0).astype(np.float64)
+            Y = (R_iX.astype(np.float64) - self.mean_X[None, :])
+        else:
+            self.mean_X = np.zeros((self.D,), dtype=np.float64)
+            Y = R_iX.astype(np.float64)
+
+        # ---- latent whitening (optional)
+        Ztrain = R_ix.astype(np.float64)
+        if self.whiten_latent:
+            self.lat_mean_x = Ztrain.mean(axis=0)
+            self.lat_std_x = np.maximum(Ztrain.std(axis=0), 1e-12)
+            Ztrain_w = (Ztrain - self.lat_mean_x) / self.lat_std_x
+        else:
+            self.lat_mean_x = np.zeros((self.d_lat,), dtype=np.float64)
+            self.lat_std_x = np.ones((self.d_lat,), dtype=np.float64)
+            Ztrain_w = Ztrain
+        self.R_ix_w = Ztrain_w  # (N,d) float64
+
+        # ---- ANN on training latents (eps + medoids snapping)
+        self.ann_train, _ = make_ann(self.ann_backend, ann_params=self.ann_params, n_jobs=self.n_jobs)
+        self.ann_train.build(self.R_ix_w.astype(self.dtype, copy=False))
+
+        # ---- eps via kNN distances on latents
+        if eps is None:
+            k_eps = int(min(max(8, int(k_eps)), self.N - 1))
+            j_iK1, D2_iK1 = self.ann_train.search(self.R_ix_w.astype(self.dtype, copy=False), k_eps + 1)
+
+            i = np.arange(self.N)[:, None]
+            is_self = (j_iK1 == i)
+
+            if np.any(is_self):
+                D2_iK = np.empty((self.N, k_eps), dtype=np.float64)
+                for ii in range(self.N):
+                    keep = (j_iK1[ii] != ii)
+                    D2_iK[ii] = D2_iK1[ii][keep][:k_eps]
+            else:
+                D2_iK = D2_iK1[:, :k_eps].astype(np.float64, copy=False)
+
+            eps_hat = median_eps_from_knn_d2(D2_iK, use_kth=bool(eps_use_kth))
+        else:
+            eps_hat = float(eps)
+
+        eps_hat *= float(eps_mul)
+        if eps_hat <= 0:
+            raise ValueError("eps must be > 0.")
+        self.eps = float(eps_hat)
+        self.ε = self.eps
+
+        # ---- choose inducing points (in whitened latent space)
+        rng = np.random.default_rng(self.seed)
+        m_eff = int(min(max(1, int(m)), self.N))
+
+        if Z_mx is not None:
+            Zm = np.asarray(Z_mx, dtype=np.float64)
+            if Zm.ndim != 2 or Zm.shape[1] != self.d_lat:
+                raise ValueError("Z_mx must be (m, d_lat).")
+            Zm_w = (Zm - self.lat_mean_x) / self.lat_std_x
+        else:
+            if inducing == "random_subset":
+                idx = rng.choice(self.N, size=m_eff, replace=False)
+                Zm_w = self.R_ix_w[idx]
+            elif inducing == "fps":
+                idx = fps_indices(self.R_ix_w, m=m_eff, seed=self.seed) # fps_indices is not defined
+                # For now, fallback to random_subset if fps_indices is not available
+                # idx = rng.choice(self.N, size=m_eff, replace=False)
+                Zm_w = self.R_ix_w[idx]
+            elif inducing == "kmeans_medoids":
+                C = _kmeans2_safe(self.R_ix_w, m_eff, seed=self.seed).astype(np.float64, copy=False)
+                j_cm, _ = self.ann_train.search(C.astype(self.dtype, copy=False), 1)
+                idx = j_cm.reshape(-1).astype(np.int64)
+
+                # de-duplicate and refill if needed
+                idx_u = np.unique(idx)
+                if idx_u.size < m_eff:
+                    needed = m_eff - idx_u.size
+                    pool = np.setdiff1d(np.arange(self.N), idx_u, assume_unique=False)
+                    extra = rng.choice(pool, size=needed, replace=False) if pool.size >= needed else rng.choice(self.N, size=needed, replace=True)
+                    idx = np.concatenate([idx_u, extra])
+                else:
+                    idx = idx_u[:m_eff]
+                Zm_w = self.R_ix_w[idx]
+            elif inducing == "given":
+                raise ValueError("Provide Z_mx when inducing='given'.")
+            else:
+                raise ValueError(f"Unknown inducing mode: {inducing!r}")
+
+        self.Z_mx_w = np.ascontiguousarray(Zm_w.astype(np.float64, copy=False))
+        self.m = int(self.Z_mx_w.shape[0])
+
+        # store raw inducing points (unwhitened) for convenience
+        self.Z_mx = (self.Z_mx_w * self.lat_std_x[None, :]) + self.lat_mean_x[None, :]
+
+        # ---- ANN on inducing points for fast prediction
+        self.ann_Z, _ = make_ann(self.ann_backend, ann_params=self.ann_params, n_jobs=self.n_jobs)
+        self.ann_Z.build(self.Z_mx_w.astype(self.dtype, copy=False))
+
+        # pred_k
+        if pred_k is None:
+            self.pred_k = None
+        else:
+            self.pred_k = int(min(max(1, int(pred_k)), self.m))
+        self.pred_κ = self.pred_k  # unicode alias
+
+        # ---- W_mm and reduced solve
+        D2_mm = sqdist_ab(self.Z_mx_w, self.Z_mx_w)
+        W_mm = np.exp(-self.beta * (D2_mm.astype(np.float64) / self.eps))
+        W_mm.flat[:: self.m + 1] += self.jitter
+        self.W_mm = W_mm  # (m,m)
+
+        G_mm = np.zeros((self.m, self.m), dtype=np.float64)
+        B_mX = np.zeros((self.m, self.D), dtype=np.float64)
+        bs = int(self.fit_block)
+
+        for i0 in range(0, self.N, bs):
+            i1 = min(self.N, i0 + bs)
+            Zi = self.R_ix_w[i0:i1]  # (b,d)
+            D2_im = sqdist_ab(Zi, self.Z_mx_w)
+            C_im = np.exp(-self.beta * (D2_im.astype(np.float64) / self.eps))
+            G_mm += C_im.T @ C_im
+            B_mX += C_im.T @ Y[i0:i1]
+
+        A_mm = G_mm + self.sigma2 * W_mm
+        A_mm.flat[:: self.m + 1] += self.jitter
+
+        cF = la.cho_factor(A_mm, lower=True, check_finite=False)
+        self.M_mX = la.cho_solve(cF, B_mX, check_finite=False)  # (m,D)
+
+    # ---------------------------
+    # Inference
+    # ---------------------------
+    def __call__(self, R_ax: Union[np.ndarray, list], *, batch_size: Optional[int] = None) -> np.ndarray:
+        R_ax = np.asarray(R_ax)
+        single = (R_ax.ndim == 1)
+        if single:
+            R_ax = R_ax[None, :]
+        R_ax = np.ascontiguousarray(R_ax.astype(self.dtype, copy=False))
+
+        if batch_size is None:
+            Y = self._decode(R_ax)
+        else:
+            bs = int(batch_size)
+            out = []
+            for s in range(0, R_ax.shape[0], bs):
+                out.append(self._decode(R_ax[s : s + bs]))
+            Y = np.vstack(out)
+
+        return Y[0] if single else Y
+
+    def _decode(self, R_ax: np.ndarray) -> np.ndarray:
+        Za = R_ax.astype(np.float64, copy=False)
+        Za_w = (Za - self.lat_mean_x) / self.lat_std_x
+
+        if self.pred_k is None or self.pred_k == self.m:
+            D2_am = sqdist_ab(Za_w, self.Z_mx_w)
+            C_am = np.exp(-self.beta * (D2_am.astype(np.float64) / self.eps))
+            Y = C_am @ self.M_mX
+        else:
+            j_aK, D2_aK = self.ann_Z.search(Za_w.astype(self.dtype, copy=False), self.pred_k)
+            W = np.exp(-self.beta * (D2_aK.astype(np.float64) / self.eps))  # (a,k)
+            M = self.M_mX[j_aK]  # (a,k,D)
+            Y = np.sum(W[:, :, None] * M, axis=1)  # (a,D)
+
+        return Y + self.mean_X[None, :]
+
+    # ============================================================
+    # Geodesic flow on pullback manifold (no C_mm storage)
+    # ============================================================
+    def _rbf_cache_single(self, r_x: np.ndarray, *, idx_m: Optional[np.ndarray]) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """
+        Cache kernel terms at position r (single point).
+        """
+        c = self.beta / self.eps
+        inv_std = 1.0 / self.lat_std_x
+
+        r = r_x.astype(np.float64, copy=False)
+        rw = (r - self.lat_mean_x) / self.lat_std_x
+
+        Zw = self.Z_mx_w if idx_m is None else self.Z_mx_w[idx_m]
+        Dw = rw[None, :] - Zw
+        D2 = np.sum(Dw * Dw, axis=1)
+
+        k_m = np.exp(-c * D2)  # (k,)
+        Dw_over = Dw * inv_std[None, :]  # (k,d) = (r-z)/std^2
+        grad_k = -(2.0 * c) * (k_m[:, None] * Dw_over)  # (k,d)
+
+        return k_m, Dw_over, grad_k
+
+    def _metric_from_gradM(
+        self,
+        grad_k: np.ndarray,  # (k,d)
+        M_kX: np.ndarray,    # (k,D)
+        *,
+        D_block: int = 8192,
+        lam: float = 1e-10,
+    ) -> Tuple[np.ndarray, Tuple[np.ndarray, bool]]:
+        """
+        Compute pullback metric g = J^T J without forming C_mm.
+        """
+        k, d = grad_k.shape
+        D = M_kX.shape[1]
+
+        g = np.zeros((d, d), dtype=np.float64)
+        for j0 in range(0, D, int(D_block)):
+            j1 = min(D, j0 + int(D_block))
+            Mb = M_kX[:, j0:j1]     # (k,block)
+            Jb = Mb.T @ grad_k      # (block,d)
+            g += Jb.T @ Jb          # (d,d)
+
+        g = 0.5 * (g + g.T)
+        g.flat[:: d + 1] += float(lam)
+
+        cF = la.cho_factor(g, lower=True, check_finite=False)
+        return g, cF
+
+    def _force_from_cache_noC(
+        self,
+        k_m: np.ndarray,       # (k,)
+        Dw_over: np.ndarray,   # (k,d)
+        v_x: np.ndarray,       # (d,)
+        M_kX: np.ndarray,      # (k,D)
+        *,
+        D_block: int = 8192,
+    ) -> np.ndarray:
+        """
+        Geodesic momentum force without storing C_mm, using blocked contractions over ambient dim D.
+        """
+        c = self.beta / self.eps
+        v = v_x.astype(np.float64, copy=False)
+
+        inv_std2 = (1.0 / self.lat_std_x) ** 2  # (d,)
+        s_m = Dw_over @ v                       # (k,)
+        S = -(2.0 * c) * (k_m * s_m)            # (k,)
+
+        term1 = (4.0 * c * c) * (k_m * s_m)[:, None] * Dw_over
+        term2 = (2.0 * c) * k_m[:, None] * (v[None, :] * inv_std2[None, :])
+        T = (term1 - term2).T                   # (d,k)
+
+        d = v.shape[0]
+        D = M_kX.shape[1]
+        f = np.zeros((d,), dtype=np.float64)
+
+        for j0 in range(0, D, int(D_block)):
+            j1 = min(D, j0 + int(D_block))
+            Mb = M_kX[:, j0:j1]   # (k,block)
+            Jv_b = S @ Mb         # (block,)
+            Hv_b = T @ Mb         # (d,block)
+            f += Hv_b @ Jv_b      # (d,)
+
+        return f
+
+    def flow(
+        self,
+        R_ax: Union[np.ndarray, list],
+        v_ax: Union[np.ndarray, list],
+        *,
+        eps: float = 1e-2,
+        K_p: int = 5,
+        K_q: int = 5,
+        D_block: int = 8192,
+        lam: float = 1e-10,
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        """
+        One generalized-leapfrog step for geodesic flow on the pullback manifold.
+        """
+        R = np.asarray(R_ax, dtype=np.float64)
+        v = np.asarray(v_ax, dtype=np.float64)
+        single = (R.ndim == 1)
+
+        if single:
+            R = R[None, :]
+            v = v[None, :]
+
+        if R.ndim != 2 or v.ndim != 2 or R.shape != v.shape or R.shape[1] != self.d_lat:
+            raise ValueError(f"Expected R_ax and v_ax shape (A,{self.d_lat}) (or ({self.d_lat},)).")
+
+        A, _d = R.shape
+
+        # Optional inducing subset indices per point
+        if self.pred_k is None or self.pred_k == self.m:
+            idx_aK = None
+        else:
+            Rw = (R - self.lat_mean_x[None, :]) / self.lat_std_x[None, :]
+            idx_aK, _ = self.ann_Z.search(Rw.astype(self.dtype, copy=False), self.pred_k)
+            idx_aK = idx_aK.astype(np.int64, copy=False)
+
+        R_next = np.empty_like(R)
+        v_next = np.empty_like(v)
+
+        for a in range(A):
+            idx = None if idx_aK is None else idx_aK[a]
+            M_kX = self.M_mX if idx is None else self.M_mX[idx]
+
+            r_n = R[a]
+            v_n = v[a]
+
+            # geometry at r_n
+            k_m_n, Dw_over_n, grad_k_n = self._rbf_cache_single(r_n, idx_m=idx)
+            g_n, cF_n = self._metric_from_gradM(grad_k_n, M_kX, D_block=D_block, lam=lam)
+
+            # momentum p_n = g(r_n) v_n
+            p_n = g_n @ v_n
+
+            # (1) implicit half-step in momentum
+            p = p_n.copy()
+            for _ in range(int(K_p)):
+                v_k = la.cho_solve(cF_n, p, check_finite=False)
+                f_k = self._force_from_cache_noC(k_m_n, Dw_over_n, v_k, M_kX, D_block=D_block)
+                p = p_n + 0.5 * float(eps) * f_k
+            p_half = p
+
+            # (2) implicit position update
+            v_half_n = la.cho_solve(cF_n, p_half, check_finite=False)
+            r = r_n + float(eps) * v_half_n
+            for _ in range(int(K_q)):
+                k_m_r, Dw_over_r, grad_k_r = self._rbf_cache_single(r, idx_m=idx)
+                g_r, cF_r = self._metric_from_gradM(grad_k_r, M_kX, D_block=D_block, lam=lam)
+                v_half_r = la.cho_solve(cF_r, p_half, check_finite=False)
+                r = r_n + 0.5 * float(eps) * (v_half_n + v_half_r)
+            r_np1 = r
+
+            # (3) explicit half-step in momentum at r_{n+1}
+            k_m_np1, Dw_over_np1, grad_k_np1 = self._rbf_cache_single(r_np1, idx_m=idx)
+            g_np1, cF_np1 = self._metric_from_gradM(grad_k_np1, M_kX, D_block=D_block, lam=lam)
+            v_mid = la.cho_solve(cF_np1, p_half, check_finite=False)
+            f_np1 = self._force_from_cache_noC(k_m_np1, Dw_over_np1, v_mid, M_kX, D_block=D_block)
+
+            p_np1 = p_half + 0.5 * float(eps) * f_np1
+            v_np1 = la.cho_solve(cF_np1, p_np1, check_finite=False)
+
+            R_next[a] = r_np1
+            v_next[a] = v_np1
+
+        if single:
+            return R_next[0], v_next[0]
+        return R_next, v_next
+
+    # ============================================================
+    # (NEW) Serialization + Hugging Face Hub
+    # ============================================================
+    def config_dict(self) -> Dict[str, Any]:
+        """
+        Config mínima (sin arrays grandes) para inspección/reproducibilidad.
+        """
+        return {
+            "class": "GPLM",
+            "beta": float(self.beta),
+            "eps": float(self.eps),
+            "sigma2": float(self.sigma2),
+            "jitter": float(self.jitter),
+            "seed": int(self.seed),
+            "center_X": bool(self.center_X),
+            "whiten_latent": bool(self.whiten_latent),
+            "inducing": str(self.inducing),
+            "fit_block": int(self.fit_block),
+            "pred_k": None if self.pred_k is None else int(self.pred_k),
+            "ann_backend": str(self.ann_backend),
+            "ann_params": None if self.ann_params is None else _to_builtin(self.ann_params),
+            "n_jobs": int(self.n_jobs),
+            "dtype": str(self.dtype_str),
+            "d_lat": int(self.d_lat),
+            "D": int(self.D),
+            "m": int(self.m),
+        }
+
+    def state_dict(self) -> Dict[str, Any]:
+        """
+        Estado completo necesario para inferencia/flow (sin datos de entrenamiento completos).
+        """
+        st: Dict[str, Any] = {
+            "config": self.config_dict(),
+            "M_mX": np.asarray(self.M_mX, dtype=np.float64),
+            "Z_mx_w": np.asarray(self.Z_mx_w, dtype=np.float64),
+            "mean_X": np.asarray(self.mean_X, dtype=np.float64),
+            "lat_mean_x": np.asarray(self.lat_mean_x, dtype=np.float64),
+            "lat_std_x": np.asarray(self.lat_std_x, dtype=np.float64),
+        }
+        return st
+
+    @classmethod
+    def from_state_dict(cls, st: Dict[str, Any]) -> "GPLM":
+        """
+        Reconstruye un objeto GPLM entrenado SIN re-entrenar.
+        """
+        if "config" not in st:
+            raise ValueError("state_dict inválido: falta 'config'.")
+
+        cfg = st["config"]
+        # construir instancia vacía
+        obj = cls.__new__(cls)
+
+        # meta/config
+        obj.beta = float(cfg["beta"])
+        obj.β = obj.beta
+        obj.eps = float(cfg["eps"])
+        obj.ε = obj.eps
+        obj.sigma2 = float(cfg["sigma2"])
+        obj.σ2 = obj.sigma2
+        obj.jitter = float(cfg["jitter"])
+        obj.seed = int(cfg["seed"])
+        obj.center_X = bool(cfg["center_X"])
+        obj.whiten_latent = bool(cfg["whiten_latent"])
+        obj.inducing = cfg.get("inducing", "given")
+        obj.fit_block = int(cfg["fit_block"])
+        obj.pred_k = cfg["pred_k"] if cfg["pred_k"] is None else int(cfg["pred_k"])
+        obj.pred_κ = obj.pred_k
+        obj.ann_backend = cfg.get("ann_backend", "auto")
+        obj.ann_params = cfg.get("ann_params", None)
+        obj.n_jobs = int(cfg.get("n_jobs", -1))
+
+        obj.dtype_str = str(cfg.get("dtype", "float32"))
+        obj.dtype = np.dtype(obj.dtype_str).type
+
+        obj.d_lat = int(cfg["d_lat"])
+        obj.D = int(cfg["D"])
+        obj.m = int(cfg["m"])
+
+        # arrays
+        obj.M_mX = np.asarray(st["M_mX"], dtype=np.float64)
+        obj.Z_mx_w = np.asarray(st["Z_mx_w"], dtype=np.float64)
+        obj.mean_X = np.asarray(st["mean_X"], dtype=np.float64)
+        obj.lat_mean_x = np.asarray(st["lat_mean_x"], dtype=np.float64)
+        obj.lat_std_x = np.asarray(st["lat_std_x"], dtype=np.float64)
+
+        # derived
+        obj.Z_mx = (obj.Z_mx_w * obj.lat_std_x[None, :]) + obj.lat_mean_x[None, :]
+
+        # ANN on inducing points (needed if pred_k is used)
+        obj.ann_Z, _ = make_ann(obj.ann_backend, ann_params=obj.ann_params, n_jobs=obj.n_jobs)
+        obj.ann_Z.build(obj.Z_mx_w.astype(obj.dtype, copy=False))
+
+        # no training data kept
+        obj.R_ix = None
+        obj.R_iX = None
+        obj.R_ix_w = None
+        obj.ann_train = None
+        obj.N = 0  # unknown/not needed for inference
+
+        return obj
+
+    def save_local(self, weights_file: str = "gplm.msgpack", config_file: str = "gplm_config.json") -> None:
+        """
+        Guarda:
+          - weights_file: msgpack con state_dict (arrays + config)
+          - config_file: JSON legible con la config
+        """
+        st = self.state_dict()
+        cfg = st["config"]
+
+        with open(config_file, "w", encoding="utf-8") as f:
+            json.dump(cfg, f, indent=2, ensure_ascii=False)
+
+        # msgpack (preferentemente flax.serialization)
+        try:
+            import flax.serialization as flax_ser  # type: ignore
+        except Exception as e:
+            raise RuntimeError("flax.serialization no está disponible; instale flax o use un backend alternativo.") from e
+
+        blob = flax_ser.msgpack_serialize(st)
+        with open(weights_file, "wb") as f:
+            f.write(blob)
+
+    @classmethod
+    def load_local(cls, weights_file: str = "gplm.msgpack") -> "GPLM":
+        """
+        Carga desde msgpack (state_dict completo) y reconstruye el objeto.
+        """
+        try:
+            import flax.serialization as flax_ser  # type: ignore
+        except Exception as e:
+            raise RuntimeError("flax.serialization no está disponible; instale flax o use un backend alternativo.") from e
+
+        with open(weights_file, "rb") as f:
+            blob = f.read()
+
+        st = flax_ser.msgpack_restore(blob)
+        return cls.from_state_dict(st)
+
+    def upload_to_huggingface(
+        self,
+        repo_id: str,
+        *,
+        token: Optional[str] = None,
+        weights_file: str = "gplm.msgpack",
+        config_file: str = "gplm_config.json",
+        repo_type: str = "model",
+        revision: Optional[str] = None,
+    ) -> None:
+        """
+        Sube (weights + config) a Hugging Face Hub.
+        """
+        try:
+            from huggingface_hub import HfApi, create_repo  # type: ignore
+        except Exception as e:
+            raise RuntimeError("huggingface_hub no está instalado. Instale con `pip install huggingface_hub`.") from e
+
+        if token is None:
+            raise ValueError("token es requerido para subir al Hub (HUGGINGFACE_TOKEN/HF_TOKEN).")
+
+        with tempfile.TemporaryDirectory() as td:
+            wpath = os.path.join(td, weights_file)
+            cpath = os.path.join(td, config_file)
+
+            self.save_local(weights_file=wpath, config_file=cpath)
+
+            create_repo(repo_id, token=token, repo_type=repo_type, exist_ok=True)
+            api = HfApi(token=token)
+
+            api.upload_file(
+                path_or_fileobj=wpath,
+                path_in_repo=weights_file,
+                repo_id=repo_id,
+                repo_type=repo_type,
+                revision=revision,
+            )
+            api.upload_file(
+                path_or_fileobj=cpath,
+                path_in_repo=config_file,
+                repo_id=repo_id,
+                repo_type=repo_type,
+                revision=revision,
+            )
+
+    @classmethod
+    def download_from_huggingface(
+        cls,
+        repo_id: str,
+        *,
+        token: Optional[str] = None,
+        weights_file: str = "gplm.msgpack",
+        repo_type: str = "model",
+        revision: Optional[str] = None,
+    ) -> "GPLM":
+        """
+        Descarga weights_file desde el Hub y reconstruye el objeto sin re-entrenar.
+        """
+        try:
+            from huggingface_hub import hf_hub_download  # type: ignore
+        except Exception as e:
+            raise RuntimeError("huggingface_hub no está instalado. Instale con `pip install huggingface_hub`.") from e
+
+        local_path = hf_hub_download(
+            repo_id=repo_id,
+            filename=weights_file,
+            repo_type=repo_type,
+            token=token,
+            revision=revision,
+        )
+        return cls.load_local(local_path)
+
+
+
+__all__ = ["GPLM", "InducingMode"]

utils.py

@@ -0,0 +1,202 @@
+# src/dima/utils.py
+from __future__ import annotations
+
+from dataclasses import asdict, is_dataclass
+from typing import Any, Dict, Iterable, Iterator, Optional, Sequence, Tuple, Union
+
+import numpy as np
+
+ArrayLike = Union[np.ndarray, Sequence[float]]
+
+
+# -------------------------
+# Basic math helpers
+# -------------------------
+def ensure_2d(X: np.ndarray) -> np.ndarray:
+    """Ensure X is 2D: (D,) -> (1,D)."""
+    X = np.asarray(X)
+    return X[None, :] if X.ndim == 1 else X
+
+
+def as_contig_f32(X: np.ndarray) -> np.ndarray:
+    """Contiguous float32 array (good default for ANN + kernels)."""
+    return np.ascontiguousarray(np.asarray(X, dtype=np.float32))
+
+
+def sqdist_ab(A: np.ndarray, B: np.ndarray) -> np.ndarray:
+    """
+    Squared Euclidean distances between rows:
+      A: (a,d), B: (b,d) -> D2: (a,b)
+    """
+    A = np.asarray(A)
+    B = np.asarray(B)
+    A2 = np.sum(A * A, axis=1, keepdims=True)          # (a,1)
+    B2 = np.sum(B * B, axis=1, keepdims=True).T        # (1,b)
+    G = A @ B.T                                        # (a,b)
+    return np.maximum(A2 + B2 - 2.0 * G, 0.0)
+
+
+def rbf_from_D2(D2: np.ndarray, *, beta: float, eps: float) -> np.ndarray:
+    """RBF kernel weights from squared distances."""
+    eps = float(eps)
+    if eps <= 0:
+        raise ValueError("eps must be > 0")
+    return np.exp(-float(beta) * (np.asarray(D2) / eps))
+
+
+# -------------------------
+# ε heuristics
+# -------------------------
+def median_eps_from_knn_d2(D2_iK: np.ndarray, *, use_kth: bool = True) -> float:
+    """
+    Median bandwidth from kNN squared distances.
+    D2_iK: (N,K) squared distances to K nearest neighbors (excluding self).
+    - use_kth=True: use the Kth neighbor distance per point, then median over points
+    - use_kth=False: use all distances, then median
+    """
+    D2_iK = np.asarray(D2_iK)
+    if D2_iK.size == 0:
+        return 1.0
+    v = D2_iK[:, -1] if use_kth else D2_iK.reshape(-1)
+    eps = float(np.median(v))
+    return max(eps, 1e-12)
+
+
+def median_eps_from_pairs(X: np.ndarray, *, max_pairs: int = 200_000, seed: int = 0) -> float:
+    """
+    Median of random-pair squared distances (rough fallback if you don't have kNN distances).
+    """
+    X = np.asarray(X)
+    N = X.shape[0]
+    if N < 2:
+        return 1.0
+
+    rng = np.random.default_rng(seed)
+    p = int(min(max_pairs, N * (N - 1) // 2))
+
+    i = rng.integers(0, N, size=p, endpoint=False)
+    j = rng.integers(0, N, size=p, endpoint=False)
+    mask = (i != j)
+    i = i[mask]
+    j = j[mask]
+    if i.size == 0:
+        return 1.0
+
+    D2 = np.sum((X[i] - X[j]) ** 2, axis=1)
+    eps = float(np.median(D2))
+    return max(eps, 1e-12)
+
+
+# -------------------------
+# Inducing / landmark selection
+# -------------------------
+def fps_indices(X: np.ndarray, m: int, *, seed: int = 0) -> np.ndarray:
+    """
+    Farthest Point Sampling indices (O(N*m)).
+    Good for space-filling inducing points / landmarks.
+    X: (N,d)
+    Returns idx: (m,)
+    """
+    X = np.asarray(X)
+    N = X.shape[0]
+    m = int(min(max(1, m), N))
+
+    rng = np.random.default_rng(seed)
+    idx = np.empty(m, dtype=np.int64)
+
+    idx[0] = int(rng.integers(0, N))
+    d2 = np.sum((X - X[idx[0]]) ** 2, axis=1)
+
+    for t in range(1, m):
+        idx[t] = int(np.argmax(d2))
+        new_d2 = np.sum((X - X[idx[t]]) ** 2, axis=1)
+        d2 = np.minimum(d2, new_d2)
+
+    return idx
+
+
+# -------------------------
+# Batching utilities
+# -------------------------
+def batched_range(n: int, batch_size: int) -> Iterator[Tuple[int, int]]:
+    """Yield (start, end) slices covering [0, n) in batches."""
+    bs = int(batch_size)
+    if bs <= 0:
+        raise ValueError("batch_size must be > 0")
+    for s in range(0, int(n), bs):
+        yield s, min(int(n), s + bs)
+
+
+def batch_iter(X: np.ndarray, batch_size: int) -> Iterator[np.ndarray]:
+    """Yield contiguous batches from X."""
+    X = np.asarray(X)
+    for s, e in batched_range(X.shape[0], batch_size):
+        yield X[s:e]
+
+
+# -------------------------
+# Metrics
+# -------------------------
+def rmse(a: np.ndarray, b: np.ndarray) -> float:
+    a = np.asarray(a)
+    b = np.asarray(b)
+    return float(np.sqrt(np.mean((a - b) ** 2)))
+
+
+def mse(a: np.ndarray, b: np.ndarray) -> float:
+    a = np.asarray(a)
+    b = np.asarray(b)
+    return float(np.mean((a - b) ** 2))
+
+
+# -------------------------
+# Device helpers (JAX optional)
+# -------------------------
+def get_jax_device(prefer: str = "auto"):
+    """
+    Safe JAX device selection.
+    prefer: "auto" | "gpu" | "cpu"
+    - returns a jax Device if jax is installed, else None
+    """
+    prefer = (prefer or "auto").lower()
+    try:
+        import jax  # local import
+    except Exception:
+        return None
+
+    devs = jax.devices()
+    gpu = [d for d in devs if d.platform == "gpu"]
+    cpu = [d for d in devs if d.platform == "cpu"]
+
+    if prefer in ("auto", "gpu"):
+        return gpu[0] if gpu else cpu[0] if cpu else devs[0]
+    if prefer == "cpu":
+        return cpu[0] if cpu else devs[0]
+    # fallback
+    return gpu[0] if gpu else cpu[0] if cpu else devs[0]
+
+
+# -------------------------
+# JSON helpers (for configs)
+# -------------------------
+def to_jsonable(x: Any) -> Any:
+    """
+    Convert common objects (numpy scalars/arrays, dataclasses) into JSON-serializable types.
+    """
+    if is_dataclass(x):
+        return {k: to_jsonable(v) for k, v in asdict(x).items()}
+
+    if isinstance(x, (np.floating, np.integer)):
+        return x.item()
+
+    if isinstance(x, np.ndarray):
+        # prefer list for small arrays; for large arrays you typically store separately
+        return x.tolist()
+
+    if isinstance(x, dict):
+        return {str(k): to_jsonable(v) for k, v in x.items()}
+
+    if isinstance(x, (list, tuple)):
+        return [to_jsonable(v) for v in x]
+
+    return x