[burn-import] ArgType Shape dosen't seem to match specs

[QuantumEntangledAndy]

Describe the bug

Shape: According to the onnx reference is defined as a vector of integers

In burn_import ArgType::Shape is defined as ArgType::Shape(usize) as in it only stores the number of dimensions not the shape.

Expected behavior

Define ArgType::Shape(usize) as ArgType::Shape(Vec<usize>)

This is needed for gather op I am working on which seems to be able to accept both tensors and shapes of a tensor

Notes

If I get the time I'll work up a PR but thought I'd post here first for comments about what it is like this

[QuantumEntangledAndy]

Looking at this close I think it would be better to make the output of the Shape convert as

ArgType::Tensor(TensorType {
            elem_type: ElementType::Int64,
            dim: 1,
            shape: tensor.shape.map(|shapes| vec![shapes.len()]),
        });

No need for a seperate ArgType as the output of Shape is a tensor of integers, which is already well described

[antimora]

@CohenAriel has implemented ONNX Gather, see #947.

@QuantumEntangledAndy let us know if this issue should still be open.

[antimora]

Closing it for now. Please re-open it if still needed.