diff --git a/.mlx_typings/mlx/nn/layers/distributed.pyi b/.mlx_typings/mlx/nn/layers/distributed.pyi
index 5be9cc4b..b23589eb 100644
--- a/.mlx_typings/mlx/nn/layers/distributed.pyi
+++ b/.mlx_typings/mlx/nn/layers/distributed.pyi
@@ -9,6 +9,12 @@ import mlx.core as mx
 from base import Module
 from mlx.nn.layers.linear import Linear
 
+def compute_shard_sizes(
+    dim: int,
+    N: int,
+    unit: int = ...,
+    weights: list[float] | None = ...,
+) -> list[int]: ...
 @lru_cache
 def sum_gradients(
     group: mx.distributed.Group,
@@ -20,6 +26,8 @@ def shard_inplace(
     *,
     segments: Union[int, list[int]] = ...,
     group: Optional[mx.distributed.Group] = ...,
+    unit: int = ...,
+    weights: list[float] | None = ...,
 ) -> None:
     """Shard a module in-place by updating its parameter dictionary with the
     sharded parameter dictionary.
@@ -51,6 +59,8 @@ def shard_linear(
     *,
     segments: Union[int, list[int]] = ...,
     group: Optional[mx.distributed.Group] = ...,
+    unit: int = ...,
+    weights: list[float] | None = ...,
 ) -> Linear:
     """Create a new linear layer that has its parameters sharded and also
     performs distributed communication either in the forward or backward
diff --git a/.mlx_typings/mlx_lm/models/switch_layers.pyi b/.mlx_typings/mlx_lm/models/switch_layers.pyi
index 235d4168..05360bfd 100644
--- a/.mlx_typings/mlx_lm/models/switch_layers.pyi
+++ b/.mlx_typings/mlx_lm/models/switch_layers.pyi
@@ -8,6 +8,9 @@ import mlx.core as mx
 import mlx.nn as nn
 
 class QuantizedSwitchLinear(nn.Module):
+    weight: mx.array
+    scales: mx.array
+
     def __init__(
         self,
         input_dims: int,
@@ -31,6 +34,8 @@ class QuantizedSwitchLinear(nn.Module):
         ...
 
 class SwitchLinear(nn.Module):
+    weight: mx.array
+
     def __init__(
         self, input_dims: int, output_dims: int, num_experts: int, bias: bool = ...
     ) -> None: ...
diff --git a/bench/exo_bench.py b/bench/exo_bench.py
index 97b6b3ac..afb52527 100644
--- a/bench/exo_bench.py
+++ b/bench/exo_bench.py
@@ -378,8 +378,21 @@ def main() -> int:
         default=1.0,
         help="System metrics polling interval in seconds (default: 1.0).",
     )
+    ap.add_argument(
+        "--tensor-strategies",
+        nargs="+",
+        default=["Naive"],
+        help="Tensor shard strategies to benchmark. Choices: Naive, Memory, Compute, Bandwidth, all. Default: Naive.",
+    )
     args = ap.parse_args()
 
+    tensor_strategies = []
+    for s in args.tensor_strategies:
+        if s.lower() == "all":
+            tensor_strategies = ["Naive", "Memory", "Compute", "Bandwidth"]
+            break
+        tensor_strategies.append(s)
+
     pp_list = parse_int_list(args.pp)
     tg_list = parse_int_list(args.tg)
     if not pp_list or not tg_list:
@@ -467,20 +480,45 @@ def main() -> int:
     all_rows: list[dict[str, Any]] = []
     all_system_metrics: dict[str, dict[str, dict[str, float]]] = {}
 
+    placement_runs: list[tuple[dict[str, Any], str]] = []
     for preview in selected:
+        sharding = str(preview["sharding"])
+        if sharding == "Tensor":
+            for ts in tensor_strategies:
+                placement_runs.append((preview, ts))
+        else:
+            placement_runs.append((preview, "Naive"))
+
+    for preview, tensor_strategy in placement_runs:
         instance = preview["instance"]
         instance_id = instance_id_from_instance(instance)
 
         sharding = str(preview["sharding"])
         instance_meta = str(preview["instance_meta"])
         n_nodes = nodes_used_in_instance(instance)
+        strategy_label = (
+            f" / strategy={tensor_strategy}" if sharding == "Tensor" else ""
+        )
 
         logger.info("=" * 80)
         logger.info(
-            f"PLACEMENT: {sharding} / {instance_meta} / nodes={n_nodes} / instance_id={instance_id}"
+            f"PLACEMENT: {sharding} / {instance_meta} / nodes={n_nodes}{strategy_label} / instance_id={instance_id}"
         )
 
-        client.request_json("POST", "/instance", body={"instance": instance})
+        if sharding == "Tensor" and tensor_strategy != "Naive":
+            client.request_json(
+                "POST",
+                "/place_instance",
+                body={
+                    "model_id": full_model_id,
+                    "sharding": sharding,
+                    "instance_meta": instance_meta,
+                    "min_nodes": n_nodes,
+                    "tensor_strategy": tensor_strategy,
+                },
+            )
+        else:
+            client.request_json("POST", "/instance", body={"instance": instance})
         try:
             wait_for_instance_ready(client, instance_id)
         except (RuntimeError, TimeoutError) as e:
@@ -541,6 +579,7 @@ def main() -> int:
                                     "placement_sharding": sharding,
                                     "placement_instance_meta": instance_meta,
                                     "placement_nodes": n_nodes,
+                                    "tensor_strategy": tensor_strategy,
                                     "instance_id": instance_id,
                                     "pp_tokens": actual_pp_tokens,
                                     "tg": tg,
@@ -595,6 +634,7 @@ def main() -> int:
                                         "placement_sharding": sharding,
                                         "placement_instance_meta": instance_meta,
                                         "placement_nodes": n_nodes,
+                                        "tensor_strategy": tensor_strategy,
                                         "instance_id": instance_id,
                                         "pp_tokens": actual_pp_tokens,
                                         "tg": tg,
@@ -656,7 +696,9 @@ def main() -> int:
         finally:
             if sampler:
                 sampler.stop()
-                placement_label = f"{sharding}/{instance_meta}/{n_nodes} nodes"
+                placement_label = (
+                    f"{sharding}/{instance_meta}/{n_nodes} nodes{strategy_label}"
+                )
                 sampler.print_summary(placement_label)
                 placement_metrics = sampler.summarize()
                 if placement_metrics:
diff --git a/dashboard/src/routes/+page.svelte b/dashboard/src/routes/+page.svelte
index 965482e6..1c121513 100644
--- a/dashboard/src/routes/+page.svelte
+++ b/dashboard/src/routes/+page.svelte
@@ -886,6 +886,8 @@
   }
 
   let selectedSharding = $state<"Pipeline" | "Tensor">("Pipeline");
+  let selectedTensorStrategy = $state<string>("Naive");
+  let showAdvancedSettings = $state(false);
   type InstanceMeta = "MlxRing" | "MlxJaccl";
 
   // Launch defaults persistence
@@ -1446,6 +1448,8 @@
             sharding: selectedSharding,
             instance_meta: selectedInstanceType,
             min_nodes: 1,
+            tensor_strategy:
+              selectedSharding === "Tensor" ? selectedTensorStrategy : "Naive",
           }),
         });
       }
@@ -5874,6 +5878,61 @@
                     </div>
                   </div>
                 </div>
+
+                <!-- Advanced Settings -->
+                {#if selectedSharding === "Tensor"}
+                  <div class="border-t border-exo-medium-gray/30 pt-3 mt-3">
+                    <button
+                      onclick={() =>
+                        (showAdvancedSettings = !showAdvancedSettings)}
+                      class="flex items-center gap-2 text-xs text-white/40 font-mono hover:text-white/60 transition-colors cursor-pointer"
+                    >
+                      <span
+                        class="transition-transform duration-200 {showAdvancedSettings
+                          ? 'rotate-90'
+                          : ''}">▶</span
+                      >
+                      Advanced
+                    </button>
+                    {#if showAdvancedSettings}
+                      <div class="mt-3 space-y-3">
+                        <div>
+                          <div class="text-xs text-white/50 font-mono mb-2">
+                            Tensor Strategy:
+                          </div>
+                          <div class="flex gap-2 flex-wrap">
+                            {#each ["Naive", "Memory", "Compute", "Bandwidth"] as strategy}
+                              <button
+                                onclick={() => {
+                                  selectedTensorStrategy = strategy;
+                                  saveLaunchDefaults();
+                                }}
+                                class="flex items-center gap-2 py-1.5 px-3 text-xs font-mono border rounded transition-all duration-200 cursor-pointer {selectedTensorStrategy ===
+                                strategy
+                                  ? 'bg-transparent text-exo-yellow border-exo-yellow'
+                                  : 'bg-transparent text-white/70 border-exo-medium-gray/50 hover:border-exo-yellow/50'}"
+                              >
+                                <span
+                                  class="w-3 h-3 rounded-full border-2 flex items-center justify-center {selectedTensorStrategy ===
+                                  strategy
+                                    ? 'border-exo-yellow'
+                                    : 'border-exo-medium-gray'}"
+                                >
+                                  {#if selectedTensorStrategy === strategy}
+                                    <span
+                                      class="w-1.5 h-1.5 rounded-full bg-exo-yellow"
+                                    ></span>
+                                  {/if}
+                                </span>
+                                {strategy}
+                              </button>
+                            {/each}
+                          </div>
+                        </div>
+                      </div>
+                    {/if}
+                  </div>
+                {/if}
               {/if}
             </div>
 
diff --git a/src/exo/api/main.py b/src/exo/api/main.py
index 5635a409..575eb215 100644
--- a/src/exo/api/main.py
+++ b/src/exo/api/main.py
@@ -383,6 +383,7 @@ class API:
             sharding=payload.sharding,
             instance_meta=payload.instance_meta,
             min_nodes=payload.min_nodes,
+            tensor_strategy=payload.tensor_strategy,
         )
         await self._send(command)
 
diff --git a/src/exo/api/types/api.py b/src/exo/api/types/api.py
index b4fa3147..e5c8e891 100644
--- a/src/exo/api/types/api.py
+++ b/src/exo/api/types/api.py
@@ -10,7 +10,7 @@ from exo.shared.types.common import CommandId, NodeId
 from exo.shared.types.memory import Memory
 from exo.shared.types.text_generation import ReasoningEffort
 from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
-from exo.shared.types.worker.shards import Sharding, ShardMetadata
+from exo.shared.types.worker.shards import Sharding, ShardMetadata, TensorShardStrategy
 from exo.utils.pydantic_ext import CamelCaseModel
 
 FinishReason = Literal[
@@ -253,6 +253,7 @@ class PlaceInstanceParams(BaseModel):
     sharding: Sharding = Sharding.Pipeline
     instance_meta: InstanceMeta = InstanceMeta.MlxRing
     min_nodes: int = 1
+    tensor_strategy: TensorShardStrategy = TensorShardStrategy.Naive
 
 
 class CreateInstanceParams(BaseModel):
diff --git a/src/exo/master/main.py b/src/exo/master/main.py
index 1511897e..438d462f 100644
--- a/src/exo/master/main.py
+++ b/src/exo/master/main.py
@@ -298,6 +298,7 @@ class Master:
                                 self.state.instances,
                                 self.state.node_memory,
                                 self.state.node_network,
+                                node_identities=self.state.node_identities,
                                 download_status=self.state.downloads,
                             )
                             transition_events = get_transition_events(
diff --git a/src/exo/master/placement.py b/src/exo/master/placement.py
index f9380693..07699dae 100644
--- a/src/exo/master/placement.py
+++ b/src/exo/master/placement.py
@@ -29,7 +29,7 @@ from exo.shared.types.events import (
     TaskStatusUpdated,
 )
 from exo.shared.types.memory import Memory
-from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo
+from exo.shared.types.profiling import MemoryUsage, NodeIdentity, NodeNetworkInfo
 from exo.shared.types.tasks import Task, TaskId, TaskStatus
 from exo.shared.types.worker.downloads import (
     DownloadCompleted,
@@ -108,6 +108,7 @@ def place_instance(
     current_instances: Mapping[InstanceId, Instance],
     node_memory: Mapping[NodeId, MemoryUsage],
     node_network: Mapping[NodeId, NodeNetworkInfo],
+    node_identities: Mapping[NodeId, NodeIdentity] | None = None,
     required_nodes: set[NodeId] | None = None,
     download_status: Mapping[NodeId, Sequence[DownloadProgress]] | None = None,
 ) -> dict[InstanceId, Instance]:
@@ -197,7 +198,12 @@ def place_instance(
         command.sharding = Sharding.Pipeline
 
     shard_assignments = get_shard_assignments(
-        command.model_card, selected_cycle, command.sharding, node_memory
+        command.model_card,
+        selected_cycle,
+        command.sharding,
+        node_memory,
+        tensor_strategy=command.tensor_strategy,
+        node_identities=node_identities,
     )
 
     cycle_digraph: Topology = topology.get_subgraph_from_nodes(selected_cycle.node_ids)
diff --git a/src/exo/master/placement_utils.py b/src/exo/master/placement_utils.py
index eb78c8fa..4c386471 100644
--- a/src/exo/master/placement_utils.py
+++ b/src/exo/master/placement_utils.py
@@ -6,7 +6,7 @@ from exo.shared.models.model_cards import ModelCard
 from exo.shared.topology import Topology
 from exo.shared.types.common import Host, NodeId
 from exo.shared.types.memory import Memory
-from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo
+from exo.shared.types.profiling import MemoryUsage, NodeIdentity, NodeNetworkInfo
 from exo.shared.types.topology import Cycle, RDMAConnection, SocketConnection
 from exo.shared.types.worker.runners import RunnerId, ShardAssignments
 from exo.shared.types.worker.shards import (
@@ -15,8 +15,54 @@ from exo.shared.types.worker.shards import (
     Sharding,
     ShardMetadata,
     TensorShardMetadata,
+    TensorShardMode,
+    TensorShardStrategy,
 )
 
+# FP32 TFLOPS by chip_id (used for Compute strategy)
+APPLE_SILICON_FLOPS: dict[str, float] = {
+    "Apple M1": 2.6,
+    "Apple M1 Pro": 5.3,
+    "Apple M1 Max": 10.4,
+    "Apple M1 Ultra": 21.0,
+    "Apple M2": 3.6,
+    "Apple M2 Pro": 6.8,
+    "Apple M2 Max": 13.6,
+    "Apple M2 Ultra": 27.2,
+    "Apple M3": 3.5,
+    "Apple M3 Pro": 5.0,
+    "Apple M3 Max": 14.2,
+    "Apple M3 Ultra": 28.4,
+    "Apple M4": 4.3,
+    "Apple M4 Pro": 9.2,
+    "Apple M4 Max": 18.4,
+    "Apple M5": 4.2,
+    "Apple M5 Pro": 8.3,
+    "Apple M5 Max": 19.9,
+}
+
+# Memory bandwidth in GB/s by chip_id (used for Bandwidth strategy)
+APPLE_SILICON_BANDWIDTH: dict[str, float] = {
+    "Apple M1": 68,
+    "Apple M1 Pro": 200,
+    "Apple M1 Max": 400,
+    "Apple M1 Ultra": 800,
+    "Apple M2": 100,
+    "Apple M2 Pro": 200,
+    "Apple M2 Max": 400,
+    "Apple M2 Ultra": 800,
+    "Apple M3": 100,
+    "Apple M3 Pro": 150,
+    "Apple M3 Max": 400,
+    "Apple M3 Ultra": 800,
+    "Apple M4": 120,
+    "Apple M4 Pro": 273,
+    "Apple M4 Max": 546,
+    "Apple M5": 154,
+    "Apple M5 Pro": 307,
+    "Apple M5 Max": 614,
+}
+
 
 def filter_cycles_by_memory(
     cycles: list[Cycle],
@@ -243,12 +289,39 @@ def _get_shard_assignments_for_pure_pipeline(
 def get_shard_assignments_for_tensor_parallel(
     model_card: ModelCard,
     cycle: Cycle,
+    node_memory: Mapping[NodeId, MemoryUsage] | None = None,
+    strategy: TensorShardStrategy = TensorShardStrategy.Naive,
+    node_identities: Mapping[NodeId, NodeIdentity] | None = None,
 ):
     total_layers = model_card.n_layers
     world_size = len(cycle)
     runner_to_shard: dict[RunnerId, ShardMetadata] = {}
     node_to_runner: dict[NodeId, RunnerId] = {}
 
+    shard_weights: list[float] | None = None
+    shard_mode = TensorShardMode.Constant
+    match strategy:
+        case TensorShardStrategy.Naive:
+            pass
+        case TensorShardStrategy.Memory:
+            if node_memory is not None:
+                shard_weights = [
+                    node_memory[node_id].ram_available.in_gb for node_id in cycle
+                ]
+            shard_mode = TensorShardMode.Greedy
+        case TensorShardStrategy.Compute:
+            if node_identities is not None:
+                shard_weights = [
+                    APPLE_SILICON_FLOPS.get(node_identities[nid].chip_id, 1.0)
+                    for nid in cycle
+                ]
+        case TensorShardStrategy.Bandwidth:
+            if node_identities is not None:
+                shard_weights = [
+                    APPLE_SILICON_BANDWIDTH.get(node_identities[nid].chip_id, 1.0)
+                    for nid in cycle
+                ]
+
     for i, node_id in enumerate(cycle):
         shard = TensorShardMetadata(
             model_card=model_card,
@@ -257,6 +330,8 @@ def get_shard_assignments_for_tensor_parallel(
             start_layer=0,
             end_layer=total_layers,
             n_layers=total_layers,
+            shard_weights=shard_weights,
+            shard_mode=shard_mode,
         )
 
         runner_id = RunnerId()
@@ -278,6 +353,8 @@ def get_shard_assignments(
     cycle: Cycle,
     sharding: Sharding,
     node_memory: Mapping[NodeId, MemoryUsage],
+    tensor_strategy: TensorShardStrategy = TensorShardStrategy.Naive,
+    node_identities: Mapping[NodeId, NodeIdentity] | None = None,
 ) -> ShardAssignments:
     match sharding:
         case Sharding.Pipeline:
@@ -290,6 +367,9 @@ def get_shard_assignments(
             return get_shard_assignments_for_tensor_parallel(
                 model_card=model_card,
                 cycle=cycle,
+                node_memory=node_memory,
+                strategy=tensor_strategy,
+                node_identities=node_identities,
             )
 
 
diff --git a/src/exo/shared/types/commands.py b/src/exo/shared/types/commands.py
index a6c988a5..5199fa3e 100644
--- a/src/exo/shared/types/commands.py
+++ b/src/exo/shared/types/commands.py
@@ -9,7 +9,7 @@ from exo.shared.types.chunks import InputImageChunk
 from exo.shared.types.common import CommandId, NodeId, SystemId
 from exo.shared.types.text_generation import TextGenerationTaskParams
 from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
-from exo.shared.types.worker.shards import Sharding, ShardMetadata
+from exo.shared.types.worker.shards import Sharding, ShardMetadata, TensorShardStrategy
 from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel
 
 
@@ -38,6 +38,7 @@ class PlaceInstance(BaseCommand):
     sharding: Sharding
     instance_meta: InstanceMeta
     min_nodes: int
+    tensor_strategy: TensorShardStrategy = TensorShardStrategy.Naive
 
 
 class CreateInstance(BaseCommand):
diff --git a/src/exo/shared/types/worker/shards.py b/src/exo/shared/types/worker/shards.py
index 60deaab7..f7d83a4a 100644
--- a/src/exo/shared/types/worker/shards.py
+++ b/src/exo/shared/types/worker/shards.py
@@ -17,6 +17,13 @@ class TensorShardMode(str, Enum):
     Constant = "Constant"
 
 
+class TensorShardStrategy(str, Enum):
+    Naive = "Naive"
+    Memory = "Memory"
+    Compute = "Compute"
+    Bandwidth = "Bandwidth"
+
+
 class BaseShardMetadata(TaggedModel):
     """
     Defines a specific shard of the model that is ready to be run on a device.
diff --git a/src/exo/worker/engines/mlx/auto_parallel.py b/src/exo/worker/engines/mlx/auto_parallel.py
index 9f2ae436..67d5fb35 100644
--- a/src/exo/worker/engines/mlx/auto_parallel.py
+++ b/src/exo/worker/engines/mlx/auto_parallel.py
@@ -14,7 +14,6 @@ from mlx.nn.layers.distributed import (
     shard_linear,
     sum_gradients,
 )
-from mlx.utils import tree_flatten
 from mlx_lm.models.base import (
     scaled_dot_product_attention,  # pyright: ignore[reportUnknownVariableType]
 )
@@ -497,9 +496,9 @@ def tensor_auto_parallel(
     on_timeout: TimeoutCallback | None,
     on_layer_loaded: LayerLoadedCallback | None,
     shard_weights: list[float] | None = None,
-    shard_mode: "TensorShardMode | None" = None,
+    shard_mode: TensorShardMode | None = None,
 ) -> nn.Module:
-    shard_mode = shard_mode or "Constant"
+    resolved_shard_mode = shard_mode or TensorShardMode.Constant
     all_to_sharded_linear = partial(
         shard_linear,
         sharding="all-to-sharded",
@@ -552,7 +551,7 @@ def tensor_auto_parallel(
             all_to_sharded_linear_in_place,
             sharded_to_all_linear_in_place,
             shard_weights=shard_weights,
-            shard_mode=shard_mode,
+            shard_mode=resolved_shard_mode,
         )
     elif isinstance(model, (DeepseekV3Model, DeepseekV32Model, KimiK25Model)):
         tensor_parallel_sharding_strategy = DeepSeekShardingStrategy(
@@ -562,7 +561,7 @@ def tensor_auto_parallel(
             all_to_sharded_linear_in_place,
             sharded_to_all_linear_in_place,
             shard_weights=shard_weights,
-            shard_mode=shard_mode,
+            shard_mode=resolved_shard_mode,
         )
     elif isinstance(model, MiniMaxModel):
         tensor_parallel_sharding_strategy = MiniMaxShardingStrategy(
@@ -572,7 +571,7 @@ def tensor_auto_parallel(
             all_to_sharded_linear_in_place,
             sharded_to_all_linear_in_place,
             shard_weights=shard_weights,
-            shard_mode=shard_mode,
+            shard_mode=resolved_shard_mode,
         )
     elif isinstance(model, GLM4MoeLiteModel):
         tensor_parallel_sharding_strategy = GLM4MoeLiteShardingStrategy(
@@ -582,7 +581,7 @@ def tensor_auto_parallel(
             all_to_sharded_linear_in_place,
             sharded_to_all_linear_in_place,
             shard_weights=shard_weights,
-            shard_mode=shard_mode,
+            shard_mode=resolved_shard_mode,
         )
     elif isinstance(model, Glm4MoeModel):
         tensor_parallel_sharding_strategy = Glm4MoeShardingStrategy(
@@ -592,7 +591,7 @@ def tensor_auto_parallel(
             all_to_sharded_linear_in_place,
             sharded_to_all_linear_in_place,
             shard_weights=shard_weights,
-            shard_mode=shard_mode,
+            shard_mode=resolved_shard_mode,
         )
     elif isinstance(
         model, (Qwen3MoeModel, Qwen3NextModel, Qwen3_5TextModel, Qwen3_5MoeModel)
@@ -604,7 +603,7 @@ def tensor_auto_parallel(
             all_to_sharded_linear_in_place,
             sharded_to_all_linear_in_place,
             shard_weights=shard_weights,
-            shard_mode=shard_mode,
+            shard_mode=resolved_shard_mode,
         )
     elif isinstance(model, GptOssModel):
         tensor_parallel_sharding_strategy = GptOssShardingStrategy(
@@ -614,7 +613,7 @@ def tensor_auto_parallel(
             all_to_sharded_linear_in_place,
             sharded_to_all_linear_in_place,
             shard_weights=shard_weights,
-            shard_mode=shard_mode,
+            shard_mode=resolved_shard_mode,
         )
     elif isinstance(model, Step35Model):
         tensor_parallel_sharding_strategy = Step35ShardingStrategy(
@@ -624,7 +623,7 @@ def tensor_auto_parallel(
             all_to_sharded_linear_in_place,
             sharded_to_all_linear_in_place,
             shard_weights=shard_weights,
-            shard_mode=shard_mode,
+            shard_mode=resolved_shard_mode,
         )
     elif isinstance(model, NemotronHModel):
         tensor_parallel_sharding_strategy = NemotronHShardingStrategy(
@@ -634,7 +633,7 @@ def tensor_auto_parallel(
             all_to_sharded_linear_in_place,
             sharded_to_all_linear_in_place,
             shard_weights=shard_weights,
-            shard_mode=shard_mode,
+            shard_mode=resolved_shard_mode,
         )
     else:
         raise ValueError(f"Unsupported model type: {type(model)}")
@@ -654,7 +653,7 @@ class TensorParallelShardingStrategy(ABC):
         all_to_sharded_linear_in_place: Callable[..., None],
         sharded_to_all_linear_in_place: Callable[..., None],
         shard_weights: list[float] | None = None,
-        shard_mode: str = "Constant",
+        shard_mode: TensorShardMode = TensorShardMode.Constant,
     ):
         self._base_all_to_sharded_linear = all_to_sharded_linear
         self._base_sharded_to_all_linear = sharded_to_all_linear
@@ -664,11 +663,13 @@ class TensorParallelShardingStrategy(ABC):
         self.shard_mode = shard_mode
         self.group = group
         self.N = group.size()
-        self._greedy_trackers: dict[str, list[list[float]]] | None = None
-        if shard_weights is not None and shard_mode == "Greedy":
+        self._greedy_trackers: dict[str, list[list[float]] | list[float]] | None = None
+        if shard_weights is not None and shard_mode == TensorShardMode.Greedy:
             self._greedy_trackers = {}
 
-    def _greedy_weights_for(self, key: str, dim: int, unit: int = 1) -> list[float] | None:
+    def _greedy_weights_for(
+        self, key: str, dim: int, unit: int = 1
+    ) -> list[float] | None:
         """Get adjusted weights for a specific projection type, and record the allocation."""
         if self.shard_weights is None or self._greedy_trackers is None:
             return self.shard_weights
@@ -677,8 +678,11 @@ class TensorParallelShardingStrategy(ABC):
         target = [dim * self.shard_weights[i] / total_w for i in range(n)]
         if key not in self._greedy_trackers:
             self._greedy_trackers[key] = [[0.0] * n, [0.0] * n, [0] * n]
-        cum_target, cum_actual, last_sizes = self._greedy_trackers[key]
-        desired = [target[i] + (cum_target[i] - cum_actual[i]) for i in range(n)]
+        tracker = cast(list[list[float]], self._greedy_trackers[key])
+        cum_target, cum_actual, last_sizes = tracker[0], tracker[1], tracker[2]
+        desired: list[float] = [
+            target[i] + (cum_target[i] - cum_actual[i]) for i in range(n)
+        ]
         min_d = min(desired)
         if min_d <= 0:
             desired = [d - min_d + 0.01 for d in desired]
@@ -687,21 +691,21 @@ class TensorParallelShardingStrategy(ABC):
             cum_target[i] += target[i]
             cum_actual[i] += actual_sizes[i]
             last_sizes[i] = actual_sizes[i]
-        self._greedy_trackers[key + "_last_weights"] = desired  # type: ignore
+        self._greedy_trackers[key + "_last_weights"] = desired
         return desired
 
     def _greedy_last_sizes(self, key: str) -> list[int]:
-        """Get the actual sizes from the last _greedy_weights_for call for this key."""
         if self._greedy_trackers is None or key not in self._greedy_trackers:
             return []
-        return self._greedy_trackers[key][2]
+        tracker = self._greedy_trackers[key]
+        assert isinstance(tracker, list) and len(tracker) == 3  # noqa: S101
+        return cast(list[int], tracker[2])
 
     def _greedy_last_weights(self, key: str) -> list[float] | None:
-        """Get the weights used in the last _greedy_weights_for call for this key."""
         if self._greedy_trackers is None:
             return self.shard_weights
-        w = self._greedy_trackers.get(key + "_last_weights")  # type: ignore
-        return w if w is not None else self.shard_weights
+        w = self._greedy_trackers.get(key + "_last_weights")
+        return cast(list[float] | None, w) if w is not None else self.shard_weights
 
     @property
     def all_to_sharded_linear(self) -> Callable[..., nn.Linear]:
@@ -750,19 +754,23 @@ class LlamaShardingStrategy(TensorParallelShardingStrategy):
             k_dim = layer.self_attn.k_proj.weight.shape[0]
             intermediate = layer.mlp.gate_proj.weight.shape[0]
             layer.self_attn.q_proj = self.all_to_sharded_linear(
-                layer.self_attn.q_proj, unit=gqa_unit,
+                layer.self_attn.q_proj,
+                unit=gqa_unit,
                 weights=self._greedy_weights_for("q", q_dim, gqa_unit),
             )
             layer.self_attn.k_proj = self.all_to_sharded_linear(
-                layer.self_attn.k_proj, unit=head_dim,
+                layer.self_attn.k_proj,
+                unit=head_dim,
                 weights=self._greedy_weights_for("k", k_dim, head_dim),
             )
             layer.self_attn.v_proj = self.all_to_sharded_linear(
-                layer.self_attn.v_proj, unit=head_dim,
+                layer.self_attn.v_proj,
+                unit=head_dim,
                 weights=self._greedy_weights_for("v", k_dim, head_dim),
             )
             layer.self_attn.o_proj = self.sharded_to_all_linear(
-                layer.self_attn.o_proj, unit=gqa_unit,
+                layer.self_attn.o_proj,
+                unit=gqa_unit,
                 weights=self._greedy_weights_for("o", q_dim, gqa_unit),
             )
             layer.self_attn.n_heads = layer.self_attn.q_proj.weight.shape[0] // head_dim
@@ -843,21 +851,30 @@ class DeepSeekShardingStrategy(TensorParallelShardingStrategy):
             o_dim = q_dim
             if layer.self_attn.q_lora_rank is None:
                 layer.self_attn.q_proj = self.all_to_sharded_linear(
-                    layer.self_attn.q_proj, unit=q_head_dim,
+                    layer.self_attn.q_proj,
+                    unit=q_head_dim,
                     weights=self._greedy_weights_for("q", q_dim, q_head_dim),
                 )
             else:
                 layer.self_attn.q_b_proj = self.all_to_sharded_linear(
-                    layer.self_attn.q_b_proj, unit=q_head_dim,
+                    layer.self_attn.q_b_proj,
+                    unit=q_head_dim,
                     weights=self._greedy_weights_for("q", q_dim, q_head_dim),
                 )
 
             layer.self_attn.o_proj = self.sharded_to_all_linear(
-                layer.self_attn.o_proj, unit=q_head_dim,
+                layer.self_attn.o_proj,
+                unit=q_head_dim,
                 weights=self._greedy_weights_for("o", o_dim, q_head_dim),
             )
             q_actual = self._greedy_last_sizes("q")
-            head_sizes = [s // q_head_dim for s in q_actual] if q_actual else compute_shard_sizes(original_num_heads, self.N, weights=self.shard_weights)
+            head_sizes = (
+                [s // q_head_dim for s in q_actual]
+                if q_actual
+                else compute_shard_sizes(
+                    original_num_heads, self.N, weights=self.shard_weights
+                )
+            )
             layer.self_attn.num_heads = head_sizes[self.group.rank()]
 
             # Logic from upstream mlx
@@ -890,23 +907,29 @@ class DeepSeekShardingStrategy(TensorParallelShardingStrategy):
             else:
                 if getattr(layer.mlp, "shared_experts", None) is not None:
                     shared_gate_dim = layer.mlp.shared_experts.gate_proj.weight.shape[0]
-                    shared_down_dim = layer.mlp.shared_experts.down_proj.weight.shape[-1]
+                    shared_down_dim = layer.mlp.shared_experts.down_proj.weight.shape[
+                        -1
+                    ]
                     shared_up_dim = layer.mlp.shared_experts.up_proj.weight.shape[0]
                     self.all_to_sharded_linear_in_place(
                         layer.mlp.shared_experts.gate_proj,
-                        weights=self._greedy_weights_for("shared_gate", shared_gate_dim),
+                        weights=self._greedy_weights_for(
+                            "shared_gate", shared_gate_dim
+                        ),
                     )
                     self.sharded_to_all_linear_in_place(
                         layer.mlp.shared_experts.down_proj,
-                        weights=self._greedy_weights_for("shared_down", shared_down_dim),
+                        weights=self._greedy_weights_for(
+                            "shared_down", shared_down_dim
+                        ),
                     )
                     self.all_to_sharded_linear_in_place(
                         layer.mlp.shared_experts.up_proj,
                         weights=self._greedy_weights_for("shared_up", shared_up_dim),
                     )
-                moe_gate_dim = layer.mlp.switch_mlp.gate_proj.weight.shape[1]
-                moe_down_dim = layer.mlp.switch_mlp.down_proj.weight.shape[-1]
-                moe_up_dim = layer.mlp.switch_mlp.up_proj.weight.shape[1]
+                moe_gate_dim = int(layer.mlp.switch_mlp.gate_proj.weight.shape[1])
+                moe_down_dim = int(layer.mlp.switch_mlp.down_proj.weight.shape[-1])
+                moe_up_dim = int(layer.mlp.switch_mlp.up_proj.weight.shape[1])
                 self.all_to_sharded_linear_in_place(
                     layer.mlp.switch_mlp.gate_proj,
                     weights=self._greedy_weights_for("moe_gate", moe_gate_dim),
@@ -962,35 +985,44 @@ class GLM4MoeLiteShardingStrategy(TensorParallelShardingStrategy):
                 timeout_seconds / total,
                 on_timeout,
             )
-            original_num_heads = layer.self_attn.num_heads  # type: ignore
+            original_num_heads = layer.self_attn.num_heads
             q_head_dim = (
                 layer.self_attn.q_b_proj.weight.shape[0] // original_num_heads
-                if layer.self_attn.q_lora_rank is not None
+                if layer.self_attn.q_lora_rank is not None  # pyright: ignore[reportUnnecessaryComparison]
                 else layer.self_attn.q_proj.weight.shape[0] // original_num_heads
-            )  # type: ignore
+            )
             q_dim = (
                 layer.self_attn.q_proj.weight.shape[0]
-                if layer.self_attn.q_lora_rank is None
+                if layer.self_attn.q_lora_rank is None  # pyright: ignore[reportUnnecessaryComparison]
                 else layer.self_attn.q_b_proj.weight.shape[0]
-            )  # type: ignore
+            )
             o_dim = q_dim
-            if layer.self_attn.q_lora_rank is None:  # type: ignore
+            if layer.self_attn.q_lora_rank is None:  # pyright: ignore[reportUnnecessaryComparison]
                 layer.self_attn.q_proj = self.all_to_sharded_linear(
-                    layer.self_attn.q_proj, unit=q_head_dim,
+                    layer.self_attn.q_proj,
+                    unit=q_head_dim,
                     weights=self._greedy_weights_for("q", q_dim, q_head_dim),
                 )
             else:
                 layer.self_attn.q_b_proj = self.all_to_sharded_linear(
-                    layer.self_attn.q_b_proj, unit=q_head_dim,
+                    layer.self_attn.q_b_proj,
+                    unit=q_head_dim,
                     weights=self._greedy_weights_for("q", q_dim, q_head_dim),
                 )
 
             layer.self_attn.o_proj = self.sharded_to_all_linear(
-                layer.self_attn.o_proj, unit=q_head_dim,
+                layer.self_attn.o_proj,
+                unit=q_head_dim,
                 weights=self._greedy_weights_for("o", o_dim, q_head_dim),
             )
             q_actual = self._greedy_last_sizes("q")
-            head_sizes = [s // q_head_dim for s in q_actual] if q_actual else compute_shard_sizes(original_num_heads, self.N, weights=self.shard_weights)
+            head_sizes = (
+                [s // q_head_dim for s in q_actual]
+                if q_actual
+                else compute_shard_sizes(
+                    original_num_heads, self.N, weights=self.shard_weights
+                )
+            )
             layer.self_attn.num_heads = head_sizes[self.group.rank()]
 
             # Logic from upstream mlx
@@ -1021,23 +1053,29 @@ class GLM4MoeLiteShardingStrategy(TensorParallelShardingStrategy):
             else:
                 if getattr(layer.mlp, "shared_experts", None) is not None:
                     shared_gate_dim = layer.mlp.shared_experts.gate_proj.weight.shape[0]
-                    shared_down_dim = layer.mlp.shared_experts.down_proj.weight.shape[-1]
+                    shared_down_dim = layer.mlp.shared_experts.down_proj.weight.shape[
+                        -1
+                    ]
                     shared_up_dim = layer.mlp.shared_experts.up_proj.weight.shape[0]
                     self.all_to_sharded_linear_in_place(
                         layer.mlp.shared_experts.gate_proj,
-                        weights=self._greedy_weights_for("shared_gate", shared_gate_dim),
+                        weights=self._greedy_weights_for(
+                            "shared_gate", shared_gate_dim
+                        ),
                     )
                     self.sharded_to_all_linear_in_place(
                         layer.mlp.shared_experts.down_proj,
-                        weights=self._greedy_weights_for("shared_down", shared_down_dim),
+                        weights=self._greedy_weights_for(
+                            "shared_down", shared_down_dim
+                        ),
                     )
                     self.all_to_sharded_linear_in_place(
                         layer.mlp.shared_experts.up_proj,
                         weights=self._greedy_weights_for("shared_up", shared_up_dim),
                     )
-                moe_gate_dim = layer.mlp.switch_mlp.gate_proj.weight.shape[1]
-                moe_down_dim = layer.mlp.switch_mlp.down_proj.weight.shape[-1]
-                moe_up_dim = layer.mlp.switch_mlp.up_proj.weight.shape[1]
+                moe_gate_dim = int(layer.mlp.switch_mlp.gate_proj.weight.shape[1])
+                moe_down_dim = int(layer.mlp.switch_mlp.down_proj.weight.shape[-1])
+                moe_up_dim = int(layer.mlp.switch_mlp.up_proj.weight.shape[1])
                 self.all_to_sharded_linear_in_place(
                     layer.mlp.switch_mlp.gate_proj,
                     weights=self._greedy_weights_for("moe_gate", moe_gate_dim),
@@ -1158,19 +1196,23 @@ class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
             q_dim = layer.self_attn.q_proj.weight.shape[0]
             k_dim = layer.self_attn.k_proj.weight.shape[0]
             layer.self_attn.q_proj = self.all_to_sharded_linear(
-                layer.self_attn.q_proj, unit=gqa_unit,
+                layer.self_attn.q_proj,
+                unit=gqa_unit,
                 weights=self._greedy_weights_for("q", q_dim, gqa_unit),
             )
             layer.self_attn.k_proj = self.all_to_sharded_linear(
-                layer.self_attn.k_proj, unit=head_dim,
+                layer.self_attn.k_proj,
+                unit=head_dim,
                 weights=self._greedy_weights_for("k", k_dim, head_dim),
             )
             layer.self_attn.v_proj = self.all_to_sharded_linear(
-                layer.self_attn.v_proj, unit=head_dim,
+                layer.self_attn.v_proj,
+                unit=head_dim,
                 weights=self._greedy_weights_for("v", k_dim, head_dim),
             )
             layer.self_attn.o_proj = self.sharded_to_all_linear(
-                layer.self_attn.o_proj, unit=gqa_unit,
+                layer.self_attn.o_proj,
+                unit=gqa_unit,
                 weights=self._greedy_weights_for("o", q_dim, gqa_unit),
             )
 
@@ -1184,9 +1226,13 @@ class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
             layer.self_attn = WrappedMiniMaxAttention(layer.self_attn, self.group)  # pyright: ignore[reportAttributeAccessIssue,reportArgumentType]
 
             # Shard the MoE.
-            moe_gate_dim = layer.block_sparse_moe.switch_mlp.gate_proj.weight.shape[1]
-            moe_down_dim = layer.block_sparse_moe.switch_mlp.down_proj.weight.shape[-1]
-            moe_up_dim = layer.block_sparse_moe.switch_mlp.up_proj.weight.shape[1]
+            moe_gate_dim = int(
+                layer.block_sparse_moe.switch_mlp.gate_proj.weight.shape[1]
+            )
+            moe_down_dim = int(
+                layer.block_sparse_moe.switch_mlp.down_proj.weight.shape[-1]
+            )
+            moe_up_dim = int(layer.block_sparse_moe.switch_mlp.up_proj.weight.shape[1])
             self.all_to_sharded_linear_in_place(
                 layer.block_sparse_moe.switch_mlp.gate_proj,
                 weights=self._greedy_weights_for("moe_gate", moe_gate_dim),
@@ -1232,19 +1278,23 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
                 q_dim = layer.self_attn.q_proj.weight.shape[0]
                 k_dim = layer.self_attn.k_proj.weight.shape[0]
                 layer.self_attn.q_proj = self.all_to_sharded_linear(
-                    layer.self_attn.q_proj, unit=gqa_unit,
+                    layer.self_attn.q_proj,
+                    unit=gqa_unit,
                     weights=self._greedy_weights_for("q", q_dim, gqa_unit),
                 )
                 layer.self_attn.k_proj = self.all_to_sharded_linear(
-                    layer.self_attn.k_proj, unit=head_dim,
+                    layer.self_attn.k_proj,
+                    unit=head_dim,
                     weights=self._greedy_weights_for("k", k_dim, head_dim),
                 )
                 layer.self_attn.v_proj = self.all_to_sharded_linear(
-                    layer.self_attn.v_proj, unit=head_dim,
+                    layer.self_attn.v_proj,
+                    unit=head_dim,
                     weights=self._greedy_weights_for("v", k_dim, head_dim),
                 )
                 layer.self_attn.o_proj = self.sharded_to_all_linear(
-                    layer.self_attn.o_proj, unit=gqa_unit,
+                    layer.self_attn.o_proj,
+                    unit=gqa_unit,
                     weights=self._greedy_weights_for("o", q_dim, gqa_unit),
                 )
                 layer.self_attn.n_heads = (
@@ -1258,8 +1308,9 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
                 if hasattr(layer, "linear_attn"):
                     linear_attn = layer.linear_attn
 
+                    k_greedy: list[float] | None = None
+                    v_greedy: list[float] | None = None
                     if isinstance(linear_attn, Qwen3NextGatedDeltaNet):
-                        # Qwen3-Next: combined projections
                         qkvz_dim = linear_attn.in_proj_qkvz.weight.shape[0]
                         ba_dim = linear_attn.in_proj_ba.weight.shape[0]
                         linear_attn.in_proj_qkvz = self.all_to_sharded_linear(
@@ -1271,19 +1322,18 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
                             weights=self._greedy_weights_for("linear_ba", ba_dim),
                         )
                     else:
-                        # Qwen3.5: separate projections
-                        # in_proj_qkv has sections [q(key_dim), k(key_dim), v(value_dim)]
-                        # that must be split section-aware, not as a contiguous block
                         head_k_dim = linear_attn.head_k_dim
                         head_v_dim = linear_attn.head_v_dim
                         key_dim = linear_attn.key_dim
                         value_dim = linear_attn.value_dim
                         b_dim = linear_attn.in_proj_b.weight.shape[0]
                         a_dim = linear_attn.in_proj_a.weight.shape[0]
-                        # Compute greedy weights ONCE per dimension — all projections
-                        # sharing the same dim must use the same weights within a layer
-                        k_greedy = self._greedy_weights_for("linear_k_dim", key_dim, head_k_dim)
-                        v_greedy = self._greedy_weights_for("linear_v_dim", value_dim, head_v_dim)
+                        k_greedy = self._greedy_weights_for(
+                            "linear_k_dim", key_dim, head_k_dim
+                        )
+                        v_greedy = self._greedy_weights_for(
+                            "linear_v_dim", value_dim, head_v_dim
+                        )
                         linear_attn.in_proj_qkv = shard_linear(
                             linear_attn.in_proj_qkv,
                             "all-to-sharded",
@@ -1293,7 +1343,8 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
                             weights=k_greedy,
                         )
                         linear_attn.in_proj_z = self.all_to_sharded_linear(
-                            linear_attn.in_proj_z, unit=head_v_dim,
+                            linear_attn.in_proj_z,
+                            unit=head_v_dim,
                             weights=v_greedy,
                         )
                         linear_attn.in_proj_b = self.all_to_sharded_linear(
@@ -1306,22 +1357,26 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
                         )
                     is_qwen3next = isinstance(linear_attn, Qwen3NextGatedDeltaNet)
                     out_dim = linear_attn.out_proj.weight.shape[-1]
-                    out_w = v_greedy if not is_qwen3next else self._greedy_weights_for("linear_out", out_dim, linear_attn.head_v_dim)  # pyright: ignore[reportPossiblyUnbound]
+                    out_w = (
+                        v_greedy
+                        if not is_qwen3next
+                        else self._greedy_weights_for(
+                            "linear_out", out_dim, linear_attn.head_v_dim
+                        )
+                    )
                     linear_attn.out_proj = self.sharded_to_all_linear(
-                        linear_attn.out_proj, unit=linear_attn.head_v_dim,
+                        linear_attn.out_proj,
+                        unit=linear_attn.head_v_dim,
                         weights=out_w,
                     )
 
-                    # Shard conv1d: depthwise conv with non-contiguous channel slicing.
-                    # Channel layout is [q(key_dim), k(key_dim), v(value_dim)].
-                    # Each rank takes its head-slice from each of the three sections.
                     rank = self.group.rank()
                     key_dim = linear_attn.key_dim
                     value_dim = linear_attn.value_dim
                     head_k_dim = linear_attn.head_k_dim
                     head_v_dim = linear_attn.head_v_dim
-                    k_w = k_greedy if not is_qwen3next else self.shard_weights  # pyright: ignore[reportPossiblyUnbound]
-                    v_w = v_greedy if not is_qwen3next else self.shard_weights  # pyright: ignore[reportPossiblyUnbound]
+                    k_w = k_greedy if not is_qwen3next else self.shard_weights
+                    v_w = v_greedy if not is_qwen3next else self.shard_weights
                     key_shard_sizes = compute_shard_sizes(
                         key_dim, self.N, unit=head_k_dim, weights=k_w
                     )
@@ -1379,22 +1434,27 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
                     )
                     q_dim = layer.self_attn.q_proj.weight.shape[0]
                     k_dim = layer.self_attn.k_proj.weight.shape[0]
-                    o_dim = layer.self_attn.o_proj.weight.shape[-1]
-                    qo_greedy = self._greedy_weights_for("qwen_qo", q_dim, kv_head_dim * 2 * gqa_repeat)
+                    qo_greedy = self._greedy_weights_for(
+                        "qwen_qo", q_dim, kv_head_dim * 2 * gqa_repeat
+                    )
                     layer.self_attn.q_proj = self.all_to_sharded_linear(
-                        layer.self_attn.q_proj, unit=kv_head_dim * 2 * gqa_repeat,
+                        layer.self_attn.q_proj,
+                        unit=kv_head_dim * 2 * gqa_repeat,
                         weights=qo_greedy,
                     )
                     layer.self_attn.k_proj = self.all_to_sharded_linear(
-                        layer.self_attn.k_proj, unit=kv_head_dim,
+                        layer.self_attn.k_proj,
+                        unit=kv_head_dim,
                         weights=self._greedy_weights_for("k", k_dim, kv_head_dim),
                     )
                     layer.self_attn.v_proj = self.all_to_sharded_linear(
-                        layer.self_attn.v_proj, unit=kv_head_dim,
+                        layer.self_attn.v_proj,
+                        unit=kv_head_dim,
                         weights=self._greedy_weights_for("v", k_dim, kv_head_dim),
                     )
                     layer.self_attn.o_proj = self.sharded_to_all_linear(
-                        layer.self_attn.o_proj, unit=kv_head_dim * gqa_repeat,
+                        layer.self_attn.o_proj,
+                        unit=kv_head_dim * gqa_repeat,
                         weights=qo_greedy,
                     )
                     layer.self_attn.num_attention_heads = (
@@ -1413,9 +1473,9 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
                     Qwen3_5SparseMoeBlock,
                 ),
             ):
-                moe_gate_dim = layer.mlp.switch_mlp.gate_proj.weight.shape[1]
-                moe_down_dim = layer.mlp.switch_mlp.down_proj.weight.shape[-1]
-                moe_up_dim = layer.mlp.switch_mlp.up_proj.weight.shape[1]
+                moe_gate_dim = int(layer.mlp.switch_mlp.gate_proj.weight.shape[1])
+                moe_down_dim = int(layer.mlp.switch_mlp.down_proj.weight.shape[-1])
+                moe_up_dim = int(layer.mlp.switch_mlp.up_proj.weight.shape[1])
                 self.all_to_sharded_linear_in_place(
                     layer.mlp.switch_mlp.gate_proj,
                     weights=self._greedy_weights_for("moe_gate", moe_gate_dim),
@@ -1436,11 +1496,15 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
                     shared_up_dim = layer.mlp.shared_expert.up_proj.weight.shape[0]
                     self.all_to_sharded_linear_in_place(
                         layer.mlp.shared_expert.gate_proj,
-                        weights=self._greedy_weights_for("shared_gate", shared_gate_dim),
+                        weights=self._greedy_weights_for(
+                            "shared_gate", shared_gate_dim
+                        ),
                     )
                     self.sharded_to_all_linear_in_place(
                         layer.mlp.shared_expert.down_proj,
-                        weights=self._greedy_weights_for("shared_down", shared_down_dim),
+                        weights=self._greedy_weights_for(
+                            "shared_down", shared_down_dim
+                        ),
                     )
                     self.all_to_sharded_linear_in_place(
                         layer.mlp.shared_expert.up_proj,
@@ -1491,19 +1555,23 @@ class Glm4MoeShardingStrategy(TensorParallelShardingStrategy):
             q_dim = layer.self_attn.q_proj.weight.shape[0]
             k_dim = layer.self_attn.k_proj.weight.shape[0]
             layer.self_attn.q_proj = self.all_to_sharded_linear(
-                layer.self_attn.q_proj, unit=gqa_unit,
+                layer.self_attn.q_proj,
+                unit=gqa_unit,
                 weights=self._greedy_weights_for("q", q_dim, gqa_unit),
             )
             layer.self_attn.k_proj = self.all_to_sharded_linear(
-                layer.self_attn.k_proj, unit=head_dim,
+                layer.self_attn.k_proj,
+                unit=head_dim,
                 weights=self._greedy_weights_for("k", k_dim, head_dim),
             )
             layer.self_attn.v_proj = self.all_to_sharded_linear(
-                layer.self_attn.v_proj, unit=head_dim,
+                layer.self_attn.v_proj,
+                unit=head_dim,
                 weights=self._greedy_weights_for("v", k_dim, head_dim),
             )
             layer.self_attn.o_proj = self.sharded_to_all_linear(
-                layer.self_attn.o_proj, unit=gqa_unit,
+                layer.self_attn.o_proj,
+                unit=gqa_unit,
                 weights=self._greedy_weights_for("o", q_dim, gqa_unit),
             )
             layer.self_attn.n_heads = layer.self_attn.q_proj.weight.shape[0] // head_dim
@@ -1512,9 +1580,9 @@ class Glm4MoeShardingStrategy(TensorParallelShardingStrategy):
             )
 
             if isinstance(layer.mlp, MoE):
-                moe_gate_dim = layer.mlp.switch_mlp.gate_proj.weight.shape[1]
-                moe_down_dim = layer.mlp.switch_mlp.down_proj.weight.shape[-1]
-                moe_up_dim = layer.mlp.switch_mlp.up_proj.weight.shape[1]
+                moe_gate_dim = int(layer.mlp.switch_mlp.gate_proj.weight.shape[1])
+                moe_down_dim = int(layer.mlp.switch_mlp.down_proj.weight.shape[-1])
+                moe_up_dim = int(layer.mlp.switch_mlp.up_proj.weight.shape[1])
                 self.all_to_sharded_linear_in_place(
                     layer.mlp.switch_mlp.gate_proj,
                     weights=self._greedy_weights_for("moe_gate", moe_gate_dim),
@@ -1529,15 +1597,21 @@ class Glm4MoeShardingStrategy(TensorParallelShardingStrategy):
                 )
                 if getattr(layer.mlp, "shared_experts", None) is not None:
                     shared_gate_dim = layer.mlp.shared_experts.gate_proj.weight.shape[0]
-                    shared_down_dim = layer.mlp.shared_experts.down_proj.weight.shape[-1]
+                    shared_down_dim = layer.mlp.shared_experts.down_proj.weight.shape[
+                        -1
+                    ]
                     shared_up_dim = layer.mlp.shared_experts.up_proj.weight.shape[0]
                     self.all_to_sharded_linear_in_place(
                         layer.mlp.shared_experts.gate_proj,
-                        weights=self._greedy_weights_for("shared_gate", shared_gate_dim),
+                        weights=self._greedy_weights_for(
+                            "shared_gate", shared_gate_dim
+                        ),
                     )
                     self.sharded_to_all_linear_in_place(
                         layer.mlp.shared_experts.down_proj,
-                        weights=self._greedy_weights_for("shared_down", shared_down_dim),
+                        weights=self._greedy_weights_for(
+                            "shared_down", shared_down_dim
+                        ),
                     )
                     self.all_to_sharded_linear_in_place(
                         layer.mlp.shared_experts.up_proj,
@@ -1589,19 +1663,23 @@ class GptOssShardingStrategy(TensorParallelShardingStrategy):
             q_dim = layer.self_attn.q_proj.weight.shape[0]
             k_dim = layer.self_attn.k_proj.weight.shape[0]
             layer.self_attn.q_proj = self.all_to_sharded_linear(
-                layer.self_attn.q_proj, unit=gqa_unit,
+                layer.self_attn.q_proj,
+                unit=gqa_unit,
                 weights=self._greedy_weights_for("q", q_dim, gqa_unit),
             )
             layer.self_attn.k_proj = self.all_to_sharded_linear(
-                layer.self_attn.k_proj, unit=head_dim,
+                layer.self_attn.k_proj,
+                unit=head_dim,
                 weights=self._greedy_weights_for("k", k_dim, head_dim),
             )
             layer.self_attn.v_proj = self.all_to_sharded_linear(
-                layer.self_attn.v_proj, unit=head_dim,
+                layer.self_attn.v_proj,
+                unit=head_dim,
                 weights=self._greedy_weights_for("v", k_dim, head_dim),
             )
             layer.self_attn.o_proj = self.sharded_to_all_linear(
-                layer.self_attn.o_proj, unit=gqa_unit,
+                layer.self_attn.o_proj,
+                unit=gqa_unit,
                 weights=self._greedy_weights_for("o", q_dim, gqa_unit),
             )
 
@@ -1618,14 +1696,23 @@ class GptOssShardingStrategy(TensorParallelShardingStrategy):
 
             rank = self.group.rank()
             q_actual = self._greedy_last_sizes("q")
-            q_head_sizes = [s // head_dim for s in q_actual] if q_actual else compute_shard_sizes(original_num_heads, self.N, unit=gqa_unit // head_dim, weights=self.shard_weights)
+            q_head_sizes = (
+                [s // head_dim for s in q_actual]
+                if q_actual
+                else compute_shard_sizes(
+                    original_num_heads,
+                    self.N,
+                    unit=gqa_unit // head_dim,
+                    weights=self.shard_weights,
+                )
+            )
             sink_start = sum(q_head_sizes[:rank])
             sink_end = sink_start + q_head_sizes[rank]
             layer.self_attn.sinks = layer.self_attn.sinks[sink_start:sink_end]
 
-            moe_gate_dim = layer.mlp.experts.gate_proj.weight.shape[1]
-            moe_down_dim = layer.mlp.experts.down_proj.weight.shape[-1]
-            moe_up_dim = layer.mlp.experts.up_proj.weight.shape[1]
+            moe_gate_dim = int(layer.mlp.experts.gate_proj.weight.shape[1])
+            moe_down_dim = int(layer.mlp.experts.down_proj.weight.shape[-1])
+            moe_up_dim = int(layer.mlp.experts.up_proj.weight.shape[1])
             self.all_to_sharded_linear_in_place(
                 layer.mlp.experts.gate_proj,
                 weights=self._greedy_weights_for("moe_gate", moe_gate_dim),
@@ -1667,19 +1754,23 @@ class Step35ShardingStrategy(TensorParallelShardingStrategy):
             q_dim = layer.self_attn.q_proj.weight.shape[0]
             k_dim = layer.self_attn.k_proj.weight.shape[0]
             layer.self_attn.q_proj = self.all_to_sharded_linear(
-                layer.self_attn.q_proj, unit=gqa_unit,
+                layer.self_attn.q_proj,
+                unit=gqa_unit,
                 weights=self._greedy_weights_for("q", q_dim, gqa_unit),
             )
             layer.self_attn.k_proj = self.all_to_sharded_linear(
-                layer.self_attn.k_proj, unit=head_dim,
+                layer.self_attn.k_proj,
+                unit=head_dim,
                 weights=self._greedy_weights_for("k", k_dim, head_dim),
             )
             layer.self_attn.v_proj = self.all_to_sharded_linear(
-                layer.self_attn.v_proj, unit=head_dim,
+                layer.self_attn.v_proj,
+                unit=head_dim,
                 weights=self._greedy_weights_for("v", k_dim, head_dim),
             )
             layer.self_attn.o_proj = self.sharded_to_all_linear(
-                layer.self_attn.o_proj, unit=gqa_unit,
+                layer.self_attn.o_proj,
+                unit=gqa_unit,
                 weights=self._greedy_weights_for("o", q_dim, gqa_unit),
             )
 
@@ -1694,7 +1785,8 @@ class Step35ShardingStrategy(TensorParallelShardingStrategy):
                 g_dim = layer.self_attn.g_proj.weight.shape[0]
                 g_unit = gqa_unit // head_dim
                 layer.self_attn.g_proj = self.all_to_sharded_linear(
-                    layer.self_attn.g_proj, unit=g_unit,
+                    layer.self_attn.g_proj,
+                    unit=g_unit,
                     weights=self._greedy_weights_for("g", g_dim, g_unit),
                 )
 
@@ -1729,9 +1821,9 @@ class Step35ShardingStrategy(TensorParallelShardingStrategy):
                     layer.mlp.share_expert.down_proj,
                     weights=self._greedy_weights_for("shared_down", shared_down_dim),
                 )
-                moe_gate_dim = layer.mlp.switch_mlp.gate_proj.weight.shape[1]
-                moe_up_dim = layer.mlp.switch_mlp.up_proj.weight.shape[1]
-                moe_down_dim = layer.mlp.switch_mlp.down_proj.weight.shape[-1]
+                moe_gate_dim = int(layer.mlp.switch_mlp.gate_proj.weight.shape[1])
+                moe_up_dim = int(layer.mlp.switch_mlp.up_proj.weight.shape[1])
+                moe_down_dim = int(layer.mlp.switch_mlp.down_proj.weight.shape[-1])
                 self.all_to_sharded_linear_in_place(
                     layer.mlp.switch_mlp.gate_proj,
                     weights=self._greedy_weights_for("moe_gate", moe_gate_dim),
@@ -1775,19 +1867,23 @@ class NemotronHShardingStrategy(TensorParallelShardingStrategy):
                 q_dim = mixer.q_proj.weight.shape[0]
                 k_dim = mixer.k_proj.weight.shape[0]
                 mixer.q_proj = self.all_to_sharded_linear(
-                    mixer.q_proj, unit=gqa_unit,
+                    mixer.q_proj,
+                    unit=gqa_unit,
                     weights=self._greedy_weights_for("q", q_dim, gqa_unit),
                 )
                 mixer.k_proj = self.all_to_sharded_linear(
-                    mixer.k_proj, unit=attn_head_dim,
+                    mixer.k_proj,
+                    unit=attn_head_dim,
                     weights=self._greedy_weights_for("k", k_dim, attn_head_dim),
                 )
                 mixer.v_proj = self.all_to_sharded_linear(
-                    mixer.v_proj, unit=attn_head_dim,
+                    mixer.v_proj,
+                    unit=attn_head_dim,
                     weights=self._greedy_weights_for("v", k_dim, attn_head_dim),
                 )
                 mixer.o_proj = self.sharded_to_all_linear(
-                    mixer.o_proj, unit=gqa_unit,
+                    mixer.o_proj,
+                    unit=gqa_unit,
                     weights=self._greedy_weights_for("o", q_dim, gqa_unit),
                 )
                 mixer.num_heads = mixer.q_proj.weight.shape[0] // attn_head_dim
@@ -1800,8 +1896,8 @@ class NemotronHShardingStrategy(TensorParallelShardingStrategy):
 
             elif isinstance(mixer, NemotronHMoE):
                 # Shard routed experts (SwitchMLP uses fc1/fc2)
-                moe_fc1_dim = mixer.switch_mlp.fc1.weight.shape[1]
-                moe_fc2_dim = mixer.switch_mlp.fc2.weight.shape[-1]
+                moe_fc1_dim = int(mixer.switch_mlp.fc1.weight.shape[1])
+                moe_fc2_dim = int(mixer.switch_mlp.fc2.weight.shape[-1])
                 self.all_to_sharded_linear_in_place(
                     mixer.switch_mlp.fc1,
                     weights=self._greedy_weights_for("moe_gate", moe_fc1_dim),
@@ -1820,7 +1916,9 @@ class NemotronHShardingStrategy(TensorParallelShardingStrategy):
                     )
                     self.sharded_to_all_linear_in_place(
                         mixer.shared_experts.down_proj,
-                        weights=self._greedy_weights_for("shared_down", shared_down_dim),
+                        weights=self._greedy_weights_for(
+                            "shared_down", shared_down_dim
+                        ),
                     )
                 mixer = ShardedMoE(mixer)  # pyright: ignore[reportArgumentType]
                 mixer.sharding_group = self.group
@@ -1844,7 +1942,8 @@ class NemotronHShardingStrategy(TensorParallelShardingStrategy):
         heads_per_group = num_heads // n_groups
         out_unit = heads_per_group * head_dim
         mixer.out_proj = self.sharded_to_all_linear(
-            mixer.out_proj, unit=out_unit,
+            mixer.out_proj,
+            unit=out_unit,
             weights=self._greedy_weights_for("mamba_out", intermediate_size, out_unit),
         )
         out_actual = self._greedy_last_sizes("mamba_out")
@@ -1852,7 +1951,9 @@ class NemotronHShardingStrategy(TensorParallelShardingStrategy):
             head_sizes = [s // head_dim for s in out_actual]
             group_sizes = [s // heads_per_group for s in head_sizes]
         else:
-            group_sizes = compute_shard_sizes(n_groups, world_size, weights=self.shard_weights)
+            group_sizes = compute_shard_sizes(
+                n_groups, world_size, weights=self.shard_weights
+            )
             head_sizes = [g * heads_per_group for g in group_sizes]
 
         groups_per_rank = group_sizes[rank]
diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py
index c87dc772..418150ea 100644
--- a/src/exo/worker/engines/mlx/utils_mlx.py
+++ b/src/exo/worker/engines/mlx/utils_mlx.py
@@ -271,7 +271,10 @@ def shard_and_load(
 
     match shard_metadata:
         case TensorShardMetadata():
-            logger.info(f"loading model from {model_path} with tensor parallelism")
+            logger.info(
+                f"loading model from {model_path} with tensor parallelism "
+                f"(weights={shard_metadata.shard_weights}, mode={shard_metadata.shard_mode})"
+            )
             model = tensor_auto_parallel(
                 model,
                 group,
diff --git a/src/exo/worker/tests/unittests/test_mlx/test_uneven_sharding.py b/src/exo/worker/tests/unittests/test_mlx/test_uneven_sharding.py
index 73858043..944214ae 100644
--- a/src/exo/worker/tests/unittests/test_mlx/test_uneven_sharding.py
+++ b/src/exo/worker/tests/unittests/test_mlx/test_uneven_sharding.py
@@ -12,6 +12,8 @@ import numpy as np
 import pytest
 from mlx.nn.layers.distributed import compute_shard_sizes
 
+from exo.shared.types.worker.shards import TensorShardMode
+
 RANDOM_SEED = 42
 INPUT_TOKENS = [1, 100, 200, 300]
 
@@ -608,7 +610,7 @@ class TestGreedyShardingTP2:
             world_size=2,
             base_port=_next_port_block(),
             shard_weights=[2.0, 1.0],
-            shard_mode="Greedy",
+            shard_mode=TensorShardMode.Greedy,
         )
 
         diff = float(np.max(np.abs(single_logits - tp2_logits)))