onnx-diagnostic 0.6.0__py3-none-any.whl → 0.6.2__py3-none-any.whl

onnx_diagnostic/reference/torch_ops/nn_ops.py

@@ -0,0 +1,196 @@
+from typing import Optional, Tuple
+import onnx
+import torch
+from ...helpers.torch_helper import onnx_dtype_to_torch_dtype
+from . import OpRunKernel, OpRunTensor
+
+
+class AveragePool_11(OpRunKernel):
+    "AveragePool"
+
+    def __init__(
+        self,
+        node: onnx.NodeProto,
+        version: Optional[int] = None,
+        verbose: int = 0,
+    ):
+        super().__init__(node, version, verbose=verbose)
+        self.auto_pad = self.get_attribute_string(node, "auto_pad", "NOTSET")
+        self.ceil_mode = bool(self.get_attribute_int(node, "ceil_mode", 0))
+        self.count_include_pad = bool(self.get_attribute_int(node, "count_include_pad", 0))
+        self.dilations = self.get_attribute_ints(node, "dilations", None)
+        self.kernel_shape: Tuple[int, ...] = (
+            self.get_attribute_ints(node, "kernel_shape") or tuple()
+        )
+        self.pads = self.get_attribute_ints(node, "pads", None)
+        self.strides = self.get_attribute_ints(node, "strides", None)
+
+    def run(self, x):
+        kernel_shape = self.kernel_shape
+        dilations = self.dilations or [1 for _ in x.shape[2:]]
+        strides = self.strides or [1 for _ in x.shape[2:]]
+        pads = self.pads or ([0 for _ in x.shape[2:]] * 2)
+        assert (
+            self.auto_pad == "NOTSET"
+        ), f"conv not implemented for auto_pad={self.auto_pad!r}"
+        assert len(set(pads)) == 1, f"conv not implemented for pads={pads}"
+        assert set(dilations) == {1}, f"conv not implemented for dilations={dilations}"
+        avg_pool = getattr(torch.nn.functional, f"avg_pool{len(kernel_shape)}d")
+        return OpRunTensor(
+            avg_pool(
+                x.tensor,
+                kernel_size=tuple(kernel_shape),
+                stride=tuple(strides),
+                padding=pads[0],
+                ceil_mode=self.ceil_mode,
+                count_include_pad=self.count_include_pad,
+                # dilation=tuple(dilations),
+            )
+        )
+
+
+class Conv_11(OpRunKernel):
+    "Conv"
+
+    def __init__(
+        self,
+        node: onnx.NodeProto,
+        version: Optional[int] = None,
+        verbose: int = 0,
+    ):
+        super().__init__(node, version, verbose=verbose)
+        self.auto_pad = self.get_attribute_string(node, "auto_pad", "NOTSET")
+        self.dilations = self.get_attribute_ints(node, "dilations", None)
+        self.group = self.get_attribute_int(node, "group", 1)
+        self.kernel_shape: Tuple[int, ...] = (
+            self.get_attribute_ints(node, "kernel_shape") or tuple()
+        )
+        self.pads = self.get_attribute_ints(node, "pads", None)
+        self.strides = self.get_attribute_ints(node, "strides", None)
+
+    def run(self, x, w, b=None):
+        kernel_shape = self.kernel_shape or w.shape[2:]
+        assert (
+            tuple(kernel_shape) == w.shape[-len(kernel_shape) :]
+        ), f"conv not implemented for kernel_shape={kernel_shape} and w.shape={w.shape}"
+        dilations = self.dilations or [1 for _ in x.shape[2:]]
+        strides = self.strides or [1 for _ in x.shape[2:]]
+
+        if self.auto_pad in {"SAME_LOWER", "SAME_UPPER"}:
+            head = []
+            tail = []
+            for i in range(len(x.shape) - 2):
+                d = x.shape[i + 2]
+                target_size = (d + strides[i] - 1) // strides[i]
+                pad_needed = (target_size - 1) * strides[i] + kernel_shape[i] - d
+                pad_head = (
+                    (pad_needed + 1) // 2 if self.auto_pad == "SAME_LOWER" else pad_needed // 2
+                )
+                pad_tail = pad_needed - pad_head
+                head.append(pad_head)
+                tail.append(pad_tail)
+            pads = head + tail
+        else:
+            pads = self.pads or ([0 for _ in x.shape[2:]] * 2)
+
+        assert len(set(pads)) == 1, (
+            f"conv not implemented for pads={pads}, "
+            f"auto_pad={self.auto_pad!r}, strides={strides}, "
+            f"x.shape={x.shape}, kernel_shape={kernel_shape}"
+        )
+
+        if b is None:
+            bias = None
+        else:
+            bias = b.tensor.squeeze()
+            if not bias.shape:
+                bias = bias.unsqueeze(0)
+        return OpRunTensor(
+            torch.nn.functional.conv2d(
+                x.tensor,
+                w.tensor,
+                bias=bias,
+                stride=tuple(strides),
+                padding=pads[0],
+                dilation=tuple(dilations),
+                groups=self.group,
+            )
+        )
+
+
+class LayerNormalization_17(OpRunKernel):
+    "LayerNormalization"
+
+    def __init__(
+        self,
+        node: onnx.NodeProto,
+        version: Optional[int] = None,
+        verbose: int = 0,
+    ):
+        super().__init__(node, version, verbose=verbose)
+        self.axis = self.get_attribute_int(node, "axis", -1)
+        self.epsilon = self.get_attribute_float(node, "epsilon", 1e-5)
+        self.stash_type = onnx_dtype_to_torch_dtype(
+            self.get_attribute_int(node, "stash_type", onnx.TensorProto.FLOAT)  # type: ignore[arg-type]
+        )
+        self.compute_std = len(node.output) > 1
+
+    def run(self, x, scale, bias=None):
+        original_dtype = x.dtype
+        if self.stash_type == torch.float32 and x.tensor.dtype != torch.float64:
+            xt = x.tensor
+            res = torch.nn.functional.layer_norm(
+                xt,
+                xt.shape[self.axis :],
+                weight=scale.tensor,
+                bias=None if bias is None else bias.tensor,
+                eps=self.epsilon,
+            )
+        else:
+            xt = x.tensor.to(self.stash_type)
+            res = torch.nn.functional.layer_norm(
+                xt,
+                xt.shape[self.axis :],
+                weight=scale.tensor.to(self.stash_type),
+                bias=None if bias is None else bias.tensor.to(self.stash_type),
+                eps=self.epsilon,
+            )
+        if not self.compute_std:
+            return OpRunTensor(res.to(original_dtype))
+        axes = tuple(range(len(xt.shape)))[self.axis :]
+        mean, var = torch.var(xt, dim=axes, keepdim=False)
+        x_inv_std_dev = torch.reciprocal(torch.sqrt(var + self.epsilon))
+        return (
+            OpRunTensor(res.to(original_dtype)),
+            OpRunTensor(mean),
+            OpRunTensor(x_inv_std_dev),
+        )
+
+
+class Softmax_13(OpRunKernel):
+    "Softmax"
+
+    def __init__(
+        self,
+        node: onnx.NodeProto,
+        version: Optional[int] = None,
+        verbose: int = 0,
+    ):
+        super().__init__(node, version, verbose=verbose)
+        self.axis = self.get_attribute_int(node, "axis", -1)
+        assert isinstance(self.axis, int), f"Unexpected value for attribute axis={self.axis!r}"
+        # this is out of spec
+        stash_type = self.get_attribute_int(node, "stash_type", None)
+        self.stash_type = None if stash_type is None else onnx_dtype_to_torch_dtype(stash_type)
+
+    def run(self, data: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(
+            torch.nn.functional.softmax(data.tensor, dim=self.axis, dtype=self.stash_type)
+        )
+
+
+class Tanh_6(OpRunKernel):
+    "Tanh"
+
+    def run(self, data: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(torch.nn.functional.tanh(data.tensor))

onnx_diagnostic/reference/torch_ops/other_ops.py

@@ -0,0 +1,106 @@
+from typing import Optional
+import onnx
+import torch
+from ...helpers.torch_helper import onnx_dtype_to_torch_dtype
+from . import OpRunKernel, OpRunTensor
+
+
+class Cast_6(OpRunKernel):
+    "Cast"
+
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None, verbose: int = 0):
+        super().__init__(node, version, verbose=verbose)
+        to = self.get_attribute_int(node, "to", 0)
+        assert isinstance(to, int), f"Unexpected value for attribute to={to!r}"
+        self.to = onnx_dtype_to_torch_dtype(to)
+        self.saturate = self.get_attribute_int(node, "saturate", 1)
+        assert self.saturate == 1, f"saturate={self.saturate} not implemented for Cast"
+
+    def run(self, data: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(data.tensor.to(self.to))
+
+
+class CastLike_15(OpRunKernel):
+    "Cast"
+
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None, verbose: int = 0):
+        super().__init__(node, version, verbose=verbose)
+        self.saturate = self.get_attribute_int(node, "saturate", 1)
+        assert self.saturate == 1, f"saturate={self.saturate} not implemented for CastLike"
+
+    def run(self, data: OpRunTensor, like: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(data.tensor.to(like.tensor.dtype))
+
+
+class Concat_1(OpRunKernel):
+    "Concat"
+
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None, verbose: int = 0):
+        super().__init__(node, version, verbose=verbose)
+        axis = self.get_attribute_int(node, "axis", 0)
+        assert isinstance(axis, int), f"Unexpected value for attribute axis={axis!r}"
+        self.axis = axis
+
+    def run(self, *data: OpRunTensor) -> OpRunTensor:
+        assert data, f"No tensor to concatenate in node name {self.name!r}"
+        devices = [d.get_device() for d in data]
+        if len(set(devices)) == 1:
+            return OpRunTensor(torch.cat([t.tensor for t in data], axis=self.axis))
+        if (
+            data[0].dtype == torch.int64
+            and self.axis == 0
+            and max(d.tensor.ndim for d in data) == 1
+            and max(d.tensor.numel() for d in data) <= 8
+        ):
+            # This is a shape
+            return OpRunTensor(torch.cat([t.tensor.cpu() for t in data], axis=self.axis))
+        index = devices.index(max(devices))
+        device = data[index].tensor.device
+        return OpRunTensor(torch.cat([t.tensor.to(device) for t in data], axis=self.axis))
+
+
+class NonZero_13(OpRunKernel):
+    "NonZero"
+
+    def run(self, x: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(torch.nonzero(x.tensor).T)
+
+
+class Tile_6(OpRunKernel):
+    "Tile"
+
+    def run(self, x: OpRunTensor, repeat: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(torch.tile(x.tensor, repeat.as_tuple_int))
+
+
+class Transpose_1(OpRunKernel):
+    "Transpose"
+
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None, verbose: int = 0):
+        super().__init__(node, version, verbose=verbose)
+        self.perm = self.get_attribute_ints(node, "perm", None)
+
+    def run(self, data: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(torch.permute(data.tensor, self.perm))
+
+
+class Trilu_14(OpRunKernel):
+    "Trilu"
+
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None, verbose: int = 0):
+        super().__init__(node, version, verbose=verbose)
+        self.upper = self.get_attribute_int(node, "upper", 1)
+
+    def run(self, data: OpRunTensor, k: Optional[OpRunTensor] = None) -> OpRunTensor:
+        diagonal = 0 if k is None else k.tensor.item()
+        if self.upper:
+            return OpRunTensor(torch.triu(data.tensor, diagonal=diagonal))
+        return OpRunTensor(torch.tril(data.tensor, diagonal=diagonal))
+
+
+class Where_9(OpRunKernel):
+    "Where"
+
+    def run(self, cond: OpRunTensor, x: OpRunTensor, y: OpRunTensor) -> OpRunTensor:
+        tcond, tx, ty = self.same_device(cond.tensor, x.tensor, y.tensor)
+        return OpRunTensor(torch.where(tcond, tx, ty))

onnx_diagnostic/reference/torch_ops/reduce_ops.py

@@ -0,0 +1,130 @@
+from typing import Optional, Tuple
+import onnx
+import torch
+from ...helpers.torch_helper import onnx_dtype_to_torch_dtype
+from . import OpRunKernel, OpRunTensor
+
+
+class ReduceOp(OpRunKernel):
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None, verbose: int = 0):
+        super().__init__(node, version, verbose=verbose)
+        self.keepdims = bool(self.get_attribute_int(node, "keepdims", 1))
+        self.noop_with_empty_axes = bool(
+            self.get_attribute_int(node, "noop_with_empty_axes", 0)
+        )
+        assert isinstance(
+            self.keepdims, bool
+        ), f"Unexpected value for attribute keepdims={self.keepdims!r}"
+        assert isinstance(self.noop_with_empty_axes, bool), (
+            f"Unexpected value for attribute "
+            f"noop_with_empty_axes={self.noop_with_empty_axes!r}"
+        )
+        assert (
+            not self.noop_with_empty_axes
+        ), f"Not implemented with noop_with_empty_axes={self.noop_with_empty_axes}"
+        # this is out of spec
+        stash_type = self.get_attribute_int(node, "stash_type", None)
+        self.stash_type = None if stash_type is None else onnx_dtype_to_torch_dtype(stash_type)
+
+
+class ReduceOpAxes(ReduceOp):
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None, verbose: int = 0):
+        super().__init__(node, version, verbose=verbose)
+        self.axes: Tuple[int, ...] = self.get_attribute_ints(node, "axes") or tuple()
+
+
+class ReduceMax_18(ReduceOp):
+    """ReduceMax"""
+
+    def run(self, x: OpRunTensor, axes: Optional[OpRunTensor] = None) -> OpRunTensor:
+        assert self.stash_type is None, f"Not implemented with stash_type={self.stash_type}"
+        if axes is None:
+            assert (
+                not self.keepdims
+            ), f"axes is Empty, keepdims={self.keepdims} for {self.__class__.__name__}"
+            return OpRunTensor(x.tensor.max())
+        taxes = axes.as_tuple_int
+        if len(taxes) == 1:
+            t = x.tensor.max(taxes[0], keepdim=self.keepdims)
+            return OpRunTensor(t.values)
+        t = x.tensor
+        for a in reversed(taxes):
+            t = t.max(a, keepdim=self.keepdims).values
+        return OpRunTensor(t)
+
+
+class ReduceMean_18(ReduceOp):
+    """ReduceMean"""
+
+    def run(self, x: OpRunTensor, axes: Optional[OpRunTensor] = None) -> OpRunTensor:
+        assert self.stash_type is None, f"Not implemented with stash_type={self.stash_type}"
+        if axes is None:
+            assert (
+                not self.keepdims
+            ), f"axes is Empty, keepdims={self.keepdims} for {self.__class__.__name__}"
+            return OpRunTensor(torch.mean(x.tensor))
+        taxes = axes.as_tuple_int
+        if len(taxes) == 1:
+            t = x.tensor.mean(taxes[0], keepdim=self.keepdims)
+            return OpRunTensor(t)
+        t = x.tensor.mean(taxes, keepdim=self.keepdims)
+        return OpRunTensor(t)
+
+
+class ReduceMin_17(ReduceOpAxes):
+    """ReduceMin"""
+
+    def run(self, x: OpRunTensor) -> OpRunTensor:
+        assert self.stash_type is None, f"Not implemented with stash_type={self.stash_type}"
+        axes = self.axes
+        if not axes:
+            assert (
+                not self.keepdims
+            ), f"axes is Empty, keepdims={self.keepdims} for {self.__class__.__name__}"
+            return OpRunTensor(x.tensor.min())
+        taxes = tuple(axes)
+        if len(taxes) == 1:
+            t = x.tensor.min(taxes[0], keepdim=self.keepdims)
+            return OpRunTensor(t.values)
+        t = x.tensor
+        for a in reversed(taxes):
+            t = t.min(a, keepdim=self.keepdims).values
+        return OpRunTensor(t)
+
+
+class ReduceMin_18(ReduceOp):
+    """ReduceMin"""
+
+    def run(self, x: OpRunTensor, axes: Optional[OpRunTensor] = None) -> OpRunTensor:
+        assert self.stash_type is None, f"Not implemented with stash_type={self.stash_type}"
+        if axes is None:
+            assert (
+                not self.keepdims
+            ), f"axes is empty, keepdims={self.keepdims} for {self.__class__.__name__}"
+            return OpRunTensor(torch.min(x.tensor))
+        taxes = axes.as_tuple_int
+        if len(taxes) == 1:
+            t = x.tensor.min(taxes[0], keepdim=self.keepdims)
+            return OpRunTensor(t.values)
+        t = x.tensor
+        for a in reversed(taxes):
+            t = t.min(a, keepdim=self.keepdims).values
+        return OpRunTensor(t)
+
+
+class ReduceSum_13(ReduceOp):
+    """ReduceSum"""
+
+    def run(self, x: OpRunTensor, axes: Optional[OpRunTensor] = None) -> OpRunTensor:
+        assert self.stash_type is None, f"Not implemented with stash_type={self.stash_type}"
+        if axes is None:
+            assert (
+                not self.keepdims
+            ), f"axes is Empty, keepdims={self.keepdims} for {self.__class__.__name__}"
+            return OpRunTensor(torch.sum(x.tensor))
+        taxes = axes.as_tuple_int
+        if len(taxes) == 1:
+            t = x.tensor.sum(taxes[0], keepdim=self.keepdims)
+            return OpRunTensor(t)
+        t = x.tensor.sum(taxes, keepdim=self.keepdims)
+        return OpRunTensor(t)

onnx_diagnostic/reference/torch_ops/sequence_ops.py

@@ -0,0 +1,65 @@
+from typing import Optional
+import onnx
+import torch
+from ...helpers.torch_helper import onnx_dtype_to_torch_dtype
+from . import OpRunKernel, OpRunSequence, OpRunTensor
+
+
+class OpRunOpSequence(OpRunKernel):
+    "Ancestor for kernel using sequences."
+
+
+class ConcatFromSequence_11(OpRunOpSequence):
+    "ConcatFromSequence"
+
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None, verbose: int = 0):
+        super().__init__(node, version, verbose=verbose)
+        axis = self.get_attribute_int(node, "axis", None)
+        assert isinstance(axis, int), f"Unexpected value for attribute axis={axis!r}"
+        self.axis = axis
+        self.new_axis = self.get_attribute_int(node, "new_axis", 0)
+
+    def run(self, input_sequence: OpRunSequence) -> OpRunTensor:
+        assert isinstance(
+            input_sequence, OpRunSequence
+        ), f"Unexpected type {type(input_sequence)} for input_sequence"
+        seq = input_sequence.sequence
+        if self.new_axis == 1:
+            if self.axis == -1:
+                seq2 = [s.unsqueeze(len(s.shape)) for s in seq]
+                res = torch.cat(seq2, axis=-1)
+            else:
+                seq2 = [s.expand(self.axis) for s in seq]
+                res = torch.cat(seq2, axis=self.axis)
+        else:
+            res = torch.cat(seq, axis=self.axis)
+        return OpRunTensor(res)
+
+
+class SequenceEmpty_11(OpRunOpSequence):
+    "SqeuenceEmpty"
+
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None, verbose: int = 0):
+        super().__init__(node, version, verbose=verbose)
+        self.dtype = onnx_dtype_to_torch_dtype(
+            self.get_attribute_int(node, "dtype", onnx.TensorProto.FLOAT)  # type: ignore[arg-type]
+        )
+
+    def run(self) -> OpRunSequence:
+        return OpRunSequence(dtype=self.dtype)
+
+
+class SequenceInsert_11(OpRunOpSequence):
+    "SqeuenceInsert"
+
+    def run(
+        self,
+        input_sequence: OpRunSequence,
+        tensor: OpRunTensor,
+        position: Optional[OpRunTensor] = None,
+    ) -> OpRunSequence:
+        assert isinstance(input_sequence, OpRunSequence), (
+            f"Unexpected type {type(input_sequence)} for input_sequence: "
+            f"{input_sequence.string_type()}"
+        )
+        return input_sequence.insert_at(tensor, position)

onnx_diagnostic/reference/torch_ops/shape_ops.py

@@ -0,0 +1,121 @@
+from typing import Optional, Tuple
+import onnx
+import torch
+from . import OpRunKernel, OpRunTensor
+
+
+class ConstantOfShape_9(OpRunKernel):
+    "ConstantOfShape"
+
+    @classmethod
+    def device_dependent(cls) -> bool:
+        """
+        Returns True if the kernel needs a device to be efficiently initialized.
+        """
+        return True
+
+    def __init__(
+        self,
+        node: onnx.NodeProto,
+        version: Optional[int] = None,
+        device: Optional[torch.device] = None,
+        verbose: int = 0,
+    ):
+        super().__init__(node, version, verbose=verbose)
+        value = self.get_attribute_tensor(node, "value")
+        if value is None:
+            value = torch.tensor([0], dtype=torch.float32)
+        self.dtype = value.dtype
+        self.device = device
+        self.value = value[0]
+
+    def run(self, shape: OpRunTensor) -> OpRunTensor:
+        # The device is unknown as shapes usually take place on CPU.
+        return OpRunTensor(
+            torch.full(
+                shape.as_tuple_int, fill_value=self.value, dtype=self.dtype, device=self.device
+            )
+        )
+
+
+class Expand_8(OpRunKernel):
+    "Expand"
+
+    def run(self, data: OpRunTensor, shape: OpRunTensor) -> OpRunTensor:
+        ishape = tuple(-1 if i == 1 else i for i in shape.as_tuple_int)
+        return OpRunTensor(data.tensor.expand(ishape))
+
+
+class Reshape_14(OpRunKernel):
+    "Reshape"
+
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None, verbose: int = 0):
+        super().__init__(node, version, verbose=verbose)
+        self.allowzero = self.get_attribute_int(node, "allowzero", 0)
+
+    def run(self, data: OpRunTensor, shape: OpRunTensor) -> OpRunTensor:
+        ishape = shape.as_tuple_int
+        assert ishape is not None, f"Unexpected return for shape={shape!r}"
+        if not self.allowzero and 0 in ishape:
+            xshape = data.tensor.shape
+            new_shape = []
+            for i, s in enumerate(ishape):
+                new_shape.append(xshape[i] if s == 0 else s)
+            return OpRunTensor(data.tensor.reshape(new_shape))
+        return OpRunTensor(data.tensor.reshape(ishape))
+
+
+class Shape_15(OpRunKernel):
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None, verbose: int = 0):
+        super().__init__(node, version, verbose=verbose)
+        self.start = self.get_attribute_int(node, "start", 0)
+        self.end = self.get_attribute_int(node, "end", None)
+
+    def run(self, data: OpRunTensor) -> OpRunTensor:
+        shape = data.shape
+        sh = shape[self.start :] if self.end is None else shape[self.start : self.end]
+        return OpRunTensor(torch.tensor(sh, dtype=torch.int64), is_constant=True)
+
+
+class Split_18(OpRunKernel):
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None, verbose: int = 0):
+        super().__init__(node, version, verbose=verbose)
+        self.axis = self.get_attribute_int(node, "axis", 0)
+        self.num_outputs = self.get_attribute_int(node, "num_outputs", None)
+
+    def run(
+        self, data: OpRunTensor, split: Optional[OpRunTensor] = None
+    ) -> Tuple[OpRunTensor, ...]:
+        if split is None:
+            assert isinstance(
+                self.num_outputs, int
+            ), f"Incompatibilities: split is None and num_outputs={self.num_outputs}"
+            size = data.tensor.shape[self.axis]
+            split_size = (
+                size // self.num_outputs
+                if size % self.num_outputs == 0
+                else size // self.num_outputs + 1
+            )
+            spl = torch.split(data.tensor, split_size, dim=self.axis)
+        else:
+            spl = torch.split(data.tensor, split.as_tuple_int, dim=self.axis)
+        return tuple(OpRunTensor(t) for t in spl)
+
+
+class Squeeze_13(OpRunKernel):
+    "Squeeze"
+
+    def run(self, data: OpRunTensor, axes: Optional[OpRunTensor] = None) -> OpRunTensor:
+        if axes is None:
+            return OpRunTensor(data.tensor.squeeze())
+        return OpRunTensor(data.tensor.squeeze(axes.as_tuple_int))
+
+
+class Unsqueeze_13(OpRunKernel):
+    "Unsqueeze"
+
+    def run(self, data: OpRunTensor, axes: OpRunTensor) -> OpRunTensor:
+        t = data.tensor
+        for i in axes.as_tuple_int:
+            t = t.unsqueeze(i)
+        return OpRunTensor(t)