Fix palettize_weights with enable_per_channel_scale=True crashing on ANE (macOS 26)

When OpPalettizerConfig is configured with enable_per_channel_scale=True,
palettize_weights wraps the constexpr_lut_to_dense output in a
constexpr_blockwise_shift_scale op (data=<dense fp16 weight>, scale=<per-channel
fp16>). On macOS 26, the MPSGraph backend lowering for that constexpr op fails
verification when targeting the Apple Neural Engine:

    'mps.dequantize' op operand #2 must be tensor of quantized values,
    but got 'tensor<1xf16>'
    ... failed assertion `original module failed verification'

The MPSGraph lowering of constexpr_blockwise_shift_scale assumes the data
operand is a quantized integer tensor (it lowers to mps.dequantize); with
enable_per_channel_scale=True, the data is the dense fp16 weight, which fails
that assumption. CPU and GPU compute units accept the wrapper and predict
correctly; only the ANE-targeted MIL -> MPSGraph dispatch is broken.

Fix: bake per_channel_scale into the LUT entries at compile time and re-emit
constexpr_lut_to_dense, instead of leaving the scale as a runtime constexpr.
Both data and scale are fp16 and the wrapper's only effect is data * scale, so
the fold is mathematically identical. The failing MPSGraph dispatch is
eliminated entirely, and CPU / GPU numerics stay bit-identical with the prior
behavior. Resulting graph also has one fewer runtime constexpr per palettized
const.

Test updated: TestPalettizeWeights::test_palettization_pcs previously asserted
that the constexpr_blockwise_shift_scale wrapper was emitted; it now asserts
the wrapper is absent (the LUT is pre-scaled). Numerical equivalence vs the
unpalettized model is verified by the existing verify_model_outputs call on
macOS 15+.

Tested:
  - test_palettization_pcs:                                    PASS
  - All 155 TestPalettizeWeights / TestJointCompressWeights:   PASS
  - Manual: Qwen3-VL 2B stateful chunk on macOS 26 + M4 ANE:
    MPSGraph verification crash gone (was reproducible at every load).

Author: john-rocky

coremltools/optimize/coreml/_quantization_passes.py

@@ -1139,12 +1139,39 @@ def transform_op(self, op: Operation):
                         "Palettization with per-channel-scale is only supported since "
                         "iOS18. Please set minimum_deployment_target accordingly."
                     )
-                new_var = mb.constexpr_blockwise_shift_scale(
-                    data=new_var,
-                    scale=per_channel_scale,
-                    offset=None,
-                    before_op=op,
+                # Bake per_channel_scale into the LUT entries at compile time
+                # and re-emit constexpr_lut_to_dense, rather than wrapping the
+                # dense fp16 weight in a runtime constexpr_blockwise_shift_scale.
+                # The wrapper makes the MPSGraph backend on macOS 26 fail
+                # MPSGraph verification at model load time on the Apple Neural
+                # Engine ("'mps.dequantize' op operand #2 must be tensor of
+                # quantized values, but got 'tensor<1xf16>'"), because the MPS
+                # lowering of that constexpr expects the data operand to be a
+                # quantized integer tensor; with per_channel_scale the data is
+                # the dense fp16 weight produced by constexpr_lut_to_dense.
+                # Folding scale into the LUT entries is mathematically identical
+                # (both data and scale are fp16, output is data * scale), so
+                # numerics are preserved while the failing runtime op is
+                # eliminated. CPU and GPU compute units are unaffected by the
+                # previous wrapper, so this also keeps their numerics
+                # bit-identical with the prior behavior.
+                lut = lut_params.lut.copy()
+                # per_channel_scale rank matches the original weight rank; LUT
+                # has additional trailing dims for [group, num_palette,
+                # vector_size]. Broadcast scale across those trailing dims.
+                pcs_bcast = per_channel_scale.reshape(
+                    per_channel_scale.shape
+                    + (1,) * (lut.ndim - per_channel_scale.ndim)
+                )
+                lut = (
+                    lut.astype(np.float32) * pcs_bcast.astype(np.float32)
+                ).astype(lut.dtype)
+                new_var = frontend_utils._construct_constexpr_lut_op(
+                    lut_params.indices,
+                    lut,
+                    lut_params.vector_axis,
                     name=op.name + "_palettized_pcs",
+                    before_op=op,
                 )
         else:
             decompressed_val = self.decompress(lut_params)

coremltools/test/optimize/coreml/test_post_training_quantization.py

@@ -1683,13 +1683,18 @@ def test_palettization_pcs(self, compute_unit, backend):
             op_type="constexpr_lut_to_dense"
         )[0]
         assert types.builtin_to_string(palettize_op.indices.dtype) == "uint4"
-        # The per-channel-scale is represented by a quant op to do scaling.
+        # The per-channel-scale is folded into the LUT entries at compile time
+        # (rather than emitted as a runtime constexpr_blockwise_shift_scale
+        # wrapper); see the comment in
+        # coremltools/optimize/coreml/_quantization_passes.py for the rationale
+        # (the wrapper was rejected by the MPSGraph lowering on Apple Neural
+        # Engine on macOS 26 because mps.dequantize expects a quantized integer
+        # data input but received the dense fp16 weight). Folding into the LUT
+        # is mathematically identical and produces a smaller graph.
         quantize_ops = mlmodel_palettized._mil_program.functions["main"].find_ops(
             op_type="constexpr_blockwise_shift_scale"
         )
-        assert len(quantize_ops) > 0
-        # Order of quant and lut op is determined by canonicalize_quantized_lut_pattern graph pass.
-        assert quantize_ops[0].outputs[0].child_ops[0].op_type == "constexpr_lut_to_dense"
+        assert len(quantize_ops) == 0
 
         if _macos_version() >= (15, 0):
             verify_model_outputs(mlmodel, mlmodel_palettized, coreml_input_values)