Repository navigation
PyTorch 'meta' device needs per-device capabilities #945
Description
PyTorch has a special device, meta, which is a dummy device with no underlying data. This is a very powerful testing tool.
However, PyTorch has not implemented special support for unknown shapes just for this device. As a result, unique etc. fail.
Proposal
Change capabilities() to capabilities(*, device=None), matching the signature of default_dtypes() and dtypes().
Change array_api_compat.torch:
def capabilities(*, device=None):
device = torch.get_default_device() if device is None else torch.device(device)
is_material = device.type != "meta"
return {
"boolean indexing": is_material,
"data-dependent shapes": is_material,
"max dimensions": 64,
}
Activity
crusaderky commentedon May 18, 2025
xref data-apis/array-api-extra#300
rgommers commentedon May 26, 2025
+1. xref #777 with a bunch of interest and discussion on this kind of functionality.
kgryte commentedon May 29, 2025
@jakevdp Would you have any objections to adding a
devicekwarg with the default beingNonefor this API? Would be good to get JAX's input on this.