fix: improve PEFT device consistency and XReg output slicing

- 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).

peft/adapters.py

@@ -110,9 +110,10 @@ class DoRALinear(nn.Module):
 
     in_f = base_linear.in_features
     out_f = base_linear.out_features
+    dev = base_linear.weight.device
 
-    self.lora_A = nn.Parameter(torch.empty(in_f, rank))
-    self.lora_B = nn.Parameter(torch.zeros(rank, out_f))
+    self.lora_A = nn.Parameter(torch.empty(in_f, rank, device=dev))
+    self.lora_B = nn.Parameter(torch.zeros(rank, out_f, device=dev))
     nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
 
     # Magnitude vector — initialised from pretrained column norms.
@@ -272,7 +273,8 @@ def load_adapter_weights(model: nn.Module, path: str) -> None:
   The model must already have adapters injected (via ``inject_adapters``)
   before calling this function.
   """
-  tensors = load_file(path, device="cpu")
+  device = str(next(model.parameters()).device)
+  tensors = load_file(path, device=device)
   trainable = {n for n, p in model.named_parameters() if p.requires_grad}
   missing = trainable - set(tensors.keys())
   if missing:

src/timesfm/utils/xreg_lib.py

@@ -509,9 +509,9 @@ class BatchedInContextXRegLinear(BatchedInContextXRegBase):
           @ x_tr_j.T
           @ y_tr_j
         )
-        outputs.append(np.array(x_te_j @ beta_hat))
+        outputs.append(np.array(x_te_j @ beta_hat)[:tel])
         if debug_info:
-          outputs_context.append(np.array(x_tr_raw_j @ beta_hat))
+          outputs_context.append(np.array(x_tr_raw_j @ beta_hat)[:trl])
 
         train_idx += trl
         test_idx += tel