Skip to content

ENH: Provide a way to disable flattening of 0d arrays to scalars #13105

@eric-wieser

Description

@eric-wieser

If only given 0d inputs, even if they are of type ndarray, ufuncs will decay their output to a scalar via [()] (as noted in #4563, #5819). While we can't change this behavior now without creating signficant pain downstream, we could add a way to opt out of it.

#13100 raises a case where np.fix resorts to calling np.asanyarray(np.ceil(x, out=out)) in order to ensure that the result is an array. Unfortunately, this has a number of draw-backs:

  • It discards duck-types, like dask arrays
  • It changes the dtype of 0d object arrays containing array-likes

Proposed implementation:

  • Create a new np.leave_wrapped sentinel object that can be passed as an out argument
  • Add support in ufunc.__call__, np.take, ... for passing out=np.leave_wrapped meaning "do not call PyArray_Return", causing the result to never be a scalar
  • Expose PyArray_Return to python as np.core.unpack_scalar
  • Implement np.fix as:
def fix(x, out=None):
    if out is None:
        do_unwrap = True
        out = np.leave_wrapped
    else:
        do_unwrap = False
    res = nx.ceil(x, out=out)
    res = nx.floor(x, out=res, where=nx.greater_equal(x, 0, out=np.leave_wrapped))

    if do_unwrap:
        res = np.unpack_scalar(res)   # PyArray_Return
    return res

Original unpack_scalars=True proposal

  • Add a new unpack_scalars=True kwarg to ufunc.__call__, ufunc.reduce. When False, the current behavior of going through PyArray_Return is disabled. Alternative names:
    • decay=True
    • unpack_0d=True
    • unpack_0d_ndarray=True (PyArray_Return already does not apply to subclasses)
  • Add a new np.unpack_scalar(arr) function to expose PyArray_Return to python code. This would not be overloadable with __array_function__, since existing uses of PyArray_Return are also not.

With these changes, the current implementation of np.fix would change from:

def fix(x, out=None):
    res = nx.asanyarray(nx.ceil(x, out=out))
    res = nx.floor(x, out=res, where=nx.greater_equal(x, 0))

    if out is None and type(res) is nx.ndarray:
        res = res[()]
    return res

to

def fix(x, out=None, *, unpack_scalars=True):
    res = nx.ceil(x, out=out, unpack_scalars=False)
    res = nx.floor(x, out=res, where=nx.greater_equal(x, 0, unpack_scalars=False), unpack_scalars=False)

    if unpack_scalars and out is None:
        res = np.unpack_scalar(res)
    return res
Details If needed, I could promote this to an NEP

Metadata

Metadata

Assignees

No one assigned

    Labels

    23 - Wish List62 - Python APIChanges or additions to the Python API. Mailing list should usually be notified.

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions