fuse_onehot_matmul_to_gather: normalize one_hot's axis against its output rank

[LeSingh1]

Problem

one_hot's axis is relative to its output rank, which is rank(indices) + 1 — one_hot.type_inference accepts axis in [-rank-1, rank]. fuse_onehot_matmul_to_gather normalizes a non-negative axis by subtracting rank(indices), which is one too few:

rank = len(onehot_op.indices.shape)
if axis >= 0:
    axis -= rank
if axis != -1:
    return False

So axis == rank(indices) - 1 — a dimension that is not last — normalizes to -1 and the pattern matches. matmul does not contract over the one-hot dimension in that case, so rewriting to gather(W, indices, axis=0) is not valid:

@mb.program(input_specs=[mb.TensorSpec(shape=(4,), dtype=types.int32)])
def prog(indices):
    oh = mb.one_hot(indices=indices, one_hot_vector_size=5, axis=0,
                    on_value=1.0, off_value=0.0)                       # (5, 4)
    return mb.matmul(x=oh, y=np.arange(12).reshape(4, 3).astype(np.float32))

apply_pass_and_basic_check(prog, "common::fuse_onehot_matmul_to_gather")
# ['one_hot', 'matmul'] -> ['gather']
# output shape (5, 3) -> (4, 3)

The original computes out[d][m] = sum over the i with indices[i] == d of W[i][m]; the rewritten program computes out[i][m] = W[indices[i]][m] — a different function, and a different output shape. It also reproduces with rank-2 indices and axis=1 ((2,4,5) becomes (2,3,5)).

The same off-by-one makes the pass miss the case it is meant to handle whenever the last axis is written non-negatively (axis == rank(indices)), which is a legal spelling of axis=-1.

common::fuse_onehot_matmul_to_gather is in the default pipeline (pass_pipeline.py).

Fix

Normalize against rank(indices) + 1, the rank axis actually refers to.

Tests

In TestFuseOnehotMatmulToGather:

• test_no_fusion_when_onehot_axis_is_not_last[rank=1,2,3] — the one-hot dimension is second to last; the pass must leave the program alone. Fails on main.
• test_fuse_onehot_matmul_to_gather_non_negative_axis[rank=1,2,3] — the last axis written as axis == rank(indices); the pass must fuse. Fails on main (the fusion is skipped).

The existing test_fuse_onehot_matmul_to_gather only uses axis=-1 and is untouched. I ran the whole class before and after the change; the pre-existing failures are identical (this environment cannot load CoreML.framework, so the assert_model_is_valid predictions fail there either way — the graph-structure assertions preceding them all run and pass).