1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162
|
from __future__ import annotations
import cupy as cp
import numpy as np
import pytest
import awkward as ak
from awkward.types import ArrayType, NumpyType, RegularType
to_list = ak.operations.to_list
@pytest.fixture(scope="function", autouse=True)
def cleanup_cuda():
yield
cp._default_memory_pool.free_all_blocks()
cp.cuda.Device().synchronize()
def test_0150_ByteMaskedArray_flatten():
content = ak.operations.from_iter(
[
[[0.0, 1.1, 2.2], [], [3.3, 4.4]],
[],
[[5.5]],
[[6.6, 7.7, 8.8, 9.9]],
[[], [10.0, 11.1, 12.2]],
],
highlevel=False,
)
mask = ak.index.Index8(np.array([0, 0, 1, 1, 0], dtype=np.int8))
array = ak.contents.ByteMaskedArray(mask, content, valid_when=False)
cuda_array = ak.to_backend(array, "cuda")
assert to_list(cuda_array) == [
[[0.0, 1.1, 2.2], [], [3.3, 4.4]],
[],
None,
None,
[[], [10.0, 11.1, 12.2]],
]
assert ak.operations.to_list(ak.operations.flatten(cuda_array, axis=1)) == [
[0.0, 1.1, 2.2],
[],
[3.3, 4.4],
[],
[10.0, 11.1, 12.2],
]
assert ak.operations.to_list(ak.operations.flatten(cuda_array, axis=-2)) == [
[0.0, 1.1, 2.2],
[],
[3.3, 4.4],
[],
[10.0, 11.1, 12.2],
]
assert ak.operations.to_list(ak.operations.flatten(cuda_array, axis=2)) == [
[0.0, 1.1, 2.2, 3.3, 4.4],
[],
None,
None,
[10.0, 11.1, 12.2],
]
assert ak.operations.to_list(ak.operations.flatten(cuda_array, axis=-1)) == [
[0.0, 1.1, 2.2, 3.3, 4.4],
[],
None,
None,
[10.0, 11.1, 12.2],
]
def test_1586_should_preserve_regulararray_numpy_regular_axis1():
a1 = ak.Array(np.array([[0.0, 1.1], [2.2, 3.3]]))
a2 = ak.from_json("[[4.4, 5.5, 6.6], [7.7, 8.8, 9.9]]")
cuda_a1 = ak.to_backend(a1, "cuda")
assert isinstance(cuda_a1.layout, ak.contents.NumpyArray)
a2 = ak.to_regular(a2, axis=1)
cuda_a2 = ak.to_backend(a2, "cuda")
c = ak.concatenate([cuda_a1, cuda_a2], axis=1)
assert c.to_list() == [[0.0, 1.1, 4.4, 5.5, 6.6], [2.2, 3.3, 7.7, 8.8, 9.9]]
assert c.type == ArrayType(RegularType(NumpyType("float64"), 5), 2)
def test_1586_should_preserve_regulararray_regular_numpy_axis1():
a1 = ak.from_json("[[0.0, 1.1], [2.2, 3.3]]")
a2 = ak.Array(np.array([[4.4, 5.5, 6.6], [7.7, 8.8, 9.9]]))
cuda_a1 = ak.to_backend(a1, "cuda")
cuda_a2 = ak.to_backend(a2, "cuda")
cuda_a1 = ak.to_regular(cuda_a1, axis=1)
assert isinstance(cuda_a2.layout, ak.contents.NumpyArray)
c = ak.concatenate([cuda_a1, cuda_a2], axis=1)
assert c.to_list() == [[0.0, 1.1, 4.4, 5.5, 6.6], [2.2, 3.3, 7.7, 8.8, 9.9]]
assert c.type == ArrayType(RegularType(NumpyType("float64"), 5), 2)
def test_1586_should_preserve_regulararray_regular_regular_axis1():
a1 = ak.from_json("[[0.0, 1.1], [2.2, 3.3]]")
a2 = ak.from_json("[[4.4, 5.5, 6.6], [7.7, 8.8, 9.9]]")
cuda_a1 = ak.to_backend(a1, "cuda")
cuda_a2 = ak.to_backend(a2, "cuda")
cuda_a1 = ak.to_regular(cuda_a1, axis=1)
cuda_a2 = ak.to_regular(cuda_a2, axis=1)
c = ak.concatenate([cuda_a1, cuda_a2], axis=1)
assert c.to_list() == [[0.0, 1.1, 4.4, 5.5, 6.6], [2.2, 3.3, 7.7, 8.8, 9.9]]
assert c.type == ArrayType(RegularType(NumpyType("float64"), 5), 2)
def test_0072_regulararray_fillna_unionarray():
content1 = ak.operations.from_iter([[], [1.1], [2.2, 2.2]], highlevel=False)
content2 = ak.operations.from_iter([["two", "two"], ["one"], []], highlevel=False)
tags = ak.index.Index8(np.array([0, 1, 0, 1, 0, 1], dtype=np.int8))
index = ak.index.Index64(np.array([0, 0, 1, 1, 2, 2], dtype=np.int64))
array = ak.contents.UnionArray(tags, index, [content1, content2])
cuda_array = ak.to_backend(array, "cuda")
assert to_list(cuda_array) == [[], ["two", "two"], [1.1], ["one"], [2.2, 2.2], []]
padded_array = ak._do.pad_none(array, 2, 1)
assert to_list(padded_array) == [
[None, None],
["two", "two"],
[1.1, None],
["one", None],
[2.2, 2.2],
[None, None],
]
value = ak.contents.NumpyArray(np.array([777]))
assert to_list(ak._do.fill_none(padded_array, value)) == [
[777, 777],
["two", "two"],
[1.1, 777],
["one", 777],
[2.2, 2.2],
[777, 777],
]
def test_0590_allow_regulararray_size_zero_ListOffsetArray_rpad_and_clip():
array = ak.highlevel.Array([[1, 2, 3], [], [4, 5]])
assert ak.operations.pad_none(array, 0, clip=True).to_list() == [
[],
[],
[],
]
array = ak.highlevel.Array([[1, 2, 3], [], [4, 5]])
cuda_array = ak.to_backend(array, "cuda")
assert ak.operations.pad_none(cuda_array, 0).to_list() == [
[1, 2, 3],
[],
[4, 5],
]
|