ENH: add bitwise_or

* draft the `bitwise_or` ufunc and turn a subset of its array API standard tests on in CI * this is extremely similar to the changes in the `bitwise_and` PR at kokkosgh-145
tylerjereddy · Jan 5, 2023 · 61608cd · 61608cd
1 parent ffd607c
commit 61608cd
Show file tree

Hide file tree

Showing 6 changed files with 542 additions and 4 deletions.
diff --git a/.github/workflows/array_api.yml b/.github/workflows/array_api.yml
@@ -30,7 +30,7 @@ jobs:
           cd /tmp
           git clone https://github.com/kokkos/pykokkos-base.git
           cd pykokkos-base
-          python setup.py install -- -DENABLE_LAYOUTS=ON -DENABLE_MEMORY_TRAITS=OFF
+          python setup.py install -- -DENABLE_LAYOUTS=ON -DENABLE_MEMORY_TRAITS=OFF -DENABLE_VIEW_RANKS=5
       - name: Install pykokkos
         run: |
           python -m pip install .
@@ -49,4 +49,4 @@ jobs:
           # for hypothesis-driven test case generation
           pytest $GITHUB_WORKSPACE/pre_compile_tools/pre_compile_ufuncs.py -s
           # only run a subset of the conformance tests to get started
-          pytest array_api_tests/meta/test_broadcasting.py array_api_tests/meta/test_equality_mapping.py array_api_tests/meta/test_signatures.py array_api_tests/meta/test_special_cases.py array_api_tests/test_constants.py array_api_tests/meta/test_utils.py array_api_tests/test_creation_functions.py::test_ones array_api_tests/test_creation_functions.py::test_ones_like array_api_tests/test_data_type_functions.py::test_result_type array_api_tests/test_operators_and_elementwise_functions.py::test_log10 array_api_tests/test_operators_and_elementwise_functions.py::test_sqrt array_api_tests/test_operators_and_elementwise_functions.py::test_isfinite array_api_tests/test_operators_and_elementwise_functions.py::test_log2 array_api_tests/test_operators_and_elementwise_functions.py::test_log1p array_api_tests/test_operators_and_elementwise_functions.py::test_isinf array_api_tests/test_operators_and_elementwise_functions.py::test_log array_api_tests/test_array_object.py::test_scalar_casting array_api_tests/test_operators_and_elementwise_functions.py::test_sign array_api_tests/test_operators_and_elementwise_functions.py::test_square array_api_tests/test_operators_and_elementwise_functions.py::test_cos array_api_tests/test_operators_and_elementwise_functions.py::test_round array_api_tests/test_operators_and_elementwise_functions.py::test_trunc array_api_tests/test_operators_and_elementwise_functions.py::test_ceil array_api_tests/test_operators_and_elementwise_functions.py::test_floor
+          pytest array_api_tests/meta/test_broadcasting.py array_api_tests/meta/test_equality_mapping.py array_api_tests/meta/test_signatures.py array_api_tests/meta/test_special_cases.py array_api_tests/test_constants.py array_api_tests/meta/test_utils.py array_api_tests/test_creation_functions.py::test_ones array_api_tests/test_creation_functions.py::test_ones_like array_api_tests/test_data_type_functions.py::test_result_type array_api_tests/test_operators_and_elementwise_functions.py::test_log10 array_api_tests/test_operators_and_elementwise_functions.py::test_sqrt array_api_tests/test_operators_and_elementwise_functions.py::test_isfinite array_api_tests/test_operators_and_elementwise_functions.py::test_log2 array_api_tests/test_operators_and_elementwise_functions.py::test_log1p array_api_tests/test_operators_and_elementwise_functions.py::test_isinf array_api_tests/test_operators_and_elementwise_functions.py::test_log array_api_tests/test_array_object.py::test_scalar_casting array_api_tests/test_operators_and_elementwise_functions.py::test_sign array_api_tests/test_operators_and_elementwise_functions.py::test_square array_api_tests/test_operators_and_elementwise_functions.py::test_cos array_api_tests/test_operators_and_elementwise_functions.py::test_round array_api_tests/test_operators_and_elementwise_functions.py::test_trunc array_api_tests/test_operators_and_elementwise_functions.py::test_ceil array_api_tests/test_operators_and_elementwise_functions.py::test_floor "array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_or[bitwise_or(x1, x2)]" "array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_or[__or__(x1, x2)]" "array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_or[__or__(x, s)]"
diff --git a/.github/workflows/main_ci.yml b/.github/workflows/main_ci.yml
@@ -30,7 +30,7 @@ jobs:
           cd /tmp
           git clone https://github.com/kokkos/pykokkos-base.git
           cd pykokkos-base
-          python setup.py install -- -DENABLE_LAYOUTS=ON -DENABLE_MEMORY_TRAITS=OFF
+          python setup.py install -- -DENABLE_LAYOUTS=ON -DENABLE_MEMORY_TRAITS=OFF -DENABLE_VIEW_RANKS=5
       - name: Install pykokkos
         run: |
           python -m pip install .

diff --git a/pykokkos/__init__.py b/pykokkos/__init__.py
@@ -61,7 +61,8 @@
                                  round,
                                  trunc,
                                  ceil,
-                                 floor)
+                                 floor,
+                                 bitwise_or)
 from pykokkos.lib.info import iinfo, finfo
 from pykokkos.lib.create import (zeros,
                                  ones,

diff --git a/pykokkos/interface/views.py b/pykokkos/interface/views.py
@@ -392,6 +392,37 @@ def __array__(self, dtype=None):
         return self.data
 
 
+    def __or__(self, other):
+        # avoid circular import with scoped import
+        from pykokkos.lib.ufuncs import bitwise_or
+        if isinstance(other, View):
+            return bitwise_or(self, other)
+        elif isinstance(other, bool):
+            other_new = pk.View([], dtype=pk.uint8)
+            other_new[:] = other
+            return bitwise_or(self, other_new)
+        elif isinstance(other, int):
+            if 0 <= other <= 255:
+                other_dtype = pk.uint8
+            elif 0 <= other <= 65535:
+                other_dtype = pk.uint16
+            elif 0 <= other <= 4294967295:
+                other_dtype = pk.uint32
+            elif 0 <= other <= 18446744073709551615:
+                other_dtype = pk.uint64
+            elif -128 <= other <= 127:
+                other_dtype = pk.int8
+            elif -32768 <= other <= 32767:
+                other_dtype = pk.int16
+            elif -2147483648 <= other <= 2147483647:
+                other_dtype = pk.int32
+            elif -9223372036854775808 <= other <= 9223372036854775807:
+                other_dtype = pk.int64
+            other_new = pk.View([], dtype=other_dtype)
+            other_new[:] = other
+            return bitwise_or(self, other_new)
+
+
     @staticmethod
     def _get_dtype_name(type_name: str) -> str:
         """