emx-onnx-cgen 0.2.0__py3-none-any.whl

emx_onnx_cgen/lowering/lstm.py

@@ -0,0 +1,355 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Iterable, Sequence
+
+from shared.scalar_types import ScalarType
+
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Node
+from .common import node_dtype, optional_name, value_dtype, value_shape
+from .registry import register_lowering
+
+ACTIVATION_KIND_BY_NAME = {
+    "Relu": 0,
+    "Tanh": 1,
+    "Sigmoid": 2,
+    "Affine": 3,
+    "LeakyRelu": 4,
+    "ThresholdedRelu": 5,
+    "ScaledTanh": 6,
+    "HardSigmoid": 7,
+    "Elu": 8,
+    "Softsign": 9,
+    "Softplus": 10,
+}
+
+DEFAULT_ACTIVATIONS = ("Sigmoid", "Tanh", "Tanh")
+
+DEFAULT_ALPHA_BY_NAME = {
+    "Affine": 1.0,
+    "LeakyRelu": 0.01,
+    "ThresholdedRelu": 1.0,
+    "ScaledTanh": 1.0,
+    "HardSigmoid": 0.2,
+    "Elu": 1.0,
+}
+
+DEFAULT_BETA_BY_NAME = {
+    "Affine": 0.0,
+    "ScaledTanh": 1.0,
+    "HardSigmoid": 0.5,
+}
+
+
+@dataclass(frozen=True)
+class LstmSpec:
+    input_x: str
+    input_w: str
+    input_r: str
+    input_b: str | None
+    input_sequence_lens: str | None
+    input_initial_h: str | None
+    input_initial_c: str | None
+    input_p: str | None
+    output_y: str | None
+    output_y_h: str | None
+    output_y_c: str | None
+    seq_length: int
+    batch_size: int
+    input_size: int
+    hidden_size: int
+    num_directions: int
+    direction: str
+    layout: int
+    input_forget: int
+    clip: float | None
+    activation_kinds: tuple[int, ...]
+    activation_alphas: tuple[float, ...]
+    activation_betas: tuple[float, ...]
+    dtype: ScalarType
+    sequence_lens_dtype: ScalarType | None
+
+
+def _normalize_activation_names(values: Iterable[object]) -> list[str]:
+    names: list[str] = []
+    for value in values:
+        if isinstance(value, bytes):
+            value = value.decode("utf-8")
+        if not isinstance(value, str):
+            raise UnsupportedOpError("LSTM activations must be strings")
+        names.append(value)
+    return names
+
+
+def _resolve_activation_params(
+    activations: Sequence[str],
+    activation_alpha: Sequence[float] | None,
+    activation_beta: Sequence[float] | None,
+) -> tuple[tuple[int, ...], tuple[float, ...], tuple[float, ...]]:
+    if activation_alpha is None:
+        activation_alpha = []
+    if activation_beta is None:
+        activation_beta = []
+    if activation_alpha and len(activation_alpha) != len(activations):
+        raise UnsupportedOpError("LSTM activation_alpha must match activations")
+    if activation_beta and len(activation_beta) != len(activations):
+        raise UnsupportedOpError("LSTM activation_beta must match activations")
+    activation_kinds: list[int] = []
+    alphas: list[float] = []
+    betas: list[float] = []
+    for idx, name in enumerate(activations):
+        kind = ACTIVATION_KIND_BY_NAME.get(name)
+        if kind is None:
+            raise UnsupportedOpError(f"Unsupported LSTM activation {name}")
+        activation_kinds.append(kind)
+        if activation_alpha:
+            alpha = float(activation_alpha[idx])
+        else:
+            alpha = DEFAULT_ALPHA_BY_NAME.get(name, 1.0)
+        if activation_beta:
+            beta = float(activation_beta[idx])
+        else:
+            beta = DEFAULT_BETA_BY_NAME.get(name, 0.0)
+        alphas.append(alpha)
+        betas.append(beta)
+    return tuple(activation_kinds), tuple(alphas), tuple(betas)
+
+
+def _resolve_activations(
+    direction: str, num_directions: int, attrs: dict[str, object]
+) -> tuple[tuple[int, ...], tuple[float, ...], tuple[float, ...]]:
+    activations_attr = attrs.get("activations")
+    if activations_attr is None:
+        activations = list(DEFAULT_ACTIVATIONS)
+    else:
+        activations = _normalize_activation_names(activations_attr)
+    if num_directions == 1:
+        if len(activations) != 3:
+            raise UnsupportedOpError("LSTM activations must have length 3")
+    else:
+        if len(activations) == 3:
+            activations = activations * 2
+        elif len(activations) != 6:
+            raise UnsupportedOpError("Bidirectional LSTM activations must be length 6")
+    activation_alpha = attrs.get("activation_alpha")
+    activation_beta = attrs.get("activation_beta")
+    return _resolve_activation_params(
+        activations,
+        activation_alpha,
+        activation_beta,
+    )
+
+
+def _expect_shape(
+    name: str, shape: tuple[int, ...], expected: tuple[int, ...]
+) -> None:
+    if shape != expected:
+        raise UnsupportedOpError(
+            f"LSTM input {name} must have shape {expected}, got {shape}"
+        )
+
+
+def _validate_direction(direction: str, num_directions: int) -> None:
+    if direction == "bidirectional" and num_directions != 2:
+        raise UnsupportedOpError(
+            "LSTM expects num_directions=2 for bidirectional models"
+        )
+    if direction in {"forward", "reverse"} and num_directions != 1:
+        raise UnsupportedOpError(
+            "LSTM expects num_directions=1 for forward/reverse models"
+        )
+    if direction not in {"forward", "reverse", "bidirectional"}:
+        raise UnsupportedOpError(f"Unsupported LSTM direction {direction}")
+
+
+def resolve_lstm_spec(graph: Graph, node: Node) -> LstmSpec:
+    if len(node.inputs) < 3 or len(node.inputs) > 8:
+        raise UnsupportedOpError("LSTM expects between 3 and 8 inputs")
+    if len(node.outputs) < 1 or len(node.outputs) > 3:
+        raise UnsupportedOpError("LSTM expects between 1 and 3 outputs")
+    input_x = node.inputs[0]
+    input_w = node.inputs[1]
+    input_r = node.inputs[2]
+    input_b = optional_name(node.inputs, 3)
+    input_sequence_lens = optional_name(node.inputs, 4)
+    input_initial_h = optional_name(node.inputs, 5)
+    input_initial_c = optional_name(node.inputs, 6)
+    input_p = optional_name(node.inputs, 7)
+    output_y = optional_name(node.outputs, 0)
+    output_y_h = optional_name(node.outputs, 1)
+    output_y_c = optional_name(node.outputs, 2)
+    if output_y is None and output_y_h is None and output_y_c is None:
+        raise UnsupportedOpError("LSTM expects at least one output")
+    op_dtype = node_dtype(
+        graph,
+        node,
+        input_x,
+        input_w,
+        input_r,
+        *(name for name in (input_b, input_initial_h, input_initial_c, input_p) if name),
+        *(name for name in (output_y, output_y_h, output_y_c) if name),
+    )
+    if not op_dtype.is_float:
+        raise UnsupportedOpError(
+            "LSTM supports float16, float, and double inputs only"
+        )
+    x_shape = value_shape(graph, input_x, node)
+    if len(x_shape) != 3:
+        raise UnsupportedOpError("LSTM input X must be rank 3")
+    layout = int(node.attrs.get("layout", 0))
+    if layout not in {0, 1}:
+        raise UnsupportedOpError("LSTM layout must be 0 or 1")
+    if layout == 0:
+        seq_length, batch_size, input_size = x_shape
+    else:
+        batch_size, seq_length, input_size = x_shape
+    w_shape = value_shape(graph, input_w, node)
+    if len(w_shape) != 3:
+        raise UnsupportedOpError("LSTM input W must be rank 3")
+    num_directions = w_shape[0]
+    hidden_size_attr = node.attrs.get("hidden_size")
+    if hidden_size_attr is None:
+        if w_shape[1] % 4 != 0:
+            raise UnsupportedOpError("LSTM W shape is not divisible by 4")
+        hidden_size = w_shape[1] // 4
+    else:
+        hidden_size = int(hidden_size_attr)
+    _validate_direction(str(node.attrs.get("direction", "forward")), num_directions)
+    direction = str(node.attrs.get("direction", "forward"))
+    expected_w_shape = (num_directions, 4 * hidden_size, input_size)
+    _expect_shape(input_w, w_shape, expected_w_shape)
+    r_shape = value_shape(graph, input_r, node)
+    expected_r_shape = (num_directions, 4 * hidden_size, hidden_size)
+    _expect_shape(input_r, r_shape, expected_r_shape)
+    if input_b is not None:
+        b_shape = value_shape(graph, input_b, node)
+        _expect_shape(input_b, b_shape, (num_directions, 8 * hidden_size))
+    if input_sequence_lens is not None:
+        seq_dtype = value_dtype(graph, input_sequence_lens, node)
+        if seq_dtype not in {ScalarType.I32, ScalarType.I64}:
+            raise UnsupportedOpError("LSTM sequence_lens must be int32 or int64")
+        seq_shape = value_shape(graph, input_sequence_lens, node)
+        if seq_shape != (batch_size,):
+            raise UnsupportedOpError(
+                "LSTM sequence_lens must match batch size"
+            )
+    state_shape = (
+        (num_directions, batch_size, hidden_size)
+        if layout == 0
+        else (batch_size, num_directions, hidden_size)
+    )
+    if input_initial_h is not None:
+        _expect_shape(
+            input_initial_h,
+            value_shape(graph, input_initial_h, node),
+            state_shape,
+        )
+    if input_initial_c is not None:
+        _expect_shape(
+            input_initial_c,
+            value_shape(graph, input_initial_c, node),
+            state_shape,
+        )
+    if input_p is not None:
+        _expect_shape(
+            input_p,
+            value_shape(graph, input_p, node),
+            (num_directions, 3 * hidden_size),
+        )
+    if output_y is not None:
+        expected_y_shape = (
+            (seq_length, num_directions, batch_size, hidden_size)
+            if layout == 0
+            else (batch_size, seq_length, num_directions, hidden_size)
+        )
+        _expect_shape(output_y, value_shape(graph, output_y, node), expected_y_shape)
+    if output_y_h is not None:
+        _expect_shape(
+            output_y_h,
+            value_shape(graph, output_y_h, node),
+            state_shape,
+        )
+    if output_y_c is not None:
+        _expect_shape(
+            output_y_c,
+            value_shape(graph, output_y_c, node),
+            state_shape,
+        )
+    input_forget = int(node.attrs.get("input_forget", 0))
+    if input_forget not in {0, 1}:
+        raise UnsupportedOpError("LSTM input_forget must be 0 or 1")
+    clip = node.attrs.get("clip")
+    if clip is not None:
+        clip = float(clip)
+        if clip < 0:
+            raise UnsupportedOpError("LSTM clip must be non-negative")
+    activation_kinds, activation_alphas, activation_betas = _resolve_activations(
+        direction, num_directions, node.attrs
+    )
+    sequence_lens_dtype = (
+        value_dtype(graph, input_sequence_lens, node)
+        if input_sequence_lens is not None
+        else None
+    )
+    return LstmSpec(
+        input_x=input_x,
+        input_w=input_w,
+        input_r=input_r,
+        input_b=input_b,
+        input_sequence_lens=input_sequence_lens,
+        input_initial_h=input_initial_h,
+        input_initial_c=input_initial_c,
+        input_p=input_p,
+        output_y=output_y,
+        output_y_h=output_y_h,
+        output_y_c=output_y_c,
+        seq_length=seq_length,
+        batch_size=batch_size,
+        input_size=input_size,
+        hidden_size=hidden_size,
+        num_directions=num_directions,
+        direction=direction,
+        layout=layout,
+        input_forget=input_forget,
+        clip=clip,
+        activation_kinds=activation_kinds,
+        activation_alphas=activation_alphas,
+        activation_betas=activation_betas,
+        dtype=op_dtype,
+        sequence_lens_dtype=sequence_lens_dtype,
+    )
+
+
+@register_lowering("LSTM")
+def lower_lstm(graph: Graph, node: Node) -> "LstmOp":
+    from ..codegen.c_emitter import LstmOp
+
+    spec = resolve_lstm_spec(graph, node)
+    return LstmOp(
+        input_x=spec.input_x,
+        input_w=spec.input_w,
+        input_r=spec.input_r,
+        input_b=spec.input_b,
+        input_sequence_lens=spec.input_sequence_lens,
+        input_initial_h=spec.input_initial_h,
+        input_initial_c=spec.input_initial_c,
+        input_p=spec.input_p,
+        output_y=spec.output_y,
+        output_y_h=spec.output_y_h,
+        output_y_c=spec.output_y_c,
+        seq_length=spec.seq_length,
+        batch_size=spec.batch_size,
+        input_size=spec.input_size,
+        hidden_size=spec.hidden_size,
+        num_directions=spec.num_directions,
+        direction=spec.direction,
+        layout=spec.layout,
+        input_forget=spec.input_forget,
+        clip=spec.clip,
+        activation_kinds=spec.activation_kinds,
+        activation_alphas=spec.activation_alphas,
+        activation_betas=spec.activation_betas,
+        dtype=spec.dtype,
+        sequence_lens_dtype=spec.sequence_lens_dtype,
+    )

emx_onnx_cgen/lowering/matmul.py

@@ -0,0 +1,120 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from ..codegen.c_emitter import MatMulOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Node
+from .common import node_dtype as _node_dtype
+from .common import value_shape as _value_shape
+from .registry import register_lowering
+
+
+@dataclass(frozen=True)
+class MatMulSpec:
+    input0_shape: tuple[int, ...]
+    input1_shape: tuple[int, ...]
+    output_shape: tuple[int, ...]
+    batch_shape: tuple[int, ...]
+    input0_batch_shape: tuple[int, ...]
+    input1_batch_shape: tuple[int, ...]
+    m: int
+    n: int
+    k: int
+    left_vector: bool
+    right_vector: bool
+
+
+def resolve_matmul_spec(graph: Graph, node: Node) -> MatMulSpec:
+    if len(node.inputs) != 2 or len(node.outputs) != 1:
+        raise UnsupportedOpError("MatMul must have 2 inputs and 1 output")
+    input0_shape = _value_shape(graph, node.inputs[0], node)
+    input1_shape = _value_shape(graph, node.inputs[1], node)
+    if len(input0_shape) < 1 or len(input1_shape) < 1:
+        raise UnsupportedOpError(
+            "MatMul inputs must be at least 1D, "
+            f"got {input0_shape} x {input1_shape}"
+        )
+    left_vector = len(input0_shape) == 1
+    right_vector = len(input1_shape) == 1
+    input0_effective = (1, input0_shape[0]) if left_vector else input0_shape
+    input1_effective = (input1_shape[0], 1) if right_vector else input1_shape
+    m, k_left = input0_effective[-2], input0_effective[-1]
+    k_right, n = input1_effective[-2], input1_effective[-1]
+    if k_left != k_right:
+        raise ShapeInferenceError(
+            f"MatMul inner dimensions must match, got {k_left} and {k_right}"
+        )
+    batch_shape, input0_batch_shape, input1_batch_shape = (
+        _broadcast_batch_shapes(
+            input0_effective[:-2], input1_effective[:-2], node
+        )
+    )
+    if left_vector and right_vector:
+        output_shape = batch_shape
+    elif left_vector:
+        output_shape = batch_shape + (n,)
+    elif right_vector:
+        output_shape = batch_shape + (m,)
+    else:
+        output_shape = batch_shape + (m, n)
+    expected_output_shape = _value_shape(graph, node.outputs[0], node)
+    if expected_output_shape != output_shape:
+        raise ShapeInferenceError(
+            "MatMul output shape must be "
+            f"{output_shape}, got {expected_output_shape}"
+        )
+    return MatMulSpec(
+        input0_shape=input0_shape,
+        input1_shape=input1_shape,
+        output_shape=output_shape,
+        batch_shape=batch_shape,
+        input0_batch_shape=input0_batch_shape,
+        input1_batch_shape=input1_batch_shape,
+        m=m,
+        n=n,
+        k=k_left,
+        left_vector=left_vector,
+        right_vector=right_vector,
+    )
+
+
+def _broadcast_batch_shapes(
+    left: tuple[int, ...], right: tuple[int, ...], node: Node
+) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
+    max_rank = max(len(left), len(right))
+    left_padded = (1,) * (max_rank - len(left)) + left
+    right_padded = (1,) * (max_rank - len(right)) + right
+    broadcast_shape = []
+    for left_dim, right_dim in zip(left_padded, right_padded):
+        if left_dim == right_dim or left_dim == 1 or right_dim == 1:
+            broadcast_shape.append(max(left_dim, right_dim))
+            continue
+        raise ShapeInferenceError(
+            "MatMul batch dimensions must be broadcastable, "
+            f"got {left} x {right}"
+        )
+    return tuple(broadcast_shape), left_padded, right_padded
+
+
+@register_lowering("MatMul")
+def lower_matmul(graph: Graph, node: Node) -> MatMulOp:
+    op_dtype = _node_dtype(graph, node, *node.inputs, *node.outputs)
+    spec = resolve_matmul_spec(graph, node)
+    return MatMulOp(
+        input0=node.inputs[0],
+        input1=node.inputs[1],
+        output=node.outputs[0],
+        input0_shape=spec.input0_shape,
+        input1_shape=spec.input1_shape,
+        output_shape=spec.output_shape,
+        batch_shape=spec.batch_shape,
+        input0_batch_shape=spec.input0_batch_shape,
+        input1_batch_shape=spec.input1_batch_shape,
+        m=spec.m,
+        n=spec.n,
+        k=spec.k,
+        left_vector=spec.left_vector,
+        right_vector=spec.right_vector,
+        dtype=op_dtype,
+    )

emx_onnx_cgen/lowering/maxpool.py

@@ -0,0 +1,195 @@
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+
+from shared.scalar_types import ScalarType
+
+from ..codegen.c_emitter import MaxPoolOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Node
+from .common import node_dtype as _node_dtype
+from .common import value_dtype as _value_dtype
+from .common import value_shape as _value_shape
+from .registry import register_lowering
+
+
+@dataclass(frozen=True)
+class MaxPoolSpec:
+    batch: int
+    channels: int
+    spatial_rank: int
+    in_spatial: tuple[int, ...]
+    out_spatial: tuple[int, ...]
+    kernel_shape: tuple[int, ...]
+    strides: tuple[int, ...]
+    pads: tuple[int, ...]
+    dilations: tuple[int, ...]
+    ceil_mode: bool
+    storage_order: int
+
+
+def resolve_maxpool_spec(graph: Graph, node: Node) -> MaxPoolSpec:
+    if len(node.inputs) != 1 or len(node.outputs) not in {1, 2}:
+        raise UnsupportedOpError("MaxPool must have 1 input and 1 or 2 outputs")
+    supported_attrs = {
+        "auto_pad",
+        "ceil_mode",
+        "dilations",
+        "kernel_shape",
+        "pads",
+        "storage_order",
+        "strides",
+    }
+    if set(node.attrs) - supported_attrs:
+        raise UnsupportedOpError("MaxPool has unsupported attributes")
+    storage_order = int(node.attrs.get("storage_order", 0))
+    if storage_order not in (0, 1):
+        raise UnsupportedOpError("MaxPool supports storage_order=0 or 1 only")
+    kernel_shape = node.attrs.get("kernel_shape")
+    if kernel_shape is None:
+        raise UnsupportedOpError("MaxPool requires kernel_shape")
+    kernel_shape = tuple(int(value) for value in kernel_shape)
+    input_shape = _value_shape(graph, node.inputs[0], node)
+    if len(input_shape) < 3:
+        raise UnsupportedOpError("MaxPool expects NCHW inputs with spatial dims")
+    spatial_rank = len(input_shape) - 2
+    if spatial_rank not in {1, 2, 3}:
+        raise UnsupportedOpError("MaxPool supports 1D/2D/3D inputs only")
+    if len(kernel_shape) != spatial_rank:
+        raise ShapeInferenceError(
+            f"MaxPool kernel_shape must have {spatial_rank} dims, got {kernel_shape}"
+        )
+    strides = tuple(
+        int(value) for value in node.attrs.get("strides", (1,) * spatial_rank)
+    )
+    if len(strides) != spatial_rank:
+        raise UnsupportedOpError("MaxPool stride rank mismatch")
+    dilations = tuple(
+        int(value) for value in node.attrs.get("dilations", (1,) * spatial_rank)
+    )
+    if len(dilations) != spatial_rank:
+        raise UnsupportedOpError("MaxPool dilation rank mismatch")
+    pads = tuple(
+        int(value)
+        for value in node.attrs.get("pads", (0,) * (2 * spatial_rank))
+    )
+    if len(pads) != 2 * spatial_rank:
+        raise UnsupportedOpError("MaxPool pads rank mismatch")
+    auto_pad = node.attrs.get("auto_pad", b"NOTSET")
+    if isinstance(auto_pad, bytes):
+        auto_pad = auto_pad.decode("utf-8", errors="ignore")
+    if auto_pad in ("", "NOTSET"):
+        pad_begin = pads[:spatial_rank]
+        pad_end = pads[spatial_rank:]
+    elif auto_pad == "VALID":
+        pad_begin = (0,) * spatial_rank
+        pad_end = (0,) * spatial_rank
+    elif auto_pad in {"SAME_UPPER", "SAME_LOWER"}:
+        pad_begin = []
+        pad_end = []
+        for dim, stride, dilation, kernel in zip(
+            input_shape[2:], strides, dilations, kernel_shape
+        ):
+            effective_kernel = dilation * (kernel - 1) + 1
+            out_dim = math.ceil(dim / stride)
+            pad_needed = max(
+                0, (out_dim - 1) * stride + effective_kernel - dim
+            )
+            if auto_pad == "SAME_UPPER":
+                pad_start = pad_needed // 2
+            else:
+                pad_start = (pad_needed + 1) // 2
+            pad_begin.append(pad_start)
+            pad_end.append(pad_needed - pad_start)
+        pad_begin = tuple(pad_begin)
+        pad_end = tuple(pad_end)
+    else:
+        raise UnsupportedOpError("MaxPool has unsupported auto_pad mode")
+    ceil_mode = int(node.attrs.get("ceil_mode", 0))
+    if ceil_mode not in (0, 1):
+        raise UnsupportedOpError("MaxPool supports ceil_mode=0 or 1 only")
+    batch, channels = input_shape[0], input_shape[1]
+    in_spatial = input_shape[2:]
+    out_spatial = []
+    for dim, stride, dilation, kernel, pad_start, pad_finish in zip(
+        in_spatial, strides, dilations, kernel_shape, pad_begin, pad_end
+    ):
+        effective_kernel = dilation * (kernel - 1) + 1
+        numerator = dim + pad_start + pad_finish - effective_kernel
+        if ceil_mode:
+            out_dim = (numerator + stride - 1) // stride + 1
+            if (out_dim - 1) * stride >= dim + pad_start:
+                out_dim -= 1
+        else:
+            out_dim = numerator // stride + 1
+        if out_dim < 0:
+            raise ShapeInferenceError(
+                "MaxPool output shape must be non-negative"
+            )
+        out_spatial.append(out_dim)
+    expected_output_shape = (batch, channels, *out_spatial)
+    output_shape = _value_shape(graph, node.outputs[0], node)
+    if output_shape != expected_output_shape:
+        raise ShapeInferenceError(
+            "MaxPool output shape must be "
+            f"{expected_output_shape}, got {output_shape}"
+        )
+    if len(node.outputs) == 2:
+        indices_shape = _value_shape(graph, node.outputs[1], node)
+        if indices_shape != expected_output_shape:
+            raise ShapeInferenceError(
+                "MaxPool indices output shape must be "
+                f"{expected_output_shape}, got {indices_shape}"
+            )
+        indices_dtype = _value_dtype(graph, node.outputs[1], node)
+        if indices_dtype != ScalarType.I64:
+            raise UnsupportedOpError("MaxPool indices output must be int64")
+    pads = (*pad_begin, *pad_end)
+    return MaxPoolSpec(
+        batch=batch,
+        channels=channels,
+        spatial_rank=spatial_rank,
+        in_spatial=in_spatial,
+        out_spatial=tuple(out_spatial),
+        kernel_shape=kernel_shape,
+        strides=strides,
+        pads=pads,
+        dilations=dilations,
+        ceil_mode=bool(ceil_mode),
+        storage_order=storage_order,
+    )
+
+
+@register_lowering("MaxPool")
+def lower_maxpool(graph: Graph, node: Node) -> MaxPoolOp:
+    if len(node.inputs) != 1 or len(node.outputs) not in {1, 2}:
+        raise UnsupportedOpError("MaxPool must have 1 input and 1 or 2 outputs")
+    op_dtype = _node_dtype(graph, node, node.inputs[0], node.outputs[0])
+    if op_dtype == ScalarType.BOOL:
+        raise UnsupportedOpError("MaxPool supports numeric inputs only")
+    spec = resolve_maxpool_spec(graph, node)
+    indices = node.outputs[1] if len(node.outputs) == 2 else None
+    indices_dtype = (
+        _value_dtype(graph, indices, node) if indices is not None else None
+    )
+    if indices_dtype is not None and indices_dtype != ScalarType.I64:
+        raise UnsupportedOpError("MaxPool indices output must be int64")
+    return MaxPoolOp(
+        input0=node.inputs[0],
+        output=node.outputs[0],
+        indices=indices,
+        batch=spec.batch,
+        channels=spec.channels,
+        spatial_rank=spec.spatial_rank,
+        in_spatial=spec.in_spatial,
+        out_spatial=spec.out_spatial,
+        kernel_shape=spec.kernel_shape,
+        strides=spec.strides,
+        pads=spec.pads,
+        dilations=spec.dilations,
+        ceil_mode=spec.ceil_mode,
+        storage_order=spec.storage_order,
+        dtype=op_dtype,
+        indices_dtype=indices_dtype,
+    )