Zero-sized tensors abort the process on GPU/ANE: a 0-length split section fails MPSCommonRuntimeCanonicalization, a width-0 output fails in MPSNDArray

[john-rocky]

What happens

A graph that contains a zero-sized tensor converts cleanly, then aborts the process when the asset is loaded or run on the GPU or the Neural Engine. A CPU-only specialization loads and runs the same asset. Two shapes of it, with two different assertions:

1. A split section of size 0 — fails a canonicalization pass at load. The empty output is never used.

a, b = x.split([x.shape[1], 0], dim=1)   # b is width-0 and dead
return a + 1.0

MPSGraphExecutable.mm:4419: failed assertion `Error: Optimize Original Module MLIR pass manager failed
Pass failed: MPSCommonRuntimeCanonicalization
Pass failed: mlir::detail::OpToOpPassAdaptor'

2. A zero-sized model output — fails at run.

return x[:, :0]

MPSNDArray.mm:893: failed assertion `[MPSNDArray, initWithBufferImpl:...] Error: buffer is not large enough.
Must be 128 bytes'

Neither is catchable. AIModel.load and the inference call abort rather than returning an error.

	CPU	GPU	Neural Engine
split with a 0-sized section	runs	abort (MPSGraphExecutable.mm:4419)	abort (MPSGraphExecutable.mm:4419)
width-0 output	runs	abort (MPSNDArray.mm:893)	abort (MPSNDArray.mm:893)

Neighbouring constructs are fine on all three: torch.cat with a width-0 operand, new_zeros(n, 0) into a cat, and indexing a width-0 tensor all convert and run.

Environment

M4 Max, macOS 27.0 (26A5416b), coreai-torch 0.4.2, coreai-core 1.0.0b2, torch 2.13.0, Python 3.11.

Reproducer

import asyncio, inspect, shutil, numpy as np, torch, coreai_torch
from pathlib import Path
from coreai_torch import TorchConverter
from coreai.runtime import AIModel, NDArray, SpecializationOptions, ComputeUnitKind


class SplitZero(torch.nn.Module):
    def forward(self, x):
        a, b = x.split([x.shape[1], 0], dim=1)  # b is width-0 and never used
        return a + 1.0


class ZeroOutput(torch.nn.Module):
    def forward(self, x):
        return x[:, :0]


async def _aw(v):
    return await v if inspect.isawaitable(v) else v


def build_and_run(mod, path, unit):
    x = torch.rand(8, 4)
    with torch.no_grad():
        ep = torch.export.export(mod.eval(), (x,)).run_decompositions(coreai_torch.get_decomp_table())
    c = TorchConverter()
    c.add_exported_program(ep, entrypoint_name="main", input_names=["x"], output_names=["o0"])
    p = c.to_coreai()
    p.optimize()
    shutil.rmtree(path, ignore_errors=True)
    p.save_asset(Path(path))

    opts = (
        SpecializationOptions.cpu_only()
        if unit == "cpu"
        else SpecializationOptions.from_preferred_compute_unit_kind(getattr(ComputeUnitKind, unit)())
    )
    loop = asyncio.new_event_loop()
    m = loop.run_until_complete(_aw(AIModel.load(Path(path), opts)))
    f = loop.run_until_complete(_aw(m.load_function("main")))
    out = loop.run_until_complete(_aw(f({"x": NDArray(x.numpy())})))
    print("ok", path, unit)


import sys

build_and_run(SplitZero(), "split_zero.aimodel", sys.argv[1])   # cpu -> ok, gpu / neural_engine -> abort
# build_and_run(ZeroOutput(), "zero_out.aimodel", sys.argv[1])  # cpu -> ok, gpu / neural_engine -> abort

Run it once per compute unit; each abort ends the process.

What I am asking for

An error instead of an abort, as in #67 — a caller can fall back or report, but only if the call returns.

Beyond that, the zero-sized split section looks like it should simply work: the empty result is dead in the graph, so nothing needs to read from a zero-byte buffer.

Where it came from

Adding a format="coreai" export to Ultralytics (ultralytics/ultralytics#25926). Their detection postprocess splits the prediction tensor into [4, num_classes, extras], where extras is 0 for plain detection models — so the width-0 section arrives without anyone writing anything unusual. The same graph with that section omitted runs on the Neural Engine and matches PyTorch to 3.1e-4 px on box coordinates.