dima / PTST.py

Update PTST.py

9042523 verified 9 months ago

24.7 kB

# PTST.py
# ============================================================
# PTST: Patch Time-Series Transformer (PatchTST-style) in JAX/Flax
#
# Design goals:
# - Same API style as your LISA/ALSA/TART:
#     model = PTST(R_tX_train, L=..., pred_len=..., ...)
#     preds = model(prefix, steps=H)
#
# - Training is automatic in __init__ (if R_tX is provided).
# - Inference is in __call__.
#
# Core ideas:
# - Patch along time: tokens are length-P patches, so seq_len reduces from L to L/P.
# - Channel-independence (default): train one univariate model shared across all channels.
#
# Forecasting:
# - Direct multi-horizon head for pred_len_train steps.
# - For steps > pred_len_train: rollout in pred_len_train-sized chunks (AR over chunks).
#
# Dependencies:
# - jax, flax, optax, numpy
# ============================================================

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Optional, Tuple

import numpy as np
import jax
import jax.numpy as jnp
import flax.linen as nn
import optax
from flax.training import train_state


# ============================================================
# Utilities
# ============================================================

def zscore(x: np.ndarray, eps: float = 1e-6) -> Tuple[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
    x = np.asarray(x, dtype=np.float32)
    mu = x.mean(axis=0, keepdims=True)
    sd = x.std(axis=0, keepdims=True) + eps
    return (x - mu) / sd, (mu, sd)


def _as_2d(X: np.ndarray) -> np.ndarray:
    X = np.asarray(X, dtype=np.float32)
    if X.ndim == 1:
        X = X[:, None]
    return X


def choose_heads(d_model: int, target_head_dim: int = 32, max_heads: int = 8) -> int:
    """Pick a reasonable number of heads given d_model."""
    n_heads = max(1, min(max_heads, d_model // target_head_dim))
    while d_model % n_heads != 0 and n_heads > 1:
        n_heads -= 1
    return int(max(1, n_heads))

def make_windows_multihorizon(X: np.ndarray, L: int, H: int):
    """
    Build multihorizon dataset:
      X_in[i]  = X[i : i+L]        (L, D)
      Y_out[i] = X[i+L : i+L+H]    (H, D)
    Returns:
      X_in : (N, L, D)
      Y_out: (N, H, D)
    """
    X = _as_2d(X)
    T, D = X.shape
    if T < L + H + 1:
        raise ValueError(f"Need T >= L+H+1. Got T={T}, L={L}, H={H}.")
    N = T - (L + H) + 1

    X_in = np.stack([X[i : i + L] for i in range(N)], axis=0)          # (N,L,D)
    Y_out = np.stack([X[i + L : i + L + H] for i in range(N)], axis=0) # (N,H,D)
    return np.ascontiguousarray(X_in), np.ascontiguousarray(Y_out)

# ============================================================
# Flax modules: Encoder-style attention
# ============================================================

class MultiHeadSelfAttentionX(nn.Module):
    d_model: int
    n_heads: int
    dropout: float = 0.0

    @nn.compact
    def __call__(self, x, deterministic: bool):
        """
        x: (B, T, d_model)
        output: (B, T, d_model)
        """
        B, T, Dm = x.shape
        assert Dm == self.d_model
        assert self.d_model % self.n_heads == 0
        head_dim = self.d_model // self.n_heads

        qkv = nn.Dense(3 * self.d_model, use_bias=False, name="qkv")(x)  # (B,T,3*Dm)
        qkv = qkv.reshape(B, T, 3, self.n_heads, head_dim)
        qkv = qkv.transpose(2, 0, 3, 1, 4)  # (3, B, H, T, hd)
        q, k, v = qkv[0], qkv[1], qkv[2]

        scale = 1.0 / math.sqrt(head_dim)
        attn_logits = jnp.einsum("bhqd,bhkd->bhqk", q, k) * scale  # (B,H,T,T)
        attn = nn.softmax(attn_logits, axis=-1)
        attn = nn.Dropout(rate=self.dropout)(attn, deterministic=deterministic)

        out = jnp.einsum("bhqk,bhkd->bhqd", attn, v)  # (B,H,T,hd)
        out = out.transpose(0, 2, 1, 3).reshape(B, T, self.d_model)
        out = nn.Dense(self.d_model, name="out_proj")(out)
        out = nn.Dropout(rate=self.dropout)(out, deterministic=deterministic)
        return out

class MultiHeadSelfAttentionXX(nn.Module):
    d_model: int
    n_heads: int
    dropout: float = 0.0

    # NEW: context-derived channel attention (your C^-)
    use_ctx_cam: bool = False
    cam_tau: float = 1.0
    cam_init_scale: float = 0.0  # gate starts at 0

    @nn.compact
    def __call__(self, x, deterministic: bool):
        """
        x: (B, T, d_model)
        """
        B, T, Dm = x.shape
        assert Dm == self.d_model
        assert self.d_model % self.n_heads == 0
        hd = self.d_model // self.n_heads

        qkv = nn.Dense(3 * self.d_model, use_bias=False, name="qkv")(x)  # (B,T,3*Dm)
        qkv = qkv.reshape(B, T, 3, self.n_heads, hd).transpose(2, 0, 3, 1, 4)  # (3,B,H,T,hd)
        q, k, v = qkv[0], qkv[1], qkv[2]  # (B,H,T,hd)

        scale = 1.0 / math.sqrt(hd)
        attn_logits = jnp.einsum("bhid,bhjd->bhij", q, k) * scale  # (B,H,T,T)
        A_plus = nn.softmax(attn_logits, axis=-1)                  # row-normalized over j
        A_plus = nn.Dropout(rate=self.dropout)(A_plus, deterministic=deterministic)

        # R'_{ihX} = A^+_{hij} R_{jhX}
        ctx = jnp.einsum("bhij,bhjd->bhid", A_plus, v)             # (B,H,T,hd)

        # unfold (hX)->x: R'_{ix}
        out = ctx.transpose(0, 2, 1, 3).reshape(B, T, self.d_model)  # (B,T,x=d_model)

        # NEW: build C^- from R' and apply R'' = R' C^-
        if self.use_ctx_cam:
            # S_{xx'} = sum_i R'_{ix} R'_{ix'}
            S = jnp.einsum("btx,bty->bxy", out, out) / jnp.maximum(T, 1)  # (B,x,x)

            # column-wise softmax (columns sum to 1): softmax over row axis
            C_minus = nn.softmax(S / self.cam_tau, axis=-2)               # (B,x,x)

            mixed = jnp.einsum("btx,bxy->bty", out, C_minus)              # (B,T,x)

            # gated residual so it starts as identity when init_scale=0
            alpha = self.param(
                "ctx_cam_scale",
                lambda k, s: jnp.array(self.cam_init_scale, jnp.float32),
                (),
            )
            out = out + alpha * (mixed - out)

        out = nn.Dense(self.d_model, name="out_proj")(out)
        out = nn.Dropout(rate=self.dropout)(out, deterministic=deterministic)
        return out

class MultiHeadSelfAttention(nn.Module):
    d_model: int
    n_heads: int
    dropout: float = 0.0

    # NEW: context-derived channel attention (your C^-)
    use_ctx_cam: bool = False
    cam_tau: float = 1.0
    cam_init_scale: float = 0.0  # gate starts at 0

    @nn.compact
    def __call__(self, x, deterministic: bool):
        """
        x: (B, T, d_model)
        """
        B, T, Dm = x.shape
        assert Dm == self.d_model
        assert self.d_model % self.n_heads == 0
        hd = self.d_model // self.n_heads

        qkv = nn.Dense(3 * self.d_model, use_bias=False, name="qkv")(x)  # (B,T,3*Dm)
        qkv = qkv.reshape(B, T, 3, self.n_heads, hd).transpose(2, 0, 3, 1, 4)  # (3,B,H,T,hd)
        q, k, v = qkv[0], qkv[1], qkv[2]  # (B,H,T,hd)

        scale = 1.0 / math.sqrt(hd)
        attn_logits = jnp.einsum("bhid,bhjd->bhij", q, k) * scale  # (B,H,T,T)
        A_plus = nn.softmax(attn_logits, axis=-1)                  # row-normalized over j
        A_plus = nn.Dropout(rate=self.dropout)(A_plus, deterministic=deterministic)

        # R'_{ihX} = A^+_{hij} R_{jhX}
        ctx = jnp.einsum("bhij,bhjd->bhid", A_plus, v)             # (B,H,T,hd)

        # unfold (hX)->x: R'_{ix}
        out = ctx.transpose(0, 2, 1, 3).reshape(B, T, self.d_model)  # (B,T,x=d_model)

        # NEW: build C^- from R' and apply R'' = R' C^-
        if self.use_ctx_cam:
            # S_{xx'} = sum_i R'_{ix} R'_{ix'}
            S  = jnp.einsum("btx,bty->bxy", out, out) / jnp.maximum(T, 1)  # (B,x,x)
            S2 = jnp.einsum("bii->bi", S)
            S  = S2[:,None,:] + S2[:,:,None] - 2*S
            S  = -1.*S

            # column-wise softmax (columns sum to 1): softmax over row axis
            C_minus = nn.softmax(S / self.cam_tau, axis=-2)               # (B,x,x)

            #mixed = jnp.einsum("btx,bxy->bty", out, C_minus)              # (B,T,x)
            # gated residual so it starts as identity when init_scale=0
            #alpha = self.param( "ctx_cam_scale", lambda k, s: jnp.array(self.cam_init_scale, jnp.float32),(),)
            #out = out + alpha * (mixed - out)
            out = jnp.einsum("btx,bxy->bty", out, C_minus)              # (B,T,x)

        out = nn.Dense(self.d_model, name="out_proj")(out)
        out = nn.Dropout(rate=self.dropout)(out, deterministic=deterministic)
        return out

class ChannelSelfAttention(nn.Module):
    d_model: int
    tau: float = 1.0
    col_norm: bool = True
    init_scale: float = 0.0

    @nn.compact
    def __call__(self, x, deterministic: bool):
        B, T, D = x.shape

        S = jnp.einsum("btd,bte->bde", x, x) / jnp.maximum(T, 1)
        C_row = nn.softmax(S / self.tau, axis=-1)
        C = jnp.swapaxes(C_row, -1, -2) if self.col_norm else C_row

        scale = self.param(
            "scale",
            lambda k, s: jnp.array(self.init_scale, jnp.float32),
            (),
        )

        # Return a delta (identity at scale=0)
        return scale * ((x @ C) - x)

class FeedForward(nn.Module):
    d_model: int
    mlp_ratio: float = 4.0
    dropout: float = 0.0

    @nn.compact
    def __call__(self, x, deterministic: bool):
        hidden = int(self.d_model * self.mlp_ratio)
        x = nn.Dense(hidden)(x)
        x = nn.gelu(x)
        x = nn.Dropout(rate=self.dropout)(x, deterministic=deterministic)
        x = nn.Dense(self.d_model)(x)
        x = nn.Dropout(rate=self.dropout)(x, deterministic=deterministic)
        return x

class TransformerEncoderBlock(nn.Module):
    d_model: int
    n_heads: int
    dropout: float = 0.0
    mlp_ratio: float = 4.0

    # toggles
    use_ctx_cam: bool = False
    use_ffn: bool = True

    # ctx-cam params
    cam_tau: float = 1.0
    cam_init_scale: float = 0.0

    @nn.compact
    def __call__(self, x, deterministic: bool):
        # --- Token self-attention (+ optional ctx channel mix inside it) ---
        h = nn.LayerNorm(name="attn_ln")(x)
        h = MultiHeadSelfAttention(
            d_model=self.d_model,
            n_heads=self.n_heads,
            dropout=self.dropout,
            use_ctx_cam=self.use_ctx_cam,
            cam_tau=self.cam_tau,
            cam_init_scale=self.cam_init_scale,
            name="mhsa",
        )(h, deterministic=deterministic)
        x = x + h

        # --- FFN (optional) ---
        if self.use_ffn:
            h2 = nn.LayerNorm(name="ffn_ln")(x)
            h2 = FeedForward(
                d_model=self.d_model, mlp_ratio=self.mlp_ratio, dropout=self.dropout
            )(h2, deterministic=deterministic)
            x = x + h2

        return x

class PTSTBackbone(nn.Module):
    """
    PatchTST-style backbone for univariate series (channel-independent training).
    Input:
      x: (B, L_eff, 1)
    Output:
      yhat: (B, pred_len)
    Options:
      - use_cam: enable context-derived channel mixing C^- inside MHSA (default False)
      - use_ffn: enable FeedForward inside each encoder block (default True)
    """
    L_eff: int
    patch_len: int
    d_model: int
    depth: int
    n_heads: int
    pred_len: int
    dropout: float = 0.0
    use_learned_pos: bool = True

    # block options
    use_cam: bool = False     # now means "use ctx-CAM inside MHSA"
    use_ffn: bool = True

    # ctx-CAM hyperparams (only used if use_cam=True)
    cam_tau: float = 1.0
    cam_init_scale: float = 0.0

    @nn.compact
    def __call__(self, x, deterministic: bool = True):
        B, L, C = x.shape
        assert C == 1
        assert L == self.L_eff
        assert self.L_eff % self.patch_len == 0

        n_patches = self.L_eff // self.patch_len

        # patchify: (B, n_patches, patch_len)
        patches = x.reshape(B, n_patches, self.patch_len)

        # patch embed (patch_len -> d_model)
        h = nn.Dense(self.d_model, name="patch_embed")(patches)  # (B, n_patches, d_model)

        # learned pos emb over patches
        if self.use_learned_pos:
            pos = self.param(
                "pos_emb",
                nn.initializers.normal(stddev=0.02),
                (n_patches, self.d_model),
            )  # (n_patches, d_model)
            h = h + pos[None, :, :]

        # encoder blocks (MHSA with optional ctx-CAM inside it; optional FFN)
        for i in range(self.depth):
            h = TransformerEncoderBlock(
                d_model=self.d_model,
                n_heads=self.n_heads,
                dropout=self.dropout,
                mlp_ratio=4.0,

                use_ctx_cam=self.use_cam,
                use_ffn=self.use_ffn,
                cam_tau=self.cam_tau,
                cam_init_scale=self.cam_init_scale,

                name=f"enc_block_{i}",
            )(h, deterministic=deterministic)

        h = nn.LayerNorm(name="final_ln")(h)

        # PatchTST-style flatten head
        h_flat = h.reshape(B, -1)  # (B, n_patches*d_model)
        yhat = nn.Dense(self.pred_len, name="head")(h_flat)  # (B, pred_len)
        return yhat


# ============================================================
# PTST: user-facing class
# ============================================================

@dataclass
class PTSTConfig:
    # data / windowing
    L: int = 128
    patch_len: int = 16
    pred_len: int = 96  # training horizon per forward pass

    # model
    d_model: int = 128
    depth: int = 4
    n_heads: Optional[int] = None
    dropout: float = 0.0
    use_learned_pos: bool = True

    # ablations
    use_ffn: bool = True          # keep/remove FFN
    use_cam: bool = False         # enable ctx-CAM inside MHSA (your C^-)

    # ctx-CAM hyperparams (only used if use_cam=True)
    cam_tau: float = 1.0
    cam_init_scale: float = 0.0   # gate; set 0.0 to start as identity

    # training
    seed: int = 0
    val_split: float = 0.1
    batch_size: int = 128
    max_epochs: int = 50
    init_lr: float = 3e-4
    min_lr: float = 1e-5
    lr_decay: float = 0.5
    tol_rel_improve: float = 1e-3
    patience: int = 5

class PTST:
    """
    PTST = PatchTST-style forecaster with TART-like API.
    Usage:
      model = PTST(F_train, L=160, patch_len=16, pred_len=120, d_model=128, depth=4)
      preds = model(prefix, steps=120)
    Notes:
    - Default is channel-independent: one univariate backbone shared across all dims.
    - Trains multi-horizon (pred_len) directly.
    - For steps > pred_len, it rolls forward in pred_len chunks.
    """

    def __init__(self, R_tX: np.ndarray, **kwargs):
        cfg = PTSTConfig(**kwargs)
        self.cfg = cfg

        R_tX = _as_2d(R_tX)
        T, D = R_tX.shape
        self.D = int(D)

        # enforce patch compatibility
        L = int(cfg.L)
        P = int(cfg.patch_len)
        if L < P:
            raise ValueError(f"L={L} must be >= patch_len={P}")
        L_eff = (L // P) * P
        if L_eff != L:
            # keep behavior explicit for reproducibility:
            # we just trim L down to nearest multiple of patch_len
            print(f"[PTST] warning: L={L} not divisible by patch_len={P}. Using L_eff={L_eff}.")
        self.L = int(L)          # requested
        self.L_eff = int(L_eff)  # actually used
        self.patch_len = int(P)

        pred_len = int(cfg.pred_len)
        if T <= self.L_eff + pred_len + 1:
            raise ValueError(f"Need T > L_eff + pred_len. Got T={T}, L_eff={self.L_eff}, pred_len={pred_len}.")

        # normalize train series channel-wise
        F_norm, (mu, sd) = zscore(R_tX)
        self._mu = mu
        self._sd = sd

        # build multihorizon dataset
        X_in, Y_out = make_windows_multihorizon(F_norm, L=self.L_eff, H=pred_len)  # (N,L,D), (N,H,D)
        N = X_in.shape[0]

        # channel-independent reshape: treat each channel as separate training sample
        # X_ci: (N*D, L_eff, 1)
        # Y_ci: (N*D, pred_len)
        # X_in: (N, L_eff, D), Y_out: (N, pred_len, D)
        if X_in.shape[1] != self.L_eff:
            raise ValueError(
                f"[PTST] X_in has wrong time length {X_in.shape[1]} (expected {self.L_eff}). "
                f"This usually means window construction is wrong."
            )
        
        # channel-independent: treat each channel as an independent training sample
        # X_ci: (N*D, L_eff, 1)
        # Y_ci: (N*D, pred_len)
        X_ci = np.transpose(X_in, (0, 2, 1))              # (N, D, L_eff)
        X_ci = X_ci.reshape(N * D, self.L_eff)[:, :, None] # (N*D, L_eff, 1)
        
        Y_ci = np.transpose(Y_out, (0, 2, 1))             # (N, D, pred_len)
        Y_ci = Y_ci.reshape(N * D, pred_len)              # (N*D, pred_len)


        # train/val split
        n_val = max(1, int(cfg.val_split * (N * D)))
        n_tr = (N * D) - n_val
        Xtr, Ytr = X_ci[:n_tr], Y_ci[:n_tr]
        Xva, Yva = X_ci[n_tr:], Y_ci[n_tr:]

        # build model
        d_model = int(max(32, cfg.d_model))
        n_heads = cfg.n_heads if cfg.n_heads is not None else choose_heads(d_model)
        self.d_model = int(d_model)
        self.n_heads = int(n_heads)
        self.depth = int(cfg.depth)
        self.pred_len = int(pred_len)

        self.model = PTSTBackbone(
            L_eff=self.L_eff,
            patch_len=self.patch_len,
            d_model=self.d_model,
            depth=self.depth,
            n_heads=self.n_heads,
            pred_len=self.pred_len,
            dropout=float(cfg.dropout),
            use_learned_pos=bool(cfg.use_learned_pos),
            
            use_cam=bool(cfg.use_cam),
            use_ffn=bool(cfg.use_ffn),
            cam_tau=float(cfg.cam_tau),
            cam_init_scale=float(cfg.cam_init_scale),

        )

        print(
            f"[PTST] Train samples: {N} windows -> {N*D} CI samples | "
            f"L_eff={self.L_eff} patch_len={self.patch_len} n_patches={self.L_eff//self.patch_len} | "
            f"D={self.D} pred_len={self.pred_len} | d_model={self.d_model} heads={self.n_heads} depth={self.depth}"
        )

        rng = jax.random.PRNGKey(int(cfg.seed))
        dummy_x = jnp.zeros((1, self.L_eff, 1), dtype=jnp.float32)
        params = self.model.init(rng, dummy_x, deterministic=True)["params"]

        def create_state(lr: float):
            tx = optax.adamw(learning_rate=lr, weight_decay=0.0)
            return train_state.TrainState.create(apply_fn=self.model.apply, params=params, tx=tx)

        state = create_state(float(cfg.init_lr))
        self._rng = rng

        @jax.jit
        def train_step(state, xb, yb, rng):
            dropout_rng, new_rng = jax.random.split(rng)

            def loss_fn(p):
                yhat = state.apply_fn({"params": p}, xb, deterministic=False, rngs={"dropout": dropout_rng})
                loss = jnp.mean((yhat - yb) ** 2)
                return loss

            loss, grads = jax.value_and_grad(loss_fn)(state.params)
            state = state.apply_gradients(grads=grads)
            return state, loss, new_rng

        @jax.jit
        def eval_step(state, xb, yb):
            yhat = state.apply_fn({"params": state.params}, xb, deterministic=True)
            return jnp.mean((yhat - yb) ** 2)

        # move to device once
        Xtr_j = jnp.asarray(Xtr)
        Ytr_j = jnp.asarray(Ytr)
        Xva_j = jnp.asarray(Xva)
        Yva_j = jnp.asarray(Yva)

        best_params = state.params
        best_val = float("inf")
        curr_lr = float(cfg.init_lr)
        epochs_no_gain = 0

        def num_batches(n: int) -> int:
            return int(np.ceil(n / cfg.batch_size))

        # training loop
        for epoch in range(1, int(cfg.max_epochs) + 1):
            idx = np.arange(n_tr)
            np.random.default_rng(epoch + cfg.seed).shuffle(idx)

            train_losses = []
            rng = self._rng

            for bi in range(num_batches(n_tr)):
                s = bi * cfg.batch_size
                e = min((bi + 1) * cfg.batch_size, n_tr)
                bidx = idx[s:e]
                xb = Xtr_j[bidx]
                yb = Ytr_j[bidx]
                state, loss, rng = train_step(state, xb, yb, rng)
                train_losses.append(float(loss))

            self._rng = rng
            tr_loss = float(np.mean(train_losses))

            val_losses = []
            for bi in range(num_batches(n_val)):
                s = bi * cfg.batch_size
                e = min((bi + 1) * cfg.batch_size, n_val)
                xb = Xva_j[s:e]
                yb = Yva_j[s:e]
                val_losses.append(float(eval_step(state, xb, yb)))
            va_loss = float(np.mean(val_losses))

            print(f"[PTST] epoch {epoch:03d} | train {tr_loss:.6e} | val {va_loss:.6e} | lr {curr_lr:.2e}")

            if va_loss + 1e-8 < best_val:
                best_val = va_loss
                best_params = state.params
                epochs_no_gain = 0
            else:
                epochs_no_gain += 1

            # plateau schedule
            if epochs_no_gain >= int(cfg.patience):
                if curr_lr > float(cfg.min_lr) * (1.0 + 1e-9):
                    curr_lr = max(float(cfg.min_lr), curr_lr * float(cfg.lr_decay))
                    print(f"[PTST] plateau → lowering LR to {curr_lr:.2e}")
                    params_now = state.params
                    state = train_state.TrainState.create(
                        apply_fn=state.apply_fn,
                        params=params_now,
                        tx=optax.adamw(learning_rate=curr_lr, weight_decay=0.0),
                    )
                    epochs_no_gain = 0
                else:
                    print(f"[PTST] early stop: lr at min and no improvement (best val {best_val:.6e})")
                    break

        # store best
        self.state = state.replace(params=best_params)
        self.params = self.state.params

        # param count
        self.param_count = int(sum(p.size for p in jax.tree_util.tree_leaves(self.params)))
        print(f"[PTST] params: {self.param_count:,}")

    # ------------------------------------------------------------
    # Internal: one-shot predict (normalized space) for a batch
    # ------------------------------------------------------------
    def _predict_block_norm(self, ctx_norm_BLD: jnp.ndarray) -> jnp.ndarray:
        """
        ctx_norm_BLD: (B, L_eff, D) normalized
        returns: (B, pred_len, D) normalized
        """
        B, L, D = ctx_norm_BLD.shape
        assert L == self.L_eff
        assert D == self.D

        # channel-independent apply:
        # reshape -> (B*D, L, 1)
        x = ctx_norm_BLD.transpose(0, 2, 1)[:, :, :, None]  # (B, D, L, 1)
        x = x.reshape(B * D, L, 1)
        yhat = self.model.apply({"params": self.params}, x, deterministic=True)  # (B*D, pred_len)
        yhat = yhat.reshape(B, D, self.pred_len)  # (B,D,H)
        yhat = yhat.transpose(0, 2, 1)            # (B,H,D)
        return yhat

    # ------------------------------------------------------------
    # Public API: forecast
    # ------------------------------------------------------------
    def __call__(self, prefix: np.ndarray, steps: int) -> np.ndarray:
        """
        Forecast `steps` time points given prefix ending at forecast start.
        prefix: (ell, D)
        returns: (steps, D)
        """
        X = _as_2d(prefix)
        ell, D = X.shape
        if D != self.D:
            raise ValueError(f"PTST expected D={self.D}, got {D}.")

        H = int(steps)
        if H <= 0:
            return np.zeros((0, self.D), dtype=np.float32)

        # take last L_eff points (pad on left if needed)
        if ell < self.L_eff:
            pad_len = self.L_eff - ell
            pad = np.repeat(X[:1], repeats=pad_len, axis=0)
            ctx = np.concatenate([pad, X], axis=0)
        else:
            ctx = X[-self.L_eff:, :]

        # normalize context
        ctx_n = (ctx - self._mu) / self._sd  # (L_eff, D)

        preds_n_all = []
        n_done = 0
        
        while n_done < H:
            ctx_n_j = jnp.asarray(ctx_n[None, :, :])  # (1,L_eff,D)
            block_n = np.array(self._predict_block_norm(ctx_n_j)[0])  # (pred_len, D)
        
            remaining = H - n_done
            take = min(self.pred_len, remaining)
        
            preds_n_all.append(block_n[:take])
            n_done += take
        
            # roll context
            ctx_n = np.concatenate([ctx_n, block_n[:take]], axis=0)
            ctx_n = ctx_n[-self.L_eff:, :]


        preds_n = np.concatenate(preds_n_all, axis=0)  # (H, D)
        preds = preds_n * self._sd + self._mu
        return preds


__all__ = ["PTST", "PTSTConfig"]

	# PTST.py
	# ============================================================
	# PTST: Patch Time-Series Transformer (PatchTST-style) in JAX/Flax
	#
	# Design goals:
	# - Same API style as your LISA/ALSA/TART:
	# model = PTST(R_tX_train, L=..., pred_len=..., ...)
	# preds = model(prefix, steps=H)
	#
	# - Training is automatic in __init__ (if R_tX is provided).
	# - Inference is in __call__.
	#
	# Core ideas:
	# - Patch along time: tokens are length-P patches, so seq_len reduces from L to L/P.
	# - Channel-independence (default): train one univariate model shared across all channels.
	#
	# Forecasting:
	# - Direct multi-horizon head for pred_len_train steps.
	# - For steps > pred_len_train: rollout in pred_len_train-sized chunks (AR over chunks).
	#
	# Dependencies:
	# - jax, flax, optax, numpy
	# ============================================================

	from __future__ import annotations

	import math
	from dataclasses import dataclass
	from typing import Optional, Tuple

	import numpy as np
	import jax
	import jax.numpy as jnp
	import flax.linen as nn
	import optax
	from flax.training import train_state


	# ============================================================
	# Utilities
	# ============================================================

	def zscore(x: np.ndarray, eps: float = 1e-6) -> Tuple[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
	x = np.asarray(x, dtype=np.float32)
	mu = x.mean(axis=0, keepdims=True)
	sd = x.std(axis=0, keepdims=True) + eps
	return (x - mu) / sd, (mu, sd)


	def _as_2d(X: np.ndarray) -> np.ndarray:
	X = np.asarray(X, dtype=np.float32)
	if X.ndim == 1:
	X = X[:, None]
	return X


	def choose_heads(d_model: int, target_head_dim: int = 32, max_heads: int = 8) -> int:
	"""Pick a reasonable number of heads given d_model."""
	n_heads = max(1, min(max_heads, d_model // target_head_dim))
	while d_model % n_heads != 0 and n_heads > 1:
	n_heads -= 1
	return int(max(1, n_heads))

	def make_windows_multihorizon(X: np.ndarray, L: int, H: int):
	"""
	Build multihorizon dataset:
	X_in[i] = X[i : i+L] (L, D)
	Y_out[i] = X[i+L : i+L+H] (H, D)

	Returns:
	X_in : (N, L, D)
	Y_out: (N, H, D)
	"""
	X = _as_2d(X)
	T, D = X.shape
	if T < L + H + 1:
	raise ValueError(f"Need T >= L+H+1. Got T={T}, L={L}, H={H}.")
	N = T - (L + H) + 1

	X_in = np.stack([X[i : i + L] for i in range(N)], axis=0) # (N,L,D)
	Y_out = np.stack([X[i + L : i + L + H] for i in range(N)], axis=0) # (N,H,D)
	return np.ascontiguousarray(X_in), np.ascontiguousarray(Y_out)

	# ============================================================
	# Flax modules: Encoder-style attention
	# ============================================================

	class MultiHeadSelfAttentionX(nn.Module):
	d_model: int
	n_heads: int
	dropout: float = 0.0

	@nn.compact
	def __call__(self, x, deterministic: bool):
	"""
	x: (B, T, d_model)
	output: (B, T, d_model)
	"""
	B, T, Dm = x.shape
	assert Dm == self.d_model
	assert self.d_model % self.n_heads == 0
	head_dim = self.d_model // self.n_heads

	qkv = nn.Dense(3 * self.d_model, use_bias=False, name="qkv")(x) # (B,T,3*Dm)
	qkv = qkv.reshape(B, T, 3, self.n_heads, head_dim)
	qkv = qkv.transpose(2, 0, 3, 1, 4) # (3, B, H, T, hd)
	q, k, v = qkv[0], qkv[1], qkv[2]

	scale = 1.0 / math.sqrt(head_dim)
	attn_logits = jnp.einsum("bhqd,bhkd->bhqk", q, k) * scale # (B,H,T,T)
	attn = nn.softmax(attn_logits, axis=-1)
	attn = nn.Dropout(rate=self.dropout)(attn, deterministic=deterministic)

	out = jnp.einsum("bhqk,bhkd->bhqd", attn, v) # (B,H,T,hd)
	out = out.transpose(0, 2, 1, 3).reshape(B, T, self.d_model)
	out = nn.Dense(self.d_model, name="out_proj")(out)
	out = nn.Dropout(rate=self.dropout)(out, deterministic=deterministic)
	return out

	class MultiHeadSelfAttentionXX(nn.Module):
	d_model: int
	n_heads: int
	dropout: float = 0.0

	# NEW: context-derived channel attention (your C^-)
	use_ctx_cam: bool = False
	cam_tau: float = 1.0
	cam_init_scale: float = 0.0 # gate starts at 0

	@nn.compact
	def __call__(self, x, deterministic: bool):
	"""
	x: (B, T, d_model)
	"""
	B, T, Dm = x.shape
	assert Dm == self.d_model
	assert self.d_model % self.n_heads == 0
	hd = self.d_model // self.n_heads

	qkv = nn.Dense(3 * self.d_model, use_bias=False, name="qkv")(x) # (B,T,3*Dm)
	qkv = qkv.reshape(B, T, 3, self.n_heads, hd).transpose(2, 0, 3, 1, 4) # (3,B,H,T,hd)
	q, k, v = qkv[0], qkv[1], qkv[2] # (B,H,T,hd)

	scale = 1.0 / math.sqrt(hd)
	attn_logits = jnp.einsum("bhid,bhjd->bhij", q, k) * scale # (B,H,T,T)
	A_plus = nn.softmax(attn_logits, axis=-1) # row-normalized over j
	A_plus = nn.Dropout(rate=self.dropout)(A_plus, deterministic=deterministic)

	# R'_{ihX} = A^+_{hij} R_{jhX}
	ctx = jnp.einsum("bhij,bhjd->bhid", A_plus, v) # (B,H,T,hd)

	# unfold (hX)->x: R'_{ix}
	out = ctx.transpose(0, 2, 1, 3).reshape(B, T, self.d_model) # (B,T,x=d_model)

	# NEW: build C^- from R' and apply R'' = R' C^-
	if self.use_ctx_cam:
	# S_{xx'} = sum_i R'_{ix} R'_{ix'}
	S = jnp.einsum("btx,bty->bxy", out, out) / jnp.maximum(T, 1) # (B,x,x)

	# column-wise softmax (columns sum to 1): softmax over row axis
	C_minus = nn.softmax(S / self.cam_tau, axis=-2) # (B,x,x)

	mixed = jnp.einsum("btx,bxy->bty", out, C_minus) # (B,T,x)

	# gated residual so it starts as identity when init_scale=0
	alpha = self.param(
	"ctx_cam_scale",
	lambda k, s: jnp.array(self.cam_init_scale, jnp.float32),
	(),
	)
	out = out + alpha * (mixed - out)

	out = nn.Dense(self.d_model, name="out_proj")(out)
	out = nn.Dropout(rate=self.dropout)(out, deterministic=deterministic)
	return out

	class MultiHeadSelfAttention(nn.Module):
	d_model: int
	n_heads: int
	dropout: float = 0.0

	# NEW: context-derived channel attention (your C^-)
	use_ctx_cam: bool = False
	cam_tau: float = 1.0
	cam_init_scale: float = 0.0 # gate starts at 0

	@nn.compact
	def __call__(self, x, deterministic: bool):
	"""
	x: (B, T, d_model)
	"""
	B, T, Dm = x.shape
	assert Dm == self.d_model
	assert self.d_model % self.n_heads == 0
	hd = self.d_model // self.n_heads

	qkv = nn.Dense(3 * self.d_model, use_bias=False, name="qkv")(x) # (B,T,3*Dm)
	qkv = qkv.reshape(B, T, 3, self.n_heads, hd).transpose(2, 0, 3, 1, 4) # (3,B,H,T,hd)
	q, k, v = qkv[0], qkv[1], qkv[2] # (B,H,T,hd)

	scale = 1.0 / math.sqrt(hd)
	attn_logits = jnp.einsum("bhid,bhjd->bhij", q, k) * scale # (B,H,T,T)
	A_plus = nn.softmax(attn_logits, axis=-1) # row-normalized over j
	A_plus = nn.Dropout(rate=self.dropout)(A_plus, deterministic=deterministic)

	# R'_{ihX} = A^+_{hij} R_{jhX}
	ctx = jnp.einsum("bhij,bhjd->bhid", A_plus, v) # (B,H,T,hd)

	# unfold (hX)->x: R'_{ix}
	out = ctx.transpose(0, 2, 1, 3).reshape(B, T, self.d_model) # (B,T,x=d_model)

	# NEW: build C^- from R' and apply R'' = R' C^-
	if self.use_ctx_cam:
	# S_{xx'} = sum_i R'_{ix} R'_{ix'}
	S = jnp.einsum("btx,bty->bxy", out, out) / jnp.maximum(T, 1) # (B,x,x)
	S2 = jnp.einsum("bii->bi", S)
	S = S2[:,None,:] + S2[:,:,None] - 2*S
	S = -1.*S

	# column-wise softmax (columns sum to 1): softmax over row axis
	C_minus = nn.softmax(S / self.cam_tau, axis=-2) # (B,x,x)

	#mixed = jnp.einsum("btx,bxy->bty", out, C_minus) # (B,T,x)
	# gated residual so it starts as identity when init_scale=0
	#alpha = self.param( "ctx_cam_scale", lambda k, s: jnp.array(self.cam_init_scale, jnp.float32),(),)
	#out = out + alpha * (mixed - out)
	out = jnp.einsum("btx,bxy->bty", out, C_minus) # (B,T,x)

	out = nn.Dense(self.d_model, name="out_proj")(out)
	out = nn.Dropout(rate=self.dropout)(out, deterministic=deterministic)
	return out

	class ChannelSelfAttention(nn.Module):
	d_model: int
	tau: float = 1.0
	col_norm: bool = True
	init_scale: float = 0.0

	@nn.compact
	def __call__(self, x, deterministic: bool):
	B, T, D = x.shape

	S = jnp.einsum("btd,bte->bde", x, x) / jnp.maximum(T, 1)
	C_row = nn.softmax(S / self.tau, axis=-1)
	C = jnp.swapaxes(C_row, -1, -2) if self.col_norm else C_row

	scale = self.param(
	"scale",
	lambda k, s: jnp.array(self.init_scale, jnp.float32),
	(),
	)

	# Return a delta (identity at scale=0)
	return scale * ((x @ C) - x)

	class FeedForward(nn.Module):
	d_model: int
	mlp_ratio: float = 4.0
	dropout: float = 0.0

	@nn.compact
	def __call__(self, x, deterministic: bool):
	hidden = int(self.d_model * self.mlp_ratio)
	x = nn.Dense(hidden)(x)
	x = nn.gelu(x)
	x = nn.Dropout(rate=self.dropout)(x, deterministic=deterministic)
	x = nn.Dense(self.d_model)(x)
	x = nn.Dropout(rate=self.dropout)(x, deterministic=deterministic)
	return x

	class TransformerEncoderBlock(nn.Module):
	d_model: int
	n_heads: int
	dropout: float = 0.0
	mlp_ratio: float = 4.0

	# toggles
	use_ctx_cam: bool = False
	use_ffn: bool = True

	# ctx-cam params
	cam_tau: float = 1.0
	cam_init_scale: float = 0.0

	@nn.compact
	def __call__(self, x, deterministic: bool):
	# --- Token self-attention (+ optional ctx channel mix inside it) ---
	h = nn.LayerNorm(name="attn_ln")(x)
	h = MultiHeadSelfAttention(
	d_model=self.d_model,
	n_heads=self.n_heads,
	dropout=self.dropout,
	use_ctx_cam=self.use_ctx_cam,
	cam_tau=self.cam_tau,
	cam_init_scale=self.cam_init_scale,
	name="mhsa",
	)(h, deterministic=deterministic)
	x = x + h

	# --- FFN (optional) ---
	if self.use_ffn:
	h2 = nn.LayerNorm(name="ffn_ln")(x)
	h2 = FeedForward(
	d_model=self.d_model, mlp_ratio=self.mlp_ratio, dropout=self.dropout
	)(h2, deterministic=deterministic)
	x = x + h2

	return x

	class PTSTBackbone(nn.Module):
	"""
	PatchTST-style backbone for univariate series (channel-independent training).

	Input:
	x: (B, L_eff, 1)

	Output:
	yhat: (B, pred_len)

	Options:
	- use_cam: enable context-derived channel mixing C^- inside MHSA (default False)
	- use_ffn: enable FeedForward inside each encoder block (default True)
	"""
	L_eff: int
	patch_len: int
	d_model: int
	depth: int
	n_heads: int
	pred_len: int
	dropout: float = 0.0
	use_learned_pos: bool = True

	# block options
	use_cam: bool = False # now means "use ctx-CAM inside MHSA"
	use_ffn: bool = True

	# ctx-CAM hyperparams (only used if use_cam=True)
	cam_tau: float = 1.0
	cam_init_scale: float = 0.0

	@nn.compact
	def __call__(self, x, deterministic: bool = True):
	B, L, C = x.shape
	assert C == 1
	assert L == self.L_eff
	assert self.L_eff % self.patch_len == 0

	n_patches = self.L_eff // self.patch_len

	# patchify: (B, n_patches, patch_len)
	patches = x.reshape(B, n_patches, self.patch_len)

	# patch embed (patch_len -> d_model)
	h = nn.Dense(self.d_model, name="patch_embed")(patches) # (B, n_patches, d_model)

	# learned pos emb over patches
	if self.use_learned_pos:
	pos = self.param(
	"pos_emb",
	nn.initializers.normal(stddev=0.02),
	(n_patches, self.d_model),
	) # (n_patches, d_model)
	h = h + pos[None, :, :]

	# encoder blocks (MHSA with optional ctx-CAM inside it; optional FFN)
	for i in range(self.depth):
	h = TransformerEncoderBlock(
	d_model=self.d_model,
	n_heads=self.n_heads,
	dropout=self.dropout,
	mlp_ratio=4.0,

	use_ctx_cam=self.use_cam,
	use_ffn=self.use_ffn,
	cam_tau=self.cam_tau,
	cam_init_scale=self.cam_init_scale,

	name=f"enc_block_{i}",
	)(h, deterministic=deterministic)

	h = nn.LayerNorm(name="final_ln")(h)

	# PatchTST-style flatten head
	h_flat = h.reshape(B, -1) # (B, n_patches*d_model)
	yhat = nn.Dense(self.pred_len, name="head")(h_flat) # (B, pred_len)
	return yhat


	# ============================================================
	# PTST: user-facing class
	# ============================================================

	@dataclass
	class PTSTConfig:
	# data / windowing
	L: int = 128
	patch_len: int = 16
	pred_len: int = 96 # training horizon per forward pass

	# model
	d_model: int = 128
	depth: int = 4
	n_heads: Optional[int] = None
	dropout: float = 0.0
	use_learned_pos: bool = True

	# ablations
	use_ffn: bool = True # keep/remove FFN
	use_cam: bool = False # enable ctx-CAM inside MHSA (your C^-)

	# ctx-CAM hyperparams (only used if use_cam=True)
	cam_tau: float = 1.0
	cam_init_scale: float = 0.0 # gate; set 0.0 to start as identity

	# training
	seed: int = 0
	val_split: float = 0.1
	batch_size: int = 128
	max_epochs: int = 50
	init_lr: float = 3e-4
	min_lr: float = 1e-5
	lr_decay: float = 0.5
	tol_rel_improve: float = 1e-3
	patience: int = 5

	class PTST:
	"""
	PTST = PatchTST-style forecaster with TART-like API.

	Usage:
	model = PTST(F_train, L=160, patch_len=16, pred_len=120, d_model=128, depth=4)
	preds = model(prefix, steps=120)

	Notes:
	- Default is channel-independent: one univariate backbone shared across all dims.
	- Trains multi-horizon (pred_len) directly.
	- For steps > pred_len, it rolls forward in pred_len chunks.
	"""

	def __init__(self, R_tX: np.ndarray, **kwargs):
	cfg = PTSTConfig(**kwargs)
	self.cfg = cfg

	R_tX = _as_2d(R_tX)
	T, D = R_tX.shape
	self.D = int(D)

	# enforce patch compatibility
	L = int(cfg.L)
	P = int(cfg.patch_len)
	if L < P:
	raise ValueError(f"L={L} must be >= patch_len={P}")
	L_eff = (L // P) * P
	if L_eff != L:
	# keep behavior explicit for reproducibility:
	# we just trim L down to nearest multiple of patch_len
	print(f"[PTST] warning: L={L} not divisible by patch_len={P}. Using L_eff={L_eff}.")
	self.L = int(L) # requested
	self.L_eff = int(L_eff) # actually used
	self.patch_len = int(P)

	pred_len = int(cfg.pred_len)
	if T <= self.L_eff + pred_len + 1:
	raise ValueError(f"Need T > L_eff + pred_len. Got T={T}, L_eff={self.L_eff}, pred_len={pred_len}.")

	# normalize train series channel-wise
	F_norm, (mu, sd) = zscore(R_tX)
	self._mu = mu
	self._sd = sd

	# build multihorizon dataset
	X_in, Y_out = make_windows_multihorizon(F_norm, L=self.L_eff, H=pred_len) # (N,L,D), (N,H,D)
	N = X_in.shape[0]

	# channel-independent reshape: treat each channel as separate training sample
	# X_ci: (N*D, L_eff, 1)
	# Y_ci: (N*D, pred_len)
	# X_in: (N, L_eff, D), Y_out: (N, pred_len, D)
	if X_in.shape[1] != self.L_eff:
	raise ValueError(
	f"[PTST] X_in has wrong time length {X_in.shape[1]} (expected {self.L_eff}). "
	f"This usually means window construction is wrong."
	)

	# channel-independent: treat each channel as an independent training sample
	# X_ci: (N*D, L_eff, 1)
	# Y_ci: (N*D, pred_len)
	X_ci = np.transpose(X_in, (0, 2, 1)) # (N, D, L_eff)
	X_ci = X_ci.reshape(N * D, self.L_eff)[:, :, None] # (N*D, L_eff, 1)

	Y_ci = np.transpose(Y_out, (0, 2, 1)) # (N, D, pred_len)
	Y_ci = Y_ci.reshape(N * D, pred_len) # (N*D, pred_len)


	# train/val split
	n_val = max(1, int(cfg.val_split * (N * D)))
	n_tr = (N * D) - n_val
	Xtr, Ytr = X_ci[:n_tr], Y_ci[:n_tr]
	Xva, Yva = X_ci[n_tr:], Y_ci[n_tr:]

	# build model
	d_model = int(max(32, cfg.d_model))
	n_heads = cfg.n_heads if cfg.n_heads is not None else choose_heads(d_model)
	self.d_model = int(d_model)
	self.n_heads = int(n_heads)
	self.depth = int(cfg.depth)
	self.pred_len = int(pred_len)

	self.model = PTSTBackbone(
	L_eff=self.L_eff,
	patch_len=self.patch_len,
	d_model=self.d_model,
	depth=self.depth,
	n_heads=self.n_heads,
	pred_len=self.pred_len,
	dropout=float(cfg.dropout),
	use_learned_pos=bool(cfg.use_learned_pos),

	use_cam=bool(cfg.use_cam),
	use_ffn=bool(cfg.use_ffn),
	cam_tau=float(cfg.cam_tau),
	cam_init_scale=float(cfg.cam_init_scale),

	)

	print(
	f"[PTST] Train samples: {N} windows -> {N*D} CI samples \| "
	f"L_eff={self.L_eff} patch_len={self.patch_len} n_patches={self.L_eff//self.patch_len} \| "
	f"D={self.D} pred_len={self.pred_len} \| d_model={self.d_model} heads={self.n_heads} depth={self.depth}"
	)

	rng = jax.random.PRNGKey(int(cfg.seed))
	dummy_x = jnp.zeros((1, self.L_eff, 1), dtype=jnp.float32)
	params = self.model.init(rng, dummy_x, deterministic=True)["params"]

	def create_state(lr: float):
	tx = optax.adamw(learning_rate=lr, weight_decay=0.0)
	return train_state.TrainState.create(apply_fn=self.model.apply, params=params, tx=tx)

	state = create_state(float(cfg.init_lr))
	self._rng = rng

	@jax.jit
	def train_step(state, xb, yb, rng):
	dropout_rng, new_rng = jax.random.split(rng)

	def loss_fn(p):
	yhat = state.apply_fn({"params": p}, xb, deterministic=False, rngs={"dropout": dropout_rng})
	loss = jnp.mean((yhat - yb) ** 2)
	return loss

	loss, grads = jax.value_and_grad(loss_fn)(state.params)
	state = state.apply_gradients(grads=grads)
	return state, loss, new_rng

	@jax.jit
	def eval_step(state, xb, yb):
	yhat = state.apply_fn({"params": state.params}, xb, deterministic=True)
	return jnp.mean((yhat - yb) ** 2)

	# move to device once
	Xtr_j = jnp.asarray(Xtr)
	Ytr_j = jnp.asarray(Ytr)
	Xva_j = jnp.asarray(Xva)
	Yva_j = jnp.asarray(Yva)

	best_params = state.params
	best_val = float("inf")
	curr_lr = float(cfg.init_lr)
	epochs_no_gain = 0

	def num_batches(n: int) -> int:
	return int(np.ceil(n / cfg.batch_size))

	# training loop
	for epoch in range(1, int(cfg.max_epochs) + 1):
	idx = np.arange(n_tr)
	np.random.default_rng(epoch + cfg.seed).shuffle(idx)

	train_losses = []
	rng = self._rng

	for bi in range(num_batches(n_tr)):
	s = bi * cfg.batch_size
	e = min((bi + 1) * cfg.batch_size, n_tr)
	bidx = idx[s:e]
	xb = Xtr_j[bidx]
	yb = Ytr_j[bidx]
	state, loss, rng = train_step(state, xb, yb, rng)
	train_losses.append(float(loss))

	self._rng = rng
	tr_loss = float(np.mean(train_losses))

	val_losses = []
	for bi in range(num_batches(n_val)):
	s = bi * cfg.batch_size
	e = min((bi + 1) * cfg.batch_size, n_val)
	xb = Xva_j[s:e]
	yb = Yva_j[s:e]
	val_losses.append(float(eval_step(state, xb, yb)))
	va_loss = float(np.mean(val_losses))

	print(f"[PTST] epoch {epoch:03d} \| train {tr_loss:.6e} \| val {va_loss:.6e} \| lr {curr_lr:.2e}")

	if va_loss + 1e-8 < best_val:
	best_val = va_loss
	best_params = state.params
	epochs_no_gain = 0
	else:
	epochs_no_gain += 1

	# plateau schedule
	if epochs_no_gain >= int(cfg.patience):
	if curr_lr > float(cfg.min_lr) * (1.0 + 1e-9):
	curr_lr = max(float(cfg.min_lr), curr_lr * float(cfg.lr_decay))
	print(f"[PTST] plateau → lowering LR to {curr_lr:.2e}")
	params_now = state.params
	state = train_state.TrainState.create(
	apply_fn=state.apply_fn,
	params=params_now,
	tx=optax.adamw(learning_rate=curr_lr, weight_decay=0.0),
	)
	epochs_no_gain = 0
	else:
	print(f"[PTST] early stop: lr at min and no improvement (best val {best_val:.6e})")
	break

	# store best
	self.state = state.replace(params=best_params)
	self.params = self.state.params

	# param count
	self.param_count = int(sum(p.size for p in jax.tree_util.tree_leaves(self.params)))
	print(f"[PTST] params: {self.param_count:,}")

	# ------------------------------------------------------------
	# Internal: one-shot predict (normalized space) for a batch
	# ------------------------------------------------------------
	def _predict_block_norm(self, ctx_norm_BLD: jnp.ndarray) -> jnp.ndarray:
	"""
	ctx_norm_BLD: (B, L_eff, D) normalized
	returns: (B, pred_len, D) normalized
	"""
	B, L, D = ctx_norm_BLD.shape
	assert L == self.L_eff
	assert D == self.D

	# channel-independent apply:
	# reshape -> (B*D, L, 1)
	x = ctx_norm_BLD.transpose(0, 2, 1)[:, :, :, None] # (B, D, L, 1)
	x = x.reshape(B * D, L, 1)
	yhat = self.model.apply({"params": self.params}, x, deterministic=True) # (B*D, pred_len)
	yhat = yhat.reshape(B, D, self.pred_len) # (B,D,H)
	yhat = yhat.transpose(0, 2, 1) # (B,H,D)
	return yhat

	# ------------------------------------------------------------
	# Public API: forecast
	# ------------------------------------------------------------
	def __call__(self, prefix: np.ndarray, steps: int) -> np.ndarray:
	"""
	Forecast `steps` time points given prefix ending at forecast start.

	prefix: (ell, D)
	returns: (steps, D)
	"""
	X = _as_2d(prefix)
	ell, D = X.shape
	if D != self.D:
	raise ValueError(f"PTST expected D={self.D}, got {D}.")

	H = int(steps)
	if H <= 0:
	return np.zeros((0, self.D), dtype=np.float32)

	# take last L_eff points (pad on left if needed)
	if ell < self.L_eff:
	pad_len = self.L_eff - ell
	pad = np.repeat(X[:1], repeats=pad_len, axis=0)
	ctx = np.concatenate([pad, X], axis=0)
	else:
	ctx = X[-self.L_eff:, :]

	# normalize context
	ctx_n = (ctx - self._mu) / self._sd # (L_eff, D)

	preds_n_all = []
	n_done = 0

	while n_done < H:
	ctx_n_j = jnp.asarray(ctx_n[None, :, :]) # (1,L_eff,D)
	block_n = np.array(self._predict_block_norm(ctx_n_j)[0]) # (pred_len, D)

	remaining = H - n_done
	take = min(self.pred_len, remaining)

	preds_n_all.append(block_n[:take])
	n_done += take

	# roll context
	ctx_n = np.concatenate([ctx_n, block_n[:take]], axis=0)
	ctx_n = ctx_n[-self.L_eff:, :]


	preds_n = np.concatenate(preds_n_all, axis=0) # (H, D)
	preds = preds_n * self._sd + self._mu
	return preds


	__all__ = ["PTST", "PTSTConfig"]