Skip to content

ufunc.at #9239

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 18 commits into from
Mar 29, 2024
Merged

ufunc.at #9239

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
18 commits
Select commit Hold shift + click to select a range
f188770
[wip] ufunc.at
guilhermeleobas Oct 12, 2023
adc2934
upd
guilhermeleobas Oct 24, 2023
806a8f7
update tests
guilhermeleobas Oct 26, 2023
c706afa
fix CI bug
guilhermeleobas Oct 26, 2023
7c91799
fix CI bug (take 2)
guilhermeleobas Oct 26, 2023
e740181
add release notes
guilhermeleobas Oct 26, 2023
d04c33f
remove comented code
guilhermeleobas Oct 30, 2023
00e4ed2
Add missing test_ufunc_at_inner_loop test
guilhermeleobas Oct 31, 2023
ab87c5d
skip unsupported typecodes
guilhermeleobas Oct 31, 2023
2db339e
enable boolean array indexing
guilhermeleobas Nov 14, 2023
e4ca86b
rewrite ufunc.at to support multi dimension indexing
guilhermeleobas Nov 14, 2023
9dcb2bf
Fix isinstance check
guilhermeleobas Nov 15, 2023
f49e05b
fix test_import error
guilhermeleobas Nov 15, 2023
7ff257e
TypeError -> TypingError
guilhermeleobas Mar 8, 2024
a5f9549
Merge remote-tracking branch 'upstream/main' into guilhermeleobas/ufu…
guilhermeleobas Mar 8, 2024
1aac0f9
dynamically compile ufunc on calling 'at'
guilhermeleobas Mar 8, 2024
137d175
adjust release notes
guilhermeleobas Mar 8, 2024
a539aa2
Update 9239.new_feature.rst
guilhermeleobas Mar 20, 2024
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
enable boolean array indexing
  • Loading branch information
@guilhermeleobas
guilhermeleobas committed Nov 14, 2023
commit 2db339ecb2e355587e41defc8e349d0da05cd787
19 changes: 13 additions & 6 deletions numba/np/ufunc/dufunc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -295,14 +295,13 @@ def ol_at(ufunc, a, indices, b=None):
msg = 'The first argument "a" must be array-like'
raise errors.NumbaTypeError(msg)

indices_int_arr = isinstance(indices, (types.Array, types.List)) \
and isinstance(indices.dtype, types.Integer)
indices_arr = isinstance(indices, (types.Array, types.List))
indices_bool_arr = indices_arr and \
isinstance(indices.dtype, types.Boolean)
indices_tuple = isinstance(indices, types.Tuple)
indices_slice = isinstance(indices, types.SliceType)
indices_scalar = not (indices_int_arr or indices_int_arr or
indices_slice or indices_tuple)
indices_empty_tuple = isinstance(indices, types.Tuple) and \
len(indices) == 0
indices_scalar = not (indices_arr or indices_slice or indices_tuple)
indices_empty_tuple = indices_tuple and len(indices) == 0
b_array = isinstance(b, (types.Array, types.Sequence, types.List,
types.Tuple))
b_scalar = not b_array
Expand Down Expand Up @@ -351,6 +350,12 @@ def impl_generic(ufunc, a, indices, b=None):
j += 1
return a

def impl_indices_bool_arr(ufunc, a, indices, b=None):
a_ = a[indices]
idx = np.indices(a_.shape)
a[indices] = ufunc.at(a_, idx, b)
return a

def impl_indices_empty_b_scalar(ufunc, a, indices, b=None):
a[()] = ufunc(a[()], b)
return a
Expand All @@ -366,6 +371,8 @@ def impl_indices_slice(ufunc, a, indices, b=None):
return impl_b_none
elif indices_slice:
return impl_indices_slice
elif indices_bool_arr:
return impl_indices_bool_arr
elif indices_empty_tuple and b_scalar:
return impl_indices_empty_b_scalar
elif indices_scalar and b_scalar:
Expand Down
3 changes: 1 addition & 2 deletions numba/tests/npyufunc/test_dufunc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -377,15 +377,14 @@ def test_ufunc_at_dtypes(self):
power_at(a, [1, 2, 3, 2], 3.5)
self.assertPreciseEqual(a, np.array([0, 1, 4414, 46, 4, 5, 6, 7, 8, 9]))

@unittest.expectedFailure
def test_ufunc_at_boolean(self):
# Test boolean indexing and boolean ufuncs
a = np.arange(10)
index = a % 2 == 0
equal_at = self._generate_jit(np.equal)
# boolean indexing not supported
equal_at(a, index, [0, 2, 4, 6, 8])
self.assertPreciseEqual(a, [1, 1, 1, 3, 1, 5, 1, 7, 1, 9])
self.assertPreciseEqual(a, np.array([1, 1, 1, 3, 1, 5, 1, 7, 1, 9]))

def test_ufunc_at_boolean2(self):
# Test unary operator
Expand Down