Decide remove_broadcast_tiles per axis so dynamic shapes keep working

[kasper0406]

remove_broadcast_tiles fires on static graphs but gives up on nearly every broadcast tile in a dynamically shaped one, leaving full-size materialisations of (b, 1) tensors in the exported model.

Root cause

Under symbolic shapes MIL mints a fresh symbol for a dynamic dimension at almost every op -- fill(shape=concat(...)) in particular -- so the two operands of one elementwise op routinely carry different symbols for the same runtime dimension. Converter output for x * rsqrt(mean(x*x, -1, keepdims=True) + 1e-6) over a (b, 8) input, with the symbols MIL assigned:

%expand_dims_0: (dim_0, 1, fp32)(Tensor) = expand_dims(x=%reduce_sum_0, axes=[1])
%fill_0:        (is1, 1, fp32)(Tensor)   = fill(shape=%concat_1, value=8.0)
%real_div_0:    (is2, 1, fp32)(Tensor)   = real_div(x=%expand_dims_0, y=%fill_0)
%add_0:         (is4, 1, fp32)(Tensor)   = add(x=%real_div_0, y=%fill_1)
%rsqrt_0:       (is4, 1, fp32)(Tensor)   = rsqrt(x=%add_0)
%tile_0:        (is4, 8, fp32)(Tensor)   = tile(x=%rsqrt_0, reps=[1, 8])
%mul_1:         (is6, 8, fp32)(Tensor)   = mul(x=%arg0, y=%tile_0)

dim_0, is1, is2, is4, is6 are all the same dimension at runtime. _consumer_output_is_unchanged decided the tile by re-broadcasting the operand shapes:

broadcast = broadcast_shapes(*operand_shapes)   # (dim_0, 8) against (is4, 1)
if broadcast is None:
    return False

and _broadcast_dims will only equate a symbolic dimension with the identical symbol, so broadcast comes back None and the tile stays. Instrumented trace of the pass on that program:

TILE tile_0 x=(is4, 1) reps=[1 8] bcast=True consumers=[('mul', (is6, 8))] -> remove=False
   consumer mul operand shapes -> [(dim_0, 8), (is4, 1)] broadcast=None out=(is6, 8)

Final models on main, running build_pass_pipeline() over a (b, 8) input: mean(x, -1, keepdims=True) * x keeps 1 tile, symbolic RMSNorm 1, symbolic LayerNorm 2, symbolic softmax 1. Each one is a full (b, 8) materialisation of a (b, 1) tensor.

Fix

The re-broadcast was answering a harder question than the pass needs. Bypassing a tile only changes an operand on the axes the tile actually replicated, so the decision splits per axis:

• reps[axis] == 1 -- the tile passes the axis through, so the operand's dimension there is literally the tile input's, whatever either is called. No comparison needed, and this is exactly where all the renamed symbols live.
• reps[axis] > 1 -- the tile replicated a size-1 axis, and bypassing it takes the operand back down to 1 there. The consumer's output only stays the same if the other operand already carries the full size on that axis. is_broadcast_tile guarantees a replicated dimension is 1 * reps[axis], i.e. a literal int, so this comparison never involves a symbol on the tile's side either.

That is per-operand reasoning, so it never has to name the consumer's output shape. The soundness argument is unchanged from before: on untouched axes the operand is identical, and on replicated axes broadcast(other, reps[axis]) and broadcast(other, 1) both come out as other == reps[axis].

After the fix the symbolic-shape finals carry no tiles at all (mean * x, RMSNorm, LayerNorm, softmax: 1, 1, 2, 1 tiles -> 0).

Static graphs are unaffected. Re-running the flax/equinox probes (LayerNorm, RMSNorm, GroupNorm, Linear, MultiHeadAttention, jnp.broadcast_to, implicit rank-broadcasting, jnp.where masks) gives identical final op counts before and after -- including flax_groupnorm's two surviving tiles, which feed a reshape rather than an elementwise op and are correctly still kept.

Known gap left as an xfail

_BROADCAST_OPS lists only elementwise ops, so a tile feeding a matmul is always kept. MIL's matmul does broadcast its batch dimensions natively -- matmul((1, 4, 8), (2, 8, 4)) type-infers to (2, 4, 4) and predicts bit-identically to numpy on the runtime -- so jnp.broadcast_to(x, (B, ...)) @ y materialises the full batch for nothing. Supporting it is not a matter of adding one name to the set: unlike the elementwise ops, only a matmul's leading axes broadcast while the trailing two are contracted, so it needs an axis rule of its own. None of the flax or equinox layers probed produce the pattern, so it is a strict=True xfail here rather than a fix.

Tests

Unit (hand-built MIL):

• test_removed_when_the_operands_carry_different_symbols -- a mul whose operands are (batch, 8) and tile((renamed_batch, 1), reps=[1, 8]), the shape the converter actually produces. Fails on main.
• test_removed_when_the_other_operand_has_a_lower_rank.
• Negatives: test_not_removed_when_the_other_operand_lacks_the_replicated_size (both operands size 1 on the replicated axis, so the output would shrink) and test_not_removed_when_the_tile_is_both_operands.
• test_batch_broadcast_ahead_of_matmul_is_removed -- strict xfail, documenting the gap above.

End-to-end:

• test_symbolic_broadcast_leaves_no_tile -- jnp.mean(x, -1, keepdims=True) * x exported with jax.export.symbolic_shape("(b, 8)"), numerics checked against JAX at two concrete batch sizes, asserts no tile survives. Fails on main.

Verification

Full suite on this branch: 430 passed, 1 skipped, 1 xfailed (python -m pytest tests/). ruff check . clean.

Also confirmed while auditing: the pass does sit before the first common::const_elimination in the assembled pipeline (indices 10 vs 14), so the module docstring's claim about tiled scalar constants holds; and re-running the pass at the end of the pipeline on the flax/equinox models changes nothing, so its position is not starving it either.

🤖 Generated with Claude Code

https://claude.ai/code/session_01Q4T3UHepKPR65o5aiXEw5E

[kasper0406]

This lines up with something I hit downstream, and I think there is a complementary half worth folding in while this file is open.

_BROADCAST_OPS deliberately excludes select, with the comment that E5RT cannot propagate shapes through a select with implicit broadcasting — so tiles feeding a select are preserved. That is correct and this PR does not regress it. But it only preserves tiles that exist, and for the most common dynamic-shape select there is no tile to preserve.

Concretely, a whole-tensor cache write jnp.where(mask, value, cache) where cache is (1, L, nkv, hd) with symbolic L and value is one row (1, 1, nkv, hd). The converter emits:

cond <- tile     [1, SYM, 1, 512]     <- tiled, preserved by this pass
a    <- concat   [1,   1, 1, 512]     <- never tiled; JAX broadcasts implicitly
b    <- input    [1, SYM, 1, 512]

E5RT then fails to load the model:

Failed to PropagateInputTensorShapes: Validation error during type
inference for select: at unknown location: Incompatible Dimension.

which is the same message the _BROADCAST_OPS comment cites. So the exclusion catches one direction and misses the other.

Two things I verified while chasing it, which may save you time:

• An explicit jnp.broadcast_to in the traced source does not help. It is folded away before reaching MIL — I compared the emitted graphs with and without it and they are byte-identical. So this cannot be fixed on the JAX side; it has to be done in MIL.
• A tile is awkward here because its reps would themselves have to be symbolic. fill_like + add works instead: fill_like takes its shape from a reference operand that already has the output shape, so no symbolic arithmetic is needed, and adding zero is exact in fp16.

I wrote that as a small pass (broadcast_select_operands) that widens any under-shaped select operand when the output has a symbolic dim, running right after remove_broadcast_tiles. It turns a --no-materialize export of a 35-layer LLM from "does not load at all" into one that loads and predicts — 6 selects widened, all of them global KV-cache writes. Implementation and rationale: kasper0406/gemma-coreml-chat@4e53341

It feels like it belongs next to this pass rather than downstream, since the two are the same problem from opposite ends: this PR stops removing tiles a select needs, that one adds the ones JAX never emitted. Happy to port it into a PR here if you want it.

One caveat on my end: I did not merge it downstream, because symbolic shapes turned out ~5.8x slower than our bucketed export (symbolic global caches cannot be Core ML states, so they revert to I/O). So the value is "RangeDim exports load again", not performance.