Skip to content

Commit

Permalink
fix more cases where the array doesn't live long enough to create a m…
Browse files Browse the repository at this point in the history
…ask with ak.Array.mask
  • Loading branch information
pfackeldey committed Dec 18, 2024
1 parent 662b8cb commit 3ac12f4
Show file tree
Hide file tree
Showing 3 changed files with 39 additions and 14 deletions.
2 changes: 1 addition & 1 deletion docs/user-guide/how-to-create-missing.md
Original file line number Diff line number Diff line change
Expand Up @@ -265,6 +265,6 @@ def faster_example():
data, mask = faster_example()
array = ak.Array(data).mask[mask]
array = ak.mask(data, mask)
array
```
3 changes: 2 additions & 1 deletion tests-cuda/test_2922a_new_cuda_kernels.py
Original file line number Diff line number Diff line change
Expand Up @@ -1003,7 +1003,8 @@ def test_2064_fill_none_record_axis_last():


def test_2064_fill_none_record_option_outside_record():
record = ak.zip({"x": [1, 4], "y": [2, 3]}).mask[[True, False]]
record = ak.zip({"x": [1, 4], "y": [2, 3]})
record = record.mask[[True, False]]

cuda_record = ak.to_backend(record, "cuda")

Expand Down
48 changes: 36 additions & 12 deletions tests-cuda/test_3136_cuda_reducers.py
Original file line number Diff line number Diff line change
Expand Up @@ -267,10 +267,12 @@ def test_2020_reduce_axis_none_sum():
ak.sum(array, axis=None, keepdims=True),
ak.to_regular(ak.Array([[63.0]], backend="cuda")),
)

arr = ak.Array([[63.0]], backend="cuda")
assert ak.almost_equal(
ak.sum(array, axis=None, keepdims=True, mask_identity=True),
ak.to_regular(
ak.Array([[63.0]], backend="cuda").mask[ak.Array([[True]], backend="cuda")]
arr.mask[ak.Array([[True]], backend="cuda")]
),
)
assert ak.sum(array[2], axis=None, mask_identity=True) is None
Expand All @@ -291,10 +293,12 @@ def test_2020_reduce_axis_none_prod():
ak.prod(array[1:], axis=None, keepdims=True),
ak.to_regular(ak.Array([[4838400.0]], backend="cuda")),
)

arr = ak.Array([[4838400.0]], backend="cuda")
assert ak.almost_equal(
ak.prod(array[1:], axis=None, keepdims=True, mask_identity=True),
ak.to_regular(
ak.Array([[4838400.0]], backend="cuda").mask[
arr.mask[
ak.Array([[True]], backend="cuda")
]
),
Expand All @@ -316,16 +320,20 @@ def test_2020_reduce_axis_none_min():
ak.min(array, axis=None, keepdims=True, initial=-100.0, mask_identity=False),
ak.to_regular(ak.Array([[-100.0]], backend="cuda")),
)

arr = ak.Array([[0.0]], backend="cuda")
assert ak.almost_equal(
ak.min(array, axis=None, keepdims=True, mask_identity=True),
ak.to_regular(
ak.Array([[0.0]], backend="cuda").mask[ak.Array([[True]], backend="cuda")]
arr.mask[ak.Array([[True]], backend="cuda")]
),
)

arr = ak.Array(ak.Array([[np.inf]], backend="cuda"))
assert ak.almost_equal(
ak.min(array[-1:], axis=None, keepdims=True, mask_identity=True),
ak.to_regular(
ak.Array(ak.Array([[np.inf]], backend="cuda")).mask[
arr.mask[
ak.Array([[False]], backend="cuda")
]
),
Expand All @@ -347,16 +355,20 @@ def test_2020_reduce_axis_none_max():
ak.max(array, axis=None, keepdims=True, initial=100.0, mask_identity=False),
ak.to_regular(ak.Array([[100.0]], backend="cuda")),
)

arr = ak.Array([[10.0]], backend="cuda")
assert ak.almost_equal(
ak.max(array, axis=None, keepdims=True, mask_identity=True),
ak.to_regular(
ak.Array([[10.0]], backend="cuda").mask[ak.Array([[True]], backend="cuda")]
arr.mask[ak.Array([[True]], backend="cuda")]
),
)

arr = ak.Array(ak.Array([[np.inf]], backend="cuda"))
assert ak.almost_equal(
ak.max(array[-1:], axis=None, keepdims=True, mask_identity=True),
ak.to_regular(
ak.Array(ak.Array([[np.inf]], backend="cuda")).mask[
arr.mask[
ak.Array([[False]], backend="cuda")
]
),
Expand All @@ -374,16 +386,20 @@ def test_2020_reduce_axis_none_count():
ak.count(array, axis=None, keepdims=True, mask_identity=False),
ak.to_regular(ak.Array([[12]], backend="cuda")),
)

arr = ak.Array([[12]], backend="cuda")
assert ak.almost_equal(
ak.count(array, axis=None, keepdims=True, mask_identity=True),
ak.to_regular(
ak.Array([[12]], backend="cuda").mask[ak.Array([[True]], backend="cuda")]
arr.mask[ak.Array([[True]], backend="cuda")]
),
)

arr = ak.Array([[0]], backend="cuda")
assert ak.almost_equal(
ak.count(array[-1:], axis=None, keepdims=True, mask_identity=True),
ak.to_regular(
ak.Array([[0]], backend="cuda").mask[ak.Array([[False]], backend="cuda")]
arr.mask[ak.Array([[False]], backend="cuda")]
),
)
assert ak.count(array[2], axis=None, mask_identity=True) is None
Expand All @@ -400,16 +416,20 @@ def test_2020_reduce_axis_none_count_nonzero():
ak.count_nonzero(array, axis=None, keepdims=True, mask_identity=False),
ak.to_regular(ak.Array([[11]], backend="cuda")),
)

arr = ak.Array([[11]], backend="cuda")
assert ak.almost_equal(
ak.count_nonzero(array, axis=None, keepdims=True, mask_identity=True),
ak.to_regular(
ak.Array([[11]], backend="cuda").mask[ak.Array([[True]], backend="cuda")]
arr.mask[ak.Array([[True]], backend="cuda")]
),
)

arr = ak.Array([[0]], backend="cuda")
assert ak.almost_equal(
ak.count_nonzero(array[-1:], axis=None, keepdims=True, mask_identity=True),
ak.to_regular(
ak.Array([[0]], backend="cuda").mask[ak.Array([[False]], backend="cuda")]
arr.mask[ak.Array([[False]], backend="cuda")]
),
)
assert ak.count_nonzero(array[2], axis=None, mask_identity=True) is None
Expand All @@ -422,8 +442,10 @@ def test_2020_reduce_axis_none_std_no_mask_axis_none():
[[0, 2, 3.0], [4, 5, 6, 7, 8], [], [9, 8, None], [10, 1], []], backend="cuda"
)
out1 = ak.std(array[-1:], axis=None, keepdims=True, mask_identity=True)

arr = ak.Array([[0.0]], backend="cuda")
out2 = ak.to_regular(
ak.Array([[0.0]], backend="cuda").mask[ak.Array([[False]], backend="cuda")]
arr.mask[ak.Array([[False]], backend="cuda")]
)
assert ak.almost_equal(out1, out2)

Expand All @@ -442,10 +464,12 @@ def test_2020_reduce_axis_none_std():
ak.std(array, axis=None, keepdims=True, mask_identity=False),
ak.to_regular([[3.139134700306227]]),
)

arr = ak.Array([[3.139134700306227]], backend="cuda")
cpt.assert_allclose(
ak.std(array, axis=None, keepdims=True, mask_identity=True),
ak.to_regular(
ak.Array([[3.139134700306227]], backend="cuda").mask[
arr.mask[
ak.Array([[True]], backend="cuda")
]
),
Expand Down

0 comments on commit 3ac12f4

Please sign in to comment.