Support broadcasting in type inference and checking

[fabio-innatera]

Right now, we assume (in both type inference and checking) that data/tensors flowing from one node to the next have the exact same shape.

PyTorch and other frameworks allow broadcasting, we should too.

ONNX documentation:
https://onnx.ai/onnx/repo-docs/Broadcasting.html

Numpy:
https://numpy.org/doc/stable/user/basics.broadcasting.html

PyTorch:
https://docs.pytorch.org/docs/stable/notes/broadcasting.html

[rowleya]

Sounds reasonable to me; this could even result in some time saving in data loading on the hardware if you do the broadcasting there directly. If this uses the numpy-compatible form, anyone who doesn't want to do this can use numpy if that works for their platform.

This would be ideal if it supported specifying parameters for a single instance of an object (such as a neuron) and then they would be copied to all other instances. I think this is a slight extension of the above, since that is more about the inputs and outputs than broadcasting the parameters - not sure if it needs to be split into another ticket therefore or can be handled here. I also don't know if this is as much of an NIR problem as an implementation detail.

[Jegp]

Excellent idea to support broadcasting in type inference. Fully behind this. How would it work in practice? Is there a simple "broadcasting extension" for the type checking we can do? Say, by permitting scalar values?

@rowleya is it correctly understood that you're suggesting that a single value (scalar) can cast to multiple neurons? If so, I wonder whether this is captured by the type declarations. NIR defines the output and input types, so if a parameter is specified by a single value, we can unambiguously assign the values to all neurons/parameters in the object.

[rowleya]

Yes, I think that is what I was suggesting; I know in many networks I have dealt with, the neurons have basically the same parameters, so it would avoid passing a list of all the same values. I guess if someone wants to specify a specific value for a few neurons, they might have to just list them all I assume (but happy to be told that there is another solution ;).

[fabio-innatera]

I implemented something like this in #157.

The current spec of NIR technically does not allow scalar values for parameters, but it is allowed because the type checking only checks if it's a numpy array - which can be an array of dim=0, i.e. a scalar.

#157 does some fancy type inference to figure out the number of neurons. Maybe it would be a better idea to allow the parameters to be scalars, but also require a num_neurons parameter on the node.