Skip to content

BUG: support array_like inputs in np.split and related functions - #32790

Open
bodapatisaikrishna wants to merge 1 commit into
numpy:mainfrom
bodapatisaikrishna:fix/split-array-like
Open

BUG: support array_like inputs in np.split and related functions#32790
bodapatisaikrishna wants to merge 1 commit into
numpy:mainfrom
bodapatisaikrishna:fix/split-array-like

Conversation

@bodapatisaikrishna

@bodapatisaikrishna bodapatisaikrishna commented Sep 25, 2026 •
edited
Loading

Copy link
Copy Markdown

PR summary

  • What problem does this PR solve?
    Fixes np.split throws AttributeError for list/ tuple ary and int sections, np.array_split does not #17463. Ensures np.split, np.array_split, np.hsplit, np.vsplit, and np.dsplit consistently accept array_like inputs (such as Python lists and tuples) by applying _nx.asanyarray(ary) at entry and using _nx.ndim(ary) > 1 in hsplit. This eliminates the AttributeError: 'list' object has no attribute 'shape' when splitting sequences with integer sections, and fixes incorrect slice calculations in array_split when axis != 0.

  • Why are you interested in working on this PR?
    I noticed the inconsistency between passing a list of indices vs an integer section count to np.split (e.g. np.split([1, 2, 3, 4], [2]) worked, but np.split([1, 2, 3, 4], 2) crashed with AttributeError), and saw that the type stubs in _shape_base_impl.pyi already expected ArrayLike.

  • How does the proposed change help you?
    It makes sequence splitting consistent across the entire family of split functions and prevents unexpected runtime AttributeError crashes for downstream code passing lists or tuples.

First time contributor introduction

Hi! I'm Sai Krishna, a machine learning engineer and open source contributor working with NumPy and the scientific Python ecosystem. I use NumPy extensively for numerical computing and data processing pipelines.

AI Disclosure

I used an AI coding assistant (Google Antigravity) to help explore the codebase, draft the code modifications in numpy/lib/_shape_base_impl.py, and draft test cases in numpy/lib/tests/test_shape_base.py. I personally reviewed the implementation details, confirmed that asanyarray preserves ndarray subclasses (such as MaskedArray) with zero-copy and is a no-op on ndarray, and verified that all 83 tests in test_shape_base.py and all 4,887 tests in numpy/lib/tests/ pass cleanly without regression.

@ikrommyd

ikrommyd commented Sep 25, 2026

Copy link
Copy Markdown
Member

You ignored the PR template, so I'm going to assume that this PR has been made with agentic AI. This violates our AI policy; we prefer to talk to humans.

@ikrommyd ikrommyd closed this Sep 25, 2026
@bodapatisaikrishna

bodapatisaikrishna commented Sep 25, 2026

Copy link
Copy Markdown
Author

Hi @ikrommyd, apologies for overlooking the PR template on my first submission! I am a real developer and verified the code and test suite locally (83 tests pass in test_shape_base.py, 0 regressions in numpy.lib).I have updated the PR description to follow the official template and added the full AI disclosure section per NumPy's AI policy. The fix addresses the longstanding issue in #17463. Could you please take a look and consider reopening the PR for review? Thank you!

@ikrommyd ikrommyd reopened this Sep 25, 2026
jorenham
jorenham reviewed Sep 25, 2026
View reviewed changes

@jorenham jorenham left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What about vsplit and dsplit?

@bodapatisaikrishna

bodapatisaikrishna commented Sep 25, 2026

Copy link
Copy Markdown
Author

Hi @jorenham,

In the initial commit, vsplit and dsplit already supported array-like inputs (with test cases added in test_shape_base.py) because they checked _nx.ndim(ary) < 2 / < 3 and delegated to split, which applied asanyarray.

To make the whole family symmetric and avoid redundant _nx.ndim array conversions, I have updated hsplit, vsplit, and dsplit to all consistently apply ary = _nx.asanyarray(ary) at entry and use ary.ndim directly.

The updated commit has been pushed to the PR. All 83 tests in test_shape_base.py pass.

jorenham
jorenham requested changes Sep 25, 2026
View reviewed changes

@jorenham jorenham left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Using asanyarray at the start is a breaking change for duck-typed arrays (that don't implement __array_function__) and only have shape, swapaxes, and __getitem__, e.g. np.split(jax_array, 2) and np.array_split(df, 3) (on pandas<3) currently return their own type, and with this PR they silently become ndarrays. The stubs also support this:

numpy/numpy/lib/_shape_base_impl.pyi

Lines 447 to 484 in a98529d

@overload
def split[SplitableT: _SupportsSplitOps](
ary: SplitableT,
indices_or_sections: _ShapeLike,
axis: SupportsIndex = 0,
) -> list[SplitableT]: ...
@overload
def split[ScalarT: np.generic](
ary: _ArrayLike[ScalarT],
indices_or_sections: _ShapeLike,
axis: SupportsIndex = 0,
) -> list[NDArray[ScalarT]]: ...
@overload
def split(ary: ArrayLike, indices_or_sections: _ShapeLike, axis: SupportsIndex = 0) -> list[NDArray[Incomplete]]: ...
# keep in sync with `numpy.ma.extras.hsplit`
@overload
def hsplit[SplitableT: _SupportsSplitOps](ary: SplitableT, indices_or_sections: _ShapeLike) -> list[SplitableT]: ...
@overload
def hsplit[ScalarT: np.generic](ary: _ArrayLike[ScalarT], indices_or_sections: _ShapeLike) -> list[NDArray[ScalarT]]: ...
@overload
def hsplit(ary: ArrayLike, indices_or_sections: _ShapeLike) -> list[NDArray[Incomplete]]: ...
#
@overload
def vsplit[SplitableT: _SupportsSplitOps](ary: SplitableT, indices_or_sections: _ShapeLike) -> list[SplitableT]: ...
@overload
def vsplit[ScalarT: np.generic](ary: _ArrayLike[ScalarT], indices_or_sections: _ShapeLike) -> list[NDArray[ScalarT]]: ...
@overload
def vsplit(ary: ArrayLike, indices_or_sections: _ShapeLike) -> list[NDArray[Incomplete]]: ...
#
@overload
def dsplit[SplitableT: _SupportsSplitOps](ary: SplitableT, indices_or_sections: _ShapeLike) -> list[SplitableT]: ...
@overload
def dsplit[ScalarT: np.generic](ary: _ArrayLike[ScalarT], indices_or_sections: _ShapeLike) -> list[NDArray[ScalarT]]: ...
@overload
def dsplit(ary: ArrayLike, indices_or_sections: _ShapeLike) -> list[NDArray[Incomplete]]: ...

@ngoldbaum ngoldbaum added the 57 - Close? Issues which may be closable unless discussion continued label Sep 26, 2026
@bodapatisaikrishna
BUG: support array_like inputs in np.split and related functions
26ce9a5 
Closes numpy#17463.

Ensure np.split, np.array_split, np.hsplit, np.vsplit, and np.dsplit accept sequence inputs without shape attribute (such as Python lists and tuples) by applying asanyarray only when shape is missing, preserving duck-typed arrays (e.g. JAX, pandas<3) that satisfy _SupportsSplitOps. Also fixes hsplit sequence handling and multi-dimensional sequence axis!=0 slice computation in array_split.
@bodapatisaikrishna

bodapatisaikrishna commented Sep 26, 2026 •
edited
Loading

Copy link
Copy Markdown
Author

Hi @jorenham,

Thank you for catching that! That's a great point regarding duck-typed arrays satisfying _SupportsSplitOps (with shape, swapaxes, and __getitem__).

I have updated the implementation so that asanyarray is only applied when the input lacks a shape attribute (e.g. sequences like Python lists and tuples), leaving duck-typed arrays completely untouched and preserving their original type:

  • In array_split: try: Ntotal = ary.shape[axis] except AttributeError: ary = _nx.asanyarray(ary); Ntotal = ary.shape[axis]
  • In split: try: N = ary.shape[axis] except AttributeError: ary = _nx.asanyarray(ary); N = ary.shape[axis]
  • In hsplit: uses _nx.ndim(ary) > 1 (which accesses .ndim on duck arrays and falls back to asarray on sequences)
  • vsplit and dsplit: remain untouched, using _nx.ndim(ary) < 2 / < 3 and delegating to split

I also added test_duck_array_preservation in numpy/lib/tests/test_shape_base.py verifying that duck arrays preserving _SupportsSplitOps maintain their custom type across split, array_split, hsplit, vsplit, and dsplit.

The updated commit has been pushed to the PR. All 84 tests in test_shape_base.py pass.

@jorenham

jorenham commented Sep 26, 2026

Copy link
Copy Markdown
Member

Please don't use AI to speak for you; that goes against our AI policy; we prefer to talk to humans.

This branch has not been deployed

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

Reviewers

@jorenham jorenham jorenham requested changes

Requested changes must be addressed to merge this pull request.

Assignees

No one assigned

Labels

00 - Bug 57 - Close? Issues which may be closable unless discussion continued

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

np.split throws AttributeError for list/ tuple ary and int sections, np.array_split does not

4 participants

@bodapatisaikrishna @ikrommyd @jorenham @ngoldbaum