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

emx_onnx_cgen/templates/dequantize_linear_op.c.j2

@@ -0,0 +1,10 @@
+static inline void {{ op_name }}({{ dim_args }}{{ params | join(', ') }}) {
+{% for dim in shape %}
+for (idx_t {{ loop_vars[loop.index0] }} = 0; {{ loop_vars[loop.index0] }} < {{ dim }}; ++{{ loop_vars[loop.index0] }}) {
+{% endfor %}
+    {{ compute_type }} dequantized = (({{ compute_type }}){{ input_expr }} - ({{ compute_type }}){{ zero_expr }}) * ({{ compute_type }}){{ scale_expr }};
+    {{ output_expr }} = ({{ output_c_type }})dequantized;
+{% for _ in shape %}
+}
+{% endfor %}
+}

emx_onnx_cgen/templates/einsum_op.c.j2

@@ -0,0 +1,55 @@
+static inline void {{ op_name }}({{ dim_args }}{{ params | join(', ') }}) {
+{% if kind == "reduce_all" %}
+    {{ acc_type }} acc = {{ zero_literal }};
+{% for var in input_loop_vars %}
+    {% for indent in range(loop.index0) %}    {% endfor %}for (idx_t {{ var }} = 0; {{ var }} < {{ input_loop_bounds[loop.index0] }}; ++{{ var }}) {
+{% endfor %}
+        {% for indent in range(input_loop_vars | length) %}    {% endfor %}acc += {{ input_expr }};
+{% for _ in input_loop_vars | reverse %}
+    {% for indent in range(loop.index0) %}    {% endfor %}}
+{% endfor %}
+    {{ output_expr }} = acc;
+{% elif kind == "sum_j" %}
+    for (idx_t {{ output_loop_vars[0] }} = 0; {{ output_loop_vars[0] }} < {{ output_loop_bounds[0] }}; ++{{ output_loop_vars[0] }}) {
+        {{ acc_type }} acc = {{ zero_literal }};
+        for (idx_t {{ reduce_loop_var }} = 0; {{ reduce_loop_var }} < {{ reduce_loop_bound }}; ++{{ reduce_loop_var }}) {
+            acc += {{ input_expr }};
+        }
+        {{ output_expr }} = acc;
+    }
+{% elif kind == "transpose" %}
+{% for var in output_loop_vars %}
+    {% for indent in range(loop.index0) %}    {% endfor %}for (idx_t {{ var }} = 0; {{ var }} < {{ output_loop_bounds[loop.index0] }}; ++{{ var }}) {
+{% endfor %}
+        {% for indent in range(output_loop_vars | length) %}    {% endfor %}{{ output_expr }} = {{ input_expr }};
+{% for _ in output_loop_vars | reverse %}
+    {% for indent in range(loop.index0) %}    {% endfor %}}
+{% endfor %}
+{% elif kind == "dot" %}
+    {{ acc_type }} acc = {{ zero_literal }};
+    for (idx_t {{ reduce_loop_var }} = 0; {{ reduce_loop_var }} < {{ reduce_loop_bound }}; ++{{ reduce_loop_var }}) {
+        acc += {{ input0_expr }} * {{ input1_expr }};
+    }
+    {{ output_expr }} = acc;
+{% elif kind == "batch_matmul" %}
+{% for var in output_loop_vars %}
+    {% for indent in range(loop.index0) %}    {% endfor %}for (idx_t {{ var }} = 0; {{ var }} < {{ output_loop_bounds[loop.index0] }}; ++{{ var }}) {
+{% endfor %}
+        {% for indent in range(output_loop_vars | length) %}    {% endfor %}{{ acc_type }} acc = {{ zero_literal }};
+        {% for indent in range(output_loop_vars | length) %}    {% endfor %}for (idx_t {{ reduce_loop_var }} = 0; {{ reduce_loop_var }} < {{ reduce_loop_bound }}; ++{{ reduce_loop_var }}) {
+            {% for indent in range(output_loop_vars | length + 1) %}    {% endfor %}acc += {{ input0_expr }} * {{ input1_expr }};
+        {% for indent in range(output_loop_vars | length) %}    {% endfor %}}
+        {% for indent in range(output_loop_vars | length) %}    {% endfor %}{{ output_expr }} = acc;
+{% for _ in output_loop_vars | reverse %}
+    {% for indent in range(loop.index0) %}    {% endfor %}}
+{% endfor %}
+{% elif kind == "batch_diagonal" %}
+{% for var in output_loop_vars %}
+    {% for indent in range(loop.index0) %}    {% endfor %}for (idx_t {{ var }} = 0; {{ var }} < {{ output_loop_bounds[loop.index0] }}; ++{{ var }}) {
+{% endfor %}
+        {% for indent in range(output_loop_vars | length) %}    {% endfor %}{{ output_expr }} = {{ input_expr }};
+{% for _ in output_loop_vars | reverse %}
+    {% for indent in range(loop.index0) %}    {% endfor %}}
+{% endfor %}
+{% endif %}
+}

emx_onnx_cgen/templates/expand_op.c.j2

@@ -0,0 +1,14 @@
+static inline void {{ op_name }}({{ dim_args }}{{ params | join(', ') }}) {
+    const {{ c_type }} *input_data = (const {{ c_type }} *){{ input0 }};
+    {{ c_type }} *output_data = ({{ c_type }} *){{ output }};
+    idx_t output_index = 0;
+{% for dim in output_shape %}
+    for (idx_t {{ loop_vars[loop.index0] }} = 0; {{ loop_vars[loop.index0] }} < {{ dim }}; ++{{ loop_vars[loop.index0] }}) {
+{% endfor %}
+        idx_t input_index = {{ input_index_expr }};
+        output_data[output_index] = input_data[input_index];
+        output_index++;
+{% for _ in output_shape %}
+    }
+{% endfor %}
+}

emx_onnx_cgen/templates/eye_like_op.c.j2

@@ -0,0 +1,27 @@
+static inline void {{ op_name }}({{ dim_args }}{{ params | join(', ') }}) {
+    (void){{ input0 }};
+    {{ c_type }} *output_data = ({{ c_type }} *){{ output }};
+    idx_t total = (idx_t){{ batch_size }} * {{ rows }} * {{ cols }};
+    for (idx_t index = 0; index < total; ++index) {
+        output_data[index] = {{ zero_literal }};
+    }
+    idx_t k = {{ k }};
+    idx_t rows = {{ rows }};
+    idx_t cols = {{ cols }};
+    idx_t row_start = k >= 0 ? 0 : -k;
+    idx_t col_start = k >= 0 ? k : 0;
+    if (row_start >= rows || col_start >= cols) {
+        return;
+    }
+    idx_t max_rows = rows - row_start;
+    idx_t max_cols = cols - col_start;
+    idx_t diag_len = max_rows < max_cols ? max_rows : max_cols;
+    for (idx_t batch = 0; batch < {{ batch_size }}; ++batch) {
+        idx_t base = batch * rows * cols;
+        for (idx_t diag = 0; diag < diag_len; ++diag) {
+            idx_t row = row_start + diag;
+            idx_t col = col_start + diag;
+            output_data[base + row * cols + col] = {{ one_literal }};
+        }
+    }
+}

emx_onnx_cgen/templates/gather_elements_op.c.j2

@@ -0,0 +1,13 @@
+static inline void {{ op_name }}({{ dim_args }}{{ params | join(', ') }}) {
+{% for dim in output_shape %}
+for (idx_t {{ loop_vars[loop.index0] }} = 0; {{ loop_vars[loop.index0] }} < {{ dim }}; ++{{ loop_vars[loop.index0] }}) {
+{% endfor %}
+idx_t gather_index = {{ indices }}{% for var in loop_vars %}[{{ var }}]{% endfor %};
+if (gather_index < 0) {
+    gather_index += {{ axis_dim }};
+}
+{{ output }}{% for var in loop_vars %}[{{ var }}]{% endfor %} = {{ data }}{% for idx in data_indices %}[{{ idx }}]{% endfor %};
+{% for _ in output_shape %}
+}
+{% endfor %}
+}

emx_onnx_cgen/templates/gather_nd_op.c.j2

@@ -0,0 +1,29 @@
+static inline void {{ op_name }}({{ dim_args }}{{ params | join(', ') }}) {
+{% if indices_prefix_shape %}
+{% for dim in indices_prefix_shape %}
+for (idx_t {{ indices_prefix_loop_vars[loop.index0] }} = 0; {{ indices_prefix_loop_vars[loop.index0] }} < {{ dim }}; ++{{ indices_prefix_loop_vars[loop.index0] }}) {
+{% endfor %}
+{% endif %}
+{% for idx in range(index_depth) %}
+    idx_t index{{ idx }} = {{ indices }}{% for var in indices_prefix_loop_vars %}[{{ var }}]{% endfor %}[{{ idx }}];
+    if (index{{ idx }} < 0) {
+        index{{ idx }} += {{ data_shape[batch_dims + idx] }};
+    }
+{% endfor %}
+{% if tail_shape %}
+{% for dim in tail_shape %}
+    for (idx_t {{ tail_loop_vars[loop.index0] }} = 0; {{ tail_loop_vars[loop.index0] }} < {{ dim }}; ++{{ tail_loop_vars[loop.index0] }}) {
+{% endfor %}
+{% endif %}
+    {{ output_index_expr }} = {{ data_index_expr }};
+{% if tail_shape %}
+{% for _ in tail_shape %}
+    }
+{% endfor %}
+{% endif %}
+{% if indices_prefix_shape %}
+{% for _ in indices_prefix_shape %}
+}
+{% endfor %}
+{% endif %}
+}

emx_onnx_cgen/templates/gather_op.c.j2

@@ -0,0 +1,13 @@
+static inline void {{ op_name }}({{ dim_args }}{{ params | join(', ') }}) {
+{% for dim in output_shape %}
+for (idx_t {{ loop_vars[loop.index0] }} = 0; {{ loop_vars[loop.index0] }} < {{ dim }}; ++{{ loop_vars[loop.index0] }}) {
+{% endfor %}
+idx_t gather_index = {{ indices }}{% for var in indices_indices %}[{{ var }}]{% endfor %};
+if (gather_index < 0) {
+    gather_index += {{ axis_dim }};
+}
+{{ output }}{% for var in loop_vars %}[{{ var }}]{% endfor %} = {{ data }}{% for idx in data_indices %}[{{ idx }}]{% endfor %};
+{% for _ in output_shape %}
+}
+{% endfor %}
+}

emx_onnx_cgen/templates/gemm_op.c.j2

@@ -0,0 +1,35 @@
+static inline void {{ op_name }}({{ dim_args }}{{ params | join(', ') }}) {
+    for (idx_t i = 0; i < {{ m }}; ++i) {
+        for (idx_t j = 0; j < {{ n }}; ++j) {
+            {{ acc_type }} acc = {{ zero_literal }};
+            for (idx_t k = 0; k < {{ k }}; ++k) {
+                {% if trans_a %}
+                const {{ c_type }} a_val = {{ input_a }}[k][i];
+                {% else %}
+                const {{ c_type }} a_val = {{ input_a }}[i][k];
+                {% endif %}
+                {% if trans_b %}
+                const {{ c_type }} b_val = {{ input_b }}[j][k];
+                {% else %}
+                const {{ c_type }} b_val = {{ input_b }}[k][j];
+                {% endif %}
+                acc += a_val * b_val;
+            }
+            {% if input_c %}
+            {% if c_rank == 2 %}
+            idx_t c_i = {% if c_dim0 == 1 %}0{% else %}i{% endif %};
+            idx_t c_j = {% if c_dim1 == 1 %}0{% else %}j{% endif %};
+            const {{ c_type }} bias = {{ input_c }}[c_i][c_j];
+            {% elif c_rank == 1 %}
+            idx_t c_j = {% if c_dim1 == 1 %}0{% else %}j{% endif %};
+            const {{ c_type }} bias = {{ input_c }}[c_j];
+            {% else %}
+            const {{ c_type }} bias = {{ input_c }}[0];
+            {% endif %}
+            {{ output }}[i][j] = acc * {{ alpha_literal }} + bias * {{ beta_literal }};
+            {% else %}
+            {{ output }}[i][j] = acc * {{ alpha_literal }};
+            {% endif %}
+        }
+    }
+}

emx_onnx_cgen/templates/grid_sample_op.c.j2

@@ -0,0 +1,184 @@
+static inline double {{ op_name }}_reflect(double value, double x_min, double x_max) {
+    const double range = x_max - x_min;
+    if (range == 0.0) {
+        return x_min;
+    }
+    if (value < x_min) {
+        const double dx = x_min - value;
+        const int n = (int)(dx / range);
+        const double r = dx - (double)n * range;
+        return (n % 2 == 0) ? (x_min + r) : (x_max - r);
+    }
+    if (value > x_max) {
+        const double dx = value - x_max;
+        const int n = (int)(dx / range);
+        const double r = dx - (double)n * range;
+        return (n % 2 == 0) ? (x_max - r) : (x_min + r);
+    }
+    return value;
+}
+{% if mode == "cubic" %}
+
+static inline void {{ op_name }}_cubic_coeffs(double x, double coeffs[4]) {
+    const double alpha = -0.75;
+    const double abs_x = fabs(x);
+    const double inv_x = 1.0 - abs_x;
+    const double span = 2.0 - abs_x;
+    coeffs[0] = ((alpha * (abs_x + 1.0) - 5.0 * alpha) * (abs_x + 1.0) + 8.0 * alpha) * (abs_x + 1.0) - 4.0 * alpha;
+    coeffs[1] = ((alpha + 2.0) * abs_x - (alpha + 3.0)) * abs_x * abs_x + 1.0;
+    coeffs[2] = ((alpha + 2.0) * inv_x - (alpha + 3.0)) * inv_x * inv_x + 1.0;
+    coeffs[3] = ((alpha * span - 5.0 * alpha) * span + 8.0 * alpha) * span - 4.0 * alpha;
+}
+{% endif %}
+
+static inline void {{ op_name }}({{ dim_args }}{{ params | join(', ') }}) {
+    const int input_spatial[{{ spatial_rank }}] = { {% for dim in input_spatial %}{{ dim }}{% if not loop.last %}, {% endif %}{% endfor %} };
+{% if padding_mode != "zeros" %}
+    const double border_min[{{ spatial_rank }}] = { {% for dim in input_spatial %}{% if align_corners %}0.0{% else %}-0.5{% endif %}{% if not loop.last %}, {% endif %}{% endfor %} };
+    const double border_max[{{ spatial_rank }}] = { {% for dim in input_spatial %}{% if align_corners %}{{ dim - 1 }}.0{% else %}{{ dim - 0.5 }}{% endif %}{% if not loop.last %}, {% endif %}{% endfor %} };
+{% endif %}
+{% set n_var = output_loop_vars[0] %}
+{% set c_var = output_loop_vars[1] %}
+{% set spatial_vars = output_loop_vars[2:] %}
+    for (idx_t {{ n_var }} = 0; {{ n_var }} < {{ output_shape[0] }}; ++{{ n_var }}) {
+        for (idx_t {{ c_var }} = 0; {{ c_var }} < {{ output_shape[1] }}; ++{{ c_var }}) {
+{% for index in range(spatial_rank) %}
+            for (idx_t {{ spatial_vars[index] }} = 0; {{ spatial_vars[index] }} < {{ output_spatial[index] }}; ++{{ spatial_vars[index] }}) {
+{% endfor %}
+                double coords[{{ spatial_rank }}];
+{% for dim in range(spatial_rank) %}
+                const double grid_{{ dim }} = (double){{ grid }}[{{ n_var }}]{% for var in spatial_vars %}[{{ var }}]{% endfor %}[{{ spatial_rank - 1 - dim }}];
+{% if align_corners %}
+                coords[{{ dim }}] = (grid_{{ dim }} + 1.0) * (double)(input_spatial[{{ dim }}] - 1) / 2.0;
+{% else %}
+                coords[{{ dim }}] = ((grid_{{ dim }} + 1.0) * (double)input_spatial[{{ dim }}] - 1.0) / 2.0;
+{% endif %}
+{% endfor %}
+{% if padding_mode != "zeros" %}
+{% for dim in range(spatial_rank) %}
+                if (coords[{{ dim }}] < border_min[{{ dim }}] || coords[{{ dim }}] > border_max[{{ dim }}]) {
+{% if padding_mode == "border" %}
+                    if (coords[{{ dim }}] < 0.0) {
+                        coords[{{ dim }}] = 0.0;
+                    } else if (coords[{{ dim }}] > (double)(input_spatial[{{ dim }}] - 1)) {
+                        coords[{{ dim }}] = (double)(input_spatial[{{ dim }}] - 1);
+                    }
+{% else %}
+                    coords[{{ dim }}] = {{ op_name }}_reflect(coords[{{ dim }}], border_min[{{ dim }}], border_max[{{ dim }}]);
+{% endif %}
+                }
+{% endfor %}
+{% endif %}
+{% if mode == "nearest" %}
+                int in_bounds = 1;
+{% for dim in range(spatial_rank) %}
+                int idx{{ dim }} = (int)nearbyint(coords[{{ dim }}]);
+{% if padding_mode == "zeros" %}
+                if (idx{{ dim }} < 0 || idx{{ dim }} >= input_spatial[{{ dim }}]) {
+                    in_bounds = 0;
+                }
+{% elif padding_mode == "border" %}
+                if (idx{{ dim }} < 0) {
+                    idx{{ dim }} = 0;
+                } else if (idx{{ dim }} >= input_spatial[{{ dim }}]) {
+                    idx{{ dim }} = input_spatial[{{ dim }}] - 1;
+                }
+{% else %}
+                idx{{ dim }} = (int){{ op_name }}_reflect((double)idx{{ dim }}, border_min[{{ dim }}], border_max[{{ dim }}]);
+{% endif %}
+{% endfor %}
+                if (!in_bounds) {
+                    {{ output }}[{{ n_var }}][{{ c_var }}]{% for var in spatial_vars %}[{{ var }}]{% endfor %} = ({{ c_type }})0;
+                } else {
+                    {{ output }}[{{ n_var }}][{{ c_var }}]{% for var in spatial_vars %}[{{ var }}]{% endfor %} = {{ input0 }}[{{ n_var }}][{{ c_var }}]{% for dim in range(spatial_rank) %}[idx{{ dim }}]{% endfor %};
+                }
+{% elif mode == "linear" %}
+                int base[{{ spatial_rank }}];
+                int upper[{{ spatial_rank }}];
+                double w0[{{ spatial_rank }}];
+                double w1[{{ spatial_rank }}];
+{% for dim in range(spatial_rank) %}
+                base[{{ dim }}] = (int)floor(coords[{{ dim }}]);
+                upper[{{ dim }}] = base[{{ dim }}] + 1;
+                w1[{{ dim }}] = coords[{{ dim }}] - (double)base[{{ dim }}];
+                w0[{{ dim }}] = 1.0 - w1[{{ dim }}];
+{% endfor %}
+                double acc = 0.0;
+{% for offsets in linear_offsets %}
+                {
+                    double weight = 1.0;
+{% for dim in range(spatial_rank) %}
+                    int idx{{ dim }} = {% if offsets[dim] == 0 %}base[{{ dim }}]{% else %}upper[{{ dim }}]{% endif %};
+                    weight *= {% if offsets[dim] == 0 %}w0[{{ dim }}]{% else %}w1[{{ dim }}]{% endif %};
+{% endfor %}
+                    int in_bounds = 1;
+{% for dim in range(spatial_rank) %}
+{% if padding_mode == "zeros" %}
+                    if (idx{{ dim }} < 0 || idx{{ dim }} >= input_spatial[{{ dim }}]) {
+                        in_bounds = 0;
+                    }
+{% elif padding_mode == "border" %}
+                    if (idx{{ dim }} < 0) {
+                        idx{{ dim }} = 0;
+                    } else if (idx{{ dim }} >= input_spatial[{{ dim }}]) {
+                        idx{{ dim }} = input_spatial[{{ dim }}] - 1;
+                    }
+{% else %}
+                    idx{{ dim }} = (int){{ op_name }}_reflect((double)idx{{ dim }}, border_min[{{ dim }}], border_max[{{ dim }}]);
+{% endif %}
+{% endfor %}
+                    if (in_bounds) {
+                        acc += weight * (double){{ input0 }}[{{ n_var }}][{{ c_var }}]{% for dim in range(spatial_rank) %}[idx{{ dim }}]{% endfor %};
+                    }
+                }
+{% endfor %}
+                {{ output }}[{{ n_var }}][{{ c_var }}]{% for var in spatial_vars %}[{{ var }}]{% endfor %} = ({{ c_type }})acc;
+{% else %}
+                int base[{{ spatial_rank }}];
+                int idxs[{{ spatial_rank }}][4];
+                double coeffs[{{ spatial_rank }}][4];
+{% for dim in range(spatial_rank) %}
+                base[{{ dim }}] = (int)floor(coords[{{ dim }}]);
+                idxs[{{ dim }}][0] = base[{{ dim }}] - 1;
+                idxs[{{ dim }}][1] = base[{{ dim }}];
+                idxs[{{ dim }}][2] = base[{{ dim }}] + 1;
+                idxs[{{ dim }}][3] = base[{{ dim }}] + 2;
+                {{ op_name }}_cubic_coeffs(coords[{{ dim }}] - (double)base[{{ dim }}], coeffs[{{ dim }}]);
+{% endfor %}
+                double acc = 0.0;
+{% for offsets in cubic_offsets %}
+                {
+                    double weight = 1.0;
+{% for dim in range(spatial_rank) %}
+                    int idx{{ dim }} = idxs[{{ dim }}][{{ offsets[dim] }}];
+                    weight *= coeffs[{{ dim }}][{{ offsets[dim] }}];
+{% endfor %}
+                    int in_bounds = 1;
+{% for dim in range(spatial_rank) %}
+{% if padding_mode == "zeros" %}
+                    if (idx{{ dim }} < 0 || idx{{ dim }} >= input_spatial[{{ dim }}]) {
+                        in_bounds = 0;
+                    }
+{% elif padding_mode == "border" %}
+                    if (idx{{ dim }} < 0) {
+                        idx{{ dim }} = 0;
+                    } else if (idx{{ dim }} >= input_spatial[{{ dim }}]) {
+                        idx{{ dim }} = input_spatial[{{ dim }}] - 1;
+                    }
+{% else %}
+                    idx{{ dim }} = (int){{ op_name }}_reflect((double)idx{{ dim }}, border_min[{{ dim }}], border_max[{{ dim }}]);
+{% endif %}
+{% endfor %}
+                    if (in_bounds) {
+                        acc += weight * (double){{ input0 }}[{{ n_var }}][{{ c_var }}]{% for dim in range(spatial_rank) %}[idx{{ dim }}]{% endfor %};
+                    }
+                }
+{% endfor %}
+                {{ output }}[{{ n_var }}][{{ c_var }}]{% for var in spatial_vars %}[{{ var }}]{% endfor %} = ({{ c_type }})acc;
+{% endif %}
+{% for index in range(spatial_rank) %}
+            }
+{% endfor %}
+        }
+    }
+}

emx_onnx_cgen/templates/group_normalization_op.c.j2

@@ -0,0 +1,46 @@
+static inline void {{ op_name }}({{ dim_args }}{{ params | join(', ') }}) {
+{% for dim in shape[:1] %}
+for (idx_t {{ loop_vars[0] }} = 0; {{ loop_vars[0] }} < {{ dim }}; ++{{ loop_vars[0] }}) {
+{% endfor %}
+    for (idx_t g = 0; g < {{ num_groups }}; ++g) {
+        {{ c_type }} sum = {{ zero_literal }};
+        for (idx_t c_in_group = 0; c_in_group < {{ group_size }}; ++c_in_group) {
+            idx_t c = g * {{ group_size }} + c_in_group;
+{% for dim in shape[2:] %}
+            for (idx_t {{ loop_vars[loop.index0 + 2] }} = 0; {{ loop_vars[loop.index0 + 2] }} < {{ dim }}; ++{{ loop_vars[loop.index0 + 2] }}) {
+{% endfor %}
+                sum += {{ input0 }}[{{ loop_vars[0] }}][c]{% for var in loop_vars[2:] %}[{{ var }}]{% endfor %};
+{% for _ in shape[2:] %}
+            }
+{% endfor %}
+        }
+        {{ c_type }} mean = sum / ({{ group_size }} * {{ spatial_size }});
+        {{ c_type }} var = {{ zero_literal }};
+        for (idx_t c_in_group = 0; c_in_group < {{ group_size }}; ++c_in_group) {
+            idx_t c = g * {{ group_size }} + c_in_group;
+{% for dim in shape[2:] %}
+            for (idx_t {{ loop_vars[loop.index0 + 2] }} = 0; {{ loop_vars[loop.index0 + 2] }} < {{ dim }}; ++{{ loop_vars[loop.index0 + 2] }}) {
+{% endfor %}
+                {{ c_type }} diff = {{ input0 }}[{{ loop_vars[0] }}][c]{% for var in loop_vars[2:] %}[{{ var }}]{% endfor %} - mean;
+                var += diff * diff;
+{% for _ in shape[2:] %}
+            }
+{% endfor %}
+        }
+        {{ c_type }} denom = {{ sqrt_fn }}(var / ({{ group_size }} * {{ spatial_size }}) + {{ epsilon_literal }});
+        for (idx_t c_in_group = 0; c_in_group < {{ group_size }}; ++c_in_group) {
+            idx_t c = g * {{ group_size }} + c_in_group;
+{% for dim in shape[2:] %}
+            for (idx_t {{ loop_vars[loop.index0 + 2] }} = 0; {{ loop_vars[loop.index0 + 2] }} < {{ dim }}; ++{{ loop_vars[loop.index0 + 2] }}) {
+{% endfor %}
+                {{ output }}[{{ loop_vars[0] }}][c]{% for var in loop_vars[2:] %}[{{ var }}]{% endfor %} =
+                    ({{ input0 }}[{{ loop_vars[0] }}][c]{% for var in loop_vars[2:] %}[{{ var }}]{% endfor %} - mean) / denom * {{ scale }}[c] + {{ bias }}[c];
+{% for _ in shape[2:] %}
+            }
+{% endfor %}
+        }
+    }
+{% for _ in shape[:1] %}
+}
+{% endfor %}
+}

emx_onnx_cgen/templates/gru_op.c.j2

@@ -0,0 +1,152 @@
+static inline void {{ op_name }}({{ dim_args }}{{ params | join(', ') }}) {
+    {% for dir in range(num_directions) %}
+    {
+        const int dir = {{ dir }};
+        const int reverse = {% if direction == "reverse" %}1{% elif direction == "bidirectional" and dir == 1 %}1{% else %}0{% endif %};
+        {% set act_f = activation_functions[dir * 2] %}
+        {% set act_g = activation_functions[dir * 2 + 1] %}
+        {{ c_type }} H_prev[{{ batch_size }}][{{ hidden_size }}];
+        for (int b = 0; b < {{ batch_size }}; ++b) {
+            for (int h = 0; h < {{ hidden_size }}; ++h) {
+                {% if input_initial_h %}
+                {% if layout == 0 %}
+                H_prev[b][h] = {{ input_initial_h }}[dir][b][h];
+                {% else %}
+                H_prev[b][h] = {{ input_initial_h }}[b][dir][h];
+                {% endif %}
+                {% else %}
+                H_prev[b][h] = {{ zero_literal }};
+                {% endif %}
+            }
+        }
+        for (int b = 0; b < {{ batch_size }}; ++b) {
+            int seq_limit = {{ seq_length }};
+            {% if input_sequence_lens %}
+            seq_limit = (int){{ input_sequence_lens }}[b];
+            if (seq_limit < 0) {
+                seq_limit = 0;
+            }
+            if (seq_limit > {{ seq_length }}) {
+                seq_limit = {{ seq_length }};
+            }
+            {% endif %}
+            for (int step = 0; step < seq_limit; ++step) {
+                int t = reverse ? (seq_limit - 1 - step) : step;
+                {{ c_type }} Z_gate[{{ hidden_size }}];
+                {{ c_type }} R_gate[{{ hidden_size }}];
+                {{ c_type }} H_next[{{ hidden_size }}];
+                for (int h = 0; h < {{ hidden_size }}; ++h) {
+                    {{ c_type }} gate_z = {{ zero_literal }};
+                    {{ c_type }} gate_r = {{ zero_literal }};
+                    for (int i = 0; i < {{ input_size }}; ++i) {
+                        {% if layout == 0 %}
+                        {{ c_type }} x_val = {{ input_x }}[t][b][i];
+                        {% else %}
+                        {{ c_type }} x_val = {{ input_x }}[b][t][i];
+                        {% endif %}
+                        gate_z += x_val * {{ input_w }}[dir][h][i];
+                        gate_r += x_val * {{ input_w }}[dir][{{ hidden_size }} + h][i];
+                    }
+                    for (int i = 0; i < {{ hidden_size }}; ++i) {
+                        {{ c_type }} h_val = H_prev[b][i];
+                        gate_z += h_val * {{ input_r }}[dir][h][i];
+                        gate_r += h_val * {{ input_r }}[dir][{{ hidden_size }} + h][i];
+                    }
+                    {% if input_b %}
+                    gate_z += {{ input_b }}[dir][h] + {{ input_b }}[dir][{{ hidden_size }} * 3 + h];
+                    gate_r += {{ input_b }}[dir][{{ hidden_size }} + h] + {{ input_b }}[dir][{{ hidden_size }} * 4 + h];
+                    {% endif %}
+                    {% if use_clip %}
+                    if (gate_z > {{ clip_literal }}) {
+                        gate_z = {{ clip_literal }};
+                    } else if (gate_z < -{{ clip_literal }}) {
+                        gate_z = -{{ clip_literal }};
+                    }
+                    if (gate_r > {{ clip_literal }}) {
+                        gate_r = {{ clip_literal }};
+                    } else if (gate_r < -{{ clip_literal }}) {
+                        gate_r = -{{ clip_literal }};
+                    }
+                    {% endif %}
+                    Z_gate[h] = {{ act_f }}(gate_z);
+                    R_gate[h] = {{ act_f }}(gate_r);
+                }
+                for (int h = 0; h < {{ hidden_size }}; ++h) {
+                    {{ c_type }} gate_h = {{ zero_literal }};
+                    for (int i = 0; i < {{ input_size }}; ++i) {
+                        {% if layout == 0 %}
+                        {{ c_type }} x_val = {{ input_x }}[t][b][i];
+                        {% else %}
+                        {{ c_type }} x_val = {{ input_x }}[b][t][i];
+                        {% endif %}
+                        gate_h += x_val * {{ input_w }}[dir][{{ hidden_size }} * 2 + h][i];
+                    }
+                    {% if linear_before_reset %}
+                    {{ c_type }} recur = {{ zero_literal }};
+                    for (int i = 0; i < {{ hidden_size }}; ++i) {
+                        recur += H_prev[b][i] * {{ input_r }}[dir][{{ hidden_size }} * 2 + h][i];
+                    }
+                    {% if input_b %}
+                    recur += {{ input_b }}[dir][{{ hidden_size }} * 5 + h];
+                    {% endif %}
+                    gate_h += R_gate[h] * recur;
+                    {% if input_b %}
+                    gate_h += {{ input_b }}[dir][{{ hidden_size }} * 2 + h];
+                    {% endif %}
+                    {% else %}
+                    for (int i = 0; i < {{ hidden_size }}; ++i) {
+                        gate_h += (R_gate[i] * H_prev[b][i]) * {{ input_r }}[dir][{{ hidden_size }} * 2 + h][i];
+                    }
+                    {% if input_b %}
+                    gate_h += {{ input_b }}[dir][{{ hidden_size }} * 2 + h] + {{ input_b }}[dir][{{ hidden_size }} * 5 + h];
+                    {% endif %}
+                    {% endif %}
+                    {% if use_clip %}
+                    if (gate_h > {{ clip_literal }}) {
+                        gate_h = {{ clip_literal }};
+                    } else if (gate_h < -{{ clip_literal }}) {
+                        gate_h = -{{ clip_literal }};
+                    }
+                    {% endif %}
+                    {{ c_type }} h_candidate = {{ act_g }}(gate_h);
+                    {{ c_type }} h_prev = H_prev[b][h];
+                    {{ c_type }} h_new = ({{ one_literal }} - Z_gate[h]) * h_candidate + Z_gate[h] * h_prev;
+                    H_next[h] = h_new;
+                    {% if output_y %}
+                    {% if layout == 0 %}
+                    {{ output_y }}[step][dir][b][h] = h_new;
+                    {% else %}
+                    {{ output_y }}[b][step][dir][h] = h_new;
+                    {% endif %}
+                    {% endif %}
+                }
+                for (int h = 0; h < {{ hidden_size }}; ++h) {
+                    H_prev[b][h] = H_next[h];
+                }
+            }
+            {% if output_y %}
+            for (int step = seq_limit; step < {{ seq_length }}; ++step) {
+                for (int h = 0; h < {{ hidden_size }}; ++h) {
+                    {% if layout == 0 %}
+                    {{ output_y }}[step][dir][b][h] = {{ zero_literal }};
+                    {% else %}
+                    {{ output_y }}[b][step][dir][h] = {{ zero_literal }};
+                    {% endif %}
+                }
+            }
+            {% endif %}
+        }
+        {% if output_y_h %}
+        for (int b = 0; b < {{ batch_size }}; ++b) {
+            for (int h = 0; h < {{ hidden_size }}; ++h) {
+                {% if layout == 0 %}
+                {{ output_y_h }}[dir][b][h] = H_prev[b][h];
+                {% else %}
+                {{ output_y_h }}[b][dir][h] = H_prev[b][h];
+                {% endif %}
+            }
+        }
+        {% endif %}
+    }
+    {% endfor %}
+}

emx_onnx_cgen/templates/hamming_window_op.c.j2

@@ -0,0 +1,12 @@
+static inline void {{ op_name }}({{ dim_args }}{{ params | join(', ') }}) {
+    const double size_value = (double){{ size }}[0];
+    const double denom = {{ periodic_literal }} ? size_value : (size_value - 1.0);
+    const double alpha = 25.0 / 46.0;
+    const double beta = 1.0 - alpha;
+    const double pi = 3.14159265358979323846;
+    for (idx_t idx = 0; idx < {{ length }}; ++idx) {
+        const double phase = (2.0 * pi * (double)idx) / denom;
+        const double value = alpha - beta * cos(phase);
+        {{ output }}[idx] = ({{ c_type }})value;
+    }
+}

emx_onnx_cgen/templates/hardmax_op.c.j2

@@ -0,0 +1,24 @@
+static inline void {{ op_name }}({{ dim_args }}{{ params | join(', ') }}) {
+    const {{ c_type }} *input_flat = (const {{ c_type }} *){{ input0 }};
+    {{ c_type }} *output_flat = ({{ c_type }} *){{ output }};
+    const idx_t outer = {{ outer }};
+    const idx_t axis_size = {{ axis_size }};
+    const idx_t inner = {{ inner }};
+    for (idx_t outer_idx = 0; outer_idx < outer; ++outer_idx) {
+        for (idx_t inner_idx = 0; inner_idx < inner; ++inner_idx) {
+            idx_t base = (outer_idx * axis_size * inner) + inner_idx;
+            {{ c_type }} max_value = input_flat[base];
+            idx_t max_index = 0;
+            for (idx_t axis_idx = 1; axis_idx < axis_size; ++axis_idx) {
+                {{ c_type }} value = input_flat[base + axis_idx * inner];
+                const {{ c_type }} prev_max = max_value;
+                max_value = {{ max_fn }}(max_value, value);
+                max_index = (value > prev_max) ? axis_idx : max_index;
+            }
+            for (idx_t axis_idx = 0; axis_idx < axis_size; ++axis_idx) {
+                output_flat[base + axis_idx * inner] =
+                    axis_idx == max_index ? {{ one_literal }} : {{ zero_literal }};
+            }
+        }
+    }
+}

emx_onnx_cgen/templates/identity_op.c.j2

@@ -0,0 +1,9 @@
+static inline void {{ op_name }}({{ dim_args }}{{ params | join(', ') }}) {
+{% for dim in shape %}
+    for (idx_t {{ loop_vars[loop.index0] }} = 0; {{ loop_vars[loop.index0] }} < {{ dim }}; ++{{ loop_vars[loop.index0] }}) {
+{% endfor %}
+    {{ output }}{% for var in loop_vars %}[{{ var }}]{% endfor %} = {{ input0 }}{% for var in loop_vars %}[{{ var }}]{% endfor %};
+{% for _ in shape %}
+    }
+{% endfor %}
+}