def _help_move_scale(
        block: Block, lut_op: Operation, scale_op: Operation, output_op: Operation
    ):
        """Move the scale from `lut_op -> scale_op -> output_op` to `lut_op -> output_op -> mul`."""
        scale_val: np.ndarray = scale_op.scale.val
        inputs = output_op.inputs
        if output_op.op_type == "linear":
            scale_val = scale_val.T
            inputs["weight"] = lut_op.outputs[0]
            if getattr(output_op, "bias", None) and output_op.bias.val is not None:
                original_bias = output_op.bias.val
                new_bias = (original_bias / np.squeeze(scale_val)).astype(original_bias.dtype)
                inputs["bias"] = new_bias
        elif output_op.op_type == "matmul":
            # Determine if the scaled weight is used by `x` or `y` in matmul.
            if output_op.y == scale_op.outputs[0]:
                if output_op.transpose_y.val is True:
                    scale_val = scale_val.T
                inputs["y"] = lut_op.outputs[0]
            else:
                if output_op.transpose_x.val is True:
                    scale_val = scale_val.T
                inputs["x"] = lut_op.outputs[0]
        else:
            if output_op.op_type != "conv":
                raise AssertionError(
                    "The scale could only be moved for linear/matmul/conv, "
                    f"but got {output_op.op_type}"
                )
            # The weight of conv has C_out at axis=0, but in output the C_out is at axis=1
            scale_val = np.squeeze(scale_val)
            if len(scale_val.shape) > 1:
                # The per-channel-scale should only have one axis with larger than 1 dim size.
                return
            channel_size = 1 if len(scale_val.shape) == 0 else scale_val.shape[0]
            scale_val = scale_val.reshape((1, channel_size, 1, 1))
            inputs["weight"] = lut_op.outputs[0]
            if getattr(output_op, "bias", None) and output_op.bias.val is not None:
                original_bias = output_op.bias.val
                new_bias = (original_bias / np.squeeze(scale_val)).astype(original_bias.dtype)
                inputs["bias"] = new_bias

        # Reconstruct the unscaled output which uses lut output as weight (skip the original scale).
        unscaled_output = getattr(mb, output_op.op_type)(**inputs, before_op=output_op)
        scaled_output = mb.mul(x=unscaled_output, y=scale_val, before_op=output_op)

        # Now the order is lut_op -> unscaled_output -> scaled_output.
        block.replace_uses_of_var_after_op(
            anchor_op=output_op,
            old_var=output_op.outputs[0],
            new_var=scaled_output,
            force_replace=True,  # Need to force replace because it involves replacing constexpr op.
        )
Move the scale from `lut_op -> scale_op -> output_op` to `lut_op -> output_op -> mul`.
coremltools/converters/mil/mil/passes/defs/optimize_quantization.py