- Initialize LoRA parameters on the same device as the base linear layer
- Load adapter weights directly to the model device instead of hardcoded CPU
- Slice XReg linear regression outputs to match the specified sequence lengths
- Replace batch-wide covariate normalization with per-input normalization
in create_covariate_matrix to prevent cross-input scale leakage.
- Refactor BatchedInContextXRegLinear.fit to solve ridge regression
per instance rather than as a single global matrix solve, avoiding
cross-contamination between batched inputs.
- Truncate JAX regression outputs to the actual train/test lengths
after the padded matrix multiply, fixing shape mismatches for
non-power-of-2 horizons (e.g. horizon=24 was returning 32 elements).
PEFTTrainer with production-grade training loop:
- PyTorch DDP multi-GPU via torchrun
- Mixed-precision training (fp16/bf16) with GradScaler
- Gradient checkpointing for long contexts
- Cosine-with-warmup LR schedule
- MSE loss + optional pinball quantile loss (9 channels)
- Early stopping on validation loss
- Adapter-only checkpointing (safetensors)
- W&B logging (rank-0 only)
- Differentiable training forward that replicates the 2.5
patch -> RevIN -> transformer -> output-head -> un-RevIN path
Sliding-window dataset that produces (context, mask, target) tuples:
- Accepts list of arrays, long-format, or wide-format DataFrames
- Context length auto-rounded to multiple of patch_len (32)
- Left-pads short series with proper masking
- Configurable stride for window overlap
darkpowerxo