fixes from adapting tvm tests (#8570)

extra/onnx.py

@@ -56,10 +56,11 @@ def buffer_parse(onnx_tensor: TensorProto) -> Tensor:
              list(onnx_tensor.uint64_data):
     if len(data) == 1: return Tensor(data[0], dtype=dtype).reshape(shape)
     return Tensor(data, dtype=dtype).reshape(shape).realize()
-  assert onnx_tensor.HasField("raw_data")
-  np_buffer = np.frombuffer(onnx_tensor.raw_data, dtype=helper.tensor_dtype_to_np_dtype(onnx_tensor.data_type)).copy().reshape(shape)
-  if np_buffer.size == 1: return Tensor(np_buffer.item(), dtype=dtype).reshape(shape)
-  return Tensor(np_buffer, dtype=dtype)
+  if onnx_tensor.HasField("raw_data"):
+    np_buffer = np.frombuffer(onnx_tensor.raw_data, dtype=helper.tensor_dtype_to_np_dtype(onnx_tensor.data_type)).copy().reshape(shape)
+    if np_buffer.size == 1: return Tensor(np_buffer.item(), dtype=dtype).reshape(shape)
+    return Tensor(np_buffer, dtype=dtype)
+  return Tensor(None)
 
 onnx_ops = importlib.import_module('extra.onnx_ops')
 ONNXLIMIT = getenv("ONNXLIMIT", -1)
@@ -98,7 +99,7 @@ def get_run_onnx(onnx_model: ModelProto):
   def prepare_input(user_input:Any, model_input:ValueInfoProto):
     type_proto = model_input.type
     if type_proto.HasField("optional_type"):
-      if user_input is None: return Tensor(None)
+      if user_input is None: return None
       type_proto = type_proto.optional_type.elem_type
     if type_proto.HasField("sequence_type"):
       if not isinstance(user_input, Sequence): raise RuntimeError(f"{model_input.name} received {user_input}, expected sequence type")
@@ -116,7 +117,7 @@ def get_run_onnx(onnx_model: ModelProto):
       # if dtype is not tensor.dtype: raise RuntimeError(f"{model_input.name} has mismatch for dtype. Expected {dtype}, received {tensor.dtype}.")
       for dim, onnx_dim in enumerate(type_proto.tensor_type.shape.dim):
         dim_param, dim_value = onnx_dim.dim_param, onnx_dim.dim_value
-        user_dim_input = user_input.shape[dim]
+        user_dim_input = tensor.shape[dim]
         if dim_param: dim_value = variable_dims[dim_param] if dim_param in variable_dims else variable_dims.setdefault(dim_param, user_dim_input)
         if user_dim_input != dim_value:
           raise RuntimeError(f"{model_input.name} has mismatch for dim={dim_param or dim}. Expected {dim_value}, received {user_dim_input}.")

extra/onnx_ops.py

@@ -141,8 +141,8 @@ def Gemm(A:Tensor, B:Tensor, C:Tensor|None=None, alpha:float=1.0, beta:float=1.0
 
 def Einsum(*Inputs:list[Tensor], equation:str): return Tensor.einsum(equation, *Inputs)
 
-def CumSum(X:Tensor, axis:int, exclusive:int=0, reverse:int=0):
-  axis = X._resolve_dim(axis)
+def CumSum(X:Tensor, axis:int|list, exclusive:int=0, reverse:int=0):
+  axis = X._resolve_dim(axis[0] if isinstance(axis, list) else 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)))
@@ -342,23 +342,22 @@ def Resize(X:Tensor, roi:list[float]|None=None, scales:list[float]|None=None, si
     else: raise ValueError(f"invalid {coordinate_transformation_mode=}")
     return index.clip(0, input_dim-1)
 
-  scales, sizes = (None if scales is None else scales[-2:]), (None if sizes is None else sizes[-2:])
+  scales, sizes = (None if scales is None else scales[2-(X.ndim-len(scales)):]), (None if sizes is None else sizes[2-(X.ndim-len(sizes)):])
   # we pre permute the axes and permute back after resize
   axes, input_shape, = (axes or list(range(X.ndim))), cast(tuple[int, ...], X.shape[2:]),
   perm = [a for a in range(len(X.shape)) if a not in axes] + list(axes)
   X = X.permute(*perm)
 
   if sizes is not None:
-    assert scales is None
     if keep_aspect_ratio_policy in ["not_larger", "not_smaller"]:
       scale_fxn = min if keep_aspect_ratio_policy == "not_larger" else max
-      scales = [scale_fxn([sizes[i] / input_shape[i] for i in range(X.ndim-2) if i+2 in axes])] * 2
+      scales = [scale_fxn([sizes[i] / input_shape[i] for i in range(len(input_shape)) if i+2 in axes])] * 2
       sizes = [int((scales[0] * input_shape[i]) + 0.5) if i+2 in axes else input_shape[i] for i in range(X.ndim-2)]
-    else: scales = [sizes[-2] / cast(int, X.size(-2)), sizes[-1] / cast(int, X.size(-1))]
+    else:
+      scales = [size / input_shape for size, input_shape in zip(sizes, input_shape)]
   else:
-    assert scales is not None
     sizes = [int(sc*sh) for sc, sh in zip(scales, input_shape)]
-  regions = [[st, ed] for st, ed in zip(roi, roi[len(roi)//2:])] if isinstance(roi, list) else [[0.0, 0.0]] * (X.ndim-2)
+  regions = [[st, ed] for st, ed in zip(roi, roi[len(roi)//2:])] if isinstance(roi, list) and roi else [[0.0, 0.0]] * (X.ndim-2)
 
   # NOTE: this transformation makes it so that we can't just call Tensor.interpolate
   # in Tensor.interpolate, we use indexes without any transformation
@@ -388,9 +387,9 @@ def CenterCropPad(t:Tensor, shape:list[int], axes:list[int]|None=None):
     elif s > tx: pad_arg[x] = ((s-tx)//2, (s-tx+1)//2)
   return t.shrink(tuple(shrink_arg)).pad(tuple(pad_arg))
 
-def OneHot(indices:Tensor, depth:float|int, values:Tensor, axis:int=-1):
+def OneHot(indices:Tensor, depth:float|int|list, values:Tensor, axis:int=-1):
   # Scalar or Rank 1 tensor containing exactly one element
-  depth = int(depth)
+  depth = int(depth[0] if isinstance(depth, list) else depth)
   indices = (indices < 0).where(indices+depth, indices)
   return indices[:, None]._one_hot_along_dim(depth, dim=axis).where(values[1], values[0])