Skip to content

array_api_compat.torch.take with 0-d indices? #460

Description

@mdhaber

array_api_compat.torch.take with 0-d indices produces a result with the same number of dimensions as the input.

from array_api_compat import torch as xp
x = xp.asarray([1])
xp.take(x, xp.asarray(0))
# tensor([1])

For non-wrapped torch.take (and all other array libraries) the output has one fewer dimension than the input.

import torch
x = torch.tensor([1])
torch.take(x, torch.asarray(0))
tensor(1)

from array_api_compat import numpy as xp
x = xp.asarray([1])
xp.take(x, xp.asarray(0))
# np.int64(1)

Technically, the standard of take requires indices to have exactly one dimension.
However, it also says:

Conceptually, take(x, indices, axis=3) is equivalent to x[:,:,:,indices,...];...

This is up for discussion in data-apis/array-api#1010.

But I'd think that either all libraries should follow the letter of the standard as closely as possible in this situation ("The output array must have the same rank (i.e., number of dimensions) as x") or array_api_compat.torch.take should be consistent with torch.take.

Activity

  1. ev-br commented on Aug 6, 2026

    @ev-br
    Member

    The dreaded spec mandates behavior more general than what torch implements strikes back again:

    • the spec requires the axis argument and negative indices;
    • torch.take does support either of these;
    • the compat layer works around it.
    • the 0D behavior just falls out of the current workaround.

    The 0D behavior is indeed not mandated by the spec, so what compat layer does is an extension to both torch and the spec behaviors.
    A possible argument could be consistency with other major libraries in how they extend the spec:
    data-apis/array-api#1010

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions