feat: Enable WebGPU Shape op support, remove EliminateShape rewrite rules (#132)

WebGPU now supports the ONNX `Shape` operator natively. This PR removes
the workaround that existed for the old constraint.

## What was removed

| Item | Reason |
|------|--------|
| `_eliminate_shape.py` | WebGPU-only workaround — replaces
`Shape(attention_mask)` with `ReduceSum+ReduceMax`. No longer needed. |
| `_eliminate_shape_test.py` | Tests for the removed module |
| `supports_shape` field from `EpCapabilities` | Only WebGPU used it
with `False`; removing the field cleans up the API |
| `EliminateShape` lowering pass in `_optimizations.py` | Nothing calls
it anymore |
| `eliminate_shape_rules` from `rewrite_rules` public API | Removed with
the implementation |

## What was enabled

The `webgpu` EP entry in `_execution_providers.py` no longer sets
`supports_shape=False` (field removed). WebGPU models now retain their
`Shape` ops, matching all other EPs.

## Test update

`test_webgpu_no_shape_nodes` → `test_webgpu_supports_shape_nodes`:
assertion flipped from `== 0` to `> 0`, confirming Shape nodes are
preserved in WebGPU graphs.

## Stats
- 416 lines deleted, 10 lines added
- 2317 tests pass

---------

Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: justinchuby

CHANGELOG.md

@@ -7,6 +7,24 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### WebGPU Shape Op Support
+
+#### Changed
+
+- WebGPU now supports the ONNX `Shape` operator natively. The `EliminateShape`
+  rewrite pass (which replaced `Shape(attention_mask)` with `ReduceSum` +
+  `ReduceMax`) has been removed.
+
+#### Removed
+
+- **Breaking**: `EpCapabilities.supports_shape` field removed. Custom EPs that
+  passed `supports_shape=True` or `supports_shape=False` to `EpCapabilities(...)`
+  will get a `TypeError`. Remove the argument — `Shape` is now universally
+  supported across all EPs.
+- `mobius.rewrite_rules.eliminate_shape_rules` removed from the public API.
+
+---
+
 ### Mistral-3 / Pixtral VLM Support
 
 #### Added

docs/ep_quickstart.md

@@ -120,7 +120,6 @@ register_ep(EpCapabilities(
     gqa_dtypes=frozenset({ir.DataType.FLOAT16}),
     qkv_pack_dtypes=frozenset({ir.DataType.FLOAT16}),
     supports_fused_rope=True,
-    supports_shape=True,
     provider_options={"some_option": "value"},
 ))
 

docs/execution_providers.md

@@ -99,7 +99,6 @@ Some EPs cannot execute certain ops. These are handled by **lowering rules**
 |---|---|---|
 | No fused RoPE inside GQA | DML | `SeparateRoPE` rewrite: GQA `do_rotary=1` → explicit `RotaryEmbedding` + GQA `do_rotary=0` |
 | No packed QKV in GQA | DML | `UnpackQKV` rewrite: packed GQA → 3 separate `MatMul` projections |
-| No `Shape` operator | WebGPU | `EliminateShape` rewrite: `Shape(attention_mask)` → `ReduceSum` + `ReduceMax` |
 | No `SkipLayerNorm` kernel | TRT-RTX, onnx-standard | `InlinePass`: expands fused ops using their registered `ir.Function` bodies |
 | No `FusedMatMul` kernel | onnx-standard | `InlinePass`: expands to `Transpose + MatMul` |
 | No `PackedMultiHeadAttention` kernel | onnx-standard | `InlinePass`: expands to block-diagonal attention bias + standard `Attention` |
@@ -119,7 +118,6 @@ class EpCapabilities:
     gqa_dtypes: frozenset[ir.DataType]               # dtypes where GQA fusion fires
     qkv_pack_dtypes: frozenset[ir.DataType]          # dtypes where PackQKV fusion fires
     supports_fused_rope: bool = True                 # False → SeparateRoPE + UnpackQKV
-    supports_shape: bool = True                      # False → EliminateShape lowering
     supports_skip_layer_norm: bool = True            # False → InlinePass expansion
     supports_fused_matmul: bool = True               # False → Transpose + MatMul via InlinePass
     supports_fused_moe: bool = True                  # False → decompose fused MoE ops
@@ -141,7 +139,7 @@ EpCapabilities(name="cuda",    gqa_dtypes={FLOAT16, BFLOAT16},
 EpCapabilities(name="dml",     gqa_dtypes={FLOAT16},
                supports_fused_rope=False)
 EpCapabilities(name="webgpu",  gqa_dtypes={FLOAT, FLOAT16},
-               supports_shape=False, default_int4_accuracy_level=4,
+               default_int4_accuracy_level=4,
                provider_options={"enableGraphCapture": "0", ...})
 EpCapabilities(name="trt-rtx", gqa_dtypes={FLOAT16, BFLOAT16},
                supports_skip_layer_norm=False, enable_graph_capture=True,
@@ -187,7 +185,7 @@ Stage 2b: InlinePass   EP-gated. Expands custom ops the EP cannot execute
             decomposition for TRT-RTX).
 
 Stage 3:  Lowering     EP-gated. Structural rewrites for EP constraints.
-          ↓ SeparateRoPE, UnpackQKV, EliminateShape
+          ↓ SeparateRoPE, UnpackQKV
             (each only fires if the EP's capabilities require it)
 
 Stage 4:  Fold         EP-agnostic. Always applied.

src/mobius/__main__.py

@@ -278,8 +278,6 @@ def _cmd_list(args: argparse.Namespace) -> None:
             extras = []
             if not caps.supports_fused_rope:
                 extras.append("no-fused-rope")
-            if not caps.supports_shape:
-                extras.append("no-shape")
             if not caps.supports_skip_layer_norm:
                 extras.append("no-skip-layer-norm")
             flags = f"  [{', '.join(extras)}]" if extras else ""

src/mobius/_execution_providers.py

@@ -60,7 +60,6 @@ class EpCapabilities:
             stage).
         supports_fused_rope: ``False`` triggers SeparateRoPE + UnpackQKV
             lowering (DML).
-        supports_shape: ``False`` triggers EliminateShape lowering (WebGPU).
         supports_skip_layer_norm: ``False`` expands SkipLayerNormalization /
             SkipSimplifiedLayerNormalization via InlinePass (TRT-RTX).
         supports_fused_moe: ``False`` decomposes fused MoE ops.
@@ -78,7 +77,6 @@ class EpCapabilities:
     gqa_dtypes: frozenset[ir.DataType] = dataclasses.field(default_factory=frozenset)
     qkv_pack_dtypes: frozenset[ir.DataType] = dataclasses.field(default_factory=frozenset)
     supports_fused_rope: bool = True
-    supports_shape: bool = True
     supports_skip_layer_norm: bool = True
     supports_fused_moe: bool = True
     supports_packed_multi_head_attention: bool = False
@@ -221,7 +219,6 @@ def _register_builtins() -> None:
             name="webgpu",
             gqa_dtypes=frozenset({ir.DataType.FLOAT, ir.DataType.FLOAT16}),
             qkv_pack_dtypes=frozenset({ir.DataType.FLOAT, ir.DataType.FLOAT16}),
-            supports_shape=False,
             default_int4_accuracy_level=4,
             provider_options={"enableGraphCapture": "0", "validationMode": "basic"},
         ),
@@ -245,7 +242,6 @@ def _register_builtins() -> None:
             gqa_dtypes=frozenset(),  # no GroupQueryAttention
             qkv_pack_dtypes=frozenset(),  # no PackQKV
             supports_fused_rope=False,  # no fused RoPE inside GQA (GQA not supported)
-            supports_shape=True,  # Shape is a standard ONNX op — no elimination needed
             supports_skip_layer_norm=False,  # inline SkipLayerNorm
             supports_packed_multi_head_attention=False,  # inline PackedMHA
         ),

src/mobius/_optimizations.py

@@ -11,7 +11,7 @@
 2. **Fusion** — promote standard ops to EP-supported fused ops
    (GQA, SkipNorm, GeluFusion). Gated by ``(ep, dtype)`` and ``model_role``.
 3. **Lowering** — decompose ops the EP cannot execute
-   (SeparateRoPE, EliminateShape).
+   (SeparateRoPE).
 4. **Fold** — final dead-node removal and constant folding.
 
 All EP knowledge is encoded in :class:`~mobius._execution_providers.EpCapabilities`
@@ -57,7 +57,6 @@
 from mobius._passes import FoldConcatInitializersPass, FoldTransposedInitializerPass
 from mobius.functions import register_function_bodies
 from mobius.rewrite_rules import (
-    eliminate_shape_rules,
     gelu_fusion_rules,
     group_query_attention_rules,
     pack_qkv_for_gqa_rules,
@@ -298,9 +297,6 @@ def _get_optimization_passes(
         lower.append(("SeparateRoPE", list(separate_rope_rules())))
         lower.append(("UnpackQKV", list(unpack_qkv_rules())))
 
-    if not caps.supports_shape:
-        lower.append(("EliminateShape", list(eliminate_shape_rules())))
-
     return fuse, lower
 
 
@@ -325,7 +321,7 @@ def optimize_model(
     2. **Fusion** — promote standard ops to EP-supported fused ops
        (e.g. GQA, SkipNorm, GeluFusion). Gated by ``(ep, dtype)`` and role.
     3. **Lowering** — decompose ops the EP cannot execute
-       (e.g. SeparateRoPE for DML, EliminateShape for WebGPU).
+       (e.g. SeparateRoPE for DML).
     4. **Fold** — final dead-node removal and constant folding.
 
     After fusion, if GQA was expected for ``(ep, dtype)`` but zero

src/mobius/components/_common.py

@@ -191,11 +191,9 @@ def create_attention_bias(
     # Get query_length and total_length from shapes.
     # query_length comes from input_ids dim 1 (the query; e.g. 1 during decode).
     # total_length comes from attention_mask dim 1 (past + current tokens).
-    # Using attention_mask for total_length lets the EliminateShape WebGPU rule
-    # eliminate that Shape op.  Using input_ids for query_length is semantically
-    # correct: during decode input_ids is (batch, 1), so query_length=1 and
-    # start = total_length - 1, giving the last row of q_indices.
-    # On WebGPU (concrete dims), Shape(input_ids, 1) is constant-folded away.
+    # Using input_ids for query_length is semantically correct: during decode
+    # input_ids is (batch, 1), so query_length=1 and start = total_length - 1,
+    # giving the last row of q_indices.
     query_length = op.Shape(input_ids, start=1, end=2)  # 1-D [1]
     total_length = op.Shape(attention_mask, start=1, end=2)  # 1-D [1]
     start = op.Sub(total_length, query_length)

src/mobius/rewrite_rules/__init__.py

@@ -30,7 +30,6 @@
 
 __all__ = [
     "bias_gelu_rules",
-    "eliminate_shape_rules",
     "gelu_fusion_rules",
     "group_query_attention_rules",
     "layer_norm_fusion_rules",
@@ -43,7 +42,6 @@
 ]
 
 from mobius.rewrite_rules._bias_gelu import bias_gelu_rules
-from mobius.rewrite_rules._eliminate_shape import eliminate_shape_rules
 from mobius.rewrite_rules._gelu_fusion import gelu_fusion_rules
 from mobius.rewrite_rules._group_query_attention import (
     group_query_attention_rules,