emx-onnx-cgen 0.3.8__py3-none-any.whl → 0.4.1.dev0__py3-none-any.whl

emx_onnx_cgen/lowering/gru.py

@@ -0,0 +1,323 @@
+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")
+
+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 GruSpec:
+    input_x: str
+    input_w: str
+    input_r: str
+    input_b: str | None
+    input_sequence_lens: str | None
+    input_initial_h: str | None
+    output_y: str | None
+    output_y_h: str | None
+    seq_length: int
+    batch_size: int
+    input_size: int
+    hidden_size: int
+    num_directions: int
+    direction: str
+    layout: int
+    linear_before_reset: 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("GRU 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("GRU activation_alpha must match activations")
+    if activation_beta and len(activation_beta) != len(activations):
+        raise UnsupportedOpError("GRU 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 GRU 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(
+    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) != 2:
+            raise UnsupportedOpError("GRU activations must have length 2")
+    else:
+        if len(activations) == 2:
+            activations = activations * 2
+        elif len(activations) != 4:
+            raise UnsupportedOpError("Bidirectional GRU activations must be length 4")
+    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"GRU 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(
+            "GRU expects num_directions=2 for bidirectional models"
+        )
+    if direction in {"forward", "reverse"} and num_directions != 1:
+        raise UnsupportedOpError(
+            "GRU expects num_directions=1 for forward/reverse models"
+        )
+    if direction not in {"forward", "reverse", "bidirectional"}:
+        raise UnsupportedOpError(f"Unsupported GRU direction {direction}")
+
+
+def resolve_gru_spec(graph: Graph, node: Node) -> GruSpec:
+    if len(node.inputs) < 3 or len(node.inputs) > 6:
+        raise UnsupportedOpError("GRU expects between 3 and 6 inputs")
+    if len(node.outputs) < 1 or len(node.outputs) > 2:
+        raise UnsupportedOpError("GRU expects between 1 and 2 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)
+    output_y = optional_name(node.outputs, 0)
+    output_y_h = optional_name(node.outputs, 1)
+    if output_y is None and output_y_h is None:
+        raise UnsupportedOpError("GRU 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) if name),
+        *(name for name in (output_y, output_y_h) if name),
+    )
+    if not op_dtype.is_float:
+        raise UnsupportedOpError(
+            "GRU supports float16, float, and double inputs only"
+        )
+    x_shape = value_shape(graph, input_x, node)
+    if len(x_shape) != 3:
+        raise UnsupportedOpError("GRU input X must be rank 3")
+    layout = int(node.attrs.get("layout", 0))
+    if layout not in {0, 1}:
+        raise UnsupportedOpError("GRU 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("GRU 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] % 3 != 0:
+            raise UnsupportedOpError("GRU W shape is not divisible by 3")
+        hidden_size = w_shape[1] // 3
+    else:
+        hidden_size = int(hidden_size_attr)
+    direction = str(node.attrs.get("direction", "forward"))
+    _validate_direction(direction, num_directions)
+    expected_w_shape = (num_directions, 3 * 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, 3 * 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, 6 * 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("GRU sequence_lens must be int32 or int64")
+        seq_shape = value_shape(graph, input_sequence_lens, node)
+        if seq_shape != (batch_size,):
+            raise UnsupportedOpError("GRU 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 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,
+        )
+    linear_before_reset = int(node.attrs.get("linear_before_reset", 0))
+    if linear_before_reset not in {0, 1}:
+        raise UnsupportedOpError("GRU linear_before_reset must be 0 or 1")
+    clip = node.attrs.get("clip")
+    if clip is not None:
+        clip = float(clip)
+        if clip < 0:
+            raise UnsupportedOpError("GRU clip must be non-negative")
+    activation_kinds, activation_alphas, activation_betas = _resolve_activations(
+        num_directions, node.attrs
+    )
+    sequence_lens_dtype = (
+        value_dtype(graph, input_sequence_lens, node)
+        if input_sequence_lens is not None
+        else None
+    )
+    return GruSpec(
+        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,
+        output_y=output_y,
+        output_y_h=output_y_h,
+        seq_length=seq_length,
+        batch_size=batch_size,
+        input_size=input_size,
+        hidden_size=hidden_size,
+        num_directions=num_directions,
+        direction=direction,
+        layout=layout,
+        linear_before_reset=linear_before_reset,
+        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("GRU")
+def lower_gru(graph: Graph, node: Node) -> "GruOp":
+    from ..ir.ops import GruOp
+
+    spec = resolve_gru_spec(graph, node)
+    return GruOp(
+        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,
+        output_y=spec.output_y,
+        output_y_h=spec.output_y_h,
+        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,
+        linear_before_reset=spec.linear_before_reset,
+        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/hamming_window.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+import numpy as np
+
+from shared.scalar_types import ScalarType
+
+from ..dtypes import dtype_info
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Initializer, Node
+from ..ir.ops import HammingWindowOp
+from ..lowering.common import value_dtype, value_shape
+from .registry import register_lowering
+
+
+_SUPPORTED_INPUT_DTYPES = {ScalarType.I32, ScalarType.I64}
+_SUPPORTED_OUTPUT_DTYPES = {
+    ScalarType.U8,
+    ScalarType.U16,
+    ScalarType.U32,
+    ScalarType.U64,
+    ScalarType.I8,
+    ScalarType.I16,
+    ScalarType.I32,
+    ScalarType.I64,
+    ScalarType.F16,
+    ScalarType.F32,
+    ScalarType.F64,
+}
+
+
+def _find_initializer(graph: Graph, name: str) -> Initializer | None:
+    for initializer in graph.initializers:
+        if initializer.name == name:
+            return initializer
+    return None
+
+
+def _read_scalar_initializer(
+    graph: Graph, name: str, node: Node
+) -> int | None:
+    initializer = _find_initializer(graph, name)
+    if initializer is None:
+        return None
+    data = np.array(initializer.data)
+    if data.size != 1:
+        raise UnsupportedOpError(
+            f"{node.op_type} size input must be a scalar"
+        )
+    return int(data.reshape(-1)[0].item())
+
+
+def _is_scalar_shape(shape: tuple[int, ...]) -> bool:
+    return shape == () or shape == (1,)
+
+
+@register_lowering("HammingWindow")
+def lower_hamming_window(graph: Graph, node: Node) -> HammingWindowOp:
+    if len(node.inputs) != 1 or len(node.outputs) != 1:
+        raise UnsupportedOpError("HammingWindow must have 1 input and 1 output")
+    size_shape = value_shape(graph, node.inputs[0], node)
+    if not _is_scalar_shape(size_shape):
+        raise UnsupportedOpError("HammingWindow size input must be a scalar")
+    input_dtype = value_dtype(graph, node.inputs[0], node)
+    if input_dtype not in _SUPPORTED_INPUT_DTYPES:
+        raise UnsupportedOpError(
+            f"HammingWindow size input must be int32 or int64, got {input_dtype.onnx_name}"
+        )
+    output_shape = value_shape(graph, node.outputs[0], node)
+    if len(output_shape) != 1:
+        raise ShapeInferenceError("HammingWindow output must be 1D")
+    output_dtype = value_dtype(graph, node.outputs[0], node)
+    if output_dtype not in _SUPPORTED_OUTPUT_DTYPES:
+        raise UnsupportedOpError(
+            "HammingWindow output dtype must be numeric, "
+            f"got {output_dtype.onnx_name}"
+        )
+    output_datatype = node.attrs.get("output_datatype")
+    if output_datatype is not None:
+        attr_dtype = dtype_info(int(output_datatype))
+        if attr_dtype != output_dtype:
+            raise UnsupportedOpError(
+                "HammingWindow output_datatype does not match output dtype"
+            )
+    periodic = int(node.attrs.get("periodic", 1))
+    if periodic not in {0, 1}:
+        raise UnsupportedOpError("HammingWindow periodic must be 0 or 1")
+    size_value = _read_scalar_initializer(graph, node.inputs[0], node)
+    if size_value is not None:
+        if size_value < 0:
+            raise ShapeInferenceError(
+                "HammingWindow size must be non-negative"
+            )
+        if output_shape[0] != size_value:
+            raise ShapeInferenceError(
+                "HammingWindow output length does not match size input"
+            )
+    return HammingWindowOp(
+        size=node.inputs[0],
+        output=node.outputs[0],
+        output_shape=output_shape,
+        periodic=periodic == 1,
+        dtype=output_dtype,
+        input_dtype=input_dtype,
+    )

emx_onnx_cgen/lowering/hardmax.py

@@ -1,53 +1,17 @@
 from __future__ import annotations
 
-from shared.scalar_types import ScalarType
-
 from ..ir.ops import HardmaxOp
 from ..errors import UnsupportedOpError
 from ..ir.model import Graph, Node
-from .common import node_dtype as _node_dtype
-from .common import onnx_opset_version as _onnx_opset_version
-from .common import shape_product as _shape_product
-from .common import value_shape as _value_shape
 from .registry import register_lowering
-from ..validation import ensure_output_shape_matches_input
-from ..validation import normalize_axis as _normalize_axis
 
 
 @register_lowering("Hardmax")
 def lower_hardmax(graph: Graph, node: Node) -> HardmaxOp:
     if len(node.inputs) != 1 or len(node.outputs) != 1:
         raise UnsupportedOpError("Hardmax must have 1 input and 1 output")
-    op_dtype = _node_dtype(graph, node, *node.inputs, *node.outputs)
-    if op_dtype not in {ScalarType.F16, ScalarType.F32, ScalarType.F64}:
-        raise UnsupportedOpError(
-            "Hardmax supports float16, float, and double inputs only"
-        )
-    input_shape = _value_shape(graph, node.inputs[0], node)
-    output_shape = _value_shape(graph, node.outputs[0], node)
-    ensure_output_shape_matches_input(node, input_shape, output_shape)
-    opset_version = _onnx_opset_version(graph)
-    default_axis = 1 if opset_version is not None and opset_version < 13 else -1
-    axis_attr = node.attrs.get("axis", default_axis)
-    axis = _normalize_axis(
-        int(axis_attr),
-        input_shape,
-        node,
-    )
-    outer = _shape_product(input_shape[:axis]) if axis > 0 else 1
-    axis_size = input_shape[axis]
-    inner = (
-        _shape_product(input_shape[axis + 1 :])
-        if axis + 1 < len(input_shape)
-        else 1
-    )
     return HardmaxOp(
         input0=node.inputs[0],
         output=node.outputs[0],
-        outer=outer,
-        axis_size=axis_size,
-        inner=inner,
-        axis=axis,
-        shape=input_shape,
-        dtype=op_dtype,
+        axis=int(node.attrs["axis"]) if "axis" in node.attrs else None,
     )

emx_onnx_cgen/lowering/identity.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
-from ..ir.ops import IdentityOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.context import GraphContext
 from ..ir.model import Graph, Node
-from .common import value_dtype, value_shape
+from ..ir.ops import IdentityOp
+from .common import value_dtype, value_has_dim_params, value_shape
 from .registry import register_lowering
 
 
@@ -13,9 +14,10 @@ def lower_identity(graph: Graph, node: Node) -> IdentityOp:
         raise UnsupportedOpError("Identity must have 1 input and 1 output")
     input_shape = value_shape(graph, node.inputs[0], node)
     output_shape = value_shape(graph, node.outputs[0], node)
+    if value_has_dim_params(graph, node.outputs[0]) or not output_shape:
+        output_shape = ()
     input_dim_params = graph.find_value(node.inputs[0]).type.dim_params
     output_dim_params = graph.find_value(node.outputs[0]).type.dim_params
-    resolved_shape = output_shape or input_shape
     if input_shape and output_shape:
         if len(input_shape) != len(output_shape):
             raise ShapeInferenceError("Identity input and output shapes must match")
@@ -35,10 +37,9 @@ def lower_identity(graph: Graph, node: Node) -> IdentityOp:
             "Identity expects matching input/output dtypes, "
             f"got {input_dtype.onnx_name} and {output_dtype.onnx_name}"
         )
+    if isinstance(graph, GraphContext):
+        graph.set_shape(node.outputs[0], input_shape)
     return IdentityOp(
         input0=node.inputs[0],
         output=node.outputs[0],
-        shape=resolved_shape,
-        dtype=output_dtype,
-        input_dtype=input_dtype,
     )

emx_onnx_cgen/lowering/logsoftmax.py

@@ -3,49 +3,15 @@ from __future__ import annotations
 from ..ir.ops import LogSoftmaxOp
 from ..errors import UnsupportedOpError
 from ..ir.model import Graph, Node
-from .common import node_dtype as _node_dtype
-from .common import onnx_opset_version as _onnx_opset_version
-from .common import shape_product as _shape_product
-from .common import value_shape as _value_shape
 from .registry import register_lowering
-from ..validation import ensure_output_shape_matches_input
-from ..validation import normalize_axis as _normalize_axis
 
 
 @register_lowering("LogSoftmax")
 def lower_logsoftmax(graph: Graph, node: Node) -> LogSoftmaxOp:
     if len(node.inputs) != 1 or len(node.outputs) != 1:
         raise UnsupportedOpError("LogSoftmax must have 1 input and 1 output")
-    op_dtype = _node_dtype(graph, node, *node.inputs, *node.outputs)
-    if not op_dtype.is_float:
-        raise UnsupportedOpError(
-            "LogSoftmax supports float16, float, and double inputs only"
-        )
-    input_shape = _value_shape(graph, node.inputs[0], node)
-    output_shape = _value_shape(graph, node.outputs[0], node)
-    ensure_output_shape_matches_input(node, input_shape, output_shape)
-    opset_version = _onnx_opset_version(graph)
-    default_axis = 1 if opset_version is not None and opset_version < 13 else -1
-    axis_attr = node.attrs.get("axis", default_axis)
-    axis = _normalize_axis(
-        int(axis_attr),
-        input_shape,
-        node,
-    )
-    outer = _shape_product(input_shape[:axis]) if axis > 0 else 1
-    axis_size = input_shape[axis]
-    inner = (
-        _shape_product(input_shape[axis + 1 :])
-        if axis + 1 < len(input_shape)
-        else 1
-    )
     return LogSoftmaxOp(
         input0=node.inputs[0],
         output=node.outputs[0],
-        outer=outer,
-        axis_size=axis_size,
-        inner=inner,
-        axis=axis,
-        shape=input_shape,
-        dtype=op_dtype,
+        axis=int(node.attrs["axis"]) if "axis" in node.attrs else None,
     )

emx_onnx_cgen/lowering/lp_pool.py

@@ -19,6 +19,8 @@ class LpPoolSpec:
     out_w: int
     kernel_h: int
     kernel_w: int
+    dilation_h: int
+    dilation_w: int
     stride_h: int
     stride_w: int
     pad_top: int
@@ -51,8 +53,10 @@ def _resolve_lp_pool_spec(graph: Graph, node: Node) -> LpPoolSpec:
     if ceil_mode != 0:
         raise UnsupportedOpError("LpPool supports ceil_mode=0 only")
     dilations = tuple(int(value) for value in node.attrs.get("dilations", (1, 1)))
-    if any(value != 1 for value in dilations):
-        raise UnsupportedOpError("LpPool supports dilations=1 only")
+    if len(dilations) != 2:
+        raise UnsupportedOpError("LpPool expects 2D dilations")
+    if any(value < 1 for value in dilations):
+        raise UnsupportedOpError("LpPool requires dilations >= 1")
     kernel_shape = node.attrs.get("kernel_shape")
     if kernel_shape is None:
         raise UnsupportedOpError("LpPool requires kernel_shape")
@@ -75,8 +79,11 @@ def _resolve_lp_pool_spec(graph: Graph, node: Node) -> LpPoolSpec:
         raise UnsupportedOpError("LpPool supports NCHW 2D inputs only")
     batch, channels, in_h, in_w = input_shape
     stride_h, stride_w = strides
-    out_h = (in_h + pad_top + pad_bottom - kernel_h) // stride_h + 1
-    out_w = (in_w + pad_left + pad_right - kernel_w) // stride_w + 1
+    dilation_h, dilation_w = dilations
+    effective_kernel_h = dilation_h * (kernel_h - 1) + 1
+    effective_kernel_w = dilation_w * (kernel_w - 1) + 1
+    out_h = (in_h + pad_top + pad_bottom - effective_kernel_h) // stride_h + 1
+    out_w = (in_w + pad_left + pad_right - effective_kernel_w) // stride_w + 1
     if out_h < 0 or out_w < 0:
         raise ShapeInferenceError("LpPool output shape must be non-negative")
     output_shape = _value_shape(graph, node.outputs[0], node)
@@ -95,6 +102,8 @@ def _resolve_lp_pool_spec(graph: Graph, node: Node) -> LpPoolSpec:
         out_w=out_w,
         kernel_h=kernel_h,
         kernel_w=kernel_w,
+        dilation_h=dilation_h,
+        dilation_w=dilation_w,
         stride_h=stride_h,
         stride_w=stride_w,
         pad_top=pad_top,
@@ -130,6 +139,8 @@ def lower_lp_pool(graph: Graph, node: Node) -> LpPoolOp:
         out_w=spec.out_w,
         kernel_h=spec.kernel_h,
         kernel_w=spec.kernel_w,
+        dilation_h=spec.dilation_h,
+        dilation_w=spec.dilation_w,
         stride_h=spec.stride_h,
         stride_w=spec.stride_w,
         pad_top=spec.pad_top,