scripts/coreml_patch.py

"""
Compatibility shim: coremltools 9.0 Torch frontend x numpy 2.x.

coremltools' `_cast` op handler folds a constant int/bool cast with
`mb.const(val=int(x.val))`. Under numpy >= 2.0, calling `int()`/`bool()` on a
length-1 (non 0-d) ndarray raises:
    TypeError: only 0-dimensional arrays can be converted to Python scalars
YOLO26's attention block emits exactly such a cast, so conversion aborts at
`.../attn/...`. We re-register a `_cast` that coerces size-1 arrays via `.item()`
first. Behaviour is otherwise identical.

Import this module before calling `coremltools.convert(...)`.
"""
import numpy as np
from coremltools.converters.mil.frontend.torch import ops as _tops
from coremltools.converters.mil.frontend.torch.ops import _get_inputs
from coremltools.converters.mil.mil import Builder as mb


def _cast_numpy2_safe(context, node, dtype, dtype_name):
    inputs = _get_inputs(context, node, expected=1)
    x = inputs[0]
    if not (len(x.shape) == 0 or np.all([d == 1 for d in x.shape])):
        raise ValueError("input to cast must be either a scalar or a length 1 tensor")

    if x.can_be_folded_to_const():
        val = x.val
        # numpy 2.x: int()/float()/bool() on a size-1, >0-d array raises. Coerce.
        if hasattr(val, "item") and np.size(val) == 1:
            val = val.item()
        if not isinstance(x.val, dtype):
            res = mb.const(val=dtype(val), name=node.name)
        else:
            res = x
    elif len(x.shape) > 0:
        x = mb.squeeze(x=x, name=node.name + "_item")
        res = mb.cast(x=x, dtype=dtype_name, name=node.name)
    else:
        res = mb.cast(x=x, dtype=dtype_name, name=node.name)
    context.add(res, node.name)


_applied = False


def apply():
    global _applied
    if not _applied:
        _tops._cast = _cast_numpy2_safe
        _applied = True
    return _applied


# Apply on import.
apply()