It's a scary change: not sure what will happen with XLA and SimpleGo ops with they find zero-dimension tensors. This may lead to undefined behavior on the ops -- since they were not designed for that.

It would be interesting to add tests for some ops of each type for some tensor with one dimension set to 0 and see if they behave correctly.

Also, you should leave the checks for negative dimensions -- those are not allowed.

It's a scary change: not sure what will happen with XLA and SimpleGo ops with they find zero-dimension tensors. This may lead to undefined behavior on the ops -- since they were not designed for that.

It would be interesting to add tests for some ops of each type for some tensor with one dimension set to 0 and see if they behave correctly.

Also, you should leave the checks for negative dimensions -- those are not allowed.

I completely agree, changing it to allow 0-dim tensors can be risky considering that we haven't done any testing. I will say that we ran Gemma inference with the XLA backend with these changes to allow 0-dim tensors and didn't run into issues, but didn't do any formal testing with that, and also didn't do any testing at all with the simplego backend. I will put the non-negative check back in for the meantime.

Thanks for the patience Riley. It seems XLA supports 0-dimension values as well.

This change has direct impact on a couple of things that I was going to ask you to test/fix. And indirect impact on all ops, but I volunteer to write the tests for them (with 0-dimension tensors) after your PR is merged -- XLA probably will just work, but the SimpleGo engine likely will require some tweaking.

1. In package pjrt, would you create a test creating and using a 0-dimension Buffer (of different ranks) ? Both as input and output of the execution of a computation ? You can use Client.BufferFromHost to create a buffers. It can be copied, or shared (if the device is CPU). And Client.NewSharedBuffer that only works on CPUs.
2. In package xlabuilder: create a test with a 0-dimension Literal. Functions to test:
   • NewLiteralFromShape
   • NewArrayLiteral: it will need a fix where it creates a pointer to the data, since the flat data passed can be nil.
   • Literal.Data also needs fixing for the 0-dimension tensor I think, when it tries to create a pointer (the data will be nil).
   • MakeShape and MakeShapeOrError (there is another one here): they also need to change the checking from dim <= 0 to dim < 0.

Ugh ... I should have supported 0-dimension tensors from the beginning, my bad.

Awesome, thanks @riley-oh6 !