-
Notifications
You must be signed in to change notification settings - Fork 59
Description
In proposal Add support for operations needed for well-known transformers, the Gather's indices data type is defined as "uint32":
indices: an MLOperand. The indices N-D tensor of the input values to gather. The values must be of type "uint32" in the range [0, N-1] where N is the size of the input dimension indexed by options.axis.
However, as discussed in Chromium CL-5017758 review, frameworks and native ML APIs usually use signed integer and some of them support negative index values:
Frameworks
- ONNX's Gather: indices data type is int32 or int64, supports negative index values, in range [-N, N-1]
- TensorFlow's Gather: indices data type is int32 or int64, only supports non-negative index values, in range [0, N-1]
Native ML APIs
- DirectML's Gather: indices data type is int64, int32, uint64 and uint32, supports negative index values
- CoreML's Gather: indices data type is integers, support negative index values
- StableHLO's Gather: indices data type is integers, doesn't support negative index values
Because the indices is an operand, its values may only available at run-time. If WebNN only supports "uint32" and range [0, N-1], frameworks that support negative indices, like ONNX, should polyfill (resolve the negative values) the indices tensor before passing to WebNN. WebNN implementation should enforce the restriction by polyfill the indices received from framework for all backends no matter whether the backend supports negative indices or not.
The best case would be single polyfill (frameworks supporting [0, N) + any backends), such as TF + DML, TF + StableHLO. The worst case might be double polyfill (frameworks support [-N, N) + any backends), such as ONNX + DML, ONNX + StableHLO.
The proposal is to support signed integers for Gather's indices and allow the negative index values, in range [-N, N-1] where N is the size of input dimensions indexed by options.axis.
With that, frameworks should just pass the indices tensor to WebNN, no polyfill needed. WebNN implementation would only need to polyfill the indices for the backends that don't support negative indices.
The best case would be non polyfill at all (any frameworks + backends supporting [-N, N)), such as ONNX + DML, TF + CoreML. The worst case might be single polyfill (any frameworks + backends supporting [0, N-1)), such as ONNX + StableHLO, TF + StableHLO.
/cc @wchao1115 @fdwr @wacky6