Add Reshape scalar optimization and Gather scalar input support

[antimora]

(Note: ignore lots of commits; this PR is a left over from #4119)

This optimization appears in models like Silero VAD where scalar values are reshaped but should remain scalars for efficiency. Previously, Reshape(scalar, [-1]) would convert the scalar to a rank-1 tensor with one element, which is wasteful. The downstream Gather operation also needed to handle scalar inputs to complete the optimization chain.

Eventually we should generalize to mark a node no-op, see #4147

Example:

Unoptmized:

pub fn forward(
    &self,
    input: Tensor<B, 2>,
    sr: i64,
    state: Tensor<B, 3>,
) -> (Tensor<B, 2>, Tensor<B, 3>) {
    let constant32_out1 = 0i64;
    let constant33_out1 = 16000i64;
    let constant41_out1 = 1i64;
    let constant42_out1 = 2f32;
    let reshape1_out1 = Tensor::<B, 1usize, Int>::full([1], sr, &*self.device)
        .reshape([-1]);
    let gather1_out1 = {
        let indices = Tensor::<B, 1, _>::from_data([constant32_out1], &*self.device);
        let selected = Tensor::select(reshape1_out1, 0, indices);
        selected.into_scalar().elem::<i64>()
    };
    let equal1_out1 = gather1_out1 == constant33_out1;
    let (if1_out1, if1_out2) = if equal1_out1 {
        let state = state.clone();
        let input = input.clone();
        let pad3_out1 = input
            .pad((0, 64, 0, 0), burn::tensor::ops::PadMode::Reflect);

Optimized:

pub fn forward(
    &self,
    input: Tensor<B, 2>,
    sr: i64,
    state: Tensor<B, 3>,
) -> (Tensor<B, 2>, Tensor<B, 3>) {
    let constant32_out1 = 0i64;
    let constant33_out1 = 16000i64;
    let constant41_out1 = 1i64;
    let constant42_out1 = 2f32;
    let reshape1_out1 = sr;
    let gather1_out1 = reshape1_out1;
    let equal1_out1 = gather1_out1 == constant33_out1;
    let (if1_out1, if1_out2) = if equal1_out1 {
        let state = state.clone();
        let input = input.clone();
        let pad5_out1 = input
            .pad((0, 64, 0, 0), burn::tensor::ops::PadMode::Reflect);
     

Checklist

• Confirmed that cargo run-checks command has been executed.
• Made sure the book is up to date with changes in this PR.

Related Issues/PRs

#4119

Changes

• Optimize Reshape(scalar, [-1]) and Reshape(scalar, [1]) to keep output as Scalar instead of converting to rank-1 Tensor
• Add Scalar input support to Gather operator in both onnx-ir and burn-import
• Add ONNX tests for both new code paths

Testing

• Add reshape_scalar_to_scalar ONNX test: Reshape([1,1] → scalar) → Reshape(scalar, [-1]) → scalar
• Add gather_scalar_input ONNX test: Reshape → Gather(scalar, 0) → scalar
• Add unit tests for scalar input handling in both onnx-ir and burn-import

[Copilot]

Pull request overview

This PR adds an optimization to keep scalar values as scalars when reshaped with [-1] or [1], rather than unnecessarily converting them to rank-1 tensors. It also adds support for gathering from scalar inputs in the Gather operation to complete the optimization chain. This optimization appears in models like Silero VAD and avoids wasteful tensor conversions.

Key Changes

• Added scalar-to-scalar optimization in Reshape for shapes [-1] and [1]
• Extended Gather operator to handle scalar inputs in both onnx-ir and burn-import
• Added comprehensive ONNX integration tests for both optimization paths

Reviewed changes

Copilot reviewed 8 out of 10 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
crates/onnx-ir/src/node/reshape.rs	Added Case 2 logic to keep scalars as scalars when reshaped to [-1] or [1], updated determine_output_type signature
crates/onnx-ir/src/node/gather.rs	Added scalar input support in type inference and config extraction, replaced error handling with scalar-specific logic, added unit test
crates/burn-import/src/burn/node/gather.rs	Added forward_scalar_gather function for scalar pass-through, updated match to handle scalar inputs, added snapshot tests
crates/burn-import/onnx-tests/tests/reshape/reshape_scalar_to_scalar.py	Python script to generate ONNX test model for reshape scalar optimization
crates/burn-import/onnx-tests/tests/reshape/reshape_scalar_to_scalar.onnx	Generated ONNX model for reshape scalar test
crates/burn-import/onnx-tests/tests/reshape/mod.rs	Added reshape_scalar_to_scalar integration test
crates/burn-import/onnx-tests/tests/gather/gather_scalar_input.py	Python script to generate ONNX test model for gather with scalar input
crates/burn-import/onnx-tests/tests/gather/gather_scalar_input.onnx	Generated ONNX model for gather scalar test
crates/burn-import/onnx-tests/tests/gather/mod.rs	Added gather_scalar_input integration test
crates/burn-import/onnx-tests/build.rs	Added new test models to build inputs

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[codecov]

Codecov Report

❌ Patch coverage is 97.24138% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 68.43%. Comparing base (38eeb1e) to head (74f7519).
⚠️ Report is 7 commits behind head on main.

Files with missing lines	Patch %	Lines
crates/onnx-ir/src/node/reshape.rs	95.23%	3 Missing ⚠️
crates/onnx-ir/src/node/gather.rs	95.83%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #4146      +/-   ##
==========================================
+ Coverage   68.39%   68.43%   +0.03%     
==========================================
  Files        1281     1281              
  Lines      156146   156274     +128     
==========================================
+ Hits       106802   106948     +146     
+ Misses      49344    49326      -18     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[laggui]

So this is only possible when the value is static or from a constant right? No runtime values

[antimora]

Yes, exactly. The optimization only applies when the shape value is statically (static or constant) known at graph-build time.