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

emx_onnx_cgen/ir/ops/nn.py

@@ -0,0 +1,580 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import Enum
+
+from shared.scalar_functions import ScalarFunction
+from shared.scalar_types import ScalarType
+
+from ...errors import ShapeInferenceError
+from ..op_base import ConvLikeOpBase, GemmLikeOpBase, MatMulLikeOpBase, RenderableOpBase
+from ..op_context import OpContext
+
+
+class EinsumKind(str, Enum):
+    REDUCE_ALL = "reduce_all"
+    SUM_J = "sum_j"
+    TRANSPOSE = "transpose"
+    DOT = "dot"
+    BATCH_MATMUL = "batch_matmul"
+    BATCH_DIAGONAL = "batch_diagonal"
+
+
+def _shape_product(shape: tuple[int, ...]) -> int:
+    product = 1
+    for dim in shape:
+        if dim < 0:
+            raise ShapeInferenceError("Dynamic dims are not supported")
+        product *= dim
+    return product
+
+
+@dataclass(frozen=True)
+class MatMulOp(MatMulLikeOpBase):
+    input0: str
+    input1: str
+    output: str
+    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
+    dtype: ScalarType
+
+@dataclass(frozen=True)
+class QLinearMatMulOp(MatMulLikeOpBase):
+    input0: str
+    input0_scale: str
+    input0_zero_point: str
+    input1: str
+    input1_scale: str
+    input1_zero_point: str
+    output_scale: str
+    output_zero_point: str
+    output: str
+    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
+    input0_dtype: ScalarType
+    input1_dtype: ScalarType
+    dtype: ScalarType
+    input0_scale_dtype: ScalarType
+    input1_scale_dtype: ScalarType
+    output_scale_dtype: ScalarType
+    input0_scale_shape: tuple[int, ...]
+    input1_scale_shape: tuple[int, ...]
+    output_scale_shape: tuple[int, ...]
+    input0_zero_shape: tuple[int, ...]
+    input1_zero_shape: tuple[int, ...]
+    output_zero_shape: tuple[int, ...]
+
+@dataclass(frozen=True)
+class EinsumOp(MatMulLikeOpBase):
+    inputs: tuple[str, ...]
+    output: str
+    kind: EinsumKind
+    input_shapes: tuple[tuple[int, ...], ...]
+    output_shape: tuple[int, ...]
+    dtype: ScalarType
+    input_dtype: ScalarType
+
+@dataclass(frozen=True)
+class GemmOp(GemmLikeOpBase):
+    input_a: str
+    input_b: str
+    input_c: str | None
+    output: str
+    m: int
+    n: int
+    k: int
+    trans_a: bool
+    trans_b: bool
+    alpha: float | int
+    beta: float | int
+    c_shape: tuple[int, ...] | None
+    dtype: ScalarType
+
+@dataclass(frozen=True)
+class AttentionOp(RenderableOpBase):
+    input_q: str
+    input_k: str
+    input_v: str
+    input_attn_mask: str | None
+    input_past_key: str | None
+    input_past_value: str | None
+    input_nonpad_kv_seqlen: str | None
+    output: str
+    output_present_key: str | None
+    output_present_value: str | None
+    output_qk_matmul: str | None
+    batch: int
+    q_heads: int
+    kv_heads: int
+    q_seq: int
+    kv_seq: int
+    total_seq: int
+    past_seq: int
+    qk_head_size: int
+    v_head_size: int
+    q_hidden_size: int | None
+    k_hidden_size: int | None
+    v_hidden_size: int | None
+    scale: float
+    is_causal: bool
+    softcap: float
+    qk_matmul_output_mode: int
+    q_rank: int
+    k_rank: int
+    v_rank: int
+    output_rank: int
+    mask_shape: tuple[int, ...] | None
+    mask_is_bool: bool
+    mask_rank: int | None
+    mask_broadcast_batch: bool
+    mask_broadcast_heads: bool
+    mask_broadcast_q_seq: bool
+    mask_q_seq: int | None
+    mask_kv_seq: int | None
+    head_group_size: int
+    dtype: ScalarType
+
+@dataclass(frozen=True)
+class RotaryEmbeddingOp(RenderableOpBase):
+    input0: str
+    cos_cache: str
+    sin_cache: str
+    position_ids: str | None
+    output: str
+    input_shape: tuple[int, ...]
+    cos_shape: tuple[int, ...]
+    sin_shape: tuple[int, ...]
+    position_ids_shape: tuple[int, ...] | None
+    dtype: ScalarType
+    position_ids_dtype: ScalarType | None
+    rotary_dim: int
+    rotary_dim_half: int
+    head_size: int
+    num_heads: int
+    seq_len: int
+    batch: int
+    input_rank: int
+    interleaved: bool
+
+@dataclass(frozen=True)
+class ConvOp(ConvLikeOpBase):
+    input0: str
+    weights: str
+    bias: str | None
+    output: str
+    batch: int
+    in_channels: int
+    out_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, ...]
+    group: int
+    dtype: ScalarType
+
+    @property
+    def out_h(self) -> int:
+        if self.spatial_rank < 1:
+            raise ValueError("Conv output height is undefined for spatial_rank < 1")
+        return self.out_spatial[0]
+
+    @property
+    def out_w(self) -> int:
+        if self.spatial_rank < 2:
+            raise ValueError("Conv output width is undefined for spatial_rank < 2")
+        return self.out_spatial[1]
+
+@dataclass(frozen=True)
+class ConvTransposeOp(ConvLikeOpBase):
+    input0: str
+    weights: str
+    bias: str | None
+    output: str
+    batch: int
+    in_channels: int
+    out_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, ...]
+    output_padding: tuple[int, ...]
+    group: int
+    dtype: ScalarType
+
+@dataclass(frozen=True)
+class AveragePoolOp(RenderableOpBase):
+    input0: str
+    output: str
+    batch: int
+    channels: int
+    in_h: int
+    in_w: int
+    out_h: int
+    out_w: int
+    kernel_h: int
+    kernel_w: int
+    stride_h: int
+    stride_w: int
+    pad_top: int
+    pad_left: int
+    pad_bottom: int
+    pad_right: int
+    count_include_pad: bool
+    dtype: ScalarType
+
+@dataclass(frozen=True)
+class LpPoolOp(RenderableOpBase):
+    input0: str
+    output: str
+    batch: int
+    channels: int
+    in_h: int
+    in_w: int
+    out_h: int
+    out_w: int
+    kernel_h: int
+    kernel_w: int
+    stride_h: int
+    stride_w: int
+    pad_top: int
+    pad_left: int
+    pad_bottom: int
+    pad_right: int
+    p: int
+    dtype: ScalarType
+
+@dataclass(frozen=True)
+class SoftmaxOp(RenderableOpBase):
+    input0: str
+    output: str
+    outer: int
+    axis_size: int
+    inner: int
+    axis: int
+    shape: tuple[int, ...]
+    dtype: ScalarType
+
+    def infer_shapes(self, ctx: OpContext) -> None:
+        input_shape = ctx.shape(self.input0)
+        axis = self.axis
+        if axis < 0:
+            axis += len(input_shape)
+        if axis < 0 or axis >= len(input_shape):
+            raise ShapeInferenceError(
+                f"Softmax axis {self.axis} is out of bounds for shape {input_shape}"
+            )
+        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
+        )
+        ctx.set_shape(self.output, input_shape)
+        ctx.set_derived(self, "outer", outer)
+        ctx.set_derived(self, "axis_size", axis_size)
+        ctx.set_derived(self, "inner", inner)
+
+@dataclass(frozen=True)
+class LogSoftmaxOp(RenderableOpBase):
+    input0: str
+    output: str
+    outer: int
+    axis_size: int
+    inner: int
+    axis: int
+    shape: tuple[int, ...]
+    dtype: ScalarType
+
+    def infer_shapes(self, ctx: OpContext) -> None:
+        input_shape = ctx.shape(self.input0)
+        axis = self.axis
+        if axis < 0:
+            axis += len(input_shape)
+        if axis < 0 or axis >= len(input_shape):
+            raise ShapeInferenceError(
+                f"LogSoftmax axis {self.axis} is out of bounds for shape {input_shape}"
+            )
+        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
+        )
+        ctx.set_shape(self.output, input_shape)
+        ctx.set_derived(self, "outer", outer)
+        ctx.set_derived(self, "axis_size", axis_size)
+        ctx.set_derived(self, "inner", inner)
+
+@dataclass(frozen=True)
+class HardmaxOp(RenderableOpBase):
+    input0: str
+    output: str
+    outer: int
+    axis_size: int
+    inner: int
+    axis: int
+    shape: tuple[int, ...]
+    dtype: ScalarType
+
+    def infer_shapes(self, ctx: OpContext) -> None:
+        input_shape = ctx.shape(self.input0)
+        axis = self.axis
+        if axis < 0:
+            axis += len(input_shape)
+        if axis < 0 or axis >= len(input_shape):
+            raise ShapeInferenceError(
+                f"Hardmax axis {self.axis} is out of bounds for shape {input_shape}"
+            )
+        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
+        )
+        ctx.set_shape(self.output, input_shape)
+        ctx.set_derived(self, "outer", outer)
+        ctx.set_derived(self, "axis_size", axis_size)
+        ctx.set_derived(self, "inner", inner)
+
+@dataclass(frozen=True)
+class NegativeLogLikelihoodLossOp(RenderableOpBase):
+    input0: str
+    target: str
+    weight: str | None
+    output: str
+    input_shape: tuple[int, ...]
+    target_shape: tuple[int, ...]
+    output_shape: tuple[int, ...]
+    n: int
+    c: int
+    d: int
+    reduction: str
+    ignore_index: int
+    input_dtype: ScalarType
+    weight_dtype: ScalarType | None
+    weight_shape: tuple[int, ...] | None
+    dtype: ScalarType
+    target_dtype: ScalarType
+
+@dataclass(frozen=True)
+class SoftmaxCrossEntropyLossOp(RenderableOpBase):
+    input0: str
+    target: str
+    weight: str | None
+    output: str
+    log_prob: str | None
+    input_shape: tuple[int, ...]
+    target_shape: tuple[int, ...]
+    output_shape: tuple[int, ...]
+    log_prob_shape: tuple[int, ...] | None
+    n: int
+    c: int
+    d: int
+    reduction: str
+    ignore_index: int | None
+    input_dtype: ScalarType
+    weight_dtype: ScalarType | None
+    weight_shape: tuple[int, ...] | None
+    dtype: ScalarType
+    target_dtype: ScalarType
+
+@dataclass(frozen=True)
+class BatchNormOp(RenderableOpBase):
+    input0: str
+    scale: str
+    bias: str
+    mean: str
+    variance: str
+    output: str
+    shape: tuple[int, ...]
+    channels: int
+    epsilon: float
+    dtype: ScalarType
+
+@dataclass(frozen=True)
+class LpNormalizationOp(RenderableOpBase):
+    input0: str
+    output: str
+    shape: tuple[int, ...]
+    axis: int
+    p: int
+    outer: int
+    axis_size: int
+    inner: int
+    dtype: ScalarType
+
+@dataclass(frozen=True)
+class InstanceNormalizationOp(RenderableOpBase):
+    input0: str
+    scale: str
+    bias: str
+    output: str
+    shape: tuple[int, ...]
+    channels: int
+    spatial_size: int
+    epsilon: float
+    dtype: ScalarType
+
+@dataclass(frozen=True)
+class GroupNormalizationOp(RenderableOpBase):
+    input0: str
+    scale: str
+    bias: str
+    output: str
+    shape: tuple[int, ...]
+    channels: int
+    num_groups: int
+    group_size: int
+    spatial_size: int
+    epsilon: float
+    dtype: ScalarType
+
+@dataclass(frozen=True)
+class LayerNormalizationOp(RenderableOpBase):
+    input0: str
+    scale: str
+    bias: str | None
+    output: str
+    mean_output: str | None
+    invstd_output: str | None
+    shape: tuple[int, ...]
+    normalized_shape: tuple[int, ...]
+    scale_shape: tuple[int, ...]
+    bias_shape: tuple[int, ...] | None
+    outer: int
+    inner: int
+    axis: int
+    epsilon: float
+    dtype: ScalarType
+
+@dataclass(frozen=True)
+class MeanVarianceNormalizationOp(RenderableOpBase):
+    input0: str
+    output: str
+    shape: tuple[int, ...]
+    axes: tuple[int, ...]
+    non_axes: tuple[int, ...]
+    reduce_count: int
+    epsilon: float
+    dtype: ScalarType
+
+@dataclass(frozen=True)
+class RMSNormalizationOp(RenderableOpBase):
+    input0: str
+    scale: str
+    output: str
+    shape: tuple[int, ...]
+    normalized_shape: tuple[int, ...]
+    scale_shape: tuple[int, ...]
+    outer: int
+    inner: int
+    axis: int
+    epsilon: float
+    dtype: ScalarType
+
+@dataclass(frozen=True)
+class LrnOp(RenderableOpBase):
+    input0: str
+    output: str
+    shape: tuple[int, ...]
+    channels: int
+    size: int
+    half: int
+    alpha: float
+    beta: float
+    bias: float
+    dtype: ScalarType
+
+@dataclass(frozen=True)
+class LstmOp(RenderableOpBase):
+    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
+
+@dataclass(frozen=True)
+class AdagradOp(RenderableOpBase):
+    rate: str
+    timestep: str
+    inputs: tuple[str, ...]
+    gradients: tuple[str, ...]
+    accumulators: tuple[str, ...]
+    outputs: tuple[str, ...]
+    accumulator_outputs: tuple[str, ...]
+    rate_shape: tuple[int, ...]
+    timestep_shape: tuple[int, ...]
+    tensor_shapes: tuple[tuple[int, ...], ...]
+    output_shapes: tuple[tuple[int, ...], ...]
+    dtype: ScalarType
+    rate_dtype: ScalarType
+    timestep_dtype: ScalarType
+    norm_coefficient: float
+    epsilon: float
+    decay_factor: float
+
+@dataclass(frozen=True)
+class MaxPoolOp(RenderableOpBase):
+    input0: str
+    output: str
+    indices: str | None
+    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
+    dtype: ScalarType
+    indices_dtype: ScalarType | None

emx_onnx_cgen/ir/ops/reduce.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from shared.scalar_types import ScalarType
+
+from ..op_base import ReduceOpBase
+from ..op_context import OpContext
+
+
+@dataclass(frozen=True)
+class ReduceOp(ReduceOpBase):
+    input0: str
+    output: str
+    input_shape: tuple[int, ...]
+    output_shape: tuple[int, ...]
+    axes: tuple[int, ...]
+    axes_input: str | None
+    axes_input_shape: tuple[int, ...] | None
+    axes_input_dtype: ScalarType | None
+    keepdims: bool
+    noop_with_empty_axes: bool
+    reduce_kind: str
+    reduce_count: int | None
+    dtype: ScalarType
+
+    def infer_types(self, ctx: OpContext) -> None:
+        ctx.dtype(self.output)
+
+    def infer_shapes(self, ctx: OpContext) -> None:
+        input_shape = ctx.shape(self.input0)
+        if self.axes_input is None:
+            axes = self.normalize_axes(self.axes, len(input_shape))
+            output_shape = self.reduced_shape(
+                input_shape, axes, keepdims=self.keepdims
+            )
+        else:
+            axes = self.axes
+            output_shape = ctx.shape(self.output)
+        ctx.set_shape(self.output, output_shape)
+        ctx.set_derived(self, "axes", axes)
+
+
+@dataclass(frozen=True)
+class ArgReduceOp(ReduceOpBase):
+    input0: str
+    output: str
+    input_shape: tuple[int, ...]
+    output_shape: tuple[int, ...]
+    axis: int
+    keepdims: bool
+    select_last_index: bool
+    reduce_kind: str
+    input_dtype: ScalarType
+    output_dtype: ScalarType
+
+    def infer_types(self, ctx: OpContext) -> None:
+        ctx.dtype(self.input0)
+        ctx.dtype(self.output)
+
+    def infer_shapes(self, ctx: OpContext) -> None:
+        input_shape = ctx.shape(self.input0)
+        axes = self.normalize_axes((self.axis,), len(input_shape))
+        output_shape = self.reduced_shape(
+            input_shape, axes, keepdims=self.keepdims
+        )
+        ctx.set_shape(self.output, output_shape)
+        ctx.set_derived(self, "axis", axes[0])
+
+
+@dataclass(frozen=True)
+class TopKOp(ReduceOpBase):
+    input0: str
+    output_values: str
+    output_indices: str
+    input_shape: tuple[int, ...]
+    output_shape: tuple[int, ...]
+    axis: int
+    k: int
+    largest: bool
+    sorted: bool
+    input_dtype: ScalarType
+    output_values_dtype: ScalarType
+    output_indices_dtype: ScalarType
+
+    def infer_types(self, ctx: OpContext) -> None:
+        ctx.dtype(self.input0)
+        ctx.dtype(self.output_values)
+        ctx.dtype(self.output_indices)
+
+    def infer_shapes(self, ctx: OpContext) -> None:
+        input_shape = ctx.shape(self.input0)
+        output_shape = ctx.shape(self.output_values)
+        ctx.set_shape(self.output_values, output_shape)
+        ctx.set_shape(self.output_indices, ctx.shape(self.output_indices))