File size: 2,098 Bytes
3ae8466 1199c3b 3ae8466 1199c3b 3ae8466 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | 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, [])
|