tinygrad 0.10.2__py3-none-any.whl → 0.11.0__py3-none-any.whl

tinygrad/frontend/onnx.py

@@ -0,0 +1,1253 @@
+# pylint: disable=possibly-unused-variable
+from typing import Any, Sequence, cast, Literal, NamedTuple, Generator
+import dataclasses, functools, io, math, types, warnings, pathlib, sys, os, struct, enum
+from io import BufferedReader
+from tinygrad.nn.state import TensorIO
+from tinygrad.tensor import Tensor, _broadcast_shape, ReductionStr
+from tinygrad.helpers import getenv, DEBUG, all_same, prod, flatten, make_tuple, argsort, is_numpy_ndarray, get_single_element, polyN
+from tinygrad.dtype import DType, ConstType, dtypes, _from_np_dtype, truncate
+from tinygrad.device import is_dtype_supported, Device
+
+# ***** protobuf definitions ******
+class WireType(enum.IntEnum):
+  """
+  Protocol Buffer wire types for decoding fields.
+  Reference: https://github.com/protocolbuffers/protobuf/blob/main/python/google/protobuf/internal/wire_format.py#L24-L29
+  """
+  VARINT = 0; FIXED64 = 1; LENGTH_DELIMITED = 2; START_GROUP = 3; END_GROUP = 4; FIXED32 = 5 # noqa: E702
+
+class AttributeType(enum.IntEnum):
+  """
+  ONNX attribute type identifiers.
+  Reference: https://github.com/onnx/onnx/blob/rel-1.18.0/onnx/onnx.proto3#L128-L145
+  """
+  FLOAT = 1; INT = 2; STRING = 3; TENSOR = 4; FLOATS = 6; INTS = 7; STRINGS = 8 # noqa: E702
+
+  def to_field_name(self) -> str: return {1: "f", 2: "i", 3: "s", 4: "t", 6: "floats", 7: "ints", 8: "strings"}[self.value]
+
+class OnnxDataType(enum.IntEnum):
+  """
+  ONNX tensor data type identifiers.
+  Reference: https://github.com/onnx/onnx/blob/rel-1.18.0/onnx/onnx.proto3#L500-L544
+  """
+  FLOAT = 1; UINT8 = 2; INT8 = 3; UINT16 = 4; INT16 = 5; INT32 = 6; INT64 = 7; BOOL = 9; FLOAT16 = 10; DOUBLE = 11; UINT32 = 12 # noqa: E702
+  UINT64 = 13; BFLOAT16 = 16 # noqa: E702
+
+  def to_dtype(self) -> DType: return dtypes.fields()[self.name.lower()]
+
+def dtype_fallback(dtype: DType, fallback_context: str) -> DType:
+  if is_dtype_supported(dtype): return dtype
+  default_dtype = dtypes.default_int if dtypes.is_int(dtype) else dtypes.default_float
+  warnings.warn(f"dtype {dtype} on {Device.DEFAULT} from {fallback_context} is not supported, falling back to {default_dtype}")
+  assert is_dtype_supported(default_dtype), f"dtype {default_dtype} must be supported on {Device.DEFAULT}"
+  return default_dtype
+
+# ***** onnx spec definitions *****
+class Domain(enum.Enum):
+  ONNX = "ai.onnx"
+  ONNX_ML = "ai.onnx.ml"
+  AI_ONNX_TRAINING = "ai.onnx.training"
+  AI_ONNX_PREVIEW_TRAINING = "ai.onnx.preview.training"
+  MICROSOFT_CONTRIB_OPS = "com.microsoft"
+  MICROSOFT_NCHWC = "com.microsoft.nchwc"
+  MICROSOFT_EXPERIMENTAL = "com.microsoft.experimental"
+  PYTORCH_ATEN = "org.pytorch.aten"
+  @classmethod
+  def from_onnx(cls, domain: str | None) -> "Domain": return cls.ONNX if domain is None or domain == "" else cls(domain)
+
+class OpSetId(NamedTuple):
+  domain: Domain
+  version: int
+
+@dataclasses.dataclass(frozen=True)
+class OnnxValue:
+  shape: tuple[str|int, ...]
+  dtype: DType
+  is_optional: bool
+  is_sequence: bool
+
+@dataclasses.dataclass(frozen=True)
+class OnnxNode:
+  op: str
+  opset_id: OpSetId
+  inputs: tuple[str, ...]
+  outputs: tuple[str, ...]
+  opts: dict[str, Any]
+
+# ***** protobuf parsing ******
+class PBBufferedReader(BufferedReader):
+  def __init__(self, tensor: Tensor):
+    assert tensor.dtype == dtypes.uint8, tensor
+    super().__init__(TensorIO(tensor))
+    self.len = tensor.nbytes()
+
+  def decode_varint(self) -> int:
+    """Reference: https://protobuf.dev/programming-guides/encoding/#varints"""
+    result = 0
+    shift = 0
+    while True:
+      data = self.read(1)
+      if data == b"": raise EOFError("decode_varint EOF")
+      result |= (data[0] & 0x7F) << shift
+      if not (data[0] & 0x80): return result
+      shift += 7
+      if shift >= 70: raise ValueError("Varint too long")
+
+  def read_delimited(self, use_tensor=False):
+    str_len = self.decode_varint()
+    if not use_tensor: return self.read(str_len)
+    raw = self.raw
+    assert isinstance(raw, TensorIO)
+    res = raw._tensor[self.tell():(self.tell()+str_len)]
+    self.seek(str_len, os.SEEK_CUR)
+    return res
+  def read_string(self) -> str: return self.read_delimited().decode("utf-8")
+  def read_bytes(self) -> Tensor: return self.read_delimited(use_tensor=True)
+  def read_float(self) -> float: return struct.unpack("<f", self.read(4))[0]
+  def read_packed_floats(self) -> Tensor: return self.read_delimited(use_tensor=True)
+  def read_int64(self) -> int: return truncate[dtypes.int64](self.decode_varint())
+  def read_packed_int64s(self) -> list[int]:
+    total_bytes_len = self.decode_varint()
+    old_pos = self.tell()
+    values = []
+    # need copy here because packed ints are varint
+    while self.tell() < total_bytes_len + old_pos: values.append(self.read_int64())
+    return values
+
+  def skip_field(self, wire_type: WireType) -> None:
+    """Skip a field based on its wire type."""
+    match wire_type:
+      case WireType.VARINT: self.decode_varint()
+      case WireType.FIXED64: self.seek(8, os.SEEK_CUR)
+      case WireType.FIXED32: self.seek(4, os.SEEK_CUR)
+      case WireType.LENGTH_DELIMITED: self.seek(self.decode_varint(), os.SEEK_CUR)
+      case _: raise ValueError(f"Unknown wire type: {wire_type}")
+
+class OnnxPBParser:
+  """
+  ONNX protobuf parser.
+  Reference: https://github.com/onnx/onnx/blob/main/onnx/onnx.proto3
+  """
+  def __init__(self, inp: Tensor|str|pathlib.Path, load_external_data: bool=True):
+    self.file_path: pathlib.Path|None = None
+    self.load_external_data = load_external_data
+    if not isinstance(inp, Tensor):
+      self.file_path = pathlib.Path(inp)
+      self.tensor = Tensor(self.file_path)
+    else: self.tensor = inp
+    self.reader = PBBufferedReader(self.tensor)
+
+  def parse(self) -> dict:
+    """Parses the ONNX model into a nested dictionary. """
+    return self._parse_ModelProto()
+
+  def _parse_message(self, end_pos: int) -> Generator[tuple[int, WireType], None, None]:
+    while self.reader.tell() < end_pos:
+      tag = self.reader.decode_varint()
+      yield tag >> 3, WireType(tag & 0x07)
+
+  def _decode_end_pos(self) -> int:
+    str_len = self.reader.decode_varint()
+    start_pos = self.reader.tell()
+    return start_pos + str_len
+
+  def _parse_ModelProto(self) -> dict:
+    """Entry point for parsing the ONNX model."""
+    obj: dict[str, Any] = {"opset_import": []}
+    for fid, wire_type in self._parse_message(self.reader.len):
+      match fid:
+        case 4: obj["domain"] = self.reader.read_string()
+        case 5: obj["model_version"] = self.reader.read_int64()
+        case 7: obj["graph"] = self._parse_GraphProto()
+        case 8: obj["opset_import"].append(self._parse_OperatorSetIdProto())
+        case _: self.reader.skip_field(wire_type)
+
+    # update opset version
+    opset_imports = {Domain.from_onnx(x.get('domain')):x.get('version', 1) for x in obj["opset_import"]}
+    for n in obj["graph"]["node"]:
+      n_ = n["parsed_node"]
+      n["parsed_node"] = OnnxNode(n_.op, OpSetId(n_.opset_id.domain, opset_imports.get(n_.opset_id.domain, 1)), n_.inputs, n_.outputs, n_.opts)
+    return obj
+
+  def _parse_GraphProto(self) -> dict:
+    obj: dict[str, Any] = {"node": [], "initializer": [], "input": [], "output": []}
+    for fid, wire_type in self._parse_message(self._decode_end_pos()):
+      match fid:
+        case 1: obj["node"].append(self._parse_NodeProto())
+        case 2: obj["name"] = self.reader.read_string()
+        case 5: obj["initializer"].append(self._parse_TensorProto())
+        case 11: obj["input"].append(self._parse_ValueInfoProto())
+        case 12: obj["output"].append(self._parse_ValueInfoProto())
+        case _: self.reader.skip_field(wire_type)
+    return obj
+
+  def _parse_NodeProto(self) -> dict:
+    obj: dict[str, Any] = {"input": [], "output": [], "attribute": [], "domain": None}
+    for fid, wire_type in self._parse_message(self._decode_end_pos()):
+      match fid:
+        case 1: obj["input"].append(self.reader.read_string())
+        case 2: obj["output"].append(self.reader.read_string())
+        case 3: obj["name"] = self.reader.read_string()
+        case 4: obj["op_type"] = self.reader.read_string()
+        case 5: obj["attribute"].append(self._parse_AttributeProto())
+        case 6: obj["doc_string"] = self.reader.read_string()
+        case 7: obj["domain"] = self.reader.read_string()
+        case _: self.reader.skip_field(wire_type)
+
+    # parse node
+    attributes = {attr_dict["name"]: attr_dict[AttributeType(attr_dict["type"]).to_field_name()] for attr_dict in obj["attribute"]}
+    opset_id = OpSetId(Domain.from_onnx(obj.get('domain')), 1)  # default version, to be updated later in _parse_ModelProto
+    obj["parsed_node"] = OnnxNode(obj["op_type"], opset_id, tuple(obj["input"]), tuple(obj["output"]), attributes)
+    return obj
+
+  def _parse_TensorProto(self) -> dict:
+    obj: dict[str, Any] = {"dims": []}
+    for fid, wire_type in self._parse_message(self._decode_end_pos()):
+      match fid:
+        case 1: obj["dims"].append(self.reader.read_int64())
+        case 2: obj["data_type"] = self.reader.read_int64()
+        case 4: obj["float_data"] = self.reader.read_packed_floats()
+        case 5: obj["int32_data"] = self.reader.read_packed_int64s()
+        case 7: obj["int64_data"] = self.reader.read_packed_int64s()
+        case 8: obj["name"] = self.reader.read_string()
+        case 9: obj["raw_data"] = self.reader.read_bytes()
+        case 10: obj["double_data"] = self.reader.read_packed_floats()
+        case 11: obj["uint64_data"] = self.reader.read_packed_int64s()
+        case 13: obj.setdefault("external_data", []).append(self._parse_StringStringEntryProto())
+        case 14: obj["data_location"] = self.reader.read_int64()
+        case _: self.reader.skip_field(wire_type)
+
+    # load external data
+    if self.load_external_data and obj.get("data_location", 0) == 1:
+      if "external_data" not in obj: raise ValueError("no external_data")
+      location, length, offset = None, None, 0
+      for kv in obj["external_data"]:
+        if kv["key"] == "location": location = kv["value"]
+        elif kv["key"] == "offset": offset = int(kv["value"])
+        elif kv["key"] == "length": length = int(kv["value"])
+      if location is None: raise ValueError("no location in external_data")
+
+      if self.file_path is None:
+        if isinstance(self.tensor.device, str) and self.tensor.device.startswith("DISK:"):
+          self.file_path = pathlib.Path(self.tensor.device[5:])
+        else: raise ValueError("onnx external_data needs the origin file path, try passing onnx file path to onnx_load")
+      ext_path = self.file_path.parent.joinpath(location)
+      if not ext_path.exists(): raise FileNotFoundError(f"external location not exists: {ext_path}")
+
+      ext_tensor = Tensor(ext_path)
+      obj["raw_data"] = ext_tensor[offset:offset+length] if length is not None else ext_tensor[offset:]
+      obj["data_location"] = 0
+
+    # parse tensor
+    to_dtype = dtype_fallback(true_dtype := OnnxDataType(obj['data_type']).to_dtype(), "buffer parse")
+    shape = tuple(obj['dims'])
+    present_fields = [field for field in ['float_data', 'int32_data', 'int64_data', 'double_data', 'uint64_data', 'raw_data'] if field in obj]
+    assert len(present_fields) == 1, f"only 1 data field is allowed from {obj=}"
+    data = obj[present_fields[0]]
+    if not isinstance(data, Tensor):
+      obj["parsed_tensor"] = Tensor(data, dtype=to_dtype).reshape(shape)
+      return obj
+    assert isinstance(data, Tensor) and data.dtype == dtypes.uint8, data
+    data = data.bitcast(true_dtype).reshape(shape)
+    data = data.to(Device.DEFAULT) if true_dtype is to_dtype else data.to("cpu").cast(to_dtype).to(Device.DEFAULT)
+    # const folding
+    if shape == ():
+      if data.dtype == dtypes.float16 and sys.version_info < (3, 12): data = data.cast(dtypes.float32)
+      data = Tensor(data.item(), dtype=to_dtype).reshape(shape)
+    obj["parsed_tensor"] = data
+    return obj
+
+  def _parse_AttributeProto(self) -> dict:
+    obj: dict[str, Any] = {"floats": [], "ints": [], "strings": []}
+    for fid, wire_type in self._parse_message(self._decode_end_pos()):
+      match fid:
+        case 1: obj["name"] = self.reader.read_string()
+        case 2: obj["f"] = self.reader.read_float()
+        case 3: obj["i"] = self.reader.read_int64()
+        case 4: obj["s"] = self.reader.read_bytes().data().tobytes().decode("utf8")
+        case 5: obj["t"] = self._parse_TensorProto()['parsed_tensor']
+        case 7: obj["floats"].append(self.reader.read_float())
+        case 8: obj["ints"].append(self.reader.read_int64())
+        case 9: obj["strings"].append(self.reader.read_bytes().data().tobytes().decode("utf8"))
+        case 20: obj["type"] = self.reader.read_int64()
+        case _: self.reader.skip_field(wire_type)
+    obj["floats"], obj["ints"], obj["strings"] = tuple(obj["floats"]), tuple(obj["ints"]), tuple(obj["strings"])
+    return obj
+
+  def _parse_ValueInfoProto(self) -> dict:
+    obj: dict[str, Any] = {}
+    for fid, wire_type in self._parse_message(self._decode_end_pos()):
+      match fid:
+        case 1: obj["name"] = self.reader.read_string()
+        case 2: obj["type"] = self._parse_TypeProto()
+        case _: self.reader.skip_field(wire_type)
+
+    # parse type
+    if "type" not in obj: return {**obj, "parsed_type": None}
+    type_obj = obj["type"]
+    if is_optional := "optional_type" in type_obj: type_obj = type_obj["optional_type"]["elem_type"]
+    if is_sequence := "sequence_type" in type_obj: type_obj = type_obj["sequence_type"]["elem_type"]
+    assert "tensor_type" in type_obj, type_obj
+    shape_dims = type_obj['tensor_type'].get('shape', {}).get('dim', [])
+    obj['parsed_type'] = OnnxValue(tuple(d.get('dim_param') or d.get('dim_value') for d in shape_dims),
+                                   OnnxDataType(type_obj['tensor_type']['elem_type']).to_dtype(), is_optional, is_sequence)
+    return obj
+
+  def _parse_TypeProto(self) -> dict:
+    obj: dict[str, Any] = {}
+    for fid, wire_type in self._parse_message(self._decode_end_pos()):
+      match fid:
+        case 1: obj["tensor_type"] = self._parse_TypeProtoTensor()
+        case 4: obj["sequence_type"] = self._parse_TypeProtoSequence()
+        case 9: obj["optional_type"] = self._parse_TypeProtoOptional()
+        case _: self.reader.skip_field(wire_type)
+    return obj
+
+  def _parse_TypeProtoTensor(self) -> dict:
+    obj: dict[str, Any] = {}
+    for fid, wire_type in self._parse_message(self._decode_end_pos()):
+      match fid:
+        case 1: obj["elem_type"] = self.reader.read_int64()
+        case 2: obj["shape"] = self._parse_TensorShapeProto()
+        case _: self.reader.skip_field(wire_type)
+    return obj
+
+  def _parse_TypeProtoSequence(self) -> dict:
+    obj = {}
+    for fid, wire_type in self._parse_message(self._decode_end_pos()):
+      match fid:
+        case 1: obj["elem_type"] = self._parse_TypeProto()
+        case _: self.reader.skip_field(wire_type)
+    return obj
+
+  def _parse_TypeProtoOptional(self) -> dict:
+    obj = {}
+    for fid, wire_type in self._parse_message(self._decode_end_pos()):
+      match fid:
+        case 1: obj["elem_type"] = self._parse_TypeProto()
+        case _: self.reader.skip_field(wire_type)
+    return obj
+
+  def _parse_TensorShapeProto(self) -> dict:
+    obj: dict[str, Any] = {"dim": []}
+    for fid, wire_type in self._parse_message(self._decode_end_pos()):
+      match fid:
+        case 1: obj["dim"].append(self._parse_TensorShapeProtoDimension())
+        case _: self.reader.skip_field(wire_type)
+    return obj
+
+  def _parse_TensorShapeProtoDimension(self) -> dict:
+    obj: dict[str, Any] = {}
+    for fid, wire_type in self._parse_message(self._decode_end_pos()):
+      match fid:
+        case 1: obj["dim_value"] = self.reader.read_int64()
+        case 2: obj["dim_param"] = self.reader.read_string()
+        case _: self.reader.skip_field(wire_type)
+    return obj
+
+  def _parse_StringStringEntryProto(self) -> dict:
+    obj: dict[str, Any] = {}
+    for fid, wire_type in self._parse_message(self._decode_end_pos()):
+      match fid:
+        case 1: obj["key"] = self.reader.read_string()
+        case 2: obj["value"] = self.reader.read_string()
+        case _: self.reader.skip_field(wire_type)
+    return obj
+
+  def _parse_OperatorSetIdProto(self) -> dict:
+    obj: dict[str, Any] = {}
+    for fid, wire_type in self._parse_message(self._decode_end_pos()):
+      match fid:
+        case 1: obj["domain"] = self.reader.read_string()
+        case 2: obj["version"] = self.reader.read_int64()
+        case _: self.reader.skip_field(wire_type)
+    return obj
+
+# ***** python const *****
+required_input_python_consts: dict[str, tuple[int, ...]] = {
+  "Tile": (1,), "Range": (0,1,2), "Expand": (1,), "Reshape": (1,), "Squeeze": (1,), "Unsqueeze": (1,), "Trilu": (1,), "ConstantOfShape": (0,),
+  "CumSum": (1,), "TopK": (1,), "Pad": (1,2,3), "MaxUnpool": (2,), "Dropout": (1,2), "CenterCropPad": (1,), "OneHot": (1,), "Compress": (1,),
+  "ImageDecoder": (0,), "AffineGrid": (1,), "Resize": (1,2,3), "Upsample": (1,), "Split": (1,), "Slice": (1,2,3,4),
+  **{"Reduce"+r: (1,) for r in ("Max", "Min", "Sum", "Mean", "SumSquare", "Prod", "L1", "L2", "LogSum", "LogSumExp")},
+  **{optim: (1,) for optim in ("Adam", "Adagrad", "Momentum")}
+}
+
+cache_misses = 0
+@functools.cache
+def _cached_to_python_const(t:Tensor):
+  if t.dtype == dtypes.uint8: return t.data().tobytes()
+  if 0 in t.shape: return []
+  return t.tolist()
+
+# Tensor -> python value cache for parameters
+def to_python_const(t:Any, op:str, idx:int) -> list[ConstType]|ConstType|bytes:
+  if idx not in required_input_python_consts.get(op, ()) or not isinstance(t, Tensor): return t
+  global cache_misses
+  ret = _cached_to_python_const(t)
+  if (info := _cached_to_python_const.cache_info()).misses > cache_misses and DEBUG >= 3:
+    print(f"Cache miss for {t}")
+    cache_misses = info.misses
+  return ret
+
+# ***** runner ******
+debug = int(getenv("DEBUGONNX", "0"))
+limit = int(getenv("ONNXLIMIT", "-1"))
+class OnnxRunner:
+  """
+  `OnnxRunner` executes an ONNX model using Tinygrad.
+
+  Args:
+    model_path: The ONNX model, provided as a file path (a string or Path object) or a Tensor.
+  """
+  def __init__(self, model_path: Tensor | str | pathlib.Path):
+    model = OnnxPBParser(model_path, load_external_data=True).parse()
+    graph = model["graph"]
+    self.is_training = any(n['parsed_node'].opset_id.domain in {Domain.AI_ONNX_TRAINING, Domain.AI_ONNX_PREVIEW_TRAINING} for n in graph["node"])
+    self.graph_values = {"": None, **{i["name"]: i["parsed_tensor"] for i in graph["initializer"]}}
+    self.graph_inputs = {i["name"]: i["parsed_type"] for i in graph["input"] if i["name"] not in self.graph_values}
+    self.graph_outputs = tuple(o["name"] for o in graph["output"])
+    self.graph_nodes = tuple(n["parsed_node"] for n in graph["node"])
+
+    self.old_training = Tensor.training
+    Tensor.training = True if self.is_training else False
+
+    self.variable_dims: dict[str, int] = {}
+    self.onnx_ops = onnx_ops
+
+  def _parse_input(self, name: str, value: Any, spec: OnnxValue):
+    if spec.is_optional and value is None: return None
+    if spec.is_sequence:
+      if not isinstance(value, Sequence): raise RuntimeError(f"input {name} received {value}, expected a sequence type")
+      sequence = [Tensor(v, dtype=spec.dtype, requires_grad=self.is_training) if not isinstance(v, Tensor) else v for v in value]
+      if not all_same(tuple(t.shape for t in sequence)): raise RuntimeError(f"Shapes for input {name} sequence must be homogeneous")
+      if not all(t.dtype is spec.dtype for t in sequence): warnings.warn(f"Dtypes for input {name} sequence aren't all {spec.dtype}")
+      return sequence
+    dtype = _from_np_dtype(value.dtype) if is_numpy_ndarray(value) else spec.dtype
+    tensor = Tensor(value, dtype=dtype, requires_grad=self.is_training) if not isinstance(value, Tensor) else value
+    if tensor.dtype is not spec.dtype: warnings.warn(f"input {name} has mismatch on dtype. Expected {spec.dtype}, received {tensor.dtype}.")
+    for dim, (onnx_dim, user_dim_input) in enumerate(zip(spec.shape, tensor.shape, strict=True)):
+      if isinstance(onnx_dim, str):
+        onnx_dim = self.variable_dims[onnx_dim] if onnx_dim in self.variable_dims else self.variable_dims.setdefault(onnx_dim, int(user_dim_input))
+      if user_dim_input != onnx_dim: raise RuntimeError(f"input {name} has mismatch on {dim=}. Expected {onnx_dim}, received {user_dim_input}.")
+    return tensor
+
+  def _select_op(self, op:str, required_opset:OpSetId) -> types.FunctionType:
+    if op not in self.onnx_ops: raise NotImplementedError(f"{op=} is not supported")
+    # return default implementation if no opset_id is specified
+    if isinstance(impl := self.onnx_ops[op], types.FunctionType): return impl
+    # match domain and select implementation with latest compatible version
+    eligible_ops = {impl_opset.version:impl_fxn for impl_opset,impl_fxn in impl.items()
+                    if impl_opset.domain == required_opset.domain and impl_opset.version <= required_opset.version}
+    if not eligible_ops: raise NotImplementedError(f"{op=} is not supported for domain {required_opset.domain} and version {required_opset.version}")
+    return eligible_ops[max(eligible_ops.keys())]
+
+  def get_empty_input_data(self, device:str|None=None, dtype:DType|None=None) -> dict[str, Tensor]:
+    return {name:Tensor.empty(*spec.shape, device=device, dtype=dtype or spec.dtype) for name, spec in self.graph_inputs.items()}
+
+  def to(self, device:str|None):
+    self.graph_values = {k:v.to(device) if isinstance(v, Tensor) else v for k,v in self.graph_values.items()}
+    self.graph_nodes = tuple(OnnxNode(n.op, n.opset_id, tuple(n.inputs), tuple(n.outputs),
+                                      {k:v.to(device) if isinstance(v, Tensor) else v for k,v in n.opts.items()}) for n in self.graph_nodes)
+    return self
+
+  def __call__(self, inputs:dict[str, Any], debug=debug):
+    for name, input_spec in self.graph_inputs.items():
+      if name not in inputs: raise RuntimeError(f"Please provide input data for {name}")
+      self.graph_values[name] = self._parse_input(name, inputs[name], input_spec)
+
+    for num, node in enumerate(self.graph_nodes):
+      inps = [to_python_const(self.graph_values[name], node.op, i) for i,name in enumerate(node.inputs)]
+      opts = node.opts
+
+      # provide additional opts
+      if node.op == "Split" and 'num_outputs' not in opts: opts['num_outputs'] = len(node.outputs)
+      if node.op == "Gradient": opts['intermediate_tensors'] = self.graph_values
+
+      if debug >= 1: print(f"{num}: op '{node.op}' opt {opts}")
+      if debug >= 2 and node.inputs: print("\tinputs:\n" + "\n".join(f"\t\t{x} - {i!r}" for x,i in zip(node.inputs, inps)))
+      ret = self._select_op(node.op, node.opset_id)(*inps, **opts)
+      ret = ret if isinstance(ret, tuple) else (ret,)
+      if debug >= 2: print("\toutputs:\n" + "\n".join(f"\t\t{x} - {o!r}" for x,o in zip(node.outputs, ret)))
+
+      self.graph_values.update(dict(zip(node.outputs, ret[:len(node.outputs)], strict=True)))
+
+      if num == limit:
+        Tensor.training = self.old_training
+        return {name:self.graph_values[name] for name in node.outputs}
+    Tensor.training = self.old_training
+    return {name:self.graph_values[name] for name in self.graph_outputs}
+
+####################
+##### ONNX OPS #####
+####################
+def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionType]]:
+  # ***** helper functions *****
+  def _resolve_const(x: Sequence[ConstType]|ConstType): return get_single_element(x) if isinstance(x, Sequence) else x
+
+  def _axes(axes, noop_with_empty_axes): return axes or ([] if noop_with_empty_axes else None)
+
+  # (padding_top, padding_left, ..., padding_bottom, padding_right, ...) -> (padding_left, padding_right, padding_top, padding_bottom, ...)
+  def _onnx_pads_to_tiny_pads(pads): return tuple(flatten(reversed(list(zip(pads, pads[len(pads)//2:])))))
+
+  AUTO_PAD_OPTIONS = Literal["NOTSET", "SAME_UPPER", "SAME_LOWER", "VALID"]
+  # (padding_height, padding_width) -> (padding_top, padding_left, padding_bottom, padding_right)
+  def _auto_pad(pads, auto_pad: AUTO_PAD_OPTIONS):
+    if auto_pad == "SAME_UPPER": return [pads[i]//2 for i in range(len(pads))] + [pads[i]-pads[i]//2 for i in range(len(pads))]
+    return [pads[i]-pads[i]//2 for i in range(len(pads))] + [pads[i]//2 for i in range(len(pads))]
+
+  def _resolve_pool_pads(x:Tensor, p_, k_, d_, s_, auto_pad:AUTO_PAD_OPTIONS):
+    if auto_pad == "VALID": return [0]*(len(k_)*2)
+    i_, (s_,d_,p_) = x.shape[-len(k_):], (make_tuple(x, len(k_)*2) for x in (s_, d_, p_))
+    if auto_pad == "NOTSET": return _onnx_pads_to_tiny_pads(p_ if len(p_)==len(k_)*2 else p_*2)
+    o_ = [((i - (1 if auto_pad in ("SAME_UPPER", "SAME_LOWER") else k)) // s + 1) for i,k,s in zip(i_, k_, s_)]
+    return _onnx_pads_to_tiny_pads(_auto_pad([(o-1)*s+k-i for o,i,k,s in zip(o_, i_, k_, s_)], auto_pad))
+
+  def _clamp_cast(x:Tensor, dtype:DType): return x.clamp(dtypes.min(dtype), dtypes.max(dtype)).cast(dtype)
+
+  def _prepare_quantize(x:Tensor, scale:Tensor, zero_point:Tensor|int, axis=1, block_size=0):
+    if axis < 0: axis += x.ndim
+    # https://github.com/onnx/onnx/blob/main/onnx/reference/ops/op_quantize_linear.py#L31
+    def reshape(val:Tensor):
+      if val.numel() == 1: return val
+      if block_size == 0: return val.reshape([val.shape[0] if dim == axis else 1 for dim in range(x.ndim)])
+      return val.repeat_interleave(block_size, axis)
+    return (reshape(scale), reshape(zero_point) if isinstance(zero_point, Tensor) else zero_point)
+
+  def _op_integer(op, inputs:list[Tensor], zero_points:list[Tensor], **opts):
+    adjusted_inputs = [inp.int() - zp for inp, zp in zip(inputs, zero_points)]
+    return op(*adjusted_inputs, **opts)
+
+  def _qlinearop_quantized(op, inputs:list[Tensor], zero_points:list[Tensor], scales:list[Tensor], out_scale:Tensor, out_zero_point:Tensor, **opts):
+    # op execution is done in quantized int
+    out = _op_integer(op, inputs, zero_points, **opts)
+    assert dtypes.is_int(out.dtype), "quantized op should've done math in int"
+    out_quantized = (out * prod(scales) / out_scale).round() + out_zero_point
+    return _clamp_cast(out_quantized, out_zero_point.dtype)
+
+  def _qlinearop_float(op, inputs:list[Tensor], zero_points:list[Tensor], scales:list[Tensor], out_scale:Tensor, out_zero_point:Tensor, **opts):
+    # op execution is done in float32
+    dequantized_inputs = [(inp.int() - zp) * scale for inp, zp, scale in zip(inputs, zero_points, scales)]
+    out = op(*dequantized_inputs, **opts)
+    assert dtypes.is_float(out.dtype), "op should've done math in float"
+    out_quantized = (out / out_scale).round() + out_zero_point
+    return _clamp_cast(out_quantized, out_zero_point.dtype)
+
+  def _onnx_training(input_group_size):
+    def __decorator(func):
+      def ___wrapper(R:Tensor, T:int, *inputs:Tensor, **kwargs):
+        R = R.detach()
+        groups = len(inputs) // input_group_size
+        ret = [func(R, T, *inps, **kwargs) for inps in (inputs[i::groups] for i in range(groups))]
+        return tuple(flatten(zip(*ret)))
+      return ___wrapper
+    return __decorator
+
+  # ***** Property/Graph Ops *****
+  def Identity(x:Tensor): return x
+  def Constant(sparse_value:Tensor|None=None, value:Tensor|None=None, value_float:float|None=None, value_floats:list[float]|None=None,
+              value_int:int|None=None, value_ints:list[int]|None=None, value_string:str|None=None, value_strings:list[str]|None=None):
+    if value is not None: return value
+    if value_float is not None: return Tensor(value_float, dtype=dtypes.float32, requires_grad=False)
+    if value_floats is not None: return Tensor(list(value_floats), dtype=dtypes.float32, requires_grad=False)
+    if value_int is not None: return Tensor(value_int, dtype=dtypes.int64, requires_grad=False)
+    if value_ints is not None: return Tensor(list(value_ints), dtype=dtypes.int64, requires_grad=False)
+    if value_string is not None or value_strings is not None or sparse_value is not None:
+      raise NotImplementedError('Constant OP not implemented for value_string, value_strings and sparse_value')
+
+  def Range(start:float|int|list[float|int], limit:float|int|list[float|int], delta:float|int|list[float|int]):
+    return Tensor.arange(start=_resolve_const(start), stop=_resolve_const(limit), step=_resolve_const(delta))
+
+  def ImageDecoder(encoded_stream:bytes, pixel_format="RGB"):
+    try: import PIL.Image
+    except ImportError as e: raise ImportError("Pillow must be installed for the ImageDecoder operator") from e
+    img = PIL.Image.open(io.BytesIO(encoded_stream))
+    if pixel_format == "BGR": return Tensor(img.tobytes(), dtype=dtypes.uint8).reshape(*img.size, 3).flip(-1)
+    if pixel_format == "RGB": return Tensor(img.tobytes(), dtype=dtypes.uint8).reshape(*img.size, 3)
+    if pixel_format == "Grayscale": return Tensor(img.convert("L").tobytes(), dtype=dtypes.uint8).reshape(*img.size, 1)
+    raise ValueError(f"pixel_format={pixel_format!r} is not supported.")
+
+  def EyeLike(x:Tensor, dtype:int|None=None, k:int=0):
+    ret = Tensor.eye(cast(int, min(x.shape)), dtype=dtype_fallback(OnnxDataType(dtype).to_dtype(), "EyeLike op") if dtype is not None else x.dtype)
+    return ret if x.size(0) == x.size(1) else ret.pad(tuple(None if d == ret.size(0) else (k, d-ret.shape[0]-k) for d in x.shape))
+
+  def OptionalHasElement(x:Tensor|None=None): return Tensor(x is not None and x.numel() > 0)
+  def OptionalGetElement(x:Tensor|None=None): return x if x is not None else Tensor([])
+  def ConstantOfShape(shape:list[int], value:Tensor|None=None):
+    if value is None: value = Tensor(0, dtype=dtypes.float32)
+    if shape == [0]: return Tensor([], dtype=value.dtype)
+    return value.expand(shape)
+
+  def Size(data:Tensor): return data.numel()
+  def Shape(data:Tensor, end:int|None=None, start:int=0): return Tensor(data.shape[start:end], dtype=dtypes.int64)
+
+  # ***** Unary Ops (math) *****
+  def Not(x:Tensor): return x.logical_not()
+  def Clip(x: Tensor, min:Tensor|None=None, max:Tensor|None=None): return x if min is None and max is None else x.clip(min, max)  # noqa: A002 # pylint: disable=redefined-builtin
+  def IsInf(x:Tensor, detect_negative:int=1, detect_positive:int=1): return x.isinf(bool(detect_positive), bool(detect_negative))
+
+  # ***** Unary Ops (activation) *****
+  def softmax_1(x:Tensor, axis:int=1): return x.softmax(axis)
+  def softmax_13(x:Tensor, axis:int=-1): return x.softmax(axis)
+  Softmax = {OpSetId(Domain.ONNX, 1):softmax_1, OpSetId(Domain.ONNX, 13):softmax_13}
+  def HardSigmoid(x:Tensor, alpha:float=0.2, beta:float=0.5): return (alpha*x + beta).clip(0, 1)
+  def Gelu(x:Tensor, approximate:str|None=None): return x.gelu() if approximate == "tanh" else 0.5 * x * (1 + (x/math.sqrt(2)).erf())
+  def BiasGelu(x: Tensor, bias: Tensor, approximate: str | None = None) -> Tensor: return Gelu(x + bias, approximate)
+  def FastGelu(x:Tensor, bias:Tensor|None=None): return (x + bias).gelu() if bias is not None else x.gelu() # this is tanh approximated
+  def PRelu(X:Tensor, slope:Tensor): return (X > 0).where(X, X * slope)
+  def LeakyRelu(X:Tensor, alpha:float=0.01): return X.leaky_relu(alpha)
+  def ThresholdedRelu(X:Tensor, alpha:float=1.0): return (X > alpha).where(X, 0)
+  def LogSoftmax(x: Tensor, axis:int=-1): return x.log_softmax(axis)
+  def Binarizer(x:Tensor, threshold:float=0.0): return (x > threshold).float()
+
+  # ***** Unary Ops (broadcasted) *****
+  def Add(x:Tensor,y:Tensor, broadcast=None, axis=None): return x + y
+  def Sub(x:Tensor|int,y:Tensor): return x - y # some test has input as int
+  def Div(x:Tensor,y:Tensor): return x.div(y, rounding_mode='trunc' if dtypes.is_int(x.dtype) else None)
+  def Less(x:Tensor,y:Tensor): return x < y
+  def LessOrEqual(x:Tensor,y:Tensor): return x <= y
+  def Greater(x:Tensor,y:Tensor): return x > y
+  def GreaterOrEqual(x:Tensor,y:Tensor): return x >= y
+  def Equal(x:Tensor,y:Tensor): return x == y
+  def And(x:Tensor,y:Tensor): return (x==y).where(x, False)
+  def Or(x:Tensor,y:Tensor): return (x==y).where(x, True)
+  def Xor(x:Tensor,y:Tensor): return x.bool().bitwise_xor(y.bool())
+  def BitwiseAnd(x:Tensor,y:Tensor): return x & y
+  def BitwiseOr(x:Tensor,y:Tensor): return x | y
+  def BitwiseXor(x:Tensor,y:Tensor): return x ^ y
+  def BitwiseNot(x:Tensor): return ~x
+  def Mod(x:Tensor,y:Tensor,fmod=0): return x - x.div(y, rounding_mode="trunc") * y if fmod else x % y
+
+  # ***** Casting Ops *****
+  # TODO: saturate
+  def Cast(x:Tensor, to:int, saturate:int=1): return x.cast(dtype_fallback(OnnxDataType(to).to_dtype(), "Cast op"))
+  def CastLike(x:Tensor, target_type:Tensor, saturate:int=1): return x.cast(target_type.dtype)
+
+  # ***** Reduce Ops *****
+  def Max(*data_0:Tensor): return functools.reduce(Tensor.maximum, data_0)
+  def Min(*data_0:Tensor): return functools.reduce(Tensor.minimum, data_0)
+  def Sum(*data_0:Tensor): return functools.reduce(Tensor.add, data_0)
+  def Mean(*data_0:Tensor): return Sum(*data_0) / len(data_0)
+  def ReduceMax(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
+    return data.max(_axes(axes, noop_with_empty_axes), keepdim=keepdims)
+  def ReduceMin(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
+    return data.min(_axes(axes, noop_with_empty_axes), keepdim=keepdims)
+  def ReduceSum(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
+    return data.sum(_axes(axes, noop_with_empty_axes), keepdim=keepdims)
+  def ReduceMean(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
+    return data.mean(_axes(axes, noop_with_empty_axes), keepdim=keepdims)
+  def ReduceSumSquare(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
+    return ReduceSum(data.square(), axes, keepdims, noop_with_empty_axes)
+  def ReduceProd(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
+    return data.prod(_axes(axes, noop_with_empty_axes), keepdim=keepdims)
+  def ReduceL1(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
+    return ReduceSum(data.abs(), axes, keepdims, noop_with_empty_axes)
+  def ReduceL2(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
+    return ReduceSumSquare(data, axes, keepdims, noop_with_empty_axes).sqrt()
+  def ReduceLogSum(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
+    return ReduceSum(data, axes, keepdims, noop_with_empty_axes).log()
+  def ReduceLogSumExp(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
+    return ReduceSum(data.exp(), axes, keepdims, noop_with_empty_axes).log()
+  def ArgMax(x:Tensor, axis:int=0, keepdims:int=1, select_last_index:int=0):
+    if select_last_index: return ((x.shape[axis]-1) - x.flip(axis).argmax(axis, keepdim=keepdims)).cast(dtypes.int64)
+    return x.argmax(axis, keepdim=keepdims).cast(dtypes.int64)
+  def ArgMin(x, axis:int=0, keepdims:int=1, select_last_index:int=0):
+    return ArgMax(-x, axis=axis, keepdims=keepdims, select_last_index=select_last_index)
+
+  # ***** Movement Ops *****
+  def Reshape(data:Tensor, shape:list[int], allowzero:int=0):
+    return data.reshape([x if x != 0 else (0 if allowzero else data.shape[i]) for i,x in enumerate(shape)])
+  def Flatten(x:Tensor, axis:int=1): return x.reshape(prod(x.shape[0:axis]), -1)
+  def Expand(x:Tensor, shape:list[int]): return x.expand(_broadcast_shape(x.shape, tuple(shape)))
+  def Shrink(x:Tensor, bias:float=0.0, lambd:float=0.5): return (x < -lambd)*(x+bias) + (x > lambd)*(x-bias)
+  def Transpose(x:Tensor, perm:list[int]|None=None): return x.permute(order=perm or list(range(x.ndim)[::-1]))
+
+  def Squeeze(data:Tensor, axes:list[int]|None=None):
+    return data.squeeze() if axes is None else functools.reduce(lambda d, dim: d.squeeze(dim), sorted(axes, reverse=True), data)
+  def Unsqueeze(data:Tensor, axes:list[int]): return functools.reduce(lambda d, dim: d.unsqueeze(dim), sorted(axes), data)
+
+  def Tile(x:Tensor, repeats:list[int]): return x.repeat(repeats)
+  def Concat(*xs:Tensor, axis:int): return Tensor.cat(*xs, dim=axis)
+  def Slice(data:Tensor, starts:list[int], ends:list[int], axes:list[int]|None=None, steps:list[int]|None=None):
+    axes = axes or list(range(data.ndim))
+    steps = steps or [1]*data.ndim
+    slices = [slice(0,x,1) for x in data.shape]
+    for i, axis in enumerate(axes): slices[axis] = slice(starts[i], ends[i], steps[i])
+    return data[tuple(slices)]
+
+  def Split(data:Tensor, split:list[int]|None=None, num_outputs:int=0, axis:int=0):
+    sz = data.shape[axis]
+    if split is None: split = [sz // num_outputs + (1 if i < sz % num_outputs else 0) for i in range(num_outputs)]
+    return data.split(split, axis)
+
+  def Pad(x:Tensor, pads:list[int], constant_value:ConstType|None=None, axes:list[int]|None=None,
+          mode:Literal["constant", "reflect", "edge", "wrap"]="constant", value=0):
+    value = constant_value or value
+    axes = axes or list(range(x.ndim))
+    real_pads = [0] * (x.ndim*2)
+    for i,axis in enumerate(axes): real_pads[axis%x.ndim], real_pads[axis%x.ndim+x.ndim] = pads[i], pads[i+len(axes)]
+    return x.pad(padding=_onnx_pads_to_tiny_pads(real_pads), mode={"edge":"replicate", "wrap":"circular"}.get(mode, mode), value=value)
+
+  def CenterCropPad(t:Tensor, shape:list[int], axes:list[int]|None=None):
+    shrink_arg:list[None|tuple[int,int]] = [None] * t.ndim
+    pad_arg:list[None|tuple[int,int]] = [None] * t.ndim
+    for s, x in zip(shape, axes or range(t.ndim)):
+      tx = t.shape[x]
+      if s < tx: shrink_arg[x] = (tx//2 - (s+1)//2, tx//2 + s//2)
+      elif s > tx: pad_arg[x] = ((s-tx)//2, (s-tx+1)//2)
+    return t.shrink(tuple(shrink_arg)).pad(tuple(pad_arg))
+
+  # ***** Processing Ops *****
+  def AveragePool(X: Tensor, kernel_shape:list[int], auto_pad:AUTO_PAD_OPTIONS="NOTSET", ceil_mode:int=0, count_include_pad:int=0,
+                  dilations:list[int]|int=1, pads:list[int]|int=0, strides:list[int]|int=1):
+    pool_pads = _resolve_pool_pads(X, pads, kernel_shape, dilations, strides, auto_pad)
+    return X.avg_pool2d(tuple(kernel_shape), strides, dilations, pool_pads, ceil_mode=ceil_mode, count_include_pad=count_include_pad)
+
+  def MaxPool(X: Tensor, kernel_shape:list[int], auto_pad:AUTO_PAD_OPTIONS="NOTSET", ceil_mode:int=0, dilations:list[int]|int=1, pads:list[int]|int=0,
+              storage_order:int=0, strides:list[int]|int=1):
+    pool_pads = _resolve_pool_pads(X, pads, kernel_shape, dilations, strides, auto_pad)
+    out = X.max_pool2d(tuple(kernel_shape), strides, dilations, pool_pads, ceil_mode=ceil_mode, return_indices=True)
+    ret, idx = cast(tuple[Tensor, Tensor], out)
+    return ret, idx.transpose(-2, -1).cast(dtypes.int64) if storage_order else idx.cast(dtypes.int64)
+
+  def Conv(X: Tensor, W: Tensor, B:Tensor|None=None, auto_pad:AUTO_PAD_OPTIONS="NOTSET", dilations:list[int]|int=1, group:int=1,
+          kernel_shape:list[int]|None=None, pads:list[int]|int=0, strides:list[int]|int=1):
+    return X.conv2d(W, B, stride=strides, groups=group, dilation=dilations,
+                    padding=_resolve_pool_pads(X, pads, kernel_shape or W.shape[2:], dilations, strides, auto_pad))
+
+  def ConvTranspose(X: Tensor, W: Tensor, B:Tensor|None=None, auto_pad:AUTO_PAD_OPTIONS="NOTSET", dilations:list[int]|int=1, group:int=1,
+                    kernel_shape:list[int]|None=None, pads:list[int]|None=None, output_shape:list[int]|None=None, output_padding:list[int]|int=0,
+                    strides:list[int]|int=1):
+    input_shape_, kernel_shape_ = X.shape[2:], (kernel_shape or W.shape[2:])
+    strides_, dilations_, output_padding_ = (make_tuple(x, len(input_shape_)) for x in (strides, dilations, output_padding))
+    if output_shape is not None: # we pad according to output_shape
+      pads = _auto_pad([s_*(i-1) + op_ + ((k_-1)*d_+1) - os for s_,i,op_,k_,d_,os in
+                        zip(strides_, input_shape_, output_padding_, kernel_shape_, dilations_, output_shape)], auto_pad)
+    if pads is None: # we generate pads
+      output_shape = output_shape or [X.shape[i+2] * strides_[i] for i in range(len(strides_))]
+      pads = [strides_[i]*(input_shape_[i]-1)+output_padding_[i]+((kernel_shape_[i]-1)*dilations_[i]+1)-output_shape[i]
+              for i in range(len(input_shape_))]
+      pads = _auto_pad(pads, auto_pad) if auto_pad != "NOTSET" else [0] * len(input_shape_) * 2
+    pads = _onnx_pads_to_tiny_pads(pads)
+    return X.conv_transpose2d(W, B, group, strides_, dilations_, pads, output_padding_)
+
+  def MaxUnpool(xT: Tensor, xI: Tensor, outshape: list[int]|None=None, kernel_shape:list[int]|None=None, pads:list[int]|int=0,
+                strides:list[int]|int=1):
+    if kernel_shape is None: kernel_shape = []
+    pads_: int | tuple[int, ...] = tuple(pads) if isinstance(pads, list) else pads
+    return Tensor.max_unpool2d(xT, xI, tuple(kernel_shape), strides, 1, pads_, outshape if outshape is None else tuple(outshape))
+
+  def GlobalAveragePool(X:Tensor): return X.mean(axis=tuple(range(2, X.ndim)), keepdim=True)
+  def GlobalMaxPool(X:Tensor): return X.max(axis=tuple(range(2, X.ndim)), keepdim=True)
+
+  def Gemm(A:Tensor, B:Tensor, C:Tensor|None=None, alpha:float=1.0, beta:float=1.0, transA:int=0, transB:int=0, broadcast=0):
+    ret = alpha * (A.transpose(transA) @ B.transpose(transB))
+    if C is not None: ret = ret + beta * (C if broadcast == 0 else C.reshape([-1 if i < len(C.shape) else 1 for i in range(ret.ndim)][::-1]))
+    return ret
+
+  def Einsum(*Inputs:list[Tensor], equation:str): return Tensor.einsum(equation, *Inputs)
+
+  def CumSum(X:Tensor, axis:int|list[int], exclusive:int=0, reverse:int=0):
+    axis = X._resolve_dim(_resolve_const(axis))
+    if reverse: X = X.flip(axis)
+    if exclusive: X = X.pad(tuple((1,0) if i == axis else None for i in range(X.ndim)))\
+                        .shrink(tuple((0,X.shape[axis]) if i == axis else None for i in range(X.ndim)))
+    return X.cumsum(axis).flip(axis) if reverse else X.cumsum(axis)
+
+  def Trilu(x:Tensor, k:int|list[int]=0, upper:int=1):
+    k_ = _resolve_const(k)
+    return x.triu(k_) if upper else x.tril(k_)
+
+  def Resize(X:Tensor, roi:list[float]|None=None, scales:list[float]|None=None, sizes:list[int]|None=None, antialias:int=0,
+        axes:list[int]|None=None, coordinate_transformation_mode:str='half_pixel', cubic_coeff_a:float=-0.75, exclude_outside:int=0,
+        extrapolation_value:float=0.0, keep_aspect_ratio_policy:str='stretch', mode:str='nearest', nearest_mode:str='round_prefer_floor'):
+    def _apply_transformation(input_sz, output_sz, scale_dim, mode):
+      index = Tensor.arange(output_sz, requires_grad=False, device=X.device)
+      if mode == "half_pixel": return (index + 0.5) / scale_dim - 0.5
+      if mode == "align_corners": return index * (input_sz - 1) / (output_sz - 1) if output_sz != 1 else Tensor.zeros_like(index)
+      if mode == "asymmetric": return index / scale_dim
+      if mode == "pytorch_half_pixel": return ((index + 0.5) / scale_dim - 0.5) if output_sz != 1 else Tensor.zeros_like(index)
+      if mode == "half_pixel_symmetric":
+        output_dim_scaled = input_sz * scale_dim
+        return (input_sz / 2) * (1 - (output_sz / output_dim_scaled)) + (index + 0.5) / scale_dim - 0.5
+      raise ValueError(f"invalid {coordinate_transformation_mode=}")
+
+    if antialias: raise NotImplementedError("antialias is not implemented")
+    axes = axes or list(range(X.ndim))
+    perm = [a for a in range(len(X.shape)) if a not in axes] + list(axes)
+    # we pre-permute the axes and permute back after resize
+    # the permute aligns X's axes to scales, sizes, and roi
+    X = X.permute(*perm)
+
+    input_shape = cast(tuple[int, ...], X.shape[2:])
+    if scales is not None: assert all(sc==1 for sc in scales[:-len(input_shape)]), "resizing batch_size dim or channel dim not supported"
+    if sizes is not None: assert tuple(sizes[:-2]) == tuple(X.shape[X.ndim-len(sizes):-2]),  "resizing batch_size dim or channel dim not supported"
+
+    scales, sizes = (None if scales is None else scales[-len(input_shape):]), (None if sizes is None else sizes[-len(input_shape):])
+    if sizes is not None:
+      if keep_aspect_ratio_policy in ["not_larger", "not_smaller"]:
+        scale_fxn = min if keep_aspect_ratio_policy == "not_larger" else max
+        scale = scale_fxn(sz / sh for sz,sh in zip(sizes, input_shape))
+        sizes, scales = [int(scale * sh + 0.5) for sh in input_shape], [scale]*len(input_shape)
+      else: scales = [sz / sh for sz, sh in zip(sizes, input_shape)]
+    else:
+      assert scales is not None, "either sizes or scales must be provided"
+      sizes = [int(sc * sh) for sc, sh in zip(scales, input_shape)]
+
+    if all(sz == sh for sz, sh in zip(sizes, input_shape)): return X.permute(*argsort(perm)) if perm else X
+
+    indexes = []
+    for input_sz, output_sz, scale in zip(input_shape, sizes, scales):
+      indexes.append(_apply_transformation(input_sz, output_sz, scale, coordinate_transformation_mode))
+
+    if mode in ["nearest", "linear"]: indexes = [idx.clip(0, sz-1) for idx, sz in zip(indexes, input_shape)]
+
+    if mode == "nearest":
+      mode_operations = {
+        "round_prefer_floor": lambda idx: (idx - 0.5).ceil(),
+        "round_prefer_ceil": lambda idx: (idx + 0.5).floor(),
+        "floor": lambda idx: idx.floor(),
+        "ceil": lambda idx: idx.ceil()
+      }
+      if nearest_mode not in mode_operations: raise ValueError(f"invalid {nearest_mode=}")
+      indexes = [mode_operations[nearest_mode](idx).int() for idx in indexes]
+      X = X[(..., *Tensor.meshgrid(*indexes))]
+
+    if mode == "linear":
+      expand = list(X.shape)
+      for i in range(-len(sizes), 0):
+        reshape, index = [1] * X.ndim, indexes[i]
+        reshape[i] = expand[i] = sizes[i]
+        low, high, perc = [y.reshape(reshape).expand(expand) for y in (index.floor().int(), index.ceil().int(), index - index.floor())]
+        X = X.gather(i, low).lerp(X.gather(i, high), perc)
+
+    if mode == "cubic":
+      A = cubic_coeff_a
+
+      # Keys weights
+      # see piecewise function in: https://en.wikipedia.org/wiki/Bicubic_interpolation#Bicubic_convolution_algorithm
+      def W0_1(x:Tensor): return polyN(x, [A + 2, -(A + 3), 0, 1])
+      def W1_2(x: Tensor): return polyN(x, [A, -5 * A, 8 * A, -4 * A])
+
+      expand = list(X.shape)
+      for i in range(-len(sizes), 0):
+        input_sz = cast(int, X.shape[i])
+        reshape, index = [1] * X.ndim, indexes[i]
+        reshape[i] = expand[i] = sizes[i]
+
+        p = index.floor().int()
+        ratio = index - p  # in [0, 1]
+
+        # Neighbor indices
+        idx0, idx1, idx2, idx3 = [p + d for d in [-1, 0, 1, 2]]
+        # Weights of distance from index and neighbor indices
+        c0, c1, c2, c3 = W1_2(ratio+1), W0_1(ratio), W0_1(-(ratio-1)), W1_2(-(ratio-2))
+
+        if exclude_outside:
+          c0 = ((idx0 >= 0) & (idx0 < input_sz)).where(c0, 0)
+          c1 = ((idx1 >= 0) & (idx1 < input_sz)).where(c1, 0)
+          c2 = ((idx2 >= 0) & (idx2 < input_sz)).where(c2, 0)
+          c3 = ((idx3 >= 0) & (idx3 < input_sz)).where(c3, 0)
+
+          total = c0 + c1 + c2 + c3
+          c0, c1, c2, c3 = c0 / (total + 1e-9), c1 / (total + 1e-9), c2 / (total + 1e-9), c3 / (total + 1e-9)
+
+        # Reshape and expand
+        expanded_indices = [y.clip(0, input_sz - 1).reshape(reshape).expand(expand) for y in [idx0, idx1, idx2, idx3]]
+        expanded_coeffs = [y.reshape(reshape).expand(expand) for y in [c0, c1, c2, c3]]
+
+        # Gather values and apply coefficients
+        gathered_values = [X.gather(i, idx) for idx in expanded_indices]
+        X = sum(v * c for v, c in zip(gathered_values, expanded_coeffs))
+    return X.permute(*argsort(perm)) if perm else X
+  def Upsample(X, scales, mode): return Resize(X=X, scales=scales, mode=mode)  # deprecated
+
+  def TopK(X:Tensor, K:int|list[int], axis:int=-1, largest:int=1, sorted:int=1):  # noqa: A002 # pylint: disable=redefined-builtin
+    val, idx = X.topk(_resolve_const(K), axis, bool(largest), bool(sorted))
+    return val, idx.cast(dtypes.int64)
+
+  # ***** Neural Network Ops *****
+  def BatchNormalization(X:Tensor, scale:Tensor, B:Tensor, input_mean:Tensor, input_var:Tensor, epsilon:float=1e-05, momentum:float=0.9,
+                        training_mode:int=0, spatial=1, is_test=0):
+    if training_mode:
+      x_detached = X.detach()
+      current_mean = x_detached.mean(axis=(0,2,3))
+      y = (x_detached - current_mean.reshape(shape=[1, -1, 1, 1]))
+      current_var = (y*y).mean(axis=(0,2,3))
+      current_invstd = current_var.add(epsilon).rsqrt()
+
+      running_mean = input_mean * momentum + current_mean * (1 - momentum)
+      running_var = input_var * momentum + current_var * (1 - momentum)
+
+      return X.batchnorm(scale, B, current_mean, current_invstd), running_mean, running_var
+    return X.batchnorm(scale, B, input_mean, (input_var + epsilon).rsqrt())
+  def GroupNormalization(x:Tensor, scale:Tensor, bias:Tensor, num_groups:int, epsilon:float=1e-05):
+    x = x.reshape(x.shape[0], num_groups, -1).layernorm(eps=epsilon).reshape(x.shape)
+    return x * scale.reshape(1, -1, *[1] * (x.ndim-2)) + bias.reshape(1, -1, *[1] * (x.ndim-2))
+  def InstanceNormalization(x:Tensor, scale:Tensor, bias:Tensor, epsilon:float=1e-05):
+    return GroupNormalization(x, scale, bias, num_groups=cast(int, x.shape[1]), epsilon=epsilon)
+  def LayerNormalization(x:Tensor, scale:Tensor, bias:Tensor, axis:int=-1, epsilon:float=1e-05, stash_type:int=1):
+    assert stash_type == 1, "only float32 is supported"
+    axes = tuple(i for i in range(axis if axis >= 0 else x.ndim + axis, x.ndim))
+    mean = x.mean(axis=axes, keepdim=True)
+    return x.layernorm(axes, epsilon).mul(scale).add(bias), mean, (x.sub(mean)).square().mean(axis=axes, keepdim=True).add(epsilon).rsqrt()
+  def SkipLayerNormalization(x:Tensor, skip:Tensor, gamma:Tensor, beta:Tensor|None=None, bias:Tensor|None=None, epsilon:float=1e-12):
+    x = x + skip
+    if bias is not None: x = x + bias
+    ret = x.layernorm(eps=epsilon) * gamma
+    if beta is not None: ret = ret + beta
+    return ret, None, None, x
+  def EmbedLayerNormalization(input_ids: Tensor, segment_ids:Tensor, word_embedding:Tensor, position_embedding:Tensor,
+                              segment_embedding:Tensor, gamma=None, beta=None, mask:Tensor|None=None,
+                              position_ids:Tensor|None=None, epsilon=1e-12, mask_index_type=0):
+    # https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.EmbedLayerNormalization
+    assert (segment_ids is None) is (segment_embedding is None)
+    assert mask is None and not mask_index_type, "functionality not supported yet"  # TODO
+    input_shape = input_ids.shape
+    seq_length = input_shape[1]
+    compute_seg_emb = (segment_embedding is not None and segment_ids is not None)
+    vocab_size, max_position_embeddings = word_embedding.shape[0], position_embedding.shape[0]
+    type_vocab_size  = (segment_embedding.shape[0] if compute_seg_emb else None)
+
+    def embedding(x:Tensor, vocab_size, weight:Tensor) -> Tensor:
+      return x.unsqueeze(-1).expand(*x.shape, vocab_size)._one_hot_along_dim(vocab_size) @ weight
+
+    # bert embedding layer
+    if position_ids is None: position_ids = Tensor.arange(seq_length, requires_grad=False).unsqueeze(0).expand(*input_shape)
+    wrd_embedding_res = embedding(input_ids, vocab_size, word_embedding)
+    pos_embedding_res = embedding(position_ids, max_position_embeddings, position_embedding)
+    seg_embedding_res = embedding(segment_ids, type_vocab_size, segment_embedding) if compute_seg_emb else None
+
+    embedding_sum = wrd_embedding_res + pos_embedding_res
+    if seg_embedding_res is not None: embedding_sum = embedding_sum + seg_embedding_res
+    out = embedding_sum.layernorm(eps=epsilon) * gamma + beta
+    return out, None, embedding_sum
+  def MeanVarianceNormalization(x:Tensor, axis:list[int]|None=None):
+    if axis is None: axis = [0,2,3]
+    return (x - x.mean(axis, keepdim=True)) / (x.std(axis, keepdim=True, correction=0) + 1e-9)
+
+  def OneHot(indices:Tensor, depth:float|int|list[int|float], values:Tensor, axis:int=-1):
+    # Scalar or Rank 1 tensor containing exactly one element
+    depth = int(_resolve_const(depth))
+    indices = indices.int()
+    indices = (indices < 0).where(indices+depth, indices)
+    return indices.unsqueeze(axis)._one_hot_along_dim(depth, dim=axis).where(values[1], values[0])
+
+  def DepthToSpace(X:Tensor, blocksize:int, mode:str="DCR"):
+    return X.rearrange("b (c h1 w1) h w -> b c (h h1) (w w1)" if mode=="CRD" else "b (h1 w1 c) h w -> b c (h h1) (w w1)", h1=blocksize, w1=blocksize)
+  def SpaceToDepth(X:Tensor, blocksize:int):
+    return X.rearrange("b c (h h1) (w w1) -> b (h1 w1 c) h w", h1=blocksize, w1=blocksize)
+
+  # Reimplemented here because you need legacy RNG for passing ONNX tests.
+  def dropout_7(data:Tensor, ratio:float=0.5, training_mode:bool=False, seed:int|None=None):
+    import numpy as np
+    if not training_mode: return data, data.full_like(True, dtype=dtypes.bool)
+    if seed is not None:
+      rand = Tensor(np.random.RandomState(seed).random(cast(tuple[int,...], data.shape)), requires_grad=False, dtype=data.dtype, device=data.device)
+    else:
+      rand = data.rand_like(requires_grad=False)
+    mask = rand >= ratio
+    return data * mask / (1.0 - ratio), mask
+  # 6 with 'is_test' needed for https://github.com/MTlab/onnx2caffe/raw/refs/heads/master/model/MobileNetV2.onnx
+  def dropout_6(data:Tensor, ratio:float=0.5, is_test=0): return dropout_7(data, ratio, training_mode=not is_test)
+  Dropout = {OpSetId(Domain.ONNX, 6):dropout_6, OpSetId(Domain.ONNX, 7):dropout_7}
+
+  def LRN(x:Tensor, size:int, alpha:float=1e-4, beta:float=0.75, bias:float=1.0):
+    pooled_x = (x**2).rearrange('b c h w -> b 1 c (h w)').pad((0,0,(size-1)//2, size//2)).avg_pool2d((size, 1), 1)
+    return x / (pooled_x.reshape(x.shape) * alpha + bias).pow(beta)
+
+  def NegativeLogLikelihoodLoss(x:Tensor, target:Tensor, weight:Tensor|None=None, ignore_index:int|None=None, reduction:ReductionStr="mean"):
+    return x.nll_loss(target, weight, ignore_index, reduction)
+  def SoftmaxCrossEntropyLoss(scores:Tensor, labels:Tensor, weights:Tensor|None=None, ignore_index:int|None=None, reduction:ReductionStr="mean"):
+    log_probs = scores.log_softmax(1)
+    return log_probs.nll_loss(labels, weights, ignore_index, reduction), log_probs
+
+  def AffineGrid(theta:Tensor, size:list[int], align_corners:int=0):
+    N, _, *spatial_dims = size
+    def generate_grid(steps):
+      if align_corners: return Tensor.linspace(-1, 1, steps, device=theta.device)
+      return Tensor.linspace(-1+1/steps, 1-1/steps, steps, device=theta.device)
+    grids = Tensor.meshgrid(*(generate_grid(d) for d in spatial_dims))
+    base_grid = Tensor.stack(*reversed(grids), Tensor.ones_like(grids[0], device=theta.device), dim=-1)
+    base_grid = base_grid.reshape(1, prod(spatial_dims), len(grids)+1).expand(N, -1, -1)
+    return (base_grid @ theta.transpose(1, 2)).reshape(N, *spatial_dims, -1)
+
+  def attention_contrib(x:Tensor, weights:Tensor, bias:Tensor|None=None, mask_index:Tensor|None=None, past:Tensor|None=None,
+                        attention_bias:Tensor|None=None, past_sequence_length:Tensor|None=None,  do_rotary:int=0, mask_filter_value:float=-10000.0,
+                        num_heads:int|None=None, past_present_share_buffer:int|None=None, qkv_hidden_sizes:list[int]|None=None,
+                        rotary_embedding_dim:int|None=None, scale:float|None=None, unidirectional:int=0):
+    assert not do_rotary and not attention_bias, "TODO"
+    if qkv_hidden_sizes is None: qkv_hidden_sizes = [weights.shape[1] // 3] * 3
+    qkv = x.linear(weights, bias)
+    q, k, v = qkv.split(qkv_hidden_sizes, dim=2)
+
+    batch_size, seq_len, _ = x.shape
+    assert num_heads is not None, "num_heads must be provided"
+    q_head_size, k_head_size, v_head_size = (sz // num_heads for sz in qkv_hidden_sizes)
+    q, k, v = (x.reshape(batch_size, seq_len, num_heads, hsz).transpose(1, 2) for x, hsz in zip((q, k, v), (q_head_size, k_head_size, v_head_size)))
+
+    present = None
+    if past is not None:
+      k, v = past[0].cat(k, dim=2), past[1].cat(v, dim=2)
+      present = k.stack(v)
+
+    if scale is None: scale = 1.0 / math.sqrt(q_head_size)
+    attn_scores = q @ k.transpose(-1, -2) * scale
+
+    if mask_index is not None:
+      assert 4 >= mask_index.ndim >= 1, f"{mask_index.ndim=}"
+      assert isinstance(batch_size, int), f"{batch_size=}"
+      if mask_index.ndim != 1: mask = mask_index.bool()
+      else:
+        if mask_index.shape[0] == batch_size:
+          mask = Tensor.arange(attn_scores.shape[-1], requires_grad=False, device=mask_index.device).unsqueeze(0) < mask_index.unsqueeze(1)
+        elif mask_index.shape[0] == 2*batch_size:
+          end_positions = mask_index[:batch_size]
+          start_positions = mask_index[batch_size:]
+          arange = Tensor.arange(seq_len).unsqueeze(0)
+          mask = (arange < end_positions.unsqueeze(1)) & (arange >= start_positions.unsqueeze(1))
+        else: raise NotImplementedError("mask_index with shape (3 * batch_size + 2) is not implemented")
+      while mask.ndim < 4: mask = mask.unsqueeze(1)
+      attn_scores = mask.where(attn_scores, mask_filter_value)
+
+    if unidirectional:
+      causal_mask = Tensor.ones((seq_len, seq_len), dtype=dtypes.bool).tril()
+      attn_scores = causal_mask.where(attn_scores, mask_filter_value)
+
+    output = attn_scores.softmax(-1) @ v
+    output = output.transpose(1, 2).reshape(batch_size, seq_len, -1)
+    return output, present
+
+  def attention_onnx(Q:Tensor, K:Tensor, V:Tensor, attn_mask:Tensor|None=None, past_key:Tensor|None=None, past_value:Tensor|None=None,
+                     is_causal:int=0, kv_num_heads:int|None=None, q_num_heads:int|None=None, qk_matmul_output_mode:int=0, scale:float|None=None,
+                     softcap:float=0.0, softmax_precision:int|None=None):
+    input_shape_len = Q.ndim
+    if input_shape_len == 3:
+      assert q_num_heads is not None and kv_num_heads is not None
+      Q = Q.reshape(Q.shape[0], q_num_heads, Q.shape[1], -1)
+      K = K.reshape(K.shape[0], kv_num_heads, K.shape[1], -1)
+      V = V.reshape(V.shape[0], kv_num_heads, V.shape[1], -1)
+
+    if past_key is not None: K = past_key.cat(K, dim=2)
+    if past_value is not None: V = past_value.cat(V, dim=2)
+    present_key, present_value = K, V
+
+    _q_heads, _kv_heads = q_num_heads or Q.shape[1], kv_num_heads or K.shape[1]
+    if _q_heads != _kv_heads:
+      K = K.repeat((1, _q_heads // _kv_heads, 1, 1))
+      V = V.repeat((1, _q_heads // _kv_heads, 1, 1))
+
+    effective_scale = scale if scale is not None else 1.0 / (cast(int, Q.shape[-1]) ** 0.5)
+    scores = (Q @ K.transpose(-1, -2)) * effective_scale
+    qk_matmul_return_val = scores
+
+    if is_causal:
+      causal_mask = Tensor.ones(Q.shape[-2], K.shape[-2], device=Q.device, dtype=dtypes.bool, requires_grad=False).tril(0)
+      scores = scores.masked_fill(causal_mask.logical_not(), -float("inf"))
+
+    if attn_mask is not None:
+      mask_to_add = attn_mask.where(0, -float("inf")) if attn_mask.dtype == dtypes.bool else attn_mask
+      scores = scores + mask_to_add
+    if qk_matmul_output_mode == 1: qk_matmul_return_val = scores
+
+    if softcap > 0.0: scores = (scores / softcap).tanh() * softcap
+    if qk_matmul_output_mode == 2: qk_matmul_return_val = scores
+
+    if softmax_precision: scores = scores.cast({1: dtypes.float32, 10: dtypes.float16, 16: dtypes.bfloat16}[softmax_precision])
+    qk_softmax = scores.softmax(-1).cast(Q.dtype)
+    if qk_matmul_output_mode == 3: qk_matmul_return_val = qk_softmax
+
+    output = (qk_softmax @ V).cast(Q.dtype)
+    if input_shape_len == 3: output = output.permute(0, 2, 1, 3).reshape(Q.shape[0], Q.shape[2], -1)
+    return output, present_key, present_value, qk_matmul_return_val
+  Attention = {OpSetId(Domain.ONNX, 1): attention_onnx, OpSetId(Domain.MICROSOFT_CONTRIB_OPS, 1): attention_contrib}
+
+  def RMSNormalization(X:Tensor, scale:Tensor, axis:int=-1, epsilon:float=1e-5):
+    norm = X.square().mean(axis=tuple(range(axis + X.ndim if axis < 0 else axis, X.ndim)), keepdim=True).add(epsilon).rsqrt()
+    return X * norm * scale
+
+  def RotaryEmbedding(X:Tensor, cos_cache:Tensor, sin_cache:Tensor, position_ids:Tensor|None=None, interleaved:int=0, num_heads:int|None=None,
+                      rotary_embedding_dim:int=0):
+    original_input_shape = X.shape
+
+    if X.ndim == 4: X = X.permute(0, 2, 1, 3)
+    elif X.ndim == 3:
+      assert num_heads is not None, "num_heads must be provided for 3D input"
+      X = X.unflatten(-1, (num_heads, X.shape[-1] // num_heads))
+
+    head_size = cast(int, X.shape[-1])
+    rot_dim = rotary_embedding_dim or head_size
+    x_rotate, x_pass = X[..., :rot_dim], X[..., rot_dim:]
+
+    cos = cos_cache[position_ids] if position_ids is not None else cos_cache[:head_size]
+    sin = sin_cache[position_ids] if position_ids is not None else sin_cache[:head_size]
+    cos = cos[..., :rot_dim//2].unsqueeze(2)
+    sin = sin[..., :rot_dim//2].unsqueeze(2)
+
+    x1, x2 = (x_rotate[..., ::2], x_rotate[..., 1::2]) if interleaved else x_rotate.chunk(2, dim=-1)
+    real = x1 * cos - x2 * sin
+    imag = x1 * sin + x2 * cos
+    x_rotated = real.stack(imag, dim=-1).flatten(start_dim=-2) if interleaved else real.cat(imag, dim=-1)
+
+    output = x_rotated.cat(x_pass, dim=-1)
+    return output.flatten(start_dim=2) if len(original_input_shape) == 3 else output.permute(0, 2, 1, 3)
+
+  # ***** Indexing Ops *****
+  def ArrayFeatureExtractor(x:Tensor, indices:Tensor): return x[..., indices]
+
+  def Gather(x:Tensor, indices:Tensor, axis:int=0):
+    if indices.numel() < 9: # NOTE lessor kernels for smaller indices but kernel number increases depending on size of indices
+      ret_shape = x.shape[:axis] + indices.shape + x.shape[axis+1:]
+      if indices.ndim > 1: indices = indices.flatten()
+      index_consts = [_cached_to_python_const(indices)] if indices.shape == () else _cached_to_python_const(indices)
+      index_consts = [x.shape[axis]+i if i<0 else i for i in index_consts]
+      args = [[(0,x) if j != axis else (i,i+1) for j, x in enumerate(x.shape)] for i in index_consts]
+      return x.shrink(arg=tuple(args[0])).cat(*[x.shrink(arg=tuple(arg)) for arg in args[1:]], dim=axis).reshape(ret_shape)
+    # NOTE faster gather, fixed number of kernels, but exceeds limited kernels for openpilot
+    return x[tuple([slice(None) if i != axis else indices for i in range(x.ndim)])]
+  def Scatter(*args, **kwargs): return ScatterElements(*args, **kwargs) # deprecated
+
+  def GatherND(x:Tensor, indices:Tensor, batch_dims:int=0):
+    if batch_dims == 0: return x[tuple(i.squeeze(-1) for i in indices.split(1, -1))]
+    x_shape, i_shape = x.shape, indices.shape
+    b = math.prod(x.shape[dim] for dim in range(batch_dims))
+    # NOTE: each batched dim of both input and indices are equal
+    x = x.reshape(b, *x.shape[batch_dims:])
+    indices = indices.reshape(b, *indices.shape[batch_dims:])
+    b_idx = Tensor.arange(b, device=x.device).reshape(b, *(1,)*(indices.ndim - 2)).expand(*indices.shape[:-1])
+    ret = x[(b_idx,) + tuple(i.squeeze(-1) for i in indices.split(1, -1))]
+    return ret.reshape(*x_shape[:batch_dims], *i_shape[batch_dims:-1], *ret.shape[indices.ndim-1:])
+  def ScatterND(x:Tensor, indices:Tensor, updates:Tensor, reduction:Literal["none", "add", "mul"]='none'):
+    assert updates.shape == indices.shape[:-1] + x.shape[cast(int, indices.shape[-1]):]
+    x = x.contiguous()
+    for index, u in zip(indices.split(1, 0), updates.split(1, 0)):
+      i = tuple(idx.squeeze(-1) for idx in index.squeeze(0).split(1, -1))
+      u = u.squeeze(0)
+      if reduction == "none": x[i] = u
+      elif reduction == "add": x[i] += u
+      elif reduction == "mul": x[i] *= u
+      else: raise NotImplementedError("reduction doesn't support max or min")
+    return x
+
+  def ScatterElements(x: Tensor, indices: Tensor, updates: Tensor, axis=0, reduction:Literal["none", "add", "mul", "min", "max"]="none"):
+    indices = (indices < 0).where(x.shape[axis], 0) + indices
+    if reduction == "none": return x.scatter(axis, indices, updates)
+    reduction_ = cast(Literal["sum", "prod", "amin", "amax"], {"add": "sum", "mul": "prod", "min": "amin", "max": "amax"}[reduction])
+    return x.scatter_reduce(axis, indices, updates, reduction_)
+  def GatherElements(x:Tensor, indices:Tensor, axis:int):
+    indices = (indices < 0).where(x.shape[axis], 0) + indices
+    return x.gather(axis, indices)
+
+  def Compress(inp:Tensor, condition:list[bool], axis:int|None=None):
+    if axis is None:
+      inp = inp.flatten()
+      axis = 0
+    axis = inp._resolve_dim(axis)
+    con = Tensor([i for i,cond in enumerate(condition) if cond]) # compress in python
+    return inp[tuple(con if i == axis else slice(None) for i in range(inp.ndim))]
+
+  # ***** Quantization Ops *****
+  def QuantizeLinear(x:Tensor, y_scale:Tensor, y_zero_point:Tensor|int=0, axis:int=1, block_size:int=0, output_dtype:int=0, saturate=1):
+    if isinstance(y_zero_point, Tensor): out_dtype = y_zero_point.dtype
+    elif output_dtype != 0: out_dtype = dtype_fallback(OnnxDataType(output_dtype).to_dtype(), "QuantizeLinear op")
+    else: out_dtype = dtypes.uint8
+    y_scale, y_zero_point = _prepare_quantize(x, y_scale, y_zero_point, axis, block_size)
+    if out_dtype == dtypes.uchar:
+      # this appears to work in practice, at least for uchar out_dtype. it folds with the quantize stuff
+      ret = _clamp_cast((x / y_scale + 0.4999999 + y_zero_point).int(), out_dtype)
+    else:
+      ret = _clamp_cast(((x / y_scale).round() + y_zero_point), out_dtype)
+    return ret.contiguous()
+
+  def DynamicQuantizeLinear(x: Tensor):
+    # only support uint8
+    qmin, qmax = dtypes.min(dtypes.uint8), dtypes.max(dtypes.uint8)
+    scale = (x.max().maximum(0) + ((-x).max()).maximum(0)) / (qmax - qmin)
+    zero_point = _clamp_cast((qmin - x.min() / scale).round(), dtypes.uint8)
+    y = _clamp_cast((x / scale).round() + zero_point, dtypes.uint8)
+    return y, scale, zero_point
+
+  def DequantizeLinear(x:Tensor, x_scale:Tensor, x_zero_point:Tensor|int=0, axis:int=1, block_size:int=0):
+    x_scale, x_zero_point = _prepare_quantize(x, x_scale, x_zero_point, axis, block_size)
+    return ((x.int() - x_zero_point) * x_scale).cast(x_scale.dtype)
+
+  def QLinearConv(x:Tensor, x_scale:Tensor, x_zero_point:Tensor, w:Tensor, w_scale:Tensor, w_zero_point:Tensor, y_scale:Tensor,
+                  y_zero_point:Tensor, B:Tensor|None=None, **opts):
+    return _qlinearop_quantized(Conv, [x,w], [x_zero_point,w_zero_point], [x_scale,w_scale], y_scale, y_zero_point, **{"B":B, **opts})
+
+  def QLinearMatMul(a:Tensor, a_scale:Tensor, a_zero_point:Tensor, b:Tensor, b_scale:Tensor, b_zero_point:Tensor, y_scale:Tensor,
+                    y_zero_point:Tensor) -> Tensor:
+    return _qlinearop_quantized(Tensor.matmul, [a,b], [a_zero_point,b_zero_point], [a_scale,b_scale], y_scale, y_zero_point)
+
+  def QLinearAdd(a:Tensor, a_scale:Tensor, a_zero_point:Tensor, b:Tensor, b_scale:Tensor, b_zero_point:Tensor, c_scale:Tensor, c_zero_point:Tensor):
+    return _qlinearop_float(Tensor.add, [a,b], [a_zero_point,b_zero_point], [a_scale,b_scale], c_scale, c_zero_point)
+
+  def QLinearMul(a:Tensor, a_scale:Tensor, a_zero_point:Tensor, b:Tensor, b_scale:Tensor, b_zero_point:Tensor, c_scale:Tensor, c_zero_point:Tensor):
+    return _qlinearop_quantized(Tensor.mul, [a,b], [a_zero_point,b_zero_point], [a_scale,b_scale], c_scale, c_zero_point)
+
+  def QLinearGlobalAveragePool(X:Tensor, x_scale:Tensor, x_zero_point:Tensor, y_scale:Tensor, y_zero_point:Tensor, channels_last:int):
+    assert channels_last == 0, "TODO NHWC"
+    return _qlinearop_float(GlobalAveragePool, [X], [x_zero_point], [x_scale], y_scale, y_zero_point)
+
+  def ConvInteger(x: Tensor, w: Tensor, x_zero_point:Tensor = Tensor(0), w_zero_point:Tensor = Tensor(0), B: Tensor | None = None, **opts) -> Tensor:
+    return _op_integer(Conv, [x,w], [x_zero_point,w_zero_point], **{"B":B, **opts})
+
+  def MatMulInteger(A: Tensor, B: Tensor, a_zero_point: Tensor = Tensor(0), b_zero_point: Tensor = Tensor(0)) -> Tensor:
+    return _op_integer(Tensor.matmul, [A,B], [a_zero_point,b_zero_point])
+
+  # ***** Training Ops *****
+  # NOTE: onnx training ops actually don't need the state for optim, all the ops work in a functional way, but we still can reuse optim.py code
+  @_onnx_training(3)
+  def Adagrad(R:Tensor, T:int, *inputs:Tensor, decay_factor:float=0.0, epsilon:float=0.0, norm_coefficient:float=0.0):
+    X, G, H = (i.detach() for i in inputs)
+    grad = norm_coefficient * X + G
+    H.assign(H + grad.square())
+    up = grad / (H.sqrt() + epsilon)
+    r = R / (1 + T * decay_factor)
+    X.assign(X.detach() - r * up)
+    return [X, H]
+
+  @_onnx_training(4)
+  def Adam(R:Tensor, T:int, *inputs:Tensor, alpha:float=0.9, beta:float=0.999, epsilon:float=0.0, norm_coefficient:float=0.0,
+          norm_coefficient_post:float=0.0):
+    from tinygrad.nn.optim import Adam as TinyAdam
+    X, G, V, H = inputs
+    G, V, H = G.detach(), V.detach(), H.detach()
+    X.grad = norm_coefficient * X.detach() + G
+    opt = TinyAdam([X], b1=alpha, b2=beta, eps=epsilon)
+    opt.m, opt.v, opt.lr = [V], [H], R
+    # need no-op for m_hat and v_hat if T == 0
+    if T == 0: opt.b1_t, opt.b2_t = opt.b1_t.zeros_like(), opt.b2_t.zeros_like()
+    else:
+      # `T-1` since it's applied again at the start of `_step`
+      opt.b1_t = Tensor([alpha**(T-1)], dtype=dtypes.float32, device=X.device, requires_grad=False)
+      opt.b2_t = Tensor([beta**(T-1)], dtype=dtypes.float32, device=X.device, requires_grad=False)
+    opt.step()
+    X = (1 - norm_coefficient_post) * X
+    return [X, V, H]
+
+  @_onnx_training(3)
+  def Momentum(R:Tensor, T:int, *inputs:Tensor, alpha:float, beta:float, mode:str, norm_coefficient:float):
+    X, G, V = (i.detach() for i in inputs)
+    grad = norm_coefficient * X + G
+    # NOTE: this beta_adjusted term makes it so we can't use SGD for nesterov
+    beta_adjusted = beta if T > 0 else 1
+    V.assign(alpha * V + grad * beta_adjusted)
+    X.assign(X - R * (V if mode == "standard" else (grad + alpha * V)))
+    return [X, V]
+
+  def Gradient(*inputs:Tensor, y:str, intermediate_tensors:dict[str, Tensor], **_):
+    intermediate_tensors[y].backward()
+    return tuple([t.grad for t in inputs])
+
+  return {
+    # Tensor ops
+    **{op: getattr(Tensor, op.lower()) for op in ("Neg", "Reciprocal", "Pow", "Sqrt", "Sign", "Abs", "Exp", "Log", "Mish", "Sin", "Cos", "Tan",
+    "Asin", "Acos", "Atan", "Relu", "Sigmoid", "MatMul", "Floor", "Ceil", "IsNaN", "Softplus", "HardSwish", "Where", "Mul", "Sinh", "Cosh",
+    "Tanh", "Softsign", "Asinh", "Acosh", "Atanh",  "Elu", "Celu", "Selu", "Round", "Erf")},
+    # Implemented ops
+    **{name:obj for name,obj in locals().items() if isinstance(obj, types.FunctionType) and not name.startswith("_") and name[0].isupper()},
+    # Version ops
+    **{name:obj for name,obj in locals().items() if isinstance(obj, dict)},
+  }
+
+onnx_ops = get_onnx_ops()