plugins/onnx/custom_ops.py

from __future__ import annotations

from typing import Dict, Iterable, Sequence

import onnx
from onnx import TensorProto, helper

DOMAIN = "com.companionforge"
VERSION = 1
PLUGIN_NAMESPACE = "companionforge"


def _attrs(attrs: Dict):
    out = dict(attrs)
    out.setdefault("plugin_namespace", PLUGIN_NAMESPACE)
    out.setdefault("plugin_version", "1")
    return out


def sparse_conv3d(inputs: Sequence[str], outputs: Sequence[str], **attrs):
    return helper.make_node("SparseConv3D", list(inputs), list(outputs), domain=DOMAIN, **_attrs(attrs))


def sparse_window_attention(inputs: Sequence[str], outputs: Sequence[str], **attrs):
    return helper.make_node("SparseWindowAttention", list(inputs), list(outputs), domain=DOMAIN, **_attrs(attrs))


def sparse_downsample(inputs: Sequence[str], outputs: Sequence[str], **attrs):
    return helper.make_node("SparseDownsample", list(inputs), list(outputs), domain=DOMAIN, **_attrs(attrs))


def sparse_upsample(inputs: Sequence[str], outputs: Sequence[str], **attrs):
    return helper.make_node("SparseUpsample", list(inputs), list(outputs), domain=DOMAIN, **_attrs(attrs))


def sparse_subdivide(inputs: Sequence[str], outputs: Sequence[str], **attrs):
    return helper.make_node("SparseSubdivide", list(inputs), list(outputs), domain=DOMAIN, **_attrs(attrs))


def mesh_topology_extract(inputs: Sequence[str], outputs: Sequence[str], **attrs):
    return helper.make_node("MeshTopologyExtract", list(inputs), list(outputs), domain=DOMAIN, **_attrs(attrs))


def model(graph: onnx.GraphProto, opset: int = 23) -> onnx.ModelProto:
    m = helper.make_model(
        graph,
        producer_name="Companion-Forge",
        producer_version="6.5-custom-sparse",
        opset_imports=[helper.make_opsetid("", opset), helper.make_opsetid(DOMAIN, VERSION)],
    )
    m.ir_version = onnx.IR_VERSION
    return m


def vi(name: str, dtype: int, shape: Iterable):
    return helper.make_tensor_value_info(name, dtype, list(shape))


def scalar(name: str, dtype: int = TensorProto.INT32):
    return helper.make_tensor_value_info(name, dtype, [])