Skip to content

PyTorch 'meta' device needs per-device capabilities #945

Open
@crusaderky

Description

@crusaderky

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,
    }

Metadata

Metadata

Assignees

No one assigned

    Labels

    API changeChanges to existing functions or objects in the API.

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions