From 07c4be157b7ca64010c7444abed8d1ccd7986c82 Mon Sep 17 00:00:00 2001
From: Ryuichi Leo Takashige <leo@exolabs.net>
Date: Wed, 28 Jan 2026 14:07:19 +0000
Subject: [PATCH] Oops

---
 src/exo/worker/engines/mlx/auto_parallel.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/exo/worker/engines/mlx/auto_parallel.py b/src/exo/worker/engines/mlx/auto_parallel.py
index 74681438..9dcf41c4 100644
--- a/src/exo/worker/engines/mlx/auto_parallel.py
+++ b/src/exo/worker/engines/mlx/auto_parallel.py
@@ -681,8 +681,8 @@ class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
         def __call__(
             self,
             x: mx.array,
-            mask: Optional[mx.array] = None,
-            cache: Optional[Any] = None,
+            mask: mx.array | Any = None,
+            cache: Any | None = None,
         ) -> mx.array:
             B, L, D = x.shape