reduce_transposes: stop passing transposes through pow and softplus_parametric

[LeSingh1]

Problem

_TransposeOptimization._UNARY_LIKE_OP_TYPES documents itself as ops with a "single non constant input" for which "Transpose(op(x)) == op(Transpose(x))". Two of its entries do not satisfy that, and common::reduce_transposes runs in the default pipeline.

pow is an elementwise binary op (class pow(elementwise_binary) in iOS15/elementwise_binary.py) — its exponent broadcasts against x. Cancelling the transposes around it re-binds the exponent to a different axis:

@mb.program(input_specs=[mb.TensorSpec(shape=(3, 3))])
def prog(x):
    t = mb.transpose(x=x, perm=[1, 0])
    p = mb.pow(x=t, y=np.array([1., 2., 3.], dtype=np.float32))
    return mb.transpose(x=p, perm=[1, 0])

apply_pass_and_basic_check(prog, "common::reduce_transposes")
# ['transpose', 'pow', 'transpose']  ->  ['pow']

The original program computes out[a][b] = x[a][b] ** y[a]; the rewritten one computes out[a][b] = x[a][b] ** y[b]. With x = np.arange(1, 10).reshape(3, 3):

	result
original	[[1, 2, 3], [16, 25, 36], [343, 512, 729]]
after the pass	[[1, 4, 27], [4, 25, 216], [7, 64, 729]]

No warning, no error. With a non-square input the rewritten program does not even build — the same program with x of shape (3, 2) raises ValueError: Incompatible dim 1 in shapes (3, 2) vs. (1, 3) from inside the pass.

softplus_parametric applies alpha[i] / beta[i] along axis 1 of its input (see iOS15/activation.py, which requires alpha.shape == (x.shape[1],)). A transpose that moves axis 1 therefore changes which channel each parameter applies to. The torch frontend emits this op for nn.Softplus with a non-default beta, so it is reachable from a normal conversion. Same repro shape:

@mb.program(input_specs=[mb.TensorSpec(shape=(1, 4, 4, 2))])
def prog(x):
    t = mb.transpose(x=x, perm=[0, 2, 1, 3])
    s = mb.softplus_parametric(x=t, alpha=np.array([1., 2., 3., 4.], np.float32),
                                    beta=np.ones(4, np.float32))
    return mb.transpose(x=s, perm=[0, 2, 1, 3])
# ['transpose', 'softplus_parametric', 'transpose'] -> ['softplus_parametric']

max|original - rewritten| here is 75.0.

Fix

• pow stays in the set but is only treated as unary-like when its exponent broadcasts identically for every ordering of the axes, i.e. all of its dims are 1 and its rank does not exceed x's. This keeps the optimization for the common pow(x, 2.0) case while making the broadcasting case fall through to _visit_materialize_op, which is the conservative/correct handling.
• softplus_parametric is removed from the set; there is no input for which passing an axis-permuting transpose through it is safe.

A possible follow-up (not done here, to keep this change small) would be to register pow as an axis-update op alongside add/mul/sub/… in _TransformAdd, which would recover the optimization for a broadcasting exponent by transposing the exponent as well.

Tests

Three tests in test_reduce_transposes_pass.py:

• test_no_fusion_pow_with_broadcasting_exponent — fails on main (transposes are removed), passes with the fix.
• test_no_fusion_softplus_parametric — fails on main, passes with the fix.
• test_fusion_pow_with_scalar_exponent — guards that the scalar case still fuses (passes both before and after).

I ran the whole of coremltools/converters/mil/mil/passes/tests/test_reduce_transposes_pass.py before and after; the set of failures is identical (my environment cannot load CoreML.framework, so the assert_model_is_valid predictions fail there either way — the graph-structure assertions that precede them all run and pass).