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

emx_onnx_cgen/runtime/evaluator.py

@@ -14,12 +14,16 @@ from ..lowering.batch_normalization import lower_batch_normalization
 from ..lowering.concat import lower_concat
 from ..lowering.constant_of_shape import lower_constant_of_shape
 from ..lowering.conv import resolve_conv_spec
+from ..lowering.conv_transpose import resolve_conv_transpose_spec
 from ..lowering.dropout import lower_dropout
 from ..lowering.cumsum import lower_cumsum
+from ..lowering.einsum import lower_einsum
 from ..lowering.flatten import lower_flatten
 from ..lowering.gemm import resolve_gemm_spec
 from ..lowering.logsoftmax import lower_logsoftmax
+from ..lowering.hardmax import lower_hardmax
 from ..lowering.lp_normalization import lower_lp_normalization
+from ..lowering.lp_pool import lower_lp_pool
 from ..lowering.grid_sample import lower_grid_sample
 from ..lowering.instance_normalization import lower_instance_normalization
 from ..lowering.group_normalization import lower_group_normalization
@@ -27,17 +31,21 @@ from ..lowering.layer_normalization import lower_layer_normalization
 from ..lowering.mean_variance_normalization import (
     lower_mean_variance_normalization,
 )
+from ..lowering.global_max_pool import lower_global_max_pool
 from ..lowering.negative_log_likelihood_loss import (
     lower_negative_log_likelihood_loss,
 )
+from ..lowering.nonzero import lower_nonzero
 from ..lowering.pad import lower_pad
 from ..lowering.expand import lower_expand
 from ..lowering.range import lower_range
+from ..lowering.one_hot import lower_onehot
 from ..lowering.split import lower_split
 from ..lowering.softmax_cross_entropy_loss import (
     lower_softmax_cross_entropy_loss,
 )
 from ..lowering.arg_reduce import lower_arg_reduce
+from ..lowering.topk import lower_topk
 from ..lowering.lstm import ACTIVATION_KIND_BY_NAME, resolve_lstm_spec
 from ..lowering.lrn import resolve_lrn_spec
 from ..lowering.matmul import lower_matmul
@@ -49,6 +57,7 @@ from ..lowering.reduce import (
     resolve_reduce_axes,
 )
 from ..lowering.reshape import lower_reshape
+from ..lowering.scatter_nd import lower_scatternd
 from ..lowering.slice import _normalize_slices
 from ..lowering.shape import lower_shape
 from ..lowering.size import lower_size
@@ -58,6 +67,7 @@ from ..lowering.squeeze import lower_squeeze
 from ..lowering.transpose import lower_transpose
 from ..lowering.unsqueeze import lower_unsqueeze
 from ..lowering.where import lower_where
+from ..lowering.quantize_linear import resolve_quantize_spec
 from ..lowering.variadic import BINARY_ONLY_OPS, VARIADIC_OP_FUNCTIONS
 from ..lowering.registry import resolve_dispatch
 from ..lowering.common import node_dtype, optional_name, value_dtype, value_shape
@@ -133,6 +143,21 @@ def _eval_matmul(evaluator: Evaluator, node: Node) -> None:
     evaluator.values[node.outputs[0]] = _apply_matmul(left, right)
 
 
+@register_evaluator("Einsum")
+def _eval_einsum(evaluator: Evaluator, node: Node) -> None:
+    lower_einsum(evaluator.graph, node)
+    equation_value = node.attrs.get("equation")
+    if equation_value is None:
+        raise UnsupportedOpError("Einsum equation attribute is required")
+    equation = (
+        equation_value.decode()
+        if isinstance(equation_value, (bytes, bytearray))
+        else str(equation_value)
+    )
+    inputs = [evaluator.values[name] for name in node.inputs]
+    evaluator.values[node.outputs[0]] = np.einsum(equation, *inputs)
+
+
 @register_evaluator("Clip")
 def _eval_clip(evaluator: Evaluator, node: Node) -> None:
     if not node.inputs or len(node.outputs) != 1:
@@ -242,10 +267,11 @@ def _eval_pad(evaluator: Evaluator, node: Node) -> None:
             pads_begin = np.zeros(rank, dtype=np.int64)
             pads_end = np.zeros(rank, dtype=np.int64)
             for axis, pad_index in enumerate(op.pads_axis_map):
-                if pad_index is None:
-                    continue
-                pads_begin[axis] = int(pads_values[pad_index])
-                pads_end[axis] = int(pads_values[pad_index + axis_count])
+                if pad_index is not None:
+                    pads_begin[axis] = int(pads_values[pad_index])
+                    pads_end[axis] = int(
+                        pads_values[pad_index + axis_count]
+                    )
             pad_width = tuple(
                 (int(pads_begin[index]), int(pads_end[index]))
                 for index in range(rank)
@@ -270,6 +296,53 @@ def _eval_pad(evaluator: Evaluator, node: Node) -> None:
     )
 
 
+@register_evaluator("ScatterND")
+def _eval_scatternd(evaluator: Evaluator, node: Node) -> None:
+    op = lower_scatternd(evaluator.graph, node)
+    data = evaluator.values[op.data]
+    indices = evaluator.values[op.indices]
+    updates = evaluator.values[op.updates]
+    output = np.array(data, copy=True)
+    index_depth = op.indices_shape[-1]
+    update_indices_shape = op.indices_shape[:-1]
+    update_count = int(np.prod(update_indices_shape)) if update_indices_shape else 1
+    flat_indices = indices.astype(np.int64, copy=False).reshape(
+        update_count, index_depth
+    )
+    tail_shape = op.data_shape[index_depth:]
+    updates_reshaped = updates.reshape((update_count,) + tail_shape)
+    for index, index_values in enumerate(flat_indices):
+        output_index: list[int | slice] = []
+        for axis, value in enumerate(index_values):
+            axis_size = op.data_shape[axis]
+            idx = int(value)
+            if idx < 0:
+                idx += axis_size
+            if idx < 0 or idx >= axis_size:
+                raise UnsupportedOpError(
+                    "ScatterND indices must be within data bounds"
+                )
+            output_index.append(idx)
+        output_index.extend([slice(None)] * len(tail_shape))
+        target = tuple(output_index)
+        update_value = updates_reshaped[index]
+        if op.reduction == "none":
+            output[target] = update_value
+        elif op.reduction == "add":
+            output[target] = output[target] + update_value
+        elif op.reduction == "mul":
+            output[target] = output[target] * update_value
+        elif op.reduction == "min":
+            output[target] = np.minimum(output[target], update_value)
+        elif op.reduction == "max":
+            output[target] = np.maximum(output[target], update_value)
+        else:
+            raise UnsupportedOpError(
+                f"Unsupported ScatterND reduction {op.reduction}"
+            )
+    evaluator.values[op.output] = output
+
+
 @register_evaluator("Celu")
 def _eval_celu(evaluator: Evaluator, node: Node) -> None:
     if len(node.inputs) != 1 or len(node.outputs) != 1:
@@ -786,6 +859,40 @@ def _eval_eye_like(evaluator: Evaluator, node: Node) -> None:
     evaluator.values[node.outputs[0]] = output
 
 
+@register_evaluator("Trilu")
+def _eval_trilu(evaluator: Evaluator, node: Node) -> None:
+    if len(node.inputs) not in {1, 2} or len(node.outputs) != 1:
+        raise UnsupportedOpError("Trilu must have 1 or 2 inputs and 1 output")
+    value = evaluator.values[node.inputs[0]]
+    if value.ndim < 2:
+        raise UnsupportedOpError("Trilu expects input rank >= 2")
+    output_dtype = value_dtype(evaluator.graph, node.outputs[0], node)
+    input_dtype = value_dtype(evaluator.graph, node.inputs[0], node)
+    if output_dtype != input_dtype:
+        raise UnsupportedOpError(
+            "Trilu expects matching input/output dtypes, "
+            f"got {input_dtype.onnx_name} and {output_dtype.onnx_name}"
+        )
+    k = 0
+    if len(node.inputs) == 2 and node.inputs[1]:
+        k_value = np.array(evaluator.values[node.inputs[1]], dtype=np.int64)
+        if k_value.size != 1:
+            raise UnsupportedOpError("Trilu k input must be scalar")
+        k = int(k_value.reshape(-1)[0])
+    upper_attr = node.attrs.get("upper", 1)
+    upper = bool(int(upper_attr))
+    rows, cols = value.shape[-2], value.shape[-1]
+    batch_shape = value.shape[:-2]
+    batch_size = int(np.prod(batch_shape)) if batch_shape else 1
+    view = value.reshape(batch_size, rows, cols)
+    if upper:
+        mask = np.triu(np.ones((rows, cols), dtype=bool), k=k)
+    else:
+        mask = np.tril(np.ones((rows, cols), dtype=bool), k=k)
+    output = np.where(mask, view, np.zeros_like(view))
+    evaluator.values[node.outputs[0]] = output.reshape(value.shape)
+
+
 @register_evaluator("Tile")
 def _eval_tile(evaluator: Evaluator, node: Node) -> None:
     if len(node.inputs) != 2 or len(node.outputs) != 1:
@@ -922,6 +1029,73 @@ def _eval_gather(evaluator: Evaluator, node: Node) -> None:
     evaluator.values[node.outputs[0]] = np.take(data, indices, axis=axis)
 
 
+@register_evaluator("GatherND")
+def _eval_gather_nd(evaluator: Evaluator, node: Node) -> None:
+    if len(node.inputs) != 2 or len(node.outputs) != 1:
+        raise UnsupportedOpError("GatherND must have 2 inputs and 1 output")
+    data = evaluator.values[node.inputs[0]]
+    indices = evaluator.values[node.inputs[1]]
+    if indices.dtype.type not in {np.int32, np.int64}:
+        raise UnsupportedOpError(
+            f"GatherND indices must be int32 or int64, got {indices.dtype}"
+        )
+    if indices.ndim < 1:
+        raise UnsupportedOpError("GatherND indices must have rank >= 1")
+    batch_dims = int(node.attrs.get("batch_dims", 0))
+    if batch_dims < 0:
+        raise UnsupportedOpError(
+            f"GatherND batch_dims must be >= 0, got {batch_dims}"
+        )
+    if batch_dims > indices.ndim - 1:
+        raise UnsupportedOpError(
+            "GatherND batch_dims must be <= indices rank - 1, "
+            f"got {batch_dims} vs {indices.ndim - 1}"
+        )
+    if batch_dims > data.ndim:
+        raise UnsupportedOpError(
+            "GatherND batch_dims must be <= data rank, "
+            f"got {batch_dims} vs {data.ndim}"
+        )
+    if tuple(data.shape[:batch_dims]) != tuple(indices.shape[:batch_dims]):
+        raise UnsupportedOpError(
+            "GatherND batch_dims must match on data/indices, "
+            f"got {data.shape} vs {indices.shape}"
+        )
+    index_depth = indices.shape[-1]
+    if index_depth <= 0:
+        raise UnsupportedOpError(
+            "GatherND indices final dimension must be >= 1"
+        )
+    if index_depth > data.ndim - batch_dims:
+        raise UnsupportedOpError(
+            "GatherND indices final dimension must be <= data rank - "
+            f"batch_dims, got {index_depth} vs {data.ndim - batch_dims}"
+        )
+    tail_shape = data.shape[batch_dims + index_depth :]
+    output_shape = indices.shape[:-1] + tail_shape
+    output = np.empty(output_shape, dtype=data.dtype)
+    indices_prefix_shape = indices.shape[:-1]
+    prefix_iter = (
+        np.ndindex(*indices_prefix_shape) if indices_prefix_shape else [()]
+    )
+    for prefix in prefix_iter:
+        raw_index = indices[prefix]
+        if index_depth == 1:
+            index_values = [int(np.asarray(raw_index).item())]
+        else:
+            index_values = [int(value) for value in raw_index]
+        for dim_index, value in enumerate(index_values):
+            if value < 0:
+                index_values[dim_index] = value + data.shape[
+                    batch_dims + dim_index
+                ]
+        data_index = list(prefix[:batch_dims]) + index_values
+        data_index.extend([slice(None)] * len(tail_shape))
+        output_index = prefix + (slice(None),) * len(tail_shape)
+        output[output_index] = data[tuple(data_index)]
+    evaluator.values[node.outputs[0]] = output
+
+
 @register_evaluator("Slice")
 def _eval_slice(evaluator: Evaluator, node: Node) -> None:
     input_value = evaluator.values[node.inputs[0]]
@@ -1101,6 +1275,28 @@ def _eval_conv(evaluator: Evaluator, node: Node) -> None:
     evaluator.values[node.outputs[0]] = _apply_conv(spec, data, weights, bias)
 
 
+@register_evaluator("ConvTranspose")
+def _eval_conv_transpose(evaluator: Evaluator, node: Node) -> None:
+    op_dtype = value_dtype(evaluator.graph, node.inputs[0], node)
+    output_dtype = value_dtype(evaluator.graph, node.outputs[0], node)
+    if op_dtype != output_dtype:
+        raise UnsupportedOpError(
+            f"{node.op_type} expects matching input/output dtypes, "
+            f"got {op_dtype.onnx_name} and {output_dtype.onnx_name}"
+        )
+    if not op_dtype.is_float:
+        raise UnsupportedOpError(
+            "ConvTranspose supports float16, float, and double inputs only"
+        )
+    spec = resolve_conv_transpose_spec(evaluator.graph, node)
+    data = evaluator.values[node.inputs[0]]
+    weights = evaluator.values[node.inputs[1]]
+    bias = evaluator.values[node.inputs[2]] if len(node.inputs) > 2 else None
+    evaluator.values[node.outputs[0]] = _apply_conv_transpose(
+        spec, data, weights, bias
+    )
+
+
 @register_evaluator("BatchNormalization")
 def _eval_batch_norm(evaluator: Evaluator, node: Node) -> None:
     op = lower_batch_normalization(evaluator.graph, node)
@@ -1133,6 +1329,59 @@ def _eval_lp_normalization(evaluator: Evaluator, node: Node) -> None:
     evaluator.values[op.output] = data / denom
 
 
+@register_evaluator("LpPool")
+def _eval_lp_pool(evaluator: Evaluator, node: Node) -> None:
+    op = lower_lp_pool(evaluator.graph, node)
+    data = evaluator.values[op.input0]
+    output = np.zeros(
+        (op.batch, op.channels, op.out_h, op.out_w), dtype=data.dtype
+    )
+    for n in range(op.batch):
+        for c in range(op.channels):
+            for out_h in range(op.out_h):
+                for out_w in range(op.out_w):
+                    h_start = out_h * op.stride_h - op.pad_top
+                    w_start = out_w * op.stride_w - op.pad_left
+                    acc = 0.0
+                    for kh in range(op.kernel_h):
+                        for kw in range(op.kernel_w):
+                            in_h = h_start + kh
+                            in_w = w_start + kw
+                            if (
+                                0 <= in_h < op.in_h
+                                and 0 <= in_w < op.in_w
+                            ):
+                                value = data[(n, c, in_h, in_w)]
+                                acc += abs(value) ** op.p
+                    output[(n, c, out_h, out_w)] = acc ** (1.0 / op.p)
+    evaluator.values[op.output] = output
+
+
+@register_evaluator("QuantizeLinear")
+def _eval_quantize_linear(evaluator: Evaluator, node: Node) -> None:
+    spec = resolve_quantize_spec(evaluator.graph, node)
+    data = evaluator.values[node.inputs[0]]
+    scale = evaluator.values[node.inputs[1]]
+    zero_point_name = optional_name(node.inputs, 2)
+    if zero_point_name is None:
+        zero_point = 0
+    else:
+        zero_point = evaluator.values[zero_point_name]
+    if spec.axis is None:
+        scaled = data / scale + zero_point
+    else:
+        shape = [1] * data.ndim
+        shape[spec.axis] = scale.shape[0]
+        scaled = data / scale.reshape(shape) + np.asarray(zero_point).reshape(
+            shape
+        )
+    rounded = np.rint(scaled)
+    info = np.iinfo(spec.output_dtype.np_dtype)
+    clipped = np.clip(rounded, info.min, info.max)
+    evaluator.values[node.outputs[0]] = clipped.astype(
+        spec.output_dtype.np_dtype, copy=False
+    )
+
 @register_evaluator("InstanceNormalization")
 def _eval_instance_normalization(evaluator: Evaluator, node: Node) -> None:
     op = lower_instance_normalization(evaluator.graph, node)
@@ -1284,6 +1533,18 @@ def _eval_maxpool(evaluator: Evaluator, node: Node) -> None:
         evaluator.values[indices_output] = indices
 
 
+@register_evaluator("GlobalMaxPool")
+def _eval_global_max_pool(evaluator: Evaluator, node: Node) -> None:
+    op = lower_global_max_pool(evaluator.graph, node)
+    value = evaluator.values[node.inputs[0]]
+    if not op.axes:
+        evaluator.values[node.outputs[0]] = value.copy()
+        return
+    evaluator.values[node.outputs[0]] = np.max(
+        value, axis=op.axes, keepdims=op.keepdims
+    )
+
+
 @register_evaluator("Softmax")
 def _eval_softmax(evaluator: Evaluator, node: Node) -> None:
     op = lower_softmax(evaluator.graph, node)
@@ -1298,6 +1559,19 @@ def _eval_logsoftmax(evaluator: Evaluator, node: Node) -> None:
     evaluator.values[node.outputs[0]] = _apply_logsoftmax(value, op.axis)
 
 
+@register_evaluator("Hardmax")
+def _eval_hardmax(evaluator: Evaluator, node: Node) -> None:
+    op = lower_hardmax(evaluator.graph, node)
+    value = evaluator.values[node.inputs[0]]
+    max_values = np.max(value, axis=op.axis, keepdims=True)
+    is_max = value == max_values
+    max_index = np.argmax(is_max, axis=op.axis)
+    output = np.zeros_like(value)
+    ones = np.array(1.0, dtype=value.dtype)
+    np.put_along_axis(output, np.expand_dims(max_index, axis=op.axis), ones, axis=op.axis)
+    evaluator.values[node.outputs[0]] = output
+
+
 @register_evaluator("NegativeLogLikelihoodLoss")
 def _eval_negative_log_likelihood_loss(
     evaluator: Evaluator, node: Node
@@ -1409,6 +1683,16 @@ def _eval_size(evaluator: Evaluator, node: Node) -> None:
     evaluator.values[op.output] = np.array(op.value, dtype=np.int64)
 
 
+@register_evaluator("NonZero")
+def _eval_nonzero(evaluator: Evaluator, node: Node) -> None:
+    op = lower_nonzero(evaluator.graph, node)
+    values = evaluator.values[op.input0]
+    indices = np.nonzero(values)
+    evaluator.values[op.output] = np.stack(indices, axis=0).astype(
+        np.int64, copy=False
+    )
+
+
 @register_evaluator("Expand")
 def _eval_expand(evaluator: Evaluator, node: Node) -> None:
     op = lower_expand(evaluator.graph, node)
@@ -1428,6 +1712,45 @@ def _eval_range(evaluator: Evaluator, node: Node) -> None:
     evaluator.values[op.output] = output
 
 
+@register_evaluator("OneHot")
+def _eval_onehot(evaluator: Evaluator, node: Node) -> None:
+    op = lower_onehot(evaluator.graph, node)
+    indices = evaluator.values[op.indices].astype(np.int64, copy=False)
+    depth_values = evaluator.values[op.depth].reshape(-1)
+    if depth_values.size != 1:
+        raise UnsupportedOpError("OneHot depth input must be a scalar")
+    depth_value = int(depth_values[0])
+    if depth_value < 0:
+        raise UnsupportedOpError("OneHot depth must be non-negative")
+    values = evaluator.values[op.values].reshape(-1)
+    if values.size != 2:
+        raise UnsupportedOpError("OneHot values input must have 2 elements")
+    off_value, on_value = values[0], values[1]
+    if depth_value == 0:
+        evaluator.values[op.output] = np.full(
+            op.output_shape, off_value, dtype=values.dtype
+        )
+        return
+    axis = op.axis
+    rank = indices.ndim
+    if axis < 0:
+        axis += rank + 1
+    depth_range = np.arange(depth_value, dtype=np.int64)
+    new_shape = (1,) * axis + (depth_value,) + (1,) * (rank - axis)
+    targets = depth_range.reshape(new_shape)
+    adjusted = np.mod(indices, depth_value) if depth_value > 0 else indices
+    values_reshaped = np.reshape(
+        adjusted, indices.shape[:axis] + (1,) + indices.shape[axis:]
+    )
+    valid_mask = (indices >= -depth_value) & (indices < depth_value)
+    valid_mask = np.reshape(
+        valid_mask, indices.shape[:axis] + (1,) + indices.shape[axis:]
+    )
+    one_hot = (targets == values_reshaped) & valid_mask
+    output = np.where(one_hot, on_value, off_value).astype(values.dtype)
+    evaluator.values[op.output] = output
+
+
 @register_evaluator("Split")
 def _eval_split(evaluator: Evaluator, node: Node) -> None:
     op = lower_split(evaluator.graph, node)
@@ -1550,6 +1873,39 @@ def _eval_arg_reduce(evaluator: Evaluator, node: Node) -> None:
     evaluator.values[op.output] = indices.astype(op.output_dtype.np_dtype)
 
 
+@register_evaluator("TopK")
+def _eval_topk(evaluator: Evaluator, node: Node) -> None:
+    op = lower_topk(evaluator.graph, node)
+    value = evaluator.values[op.input0]
+    moved = np.moveaxis(value, op.axis, -1)
+    axis_dim = moved.shape[-1]
+    flat = moved.reshape(-1, axis_dim)
+    values_out = np.empty((flat.shape[0], op.k), dtype=value.dtype)
+    indices_out = np.empty((flat.shape[0], op.k), dtype=np.int64)
+    for row_index in range(flat.shape[0]):
+        row = flat[row_index]
+        order = sorted(
+            range(axis_dim),
+            key=lambda idx: (
+                -row[idx].item() if op.largest else row[idx].item(),
+                idx,
+            ),
+        )
+        topk = order[: op.k]
+        indices_out[row_index] = topk
+        values_out[row_index] = row[topk]
+    values_out = values_out.reshape(moved.shape[:-1] + (op.k,))
+    indices_out = indices_out.reshape(moved.shape[:-1] + (op.k,))
+    values_out = np.moveaxis(values_out, -1, op.axis)
+    indices_out = np.moveaxis(indices_out, -1, op.axis)
+    evaluator.values[op.output_values] = values_out.astype(
+        op.output_values_dtype.np_dtype
+    )
+    evaluator.values[op.output_indices] = indices_out.astype(
+        op.output_indices_dtype.np_dtype
+    )
+
+
 def _eval_binary_unary(evaluator: Evaluator, node: Node) -> None:
     if node.op_type == "BitShift":
         if len(node.inputs) != 2 or len(node.outputs) != 1:
@@ -1671,9 +2027,8 @@ def _matmul_batch_broadcastable(
     left_padded = (1,) * (max_rank - len(left)) + left
     right_padded = (1,) * (max_rank - len(right)) + right
     for left_dim, right_dim in zip(left_padded, right_padded):
-        if left_dim == right_dim or left_dim == 1 or right_dim == 1:
-            continue
-        return False
+        if not (left_dim == right_dim or left_dim == 1 or right_dim == 1):
+            return False
     return True
 
 
@@ -1916,7 +2271,9 @@ def _apply_attention(
     return output, key_total, value_total, qk_output
 
 
-def _apply_conv(spec, data: np.ndarray, weights: np.ndarray, bias: np.ndarray | None) -> np.ndarray:
+def _apply_conv(
+    spec, data: np.ndarray, weights: np.ndarray, bias: np.ndarray | None
+) -> np.ndarray:
     output = np.zeros(
         (spec.batch, spec.out_channels, *spec.out_spatial),
         dtype=data.dtype,
@@ -1958,15 +2315,67 @@ def _apply_conv(spec, data: np.ndarray, weights: np.ndarray, bias: np.ndarray |
                                     valid = False
                                     break
                                 in_index.append(in_dim)
-                            if not valid:
-                                continue
-                            acc += data[(n, ic_global, *in_index)] * weights[
-                                (oc_global, ic, *kernel_index)
-                            ]
+                            if valid:
+                                acc += data[(n, ic_global, *in_index)] * weights[
+                                    (oc_global, ic, *kernel_index)
+                                ]
                     output[(n, oc_global, *out_index)] = acc
     return output
 
 
+def _apply_conv_transpose(
+    spec, data: np.ndarray, weights: np.ndarray, bias: np.ndarray | None
+) -> np.ndarray:
+    output = np.zeros(
+        (spec.batch, spec.out_channels, *spec.out_spatial), dtype=data.dtype
+    )
+    if bias is not None:
+        output += bias.reshape((1, spec.out_channels) + (1,) * spec.spatial_rank)
+    pad_begin = spec.pads[: spec.spatial_rank]
+    group_in_channels = spec.in_channels // spec.group
+    group_out_channels = spec.out_channels // spec.group
+    for n in range(spec.batch):
+        for g in range(spec.group):
+            oc_base = g * group_out_channels
+            ic_base = g * group_in_channels
+            for ic in range(group_in_channels):
+                ic_global = ic_base + ic
+                for in_index in np.ndindex(*spec.in_spatial):
+                    value = data[(n, ic_global, *in_index)]
+                    for oc in range(group_out_channels):
+                        oc_global = oc_base + oc
+                        for kernel_index in np.ndindex(*spec.kernel_shape):
+                            out_index = []
+                            valid = True
+                            for (
+                                in_dim,
+                                kernel_dim,
+                                stride,
+                                dilation,
+                                pad,
+                                out_size,
+                            ) in zip(
+                                in_index,
+                                kernel_index,
+                                spec.strides,
+                                spec.dilations,
+                                pad_begin,
+                                spec.out_spatial,
+                            ):
+                                out_dim = (
+                                    in_dim * stride + kernel_dim * dilation - pad
+                                )
+                                if out_dim < 0 or out_dim >= out_size:
+                                    valid = False
+                                    break
+                                out_index.append(out_dim)
+                            if valid:
+                                output[(n, oc_global, *out_index)] += (
+                                    value * weights[(ic_global, oc, *kernel_index)]
+                                )
+    return output
+
+
 def _apply_lrn(spec, data: np.ndarray) -> np.ndarray:
     output = np.empty_like(data)
     spatial_shape = spec.shape[2:]
@@ -2002,15 +2411,15 @@ def _apply_average_pool(op, data: np.ndarray) -> np.ndarray:
                         if ih < 0 or ih >= op.in_h:
                             if op.count_include_pad:
                                 count += op.kernel_w
-                            continue
-                        for kw in range(op.kernel_w):
-                            iw = ow * op.stride_w + kw - op.pad_left
-                            if iw < 0 or iw >= op.in_w:
-                                if op.count_include_pad:
+                        else:
+                            for kw in range(op.kernel_w):
+                                iw = ow * op.stride_w + kw - op.pad_left
+                                if iw < 0 or iw >= op.in_w:
+                                    if op.count_include_pad:
+                                        count += 1
+                                else:
+                                    acc += data[n, c, ih, iw]
                                     count += 1
-                                continue
-                            acc += data[n, c, ih, iw]
-                            count += 1
                     output[n, c, oh, ow] = 0.0 if count == 0 else acc / float(count)
     return output
 
@@ -2059,25 +2468,30 @@ def _apply_maxpool(
                             valid = False
                             break
                         in_index.append(idx)
-                    if not valid:
-                        continue
-                    value = data[(n, c, *in_index)]
-                    if value > max_value or not has_value:
-                        max_value = value
-                        has_value = True
-                        if return_indices:
-                            linear_index = n * spec.channels + c
-                            if spec.storage_order == 0:
-                                for idx, size in zip(in_index, spec.in_spatial):
-                                    linear_index = linear_index * size + idx
-                            else:
-                                spatial_index = 0
-                                spatial_stride = 1
-                                for idx, size in zip(in_index, spec.in_spatial):
-                                    spatial_index += idx * spatial_stride
-                                    spatial_stride *= size
-                                linear_index = linear_index * spatial_stride + spatial_index
-                            max_index = linear_index
+                    if valid:
+                        value = data[(n, c, *in_index)]
+                        if value > max_value or not has_value:
+                            max_value = value
+                            has_value = True
+                            if return_indices:
+                                linear_index = n * spec.channels + c
+                                if spec.storage_order == 0:
+                                    for idx, size in zip(
+                                        in_index, spec.in_spatial
+                                    ):
+                                        linear_index = linear_index * size + idx
+                                else:
+                                    spatial_index = 0
+                                    spatial_stride = 1
+                                    for idx, size in zip(
+                                        in_index, spec.in_spatial
+                                    ):
+                                        spatial_index += idx * spatial_stride
+                                        spatial_stride *= size
+                                    linear_index = (
+                                        linear_index * spatial_stride + spatial_index
+                                    )
+                                max_index = linear_index
                 output[(n, c, *out_index)] = max_value
                 if return_indices and indices is not None:
                     indices[(n, c, *out_index)] = max_index
@@ -2162,8 +2576,12 @@ def _apply_lstm(
         beta_g = spec.activation_betas[act_offset + 1]
         beta_h = spec.activation_betas[act_offset + 2]
         for step in range(seq_length):
-            t_index = step if dir_kind == "forward" else seq_length - 1 - step
-            x_t = x[t_index]
+            if dir_kind == "forward":
+                x_t = x[step]
+            else:
+                t_indices = sequence_lens - 1 - step
+                t_indices = np.clip(t_indices, 0, seq_length - 1)
+                x_t = x[t_indices, np.arange(batch_size)]
             gates = x_t @ w_dir.T + h_prev @ r_dir.T + bias
             if spec.clip is not None and spec.clip > 0:
                 gates = np.clip(gates, -spec.clip, spec.clip)