Fix indices dim check in gather_update_outputs

[laggui]

Checklist

• Confirmed that run-checks all script has been executed.

Related Issues/PRs

We missed a little something in PR #2141

Changes

Updated gather_update_outputs to accept scalar argtype when checking for indices dim.

[laggui]

@hexd0t you can also review since you contributed the last PR

[antimora]

LGTM

[hexd0t]

LGTM, although I'm not sure why the gather_scalar-Test I had added didn't fail?

Also, in the ONNX model I'm currently trying, I get Tensor(TensorType { elem_type: Int64, dim: 0, shape: Some([]) }) instead of Scalars at this point (which I understand burn doesn't really want to have anyways), stemming from an ONNX constant:

[codecov]

Codecov Report

Attention: Patch coverage is 57.14286% with 3 lines in your changes missing coverage. Please review.

Project coverage is 86.21%. Comparing base (5a0c1dc) to head (c7db6b8).
Report is 1 commits behind head on main.

Files	Patch %	Lines
crates/onnx-ir/src/dim_inference.rs	57.14%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2149      +/-   ##
==========================================
- Coverage   86.22%   86.21%   -0.01%     
==========================================
  Files         694      694              
  Lines       88754    88755       +1     
==========================================
  Hits        76524    76524              
- Misses      12230    12231       +1     

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

[laggui]

Seems like the model generated for the test actually is a TensorType with dim 0 and not a scalar type.

The error was reported to me by a user, and in his case the node received a ScalarType, hence the addition in this PR.

I didn't have time to check why in one case it remained a tensor type with dim 0 while the other was converted to a scalar type, but that might be a parsing issue 🤔

/edit: ahh it could be due to the constant as you point out - it seems they're not handled the same and we have a draft PR opened to fix that iirc.