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

emx_onnx_cgen/templates/resize_op.c.j2

@@ -0,0 +1,277 @@
+static inline void {{ op_name }}({{ dim_args }}{{ params | join(', ') }}) {
+    const int64_t input_shape[{{ rank }}] = { {% for dim in input_shape %}{{ dim }}{% if not loop.last %}, {% endif %}{% endfor %} };
+    const int64_t output_shape[{{ rank }}] EMX_UNUSED = { {% for dim in output_shape %}{{ dim }}{% if not loop.last %}, {% endif %}{% endfor %} };
+    double scales[{{ rank }}] EMX_UNUSED;
+{% if roi_input %}
+    double roi[{{ rank * 2 }}] EMX_UNUSED;
+{% if roi_axes %}
+    for (idx_t r = 0; r < {{ rank }}; ++r) {
+        roi[r] = 0.0;
+        roi[r + {{ rank }}] = 1.0;
+    }
+{% for axis in axes %}
+    roi[{{ axis }}] = (double){{ roi_input }}[{{ roi_axis_map[loop.index0] }}];
+    roi[{{ axis + rank }}] = (double){{ roi_input }}[{{ roi_axis_map[loop.index0] + (axes | length) }}];
+{% endfor %}
+{% else %}
+    for (idx_t r = 0; r < {{ rank * 2 }}; ++r) {
+        roi[r] = (double){{ roi_input }}[r];
+    }
+{% endif %}
+{% else %}
+    const double roi[{{ rank * 2 }}] EMX_UNUSED = {
+{% for _ in range(rank) %}        0.0{% if not loop.last %},{% endif %}
+{% endfor %}{% if rank > 0 %},{% endif %}
+{% for _ in range(rank) %}        1.0{% if not loop.last %},{% endif %}
+{% endfor %}
+    };
+{% endif %}
+{% if scales_input %}
+{% if scales_axes %}
+    for (idx_t r = 0; r < {{ rank }}; ++r) {
+        scales[r] = 1.0;
+    }
+{% for axis in axes %}
+    scales[{{ axis }}] = (double){{ scales_input }}[{{ scales_axis_map[loop.index0] }}];
+{% endfor %}
+{% else %}
+{% for axis in range(rank) %}
+    scales[{{ axis }}] = (double){{ scales_input }}[{{ axis }}];
+{% endfor %}
+{% endif %}
+{% elif sizes_input %}
+{% if keep_aspect_ratio_policy != "stretch" %}
+{% if keep_aspect_ratio_policy == "not_larger" %}
+    double scale_value = INFINITY;
+{% else %}
+    double scale_value = 0.0;
+{% endif %}
+{% for axis in axes %}
+    {
+        const double size_axis = (double){{ sizes_input }}[{{ sizes_axis_map[loop.index0] }}];
+        const double axis_scale = size_axis / (double)input_shape[{{ axis }}];
+{% if keep_aspect_ratio_policy == "not_larger" %}
+        if (axis_scale < scale_value) {
+            scale_value = axis_scale;
+        }
+{% else %}
+        if (axis_scale > scale_value) {
+            scale_value = axis_scale;
+        }
+{% endif %}
+    }
+{% endfor %}
+    for (idx_t r = 0; r < {{ rank }}; ++r) {
+        scales[r] = 1.0;
+    }
+{% for axis in axes %}
+    scales[{{ axis }}] = scale_value;
+{% endfor %}
+{% else %}
+    for (idx_t r = 0; r < {{ rank }}; ++r) {
+        scales[r] = 1.0;
+    }
+{% for axis in axes %}
+    scales[{{ axis }}] = (double){{ sizes_input }}[{{ sizes_axis_map[loop.index0] }}] / (double)input_shape[{{ axis }}];
+{% endfor %}
+{% endif %}
+{% else %}
+{% for axis in range(rank) %}
+    scales[{{ axis }}] = {{ scales[axis] }};
+{% endfor %}
+{% endif %}
+{% for dim in output_shape %}
+for (idx_t {{ loop_vars[loop.index0] }} = 0; {{ loop_vars[loop.index0] }} < {{ dim }}; ++{{ loop_vars[loop.index0] }}) {
+{% endfor %}
+int use_extrapolation = 0;
+{% for dim in range(rank) %}
+double x_orig{{ dim }};
+{% if coordinate_transformation_mode == "align_corners" %}
+if (output_shape[{{ dim }}] == 1) {
+    x_orig{{ dim }} = 0.0;
+} else {
+    x_orig{{ dim }} = (double){{ loop_vars[dim] }} * (input_shape[{{ dim }}] - 1) / (double)(output_shape[{{ dim }}] - 1);
+}
+{% elif coordinate_transformation_mode == "asymmetric" %}
+x_orig{{ dim }} = (double){{ loop_vars[dim] }} / scales[{{ dim }}];
+{% elif coordinate_transformation_mode == "tf_crop_and_resize" %}
+{
+    const double roi_start = roi[{{ dim }}];
+    const double roi_end = roi[{{ dim + rank }}];
+    if (output_shape[{{ dim }}] == 1) {
+        x_orig{{ dim }} = (roi_end - roi_start) * (input_shape[{{ dim }}] - 1) / 2.0;
+    } else {
+        x_orig{{ dim }} = (double){{ loop_vars[dim] }} * (roi_end - roi_start) * (input_shape[{{ dim }}] - 1)
+            / (double)(output_shape[{{ dim }}] - 1);
+    }
+    x_orig{{ dim }} += roi_start * (input_shape[{{ dim }}] - 1);
+    if (x_orig{{ dim }} < 0.0 || x_orig{{ dim }} > (double)(input_shape[{{ dim }}] - 1)) {
+        use_extrapolation = 1;
+    }
+}
+{% elif coordinate_transformation_mode == "pytorch_half_pixel" %}
+if (output_shape[{{ dim }}] == 1) {
+    x_orig{{ dim }} = -0.5;
+} else {
+    x_orig{{ dim }} = ((double){{ loop_vars[dim] }} + 0.5) / scales[{{ dim }}] - 0.5;
+}
+{% elif coordinate_transformation_mode == "half_pixel_symmetric" %}
+{
+    const double output_width = scales[{{ dim }}] * (double)input_shape[{{ dim }}];
+    const double adjustment = (double)output_shape[{{ dim }}] / output_width;
+    const double center = (double)input_shape[{{ dim }}] / 2.0;
+    const double offset = center * (1.0 - adjustment);
+    x_orig{{ dim }} = offset + ((double){{ loop_vars[dim] }} + 0.5) / scales[{{ dim }}] - 0.5;
+}
+{% else %}
+x_orig{{ dim }} = ((double){{ loop_vars[dim] }} + 0.5) / scales[{{ dim }}] - 0.5;
+{% endif %}
+{% endfor %}
+if (use_extrapolation) {
+    {{ output }}{% for var in loop_vars %}[{{ var }}]{% endfor %} = ({{ c_type }}){{ extrapolation_value }};
+} else {
+{% if mode == "nearest" %}
+{% for dim in range(rank) %}
+    const double x_val{{ dim }} = x_orig{{ dim }};
+    const double x_floor{{ dim }} EMX_UNUSED = floor(x_val{{ dim }});
+    const double x_ceil{{ dim }} EMX_UNUSED = ceil(x_val{{ dim }});
+    int idx{{ dim }};
+{% if nearest_mode == "round_prefer_floor" %}
+    idx{{ dim }} = (x_val{{ dim }} - x_floor{{ dim }} <= 0.5) ? (int)x_floor{{ dim }} : (int)x_ceil{{ dim }};
+{% elif nearest_mode == "round_prefer_ceil" %}
+    idx{{ dim }} = (x_val{{ dim }} - x_floor{{ dim }} < 0.5) ? (int)x_floor{{ dim }} : (int)x_ceil{{ dim }};
+{% elif nearest_mode == "floor" %}
+    idx{{ dim }} = (int)x_floor{{ dim }};
+{% else %}
+    idx{{ dim }} = (int)x_ceil{{ dim }};
+{% endif %}
+    if (idx{{ dim }} < 0) {
+        idx{{ dim }} = 0;
+    } else if (idx{{ dim }} >= input_shape[{{ dim }}]) {
+        idx{{ dim }} = (int)input_shape[{{ dim }}] - 1;
+    }
+{% endfor %}
+    {{ output }}{% for var in loop_vars %}[{{ var }}]{% endfor %} = {{ input0 }}{% for dim in range(rank) %}[idx{{ dim }}]{% endfor %};
+{% else %}
+    double acc = 0.0;
+{% for dim in range(rank) %}
+    double x_floor{{ dim }} = floor(x_orig{{ dim }});
+    double ratio{{ dim }} = x_orig{{ dim }} - x_floor{{ dim }};
+    const int is_integer{{ dim }} = ratio{{ dim }} == 0.0;
+    if (is_integer{{ dim }}) {
+        ratio{{ dim }} = 1.0;
+    }
+    int count{{ dim }};
+{% if antialias %}
+    double scale_clamped{{ dim }} = scales[{{ dim }}] < 1.0 ? scales[{{ dim }}] : 1.0;
+{% if mode == "linear" %}
+    int coeff_start{{ dim }} = (int)floor(-1.0 / scale_clamped{{ dim }}) + 1;
+    count{{ dim }} = 2 - 2 * coeff_start{{ dim }};
+{% else %}
+    int coeff_start{{ dim }} = (int)floor(-2.0 / scale_clamped{{ dim }}) + 1;
+    int coeff_end{{ dim }} = 2 - coeff_start{{ dim }};
+    count{{ dim }} = coeff_end{{ dim }} - coeff_start{{ dim }};
+{% endif %}
+{% else %}
+{% if mode == "linear" %}
+    count{{ dim }} = 2;
+{% else %}
+    count{{ dim }} = 4;
+{% endif %}
+{% endif %}
+    int start_index{{ dim }} = (int)x_floor{{ dim }} - (count{{ dim }} / 2);
+    if (!is_integer{{ dim }}) {
+        start_index{{ dim }} += 1;
+    }
+    int idx{{ dim }}_values[count{{ dim }}];
+    double coeff{{ dim }}[count{{ dim }}];
+{% if antialias %}
+    double coeff_sum{{ dim }} = 0.0;
+    for (int c = 0; c < count{{ dim }}; ++c) {
+        idx{{ dim }}_values[c] = start_index{{ dim }} + c;
+{% if mode == "linear" %}
+        double arg = (coeff_start{{ dim }} + c - ratio{{ dim }}) * scale_clamped{{ dim }};
+        double coeff = 1.0 - fabs(arg);
+        if (coeff < 0.0) {
+            coeff = 0.0;
+        }
+{% else %}
+        double x = scale_clamped{{ dim }} * (coeff_start{{ dim }} + c - ratio{{ dim }});
+        double ax = fabs(x);
+        double coeff;
+        if (ax <= 1.0) {
+            coeff = ({{ cubic_coeff_a }} + 2.0) * ax * ax * ax - ({{ cubic_coeff_a }} + 3.0) * ax * ax + 1.0;
+        } else if (ax < 2.0) {
+            coeff = {{ cubic_coeff_a }} * ax * ax * ax - 5.0 * {{ cubic_coeff_a }} * ax * ax + 8.0 * {{ cubic_coeff_a }} * ax - 4.0 * {{ cubic_coeff_a }};
+        } else {
+            coeff = 0.0;
+        }
+{% endif %}
+        coeff{{ dim }}[c] = coeff;
+        coeff_sum{{ dim }} += coeff;
+    }
+    if (coeff_sum{{ dim }} > 0.0) {
+        for (int c = 0; c < count{{ dim }}; ++c) {
+            coeff{{ dim }}[c] /= coeff_sum{{ dim }};
+        }
+    }
+{% else %}
+    for (int c = 0; c < count{{ dim }}; ++c) {
+        idx{{ dim }}_values[c] = start_index{{ dim }} + c;
+    }
+{% if mode == "linear" %}
+    coeff{{ dim }}[0] = 1.0 - ratio{{ dim }};
+    coeff{{ dim }}[1] = ratio{{ dim }};
+{% else %}
+    {
+        const double t = 1.0 - ratio{{ dim }};
+        coeff{{ dim }}[0] = (({{ cubic_coeff_a }} * (ratio{{ dim }} + 1.0) - 5.0 * {{ cubic_coeff_a }}) * (ratio{{ dim }} + 1.0) + 8.0 * {{ cubic_coeff_a }}) * (ratio{{ dim }} + 1.0) - 4.0 * {{ cubic_coeff_a }};
+        coeff{{ dim }}[1] = (({{ cubic_coeff_a }} + 2.0) * ratio{{ dim }} - ({{ cubic_coeff_a }} + 3.0)) * ratio{{ dim }} * ratio{{ dim }} + 1.0;
+        coeff{{ dim }}[2] = (({{ cubic_coeff_a }} + 2.0) * t - ({{ cubic_coeff_a }} + 3.0)) * t * t + 1.0;
+        coeff{{ dim }}[3] = (({{ cubic_coeff_a }} * (t + 1.0) - 5.0 * {{ cubic_coeff_a }}) * (t + 1.0) + 8.0 * {{ cubic_coeff_a }}) * (t + 1.0) - 4.0 * {{ cubic_coeff_a }};
+    }
+{% endif %}
+{% endif %}
+{% if exclude_outside %}
+    {
+        double coeff_sum{{ dim }} = 0.0;
+        for (int c = 0; c < count{{ dim }}; ++c) {
+            if (idx{{ dim }}_values[c] < 0 || idx{{ dim }}_values[c] >= input_shape[{{ dim }}]) {
+                coeff{{ dim }}[c] = 0.0;
+            }
+            coeff_sum{{ dim }} += coeff{{ dim }}[c];
+        }
+        if (coeff_sum{{ dim }} > 0.0) {
+            for (int c = 0; c < count{{ dim }}; ++c) {
+                coeff{{ dim }}[c] /= coeff_sum{{ dim }};
+            }
+        }
+    }
+{% endif %}
+{% endfor %}
+{% for dim in range(rank) %}
+    for (int n{{ dim }} = 0; n{{ dim }} < count{{ dim }}; ++n{{ dim }}) {
+{% endfor %}
+        double weight = 1.0;
+{% for dim in range(rank) %}
+        weight *= coeff{{ dim }}[n{{ dim }}];
+{% endfor %}
+{% for dim in range(rank) %}
+        int idx{{ dim }} = idx{{ dim }}_values[n{{ dim }}];
+        if (idx{{ dim }} < 0) {
+            idx{{ dim }} = 0;
+        } else if (idx{{ dim }} >= input_shape[{{ dim }}]) {
+            idx{{ dim }} = (int)input_shape[{{ dim }}] - 1;
+        }
+{% endfor %}
+        acc += weight * (double){{ input0 }}{% for dim in range(rank) %}[idx{{ dim }}]{% endfor %};
+{% for dim in range(rank) %}
+    }
+{% endfor %}
+    {{ output }}{% for var in loop_vars %}[{{ var }}]{% endfor %} = ({{ c_type }})acc;
+{% endif %}
+}
+{% for _ in output_shape %}
+}
+{% endfor %}
+}

emx_onnx_cgen/templates/rms_normalization_op.c.j2

@@ -0,0 +1,28 @@
+static inline void {{ op_name }}({{ dim_args }}{{ params | join(', ') }}) {
+{% for dim in prefix_shape %}
+for (idx_t {{ prefix_loop_vars[loop.index0] }} = 0; {{ prefix_loop_vars[loop.index0] }} < {{ dim }}; ++{{ prefix_loop_vars[loop.index0] }}) {
+{% endfor %}
+    {{ c_type }} sum = {{ zero_literal }};
+{% for dim in norm_shape %}
+    for (idx_t {{ norm_loop_vars[loop.index0] }} = 0; {{ norm_loop_vars[loop.index0] }} < {{ dim }}; ++{{ norm_loop_vars[loop.index0] }}) {
+{% endfor %}
+        {{ c_type }} value = {{ input0 }}{% for var in prefix_loop_vars %}[{{ var }}]{% endfor %}{% for var in norm_loop_vars %}[{{ var }}]{% endfor %};
+        sum += value * value;
+{% for _ in norm_shape %}
+    }
+{% endfor %}
+    {{ c_type }} mean_square = sum / {{ inner }};
+    {{ c_type }} denom = {{ sqrt_fn }}(mean_square + {{ epsilon_literal }});
+{% for dim in norm_shape %}
+    for (idx_t {{ norm_loop_vars[loop.index0] }} = 0; {{ norm_loop_vars[loop.index0] }} < {{ dim }}; ++{{ norm_loop_vars[loop.index0] }}) {
+{% endfor %}
+        {{ c_type }} value = {{ input0 }}{% for var in prefix_loop_vars %}[{{ var }}]{% endfor %}{% for var in norm_loop_vars %}[{{ var }}]{% endfor %} / denom;
+        value = value * {{ scale }}{% for var in scale_index_vars %}[{{ var }}]{% endfor %};
+        {{ output }}{% for var in prefix_loop_vars %}[{{ var }}]{% endfor %}{% for var in norm_loop_vars %}[{{ var }}]{% endfor %} = value;
+{% for _ in norm_shape %}
+    }
+{% endfor %}
+{% for _ in prefix_shape %}
+}
+{% endfor %}
+}

emx_onnx_cgen/templates/rotary_embedding_op.c.j2

@@ -0,0 +1,66 @@
+static inline void {{ op_name }}({{ dim_args }}{{ params | join(', ') }}) {
+{% if input_rank == 3 %}
+for (idx_t b = 0; b < {{ batch }}; ++b) {
+    for (idx_t s = 0; s < {{ seq_len }}; ++s) {
+        for (idx_t h = 0; h < {{ num_heads }}; ++h) {
+            idx_t head_offset = h * {{ head_size }};
+            for (idx_t d = 0; d < {{ rotary_dim_half }}; ++d) {
+{% if interleaved %}
+                idx_t idx1 = 2 * d;
+                idx_t idx2 = 2 * d + 1;
+{% else %}
+                idx_t idx1 = d;
+                idx_t idx2 = d + {{ rotary_dim_half }};
+{% endif %}
+{% if has_position_ids %}
+                idx_t pos = (idx_t){{ position_ids }}[b][s];
+                {{ c_type }} cos_val = {{ cos_cache }}[pos][d];
+                {{ c_type }} sin_val = {{ sin_cache }}[pos][d];
+{% else %}
+                {{ c_type }} cos_val = {{ cos_cache }}[b][s][d];
+                {{ c_type }} sin_val = {{ sin_cache }}[b][s][d];
+{% endif %}
+                {{ c_type }} x1 = {{ input0 }}[b][s][head_offset + idx1];
+                {{ c_type }} x2 = {{ input0 }}[b][s][head_offset + idx2];
+                {{ output }}[b][s][head_offset + idx1] = (cos_val * x1) - (sin_val * x2);
+                {{ output }}[b][s][head_offset + idx2] = (sin_val * x1) + (cos_val * x2);
+            }
+            for (idx_t d = {{ rotary_dim }}; d < {{ head_size }}; ++d) {
+                {{ output }}[b][s][head_offset + d] = {{ input0 }}[b][s][head_offset + d];
+            }
+        }
+    }
+}
+{% else %}
+for (idx_t b = 0; b < {{ batch }}; ++b) {
+    for (idx_t h = 0; h < {{ num_heads }}; ++h) {
+        for (idx_t s = 0; s < {{ seq_len }}; ++s) {
+            for (idx_t d = 0; d < {{ rotary_dim_half }}; ++d) {
+{% if interleaved %}
+                idx_t idx1 = 2 * d;
+                idx_t idx2 = 2 * d + 1;
+{% else %}
+                idx_t idx1 = d;
+                idx_t idx2 = d + {{ rotary_dim_half }};
+{% endif %}
+{% if has_position_ids %}
+                idx_t pos = (idx_t){{ position_ids }}[b][s];
+                {{ c_type }} cos_val = {{ cos_cache }}[pos][d];
+                {{ c_type }} sin_val = {{ sin_cache }}[pos][d];
+{% else %}
+                {{ c_type }} cos_val = {{ cos_cache }}[b][s][d];
+                {{ c_type }} sin_val = {{ sin_cache }}[b][s][d];
+{% endif %}
+                {{ c_type }} x1 = {{ input0 }}[b][h][s][idx1];
+                {{ c_type }} x2 = {{ input0 }}[b][h][s][idx2];
+                {{ output }}[b][h][s][idx1] = (cos_val * x1) - (sin_val * x2);
+                {{ output }}[b][h][s][idx2] = (sin_val * x1) + (cos_val * x2);
+            }
+            for (idx_t d = {{ rotary_dim }}; d < {{ head_size }}; ++d) {
+                {{ output }}[b][h][s][d] = {{ input0 }}[b][h][s][d];
+            }
+        }
+    }
+}
+{% endif %}
+}

emx_onnx_cgen/templates/scatter_nd_op.c.j2

@@ -0,0 +1,52 @@
+static inline void {{ op_name }}({{ dim_args }}{{ params | join(', ') }}) {
+{% for dim in output_shape %}
+for (idx_t {{ output_loop_vars[loop.index0] }} = 0; {{ output_loop_vars[loop.index0] }} < {{ dim }}; ++{{ output_loop_vars[loop.index0] }}) {
+{% endfor %}
+    {{ output }}{% for var in output_loop_vars %}[{{ var }}]{% endfor %} = {{ data }}{% for var in output_loop_vars %}[{{ var }}]{% endfor %};
+{% for _ in output_shape %}
+}
+{% endfor %}
+
+{% 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[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 %}
+    {{ c_type }} update_val = {{ updates_index_expr }};
+{% if reduction == "none" %}
+    {{ output_index_expr }} = update_val;
+{% elif reduction == "add" %}
+    {{ output_index_expr }} += update_val;
+{% elif reduction == "mul" %}
+    {{ output_index_expr }} *= update_val;
+{% elif reduction == "max" %}
+    if (update_val > {{ output_index_expr }}) {
+        {{ output_index_expr }} = update_val;
+    }
+{% elif reduction == "min" %}
+    if (update_val < {{ output_index_expr }}) {
+        {{ output_index_expr }} = update_val;
+    }
+{% endif %}
+{% if tail_shape %}
+{% for _ in tail_shape %}
+    }
+{% endfor %}
+{% endif %}
+{% if indices_prefix_shape %}
+{% for _ in indices_prefix_shape %}
+}
+{% endfor %}
+{% endif %}
+}

emx_onnx_cgen/templates/shape_op.c.j2

@@ -0,0 +1,6 @@
+static inline void {{ op_name }}({{ dim_args }}{{ params | join(', ') }}) {
+    (void){{ input0 }};
+{% for value in values %}
+    {{ output }}[{{ loop.index0 }}] = {{ value }};
+{% endfor %}
+}

emx_onnx_cgen/templates/size_op.c.j2

@@ -0,0 +1,4 @@
+static inline void {{ op_name }}({{ dim_args }}{{ params | join(', ') }}) {
+    (void){{ input0 }};
+    {{ output }}[0] = {{ value }};
+}

emx_onnx_cgen/templates/slice_op.c.j2

@@ -0,0 +1,9 @@
+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 %}
+{{ output }}{% for var in loop_vars %}[{{ var }}]{% endfor %} = {{ input0 }}{% for idx in input_indices %}[{{ idx }}]{% endfor %};
+{% for _ in output_shape %}
+}
+{% endfor %}
+}

emx_onnx_cgen/templates/slice_op_dynamic.c.j2

@@ -0,0 +1,70 @@
+static inline void {{ op_name }}({{ dim_args }}{{ params | join(', ') }}) {
+    const idx_t input_rank = {{ input_rank }};
+    const idx_t input_dims[] = { {% for dim in input_shape %}{{ dim }}{% if not loop.last %}, {% endif %}{% endfor %} };
+    idx_t start_indices[{{ input_rank }}];
+    idx_t step_values[{{ input_rank }}];
+    for (idx_t idx = 0; idx < input_rank; ++idx) {
+        start_indices[idx] = 0;
+        step_values[idx] = 1;
+    }
+    for (idx_t i = 0; i < {{ starts_len }}; ++i) {
+{% if axes_input %}
+        idx_t axis_value = {{ axes_input }}[i];
+{% else %}
+        idx_t axis_value = i;
+{% endif %}
+        if (axis_value < 0) {
+            axis_value += input_rank;
+        }
+        idx_t axis = axis_value;
+        idx_t dim = input_dims[axis];
+        idx_t start_value = {{ starts_input }}[i];
+        idx_t step_value = 1;
+{% if steps_input %}
+        step_value = {{ steps_input }}[i];
+{% endif %}
+        idx_t end_value = {{ ends_input }}[i];
+        if (start_value < 0) {
+            start_value += dim;
+        }
+        if (end_value < 0) {
+            end_value += dim;
+        }
+        if (step_value > 0) {
+            if (start_value < 0) {
+                start_value = 0;
+            }
+            if (start_value > dim) {
+                start_value = dim;
+            }
+            if (end_value < 0) {
+                end_value = 0;
+            }
+            if (end_value > dim) {
+                end_value = dim;
+            }
+        } else {
+            if (start_value < 0) {
+                start_value = -1;
+            }
+            if (start_value >= dim) {
+                start_value = dim - 1;
+            }
+            if (end_value < -1) {
+                end_value = -1;
+            }
+            if (end_value >= dim) {
+                end_value = dim - 1;
+            }
+        }
+        start_indices[axis] = start_value;
+        step_values[axis] = step_value;
+    }
+{% for dim in output_shape %}
+    for (idx_t {{ output_loop_vars[loop.index0] }} = 0; {{ output_loop_vars[loop.index0] }} < {{ dim }}; ++{{ output_loop_vars[loop.index0] }}) {
+{% endfor %}
+        {{ output }}{% for var in output_loop_vars %}[{{ var }}]{% endfor %} = {{ input0 }}{% for var in output_loop_vars %}[start_indices[{{ loop.index0 }}] + ({{ var }} * step_values[{{ loop.index0 }}])]{% endfor %};
+{% for _ in output_shape %}
+    }
+{% endfor %}
+}

emx_onnx_cgen/templates/softmax_cross_entropy_loss_op.c.j2

@@ -0,0 +1,105 @@
+static inline void {{ op_name }}({{ dim_args }}{{ params | join(', ') }}) {
+    const {{ c_type }} *input_flat = (const {{ c_type }} *){{ input0 }};
+    const {{ target_c_type }} *target_flat = (const {{ target_c_type }} *){{ target }};
+    {{ c_type }} *output_flat = ({{ c_type }} *){{ output }};
+{% if log_prob %}
+    {{ c_type }} *log_prob_flat = ({{ c_type }} *){{ log_prob }};
+{% endif %}
+    const idx_t n = {{ n }};
+    const idx_t c = {{ c }};
+    const idx_t d = {{ d }};
+{% if reduction != "none" %}
+    {{ acc_type }} loss_sum = {{ acc_zero_literal }};
+{% if reduction == "mean" and (weight or use_ignore_index) %}
+    {{ acc_type }} weight_sum = {{ acc_zero_literal }};
+{% endif %}
+{% endif %}
+    for (idx_t n_idx = 0; n_idx < n; ++n_idx) {
+        for (idx_t d_idx = 0; d_idx < d; ++d_idx) {
+            idx_t target_index = n_idx * d + d_idx;
+            {{ target_c_type }} target_value = target_flat[target_index];
+{% if use_ignore_index %}
+            bool ignored = ((int64_t)target_value == {{ ignore_index }});
+{% endif %}
+            idx_t class_index = (idx_t)target_value;
+            idx_t base = (n_idx * c * d) + d_idx;
+            {{ acc_type }} max_value = ({{ acc_type }})input_flat[base];
+            for (idx_t c_idx = 1; c_idx < c; ++c_idx) {
+                {{ acc_type }} value = ({{ acc_type }})input_flat[base + c_idx * d];
+                max_value = {{ max_fn }}(max_value, value);
+            }
+            {{ acc_type }} sum = {{ acc_zero_literal }};
+            for (idx_t c_idx = 0; c_idx < c; ++c_idx) {
+                {{ acc_type }} value = ({{ acc_type }})input_flat[base + c_idx * d] - max_value;
+                sum += {{ acc_exp_fn }}(value);
+            }
+            {{ acc_type }} logsum = {{ acc_log_fn }}(sum);
+{% if log_prob %}
+            {{ acc_type }} loss_value = {{ acc_zero_literal }};
+            for (idx_t c_idx = 0; c_idx < c; ++c_idx) {
+                {{ acc_type }} log_prob_value = ({{ acc_type }})input_flat[base + c_idx * d] - max_value - logsum;
+                log_prob_flat[base + c_idx * d] = ({{ c_type }})log_prob_value;
+                if (c_idx == class_index) {
+                    loss_value = -log_prob_value;
+                }
+            }
+{% if use_ignore_index %}
+            if (ignored) {
+{% if reduction == "none" %}
+                output_flat[target_index] = {{ zero_literal }};
+{% endif %}
+            }
+{% endif %}
+{% else %}
+            {{ acc_type }} loss_value = {{ acc_zero_literal }};
+{% if use_ignore_index %}
+            if (ignored) {
+{% if reduction == "none" %}
+                output_flat[target_index] = {{ zero_literal }};
+{% endif %}
+            } else {
+{% endif %}
+                {{ acc_type }} log_prob_value = ({{ acc_type }})input_flat[base + class_index * d] - max_value - logsum;
+                loss_value = -log_prob_value;
+{% if use_ignore_index %}
+            }
+{% endif %}
+{% endif %}
+{% if use_ignore_index %}
+            if (!ignored) {
+{% endif %}
+{% if weight %}
+                {{ acc_type }} sample_weight = {{ weight }}[class_index];
+                loss_value *= sample_weight;
+{% endif %}
+{% if reduction == "none" %}
+                output_flat[target_index] = loss_value;
+{% else %}
+                loss_sum += loss_value;
+{% if reduction == "mean" %}
+{% if weight %}
+                weight_sum += sample_weight;
+{% elif use_ignore_index %}
+                weight_sum += {{ acc_one_literal }};
+{% endif %}
+{% endif %}
+{% endif %}
+{% if use_ignore_index %}
+            }
+{% endif %}
+        }
+    }
+{% if reduction == "mean" %}
+{% if weight or use_ignore_index %}
+    if (weight_sum == {{ acc_zero_literal }}) {
+        output_flat[0] = {{ zero_literal }};
+    } else {
+        output_flat[0] = loss_sum / weight_sum;
+    }
+{% else %}
+    output_flat[0] = loss_sum / ({{ n }} * {{ d }});
+{% endif %}
+{% elif reduction == "sum" %}
+    output_flat[0] = loss_sum;
+{% endif %}
+}

emx_onnx_cgen/templates/softmax_op.c.j2

@@ -0,0 +1,26 @@
+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];
+            for (idx_t axis_idx = 1; axis_idx < axis_size; ++axis_idx) {
+                {{ c_type }} value = input_flat[base + axis_idx * inner];
+                max_value = {{ max_fn }}(max_value, value);
+            }
+            {{ c_type }} sum = 0;
+            for (idx_t axis_idx = 0; axis_idx < axis_size; ++axis_idx) {
+                {{ c_type }} value = {{ exp_fn }}(input_flat[base + axis_idx * inner] - max_value);
+                output_flat[base + axis_idx * inner] = value;
+                sum += value;
+            }
+            for (idx_t axis_idx = 0; axis_idx < axis_size; ++axis_idx) {
+                output_flat[base + axis_idx * inner] /= sum;
+            }
+        }
+    }
+}

emx_onnx_cgen/templates/space_to_depth_op.c.j2

@@ -0,0 +1,22 @@
+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 (idx_t n = 0; n < {{ batch }}; ++n) {
+        for (idx_t c_out = 0; c_out < {{ out_channels }}; ++c_out) {
+            idx_t c_in = c_out % {{ in_channels }};
+            idx_t temp = c_out / {{ in_channels }};
+            idx_t offset_h = temp / {{ blocksize }};
+            idx_t offset_w = temp % {{ blocksize }};
+            for (idx_t h_out = 0; h_out < {{ out_h }}; ++h_out) {
+                idx_t h_in = h_out * {{ blocksize }} + offset_h;
+                for (idx_t w_out = 0; w_out < {{ out_w }}; ++w_out) {
+                    idx_t w_in = w_out * {{ blocksize }} + offset_w;
+                    idx_t input_index = ((n * {{ in_channels }} + c_in) * {{ in_h }} + h_in) * {{ in_w }} + w_in;
+                    output_data[output_index] = input_data[input_index];
+                    output_index++;
+                }
+            }
+        }
+    }
+}