Restrict dimensions size to [0, INT32_MAX]

[philloooo]

Currently we use uint32 for dimensions, so the range is [0, UINT32_MAX].

Both CoreML and tflite backends actually use int32_t for dimensions, so the range can't exceed INT32_MAX, meaning if we pass dimensions that's between INT32_MAX and UINT32_MAX they faill on these backends.

We changed from int32 to uint32 to add the non-negative constraint, but it gaves an extra bit for the dimensions range that's not supported for two of the backends.

How can we make this better?

option 1 - change back to int32

Specifies in spec validation steps that they should be non negative.

option 2 - keep uint32

Specifies in spec validation steps that they can't exceed INT32_MAX

thoughts? @huningxin @fdwr @inexorabletash

[inexorabletash]

Those options seem about equivalent to me. Developers will see the same thing - an exception if they pass a value outside of [0, INT32_MAX]. Same amount of spec work and implementation work in either case I think. And we can relax in the future, e.g. treat negative indices as from the end, or increase to >31 bits if future backends all support it.

For implementers or future spec authors, I think it might be ever so slightly clearer to specify as (signed) int32, as that means they will probably think "what does negative mean?" and find the notes that explain it, rather than seeing uint32 and just assuming the spec actually allows 32-bit indices. But that's what really WPTs are for.

[fdwr]

🤔 I'd rather restrict the range to 2 billion (or whatever the backend's per-dimension limit might be) and keep the dimensions unsigned than introduce the implication (by adding a sign) that negative dimensions are legal/possible.

Reviewing Ningxin's table of low level API's from #279 and including StableHLO, we see TFLite and CoreML appear to be the minority (even Apple's MPS and BNNS use signless values).

native API	data type of dimension value
XNNPACK	size_t
DirectML	UINT
NNAPI	uint32_t
MLAS	size_t
OpenVINO	size_t
BNNS	size_t
MPS	NSUInteger
StableHLO	DimensionSize

Both CoreML and tflite backends actually use int32_t for dimensions

Is this specifically a TFLite limitation then, as it evidently diverges from the normal TF?

    x = tf.ones([65536, 65536])
    print("shape:", x.shape)
    print("dtype:", x.dtype)
    y = tf.reshape(x, [65536 * 65536])
    print("shape:", y.shape)
    print("dtype:", y.dtype)

shape: (65536, 65536)
dtype: <dtype: 'float32'>
shape: (4294967296,) ✅
dtype: <dtype: 'float32'>

There's also ...

option 3 - keep uint32 and report per-dimension limit

Report something like opSupportLimits or WebGPU's GPUSupportedLimits with the maximum dimension size for that backend (which could in theory be even less per-dimension than 2 billion, not that I've witnessed it). e.g. MLSupportedLimits::maximumDimensionLength.

I'm not too worried about the limitation in any case though, as it's very unlikely (I've never encountered a case) that you'd have a single tensor with a single dimension greater than 2 billion. Even more than 65'536 is highly unlikely, except during transient reshape's as a large 1D tensor.

in the future ... treat negative indices as from the end

😬 We already intentionally eliminated most of that higher level policy wiggliness, which should be resolved to actual values by the time it reaches lower level WebNN (like removing the null wildcard from reshape and negative axes). So I wouldn't want to regress that for dimensions or axes in general (with the one notable exception being gather because the data is directly in the tensors).

[huningxin]

option 2 - uint32_t with [0, INT32_MAX] range seems a good starter to me.

The spec can easily put this upper limit into valid dimension step and enable the synchronous validation at the renderer side. Then backend implementation (CoreML/TFLite) can cast dimension size to int32_t without asynchronous validation at build time.

Option 2 would also allow us moving to option 3 if needed. WDYT?