mps support added

src/tirex2/api_adapter/forecast.py

@@ -14,6 +14,27 @@
 ForecastOutputType = Literal["torch", "numpy", "gluonts", "fev"]
 
 
+def _is_oom_error(exc: BaseException) -> bool:
+    """Return whether ``exc`` is an out-of-memory error, on CUDA or MPS.
+
+    CUDA raises :class:`torch.cuda.OutOfMemoryError`; MPS currently surfaces OOM
+    as a plain :class:`RuntimeError` whose message mentions running out of memory.
+    """
+    if isinstance(exc, torch.cuda.OutOfMemoryError):
+        return True
+    if isinstance(exc, RuntimeError):
+        return "out of memory" in str(exc).lower()
+    return False
+
+
+def _empty_device_cache(device: str) -> None:
+    """Release the caching allocator for ``device`` (no-op on CPU)."""
+    if device.startswith("cuda"):
+        torch.cuda.empty_cache()
+    elif device == "mps":
+        torch.mps.empty_cache()
+
+
 def _format_output(forecasts, meta, output_type, quantile_levels):
     """Render a batch of per-series ``[V_t, Q, H]`` forecasts in the requested output format."""
     if output_type == "torch":
@@ -42,11 +63,11 @@ def _predict_adaptive(
     quantile_levels,
     **predict_kwargs,
 ):
-    """Yield formatted forecasts batch by batch, halving the batch size on CUDA OOM.
+    """Yield formatted forecasts batch by batch, halving the batch size on device OOM.
 
     Walks contiguous ``[start, end)`` windows of at most ``batch_size`` series,
     forecasting and formatting each (slicing ``meta`` alongside ``timeseries``).
-    When a window raises :class:`torch.cuda.OutOfMemoryError`, the CUDA cache is
+    When a window runs out of memory (CUDA or MPS), the device cache is
     cleared, the batch size is halved (floor of 1), and the *same* window is retried
     at the smaller size. The reduced size persists for the rest of this call, so a
     single oversized window pins it down only here - a fresh call starts again from
@@ -59,20 +80,23 @@ def _predict_adaptive(
     """
     assert batch_size >= 1, "Batch size must be >= 1"
     num_items = len(timeseries)
+    device = str(getattr(model, "device", "cpu"))
     start = 0
     current = batch_size
     while start < num_items:
         end = min(start + current, num_items)
         try:
             forecasts = model.predict(timeseries[start:end], prediction_length, **predict_kwargs)
             formatted = _format_output(forecasts, meta[start:end], output_type, quantile_levels)
-        except torch.cuda.OutOfMemoryError:
-            torch.cuda.empty_cache()
+        except (torch.cuda.OutOfMemoryError, RuntimeError) as exc:
+            if not _is_oom_error(exc):
+                raise
+            _empty_device_cache(device)
             if current == 1:
-                logger.error("CUDA OOM at batch size 1 (series index %d); cannot shrink further.", start)
+                logger.error("Device OOM at batch size 1 (series index %d); cannot shrink further.", start)
                 raise
             current = max(1, current // 2)
-            logger.warning("CUDA OOM at series index %d; halving batch size to %d and retrying.", start, current)
+            logger.warning("Device OOM at series index %d; halving batch size to %d and retrying.", start, current)
             continue
         yield formatted
         start = end

src/tirex2/base.py

@@ -60,9 +60,10 @@ def load_model(
         Local directory holding ``model-config.yaml`` and ``model.ckpt``. Values
         of the form ``hf://org/repo`` or ``org/repo`` are treated as Hugging Face
         model repo ids and downloaded with :func:`huggingface_hub.snapshot_download`.
-    device : {"cpu", "cuda"}
+    device : {"cpu", "cuda", "mps"}
         Runtime device and recurrent-kernel family to use. This overrides any
-        device/backend stored in the checkpoint config.
+        device/backend stored in the checkpoint config. ``"mps"`` runs on Apple
+        Metal using the same pure-PyTorch (native) kernels as ``"cpu"``.
     hf_kwargs : dict, optional
         Extra keyword arguments forwarded to ``snapshot_download`` for Hugging
         Face paths, e.g. ``{"revision": "main", "local_files_only": True}``.
@@ -76,6 +77,8 @@ def load_model(
     """
     if device.startswith("cuda") and not torch.cuda.is_available():
         raise RuntimeError("Execution on CUDA was requested but is not available.")
+    if device == "mps" and not torch.backends.mps.is_available():
+        raise RuntimeError("Execution on MPS was requested but is not available.")
 
     ckpt_dir = _resolve_ckpt_dir(ckpt_path, hf_kwargs=hf_kwargs)
     config_file = ckpt_dir / CONFIG_FILENAME

src/tirex2/model/component/bi_xlstm.py

@@ -47,7 +47,7 @@ class BiXLSTM(nn.Module):
         Index of this block in the overall stack.
     num_blocks : int
         Total number of blocks in the stack.
-    device : {"cpu", "cuda"}
+    device : {"cpu", "cuda", "mps"}
         Device used to choose recurrent kernels.
     dropout : float
         Dropout probability applied in the combination projection.
@@ -64,7 +64,7 @@ def __init__(
         config: xLSTMMixedConfig,
         block_idx: int,
         num_blocks: int,
-        device: Literal["cpu", "cuda"],
+        device: Literal["cpu", "cuda", "mps"],
         dropout: float = 0.0,
         share_weights: bool = True,
         cell_type: Literal["slstm", "mlstm"] = "slstm",

src/tirex2/model/component/flashrnn_slstm.py

@@ -340,17 +340,19 @@ def _bias_int2ext(bias_int: torch.Tensor):
         return bias_int
 
 
-def _flashrnn_backend(device: Literal["cpu", "cuda"]) -> str:
+def _flashrnn_backend(device: Literal["cpu", "cuda", "mps"]) -> str:
     match device:
-        case "cpu":
+        case "cpu" | "mps":
+            # The "vanilla" backend is pure PyTorch and device-agnostic, so it
+            # runs on Apple Metal; FlashRNN's fused "cuda" backend does not.
             return "vanilla"
         case "cuda":
             return "cuda"
         case _:
-            raise ValueError(f"device must be 'cpu' or 'cuda', got {device!r}.")
+            raise ValueError(f"device must be 'cpu', 'cuda', or 'mps', got {device!r}.")
 
 
-def init_cell(config: xLSTMMixedConfig, block_idx: int, num_blocks: int, device: Literal["cpu", "cuda"]):
+def init_cell(config: xLSTMMixedConfig, block_idx: int, num_blocks: int, device: Literal["cpu", "cuda", "mps"]):
     """Instantiate an sLSTM cell for the requested runtime device."""
     return sLSTMFlashRNNLayer(
         FlashRNNLayerConfig(

src/tirex2/model/component/mlstm_block.py

@@ -145,9 +145,14 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return y
 
 
-def _mlstm_backend_config(config: xLSTMMixedConfig, device: Literal["cpu", "cuda"]) -> mLSTMBackendConfig:
-    """Return the mLSTM kernel backend matching the requested runtime device."""
-    if device == "cpu":
+def _mlstm_backend_config(config: xLSTMMixedConfig, device: Literal["cpu", "cuda", "mps"]) -> mLSTMBackendConfig:
+    """Return the mLSTM kernel backend matching the requested runtime device.
+
+    ``"mps"`` shares the ``"cpu"`` configuration: the pure-PyTorch native kernels
+    are device-agnostic and run on Apple Metal, whereas the Triton kernels used by
+    ``"cuda"`` are unavailable there.
+    """
+    if device in ("cpu", "mps"):
         return mLSTMBackendConfig(
             chunkwise_kernel="chunkwise--native_autograd",
             sequence_kernel="native_sequence__native",
@@ -173,10 +178,10 @@ def _mlstm_backend_config(config: xLSTMMixedConfig, device: Literal["cpu", "cuda
             inference_state_dtype="float32",
         )
 
-    raise ValueError(f"device must be 'cpu' or 'cuda', got {device!r}.")
+    raise ValueError(f"device must be 'cpu', 'cuda', or 'mps', got {device!r}.")
 
 
-def init_cell(config: xLSTMMixedConfig, device: Literal["cpu", "cuda"]) -> mLSTMLayer:
+def init_cell(config: xLSTMMixedConfig, device: Literal["cpu", "cuda", "mps"]) -> mLSTMLayer:
     """Instantiate an mLSTM cell for the requested runtime device."""
     layer = mLSTMLayer(
         conv_mLSTMLayerConfig(

src/tirex2/model/component/variate_mixing_block.py

@@ -40,7 +40,7 @@ class TimeMixerConfig:
     num_heads: int = 4
     num_slstm_heads: int | None = None
 
-    device: Literal["cpu", "cuda"] = "cuda"
+    device: Literal["cpu", "cuda", "mps"] = "cuda"
 
 
 @dataclass

src/tirex2/model/tirex2.py

@@ -19,16 +19,16 @@
 )
 from .types import TimeseriesType
 
-Device = Literal["cpu", "cuda"]
+Device = Literal["cpu", "cuda", "mps"]
 MatmulPrecision = Literal["highest", "high", "medium"]
 
 logger = logging.getLogger(__file__)
 
 
 def _normalize_device(device: str) -> Device:
     """Normalize the public runtime device selector used by TiRex2."""
-    if device not in ("cpu", "cuda"):
-        raise ValueError(f"device must be 'cpu' or 'cuda', got {device!r}.")
+    if device not in ("cpu", "cuda", "mps"):
+        raise ValueError(f"device must be 'cpu', 'cuda', or 'mps', got {device!r}.")
     return device
 
 
@@ -90,9 +90,10 @@ def from_dict(
             Whether the attention variate mixers should RMS-normalize their
             query/key vectors. Injected into every mixer built from a dict
             template (default: True).
-        device : {"cpu", "cuda"}
+        device : {"cpu", "cuda", "mps"}
             Runtime device used to choose recurrent kernels. This overrides
-            any serialized time mixer device/backend setting.
+            any serialized time mixer device/backend setting. ``"mps"`` uses the
+            same pure-PyTorch (native) kernels as ``"cpu"``, run on Apple Metal.
         """
         templates = {}
         for name, template in config["templates"].items():