Skip to content

Initial implementation of np.take_along_axis #7202

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sklam merged 34 commits into from
Sep 16, 2021
Merged

Initial implementation of np.take_along_axis #7202

Changes from 1 commit
Commits
Show all changes
34 commits
Select commit Hold shift + click to select a range
bf1e335
Move utility method out, so it can be used by other modules.
itamarst Jun 21, 2021
4582bc8
A sketch of sorts.
itamarst Jun 30, 2021
b82dd24
A bit closer to working, maybe.
itamarst Jul 2, 2021
0cc9d7c
Even closer, maybe.
itamarst Jul 2, 2021
ada5b42
Actually works for some values.
itamarst Jul 7, 2021
372e34c
Fix off by one bug that caused segfaults for some axis.
itamarst Jul 7, 2021
d7c7985
Some documentation and comments for np.take_along_axis.
itamarst Jul 9, 2021
df2f29a
Tests for error cases of np.take_along_axis.
itamarst Jul 9, 2021
fe60ca9
Fix axis=None case.
itamarst Jul 9, 2021
d95a4bd
More test cases for np.take_along_axis().
itamarst Jul 9, 2021
0fcfb09
Tests (and corresponding bug fix) for take_along_axis() semantics.
itamarst Jul 12, 2021
dcab297
Fix flakes.
itamarst Jul 27, 2021
0dc3b76
Merge branch 'master' into np-take_along_axis
itamarst Aug 10, 2021
e759d98
Sometimes we take NumPy code from examples or other parts of the docu…
itamarst Aug 10, 2021
75ded55
Function was renamed.
itamarst Aug 10, 2021
da3b6e4
More succinct testing strategy
itamarst Aug 18, 2021
cfc28ed
Better explanation
itamarst Aug 18, 2021
7a70596
Point to stable location, in source code.
itamarst Aug 18, 2021
1ae5f53
Get rid of additional, erroneous outer loop.
itamarst Aug 18, 2021
b676024
Bounds check is up to user.
itamarst Aug 18, 2021
dbcf274
These functions will only ever be used from JITed code.
itamarst Aug 18, 2021
9ec64da
Minor optimization.
itamarst Aug 18, 2021
bd64769
Bit nicer testing.
itamarst Aug 18, 2021
7e65814
More input validation.
itamarst Aug 18, 2021
ede8d01
Flakes.
itamarst Aug 18, 2021
c86836a
A failing test.
itamarst Aug 18, 2021
f7b94d9
Get rid of combinatoric code repetition by doing a little song and da…
itamarst Aug 30, 2021
4ee6bb2
Better explanation.
itamarst Aug 30, 2021
b951a63
Merge remote-tracking branch 'origin/master' into np-take_along_axis
itamarst Aug 30, 2021
fd3c73d
Support for pre-broadcasting.
itamarst Aug 30, 2021
1d180d5
Additional tests for new broadcasting support.
itamarst Aug 30, 2021
47a32e9
Flake fixes.
itamarst Aug 30, 2021
dbedd9c
Apply suggestions from code review
itamarst Sep 13, 2021
01f9c9f
Fix flake.
itamarst Sep 13, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
These functions will only ever be used from JITed code.
  • Loading branch information
@itamarst
itamarst committed Aug 18, 2021
commit dbcf274daef32988cb7f73cb29b27c2db160393a
10 changes: 5 additions & 5 deletions numba/np/arrayobj.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@

import numpy as np

from numba import pndindex, literal_unroll
from numba import pndindex, literal_unroll, njit
from numba.core import types, utils, typing, errors, cgutils, extending
from numba.np.numpy_support import (as_dtype, carray, farray, is_contiguous,
is_fortran, check_is_integer)
Expand Down Expand Up @@ -5441,21 +5441,21 @@ def impl(arr, axis1, axis2):
return impl


@register_jitable
@njit
Comment thread
itamarst marked this conversation as resolved.
Resolved
def _take_along_axis_impl_set_ni(arr, indices, axis, Ni, Nk):
for i in range(len(Ni)):
Ni = tuple_setitem(Ni, i, arr.shape[i])
return _take_along_axis_impl(arr, indices, axis, Ni, Nk)


@register_jitable
@njit
def _take_along_axis_impl_set_nk(arr, indices, axis, Ni, Nk):
for i in range(len(Nk)):
Nk = tuple_setitem(Nk, i, arr.shape[axis + 1 + i])
return _take_along_axis_impl(arr, indices, axis, Ni, Nk)


@register_jitable
@njit
def _take_along_axis_impl_set_ni_nk(arr, indices, axis, Ni, Nk):
for i in range(len(Ni)):
Ni = tuple_setitem(Ni, i, arr.shape[i])
Expand All @@ -5464,7 +5464,7 @@ def _take_along_axis_impl_set_ni_nk(arr, indices, axis, Ni, Nk):
return _take_along_axis_impl(arr, indices, axis, Ni, Nk)


@register_jitable
@njit
def _take_along_axis_impl(arr, indices, axis, Ni, Nk):
# Based on example code in
# https://github.com/numpy/numpy/blob/623bc1fae1d47df24e7f1e29321d0c0ba2771ce0/numpy/lib/shape_base.py#L90
Comment thread
itamarst marked this conversation as resolved.
Resolved
Expand Down