Merge branch 'main' into ianna/3179-cuda-target-in-numbavectorize-not…

…-working-correctly
scikit-hep · Sep 26, 2024 · 8b0473b · 8b0473b
2 parents e9c06c6 + 704fb45
commit 8b0473b
Show file tree

Hide file tree

Showing 53 changed files with 841 additions and 144 deletions.
diff --git a/.all-contributorsrc b/.all-contributorsrc
@@ -465,6 +465,24 @@
       "contributions": [
         "code"
       ]
+    },
+    {
+      "login": "ariostas",
+      "name": "Andres Rios Tascon",
+      "avatar_url": "https://avatars.githubusercontent.com/u/7596837?v=4",
+      "profile": "http://www.ariostas.com",
+      "contributions": [
+        "code"
+      ]
+    },
+    {
+      "login": "maxymnaumchyk",
+      "name": "maxymnaumchyk",
+      "avatar_url": "https://avatars.githubusercontent.com/u/70752300?v=4",
+      "profile": "https://github.com/maxymnaumchyk",
+      "contributions": [
+        "code"
+      ]
     }
   ],
   "contributorsPerLine": 7,

diff --git a/.github/workflows/build-wheels.yml b/.github/workflows/build-wheels.yml
@@ -105,7 +105,7 @@ jobs:
     - name: Prepare build files
       run: pipx run nox -s prepare
 
-    - uses: pypa/cibuildwheel@v2.19
+    - uses: pypa/cibuildwheel@v2.20
       env:
         CIBW_BUILD: "${{ matrix.build }}*"
         CIBW_ARCHS: ${{ matrix.arch }}
@@ -136,7 +136,7 @@ jobs:
       SOURCE_DATE_EPOCH: ${{ needs.determine-source-date-epoch.outputs.source-date-epoch }}
     strategy:
       matrix:
-        python: [38, 39, 310, 311, 312]
+        python: [38, 39, 310, 311, 312, 313]
         arch: [aarch64]
     steps:
 
@@ -157,7 +157,7 @@ jobs:
 
     - uses: docker/[email protected]
 
-    - uses: pypa/cibuildwheel@v2.19
+    - uses: pypa/cibuildwheel@v2.20
       env:
         CIBW_BUILD: cp${{ matrix.python }}-*
         CIBW_ARCHS: ${{ matrix.arch }}

diff --git a/.github/workflows/deploy-cpp.yml b/.github/workflows/deploy-cpp.yml
@@ -35,8 +35,8 @@ jobs:
       run: ls -l dist/
 
     - name: Generate artifact attestation for sdist and wheel
-      uses: actions/attest-build-provenance@5e9cb68e95676991667494a6a4e59b8a2f13e1d0 # v1.3.3
+      uses: actions/attest-build-provenance@1c608d11d69870c2092266b3f9a6f3abbf17002c # v1.4.3
       with:
         subject-path: "dist/awkward*cpp-*"
 
-    - uses: pypa/gh-action-pypi-publish@v1.9.0
+    - uses: pypa/gh-action-pypi-publish@v1.10.1
diff --git a/.github/workflows/deploy.yml b/.github/workflows/deploy.yml
@@ -88,7 +88,7 @@ jobs:
       run: pipx run twine check dist/*
 
     - name: Generate artifact attestation for sdist and wheel
-      uses: actions/attest-build-provenance@5e9cb68e95676991667494a6a4e59b8a2f13e1d0 # v1.3.3
+      uses: actions/attest-build-provenance@1c608d11d69870c2092266b3f9a6f3abbf17002c # v1.4.3
       with:
         subject-path: "dist/awkward-*"
 
@@ -135,7 +135,7 @@ jobs:
         GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
       run: gh attestation verify dist/awkward-*.whl --repo ${{ github.repository }}
 
-    - uses: pypa/gh-action-pypi-publish@v1.9.0
+    - uses: pypa/gh-action-pypi-publish@v1.10.1
 
   publish-headers:
     name: "Publish header-only libraries alongside release"

diff --git a/.github/workflows/packaging-test.yml b/.github/workflows/packaging-test.yml
@@ -68,15 +68,15 @@ jobs:
     - name: Prepare build files
       run: pipx run nox -s prepare
 
-    - uses: pypa/cibuildwheel@v2.19
+    - uses: pypa/cibuildwheel@v2.20
       env:
         CIBW_ARCHS_MACOS: universal2
         CIBW_BUILD: cp39-win_amd64 cp310-manylinux_x86_64 cp38-macosx_universal2
       with:
         config-file: cibuildwheel.toml
         package-dir: awkward-cpp
 
-    - uses: pypa/cibuildwheel@v2.19
+    - uses: pypa/cibuildwheel@v2.20
       if: matrix.os == 'ubuntu-latest'
       env:
         CIBW_BUILD: cp312-manylinux_x86_64

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -39,6 +39,7 @@ jobs:
           - ubuntu-latest
           - macos-13
         python-version:
+          - '3.13'
           - '3.12'
           - '3.11'
           - '3.10'
@@ -65,6 +66,10 @@ jobs:
             python-architecture: x64
             runs-on: ubuntu-latest
             dependencies-kind: pypy
+          - python-version: '3.11'
+            python-architecture: x64
+            runs-on: ubuntu-latest
+            dependencies-kind: ml
 
     runs-on: ${{ matrix.runs-on }}
 
@@ -106,6 +111,11 @@ jobs:
           files: |
             awkward-cpp/dist/*.whl
 
+      - name: Add workaround for 3.13 + cramjam
+        if: matrix.python-version == '3.13'
+        run: echo 'PYO3_USE_ABI3_FORWARD_COMPATIBILITY=1' >> $GITHUB_ENV
+        shell: bash
+
       - name: Install awkward, awkward-cpp, and dependencies
         run: >-
           python -m pip install -v . ${{ steps.find-wheel.outputs.paths }} pytest-github-actions-annotate-failures

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -27,7 +27,7 @@ repos:
     additional_dependencies: [pyyaml]
 
 - repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.5.5
+  rev: v0.6.4
   hooks:
   - id: ruff
     args: ["--fix", "--show-fixes"]
@@ -62,20 +62,20 @@ repos:
     files: ^tests/
 
 - repo: https://github.com/python-jsonschema/check-jsonschema
-  rev: 0.29.1
+  rev: 0.29.2
   hooks:
     - id: check-github-workflows
       args: ["--verbose"]
 
 - repo: https://github.com/pre-commit/mirrors-mypy
-  rev: v1.11.0
+  rev: v1.11.2
   hooks:
     - id: mypy
       files: src
       additional_dependencies:
         - numpy>=1.24
 
 - repo: https://github.com/abravalheri/validate-pyproject
-  rev: v0.18
+  rev: v0.19
   hooks:
     - id: validate-pyproject
diff --git a/CITATION.cff b/CITATION.cff
@@ -51,6 +51,6 @@ authors:
   email: "[email protected]"
 - family-names: "Goyal"
   given-names: "Manasvi"
-  affiliation: "Delhi Technological University"
+  affiliation: "Harvard University"
   orcid: "https://orcid.org/0000-0001-6321-7491"
   email: "[email protected]"
diff --git a/README.md b/README.md
@@ -152,7 +152,7 @@ year = {2018}
 
 # Acknowledgements
 
-Support for this work was provided by NSF cooperative agreements [OAC-1836650](https://www.nsf.gov/awardsearch/showAward?AWD_ID=1836650) and [PHY-2323298](https://www.nsf.gov/awardsearch/showAward?AWD_ID=2323298) (IRIS-HEP), grant [OAC-1450377](https://nsf.gov/awardsearch/showAward?AWD_ID=1450377) (DIANA/HEP), [PHY-2121686](https://www.nsf.gov/awardsearch/showAward?AWD_ID=2121686) (US-CMS LHC Ops), and [OAC-2103945](https://www.nsf.gov/awardsearch/showAward?AWD_ID=2103945) (Awkward Array).
+Support for this work was provided by NSF cooperative agreement [OAC-1836650](https://www.nsf.gov/awardsearch/showAward?AWD_ID=1836650) (IRIS-HEP 1), [PHY-2323298](https://www.nsf.gov/awardsearch/showAward?AWD_ID=2323298) (IRIS-HEP 2), grant [OAC-1450377](https://nsf.gov/awardsearch/showAward?AWD_ID=1450377) (DIANA/HEP), [PHY-1520942](https://www.nsf.gov/awardsearch/showAward?AWD_ID=1520942) and [PHY-2121686](https://www.nsf.gov/awardsearch/showAward?AWD_ID=2121686) (US-CMS LHC Ops), and [OAC-2103945](https://www.nsf.gov/awardsearch/showAward?AWD_ID=2103945) (Awkward Array).
 
 We also thank [Erez Shinan](https://github.com/erezsh) and the developers of the [Lark standalone parser](https://github.com/lark-parser/lark), which is used to parse type strings as type objects.
 
@@ -228,6 +228,8 @@ Thanks especially to the gracious help of Awkward Array contributors (including
     </tr>
     <tr>
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/pfackeldey"><img src="https://avatars.githubusercontent.com/u/18463582?v=4?s=100" width="100px;" alt="Peter Fackeldey"/><br /><sub><b>Peter Fackeldey</b></sub></a><br /><a href="https://github.com/scikit-hep/awkward/commits?author=pfackeldey" title="Code">💻</a></td>
+      <td align="center" valign="top" width="14.28%"><a href="http://www.ariostas.com"><img src="https://avatars.githubusercontent.com/u/7596837?v=4?s=100" width="100px;" alt="Andres Rios Tascon"/><br /><sub><b>Andres Rios Tascon</b></sub></a><br /><a href="https://github.com/scikit-hep/awkward/commits?author=ariostas" title="Code">💻</a></td>
+      <td align="center" valign="top" width="14.28%"><a href="https://github.com/maxymnaumchyk"><img src="https://avatars.githubusercontent.com/u/70752300?v=4?s=100" width="100px;" alt="maxymnaumchyk"/><br /><sub><b>maxymnaumchyk</b></sub></a><br /><a href="https://github.com/scikit-hep/awkward/commits?author=maxymnaumchyk" title="Code">💻</a></td>
     </tr>
   </tbody>
 </table>

diff --git a/awkward-cpp/pyproject.toml b/awkward-cpp/pyproject.toml
@@ -1,13 +1,13 @@
 [build-system]
 requires = [
-    "scikit-build-core>=0.9",
+    "scikit-build-core>=0.10",
     "pybind11",
 ]
 build-backend = "scikit_build_core.build"
 
 [project]
 name = "awkward_cpp"
-version = "37"
+version = "38"
 dependencies = [
     "numpy>=1.18.0",
     "importlib_resources;python_version < \"3.9\""
@@ -59,7 +59,7 @@ Releases = "https://github.com/scikit-hep/awkward-1.0/releases"
 
 
 [tool.scikit-build]
-minimum-version = "0.9"
+minimum-version = "build-system.requires"
 build-dir = "build/{cache_tag}"
 sdist.reproducible = true
 sdist.include = [

diff --git a/cibuildwheel.toml b/cibuildwheel.toml
@@ -23,5 +23,5 @@ build-verbosity = 1
 PIP_ONLY_BINARY = "cmake,numpy"
 
 [[tool.cibuildwheel.overrides]]
-select = "cp312-*"
-environment.PIP_PRE = "1"
+select = "cp313*"
+environment.PYO3_USE_ABI3_FORWARD_COMPATIBILITY = "1"
diff --git a/docs/_templates/funding.html b/docs/_templates/funding.html
@@ -1 +1 @@
-<p>Support for this work was provided by NSF cooperative agreement <a href="https://www.nsf.gov/awardsearch/showAward?AWD_ID=1836650">OAC-1836650</a> (IRIS-HEP), grant <a href="https://nsf.gov/awardsearch/showAward?AWD_ID=1450377">OAC-1450377</a> (DIANA/HEP), <a href="https://www.nsf.gov/awardsearch/showAward?AWD_ID=1520942">PHY-1520942</a> (US-CMS LHC Ops), and <a href="https://www.nsf.gov/awardsearch/showAward?AWD_ID=2103945">OAC-2103945</a> (Awkward Array).</p>
+<p>Support for this work was provided by NSF cooperative agreement <a href="https://www.nsf.gov/awardsearch/showAward?AWD_ID=1836650">OAC-1836650</a> (IRIS-HEP 1), <a href="https://www.nsf.gov/awardsearch/showAward?AWD_ID=2323298">PHY-2323298</a> (IRIS-HEP 2), grant <a href="https://nsf.gov/awardsearch/showAward?AWD_ID=1450377">OAC-1450377</a> (DIANA/HEP), <a href="https://www.nsf.gov/awardsearch/showAward?AWD_ID=1520942">PHY-1520942</a> and <a href="https://www.nsf.gov/awardsearch/showAward?AWD_ID=2121686">PHY-2121686</a> (US-CMS LHC Ops), and <a href="https://www.nsf.gov/awardsearch/showAward?AWD_ID=2103945">OAC-2103945</a> (Awkward Array).</p>
diff --git a/docs/reference/toctree.txt b/docs/reference/toctree.txt
@@ -299,11 +299,12 @@
     generated/ak.backend
 
 .. toctree::
-    :caption: Approximation
+    :caption: Approximation and comparison
 
     generated/ak.round
     generated/ak.isclose
     generated/ak.almost_equal
+    generated/ak.array_equal
 
 .. toctree::
     :caption: NumPy compatibility

diff --git a/pyproject.toml b/pyproject.toml
@@ -7,7 +7,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "awkward"
-version = "2.6.7"
+version = "2.6.8"
 description = "Manipulate JSON-like data with NumPy-like idioms."
 license = { text = "BSD-3-Clause" }
 requires-python = ">=3.8"
@@ -32,6 +32,7 @@ classifiers = [
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
     "Topic :: Scientific/Engineering",
     "Topic :: Scientific/Engineering :: Information Analysis",
     "Topic :: Scientific/Engineering :: Mathematics",
@@ -40,7 +41,7 @@ classifiers = [
     "Topic :: Utilities",
 ]
 dependencies = [
-    "awkward_cpp==37",
+    "awkward_cpp==38",
     "importlib_metadata>=4.13.0;python_version < \"3.12\"",
     "numpy>=1.18.0",
     "packaging",
@@ -136,7 +137,7 @@ filterwarnings = [
 ]
 log_cli_level = "info"
 
-[tool.pylint.master]
+[tool.pylint]
 py-version = "3.8"
 jobs = "0"
 ignore-paths = [
@@ -251,7 +252,6 @@ ignore_errors = true
 ignore_missing_imports = true
 
 [tool.ruff]
-src = ["src"]
 extend-exclude = [
     "studies",
     "pybind11",
@@ -289,10 +289,6 @@ ignore = [
     "PLC1901", # x == "" can be simplified to not x (empty string is falsey)
     "ISC001",  # Conflicts with the formatter in 0.1.2
 ]
-unfixable = [
-    "T20",  # Removes print statements
-    "F841", # Removes unused variables
-]
 typing-modules = ["awkward._typing"]
 external = []
 mccabe.max-complexity = 100

diff --git a/requirements-test-full.txt b/requirements-test-full.txt
@@ -1,9 +1,9 @@
 fsspec>=2022.11.0;sys_platform != "win32"
-jax[cpu]>=0.2.15;sys_platform != "win32" and python_version < "3.12"
-numba>=0.50.0;sys_platform != "win32" and python_version < "3.12"
-numexpr>=2.7; python_version < "3.12"
-pandas>=0.24.0;sys_platform != "win32" and python_version < "3.12"
-pyarrow==16.0.0;sys_platform != "win32" and python_version < "3.12"
+jax[cpu]>=0.2.15;sys_platform != "win32" and python_version < "3.13"
+numba>=0.50.0;sys_platform != "win32" and python_version < "3.13"
+numexpr>=2.7; python_version < "3.13"
+pandas>=0.24.0;sys_platform != "win32" and python_version < "3.13"
+pyarrow==16.0.0;sys_platform != "win32" and python_version < "3.13"
 pytest>=6
 pytest-cov
 pytest-xdist

diff --git a/requirements-test-ml.txt b/requirements-test-ml.txt
@@ -0,0 +1,6 @@
+fsspec>=2022.11.0;sys_platform != "win32"
+pytest>=6
+pytest-cov
+pytest-xdist
+tensorflow >= 2.12
+torch >= 2.4.0
diff --git a/src/awkward/_broadcasting.py b/src/awkward/_broadcasting.py
@@ -701,6 +701,36 @@ def broadcast_any_list():
                 for x, p in zip(outcontent, parameters)
             )
 
+    def broadcast_any_option_all_UnmaskedArray():
+        nextinputs = []
+        nextparameters = []
+        for x in inputs:
+            if isinstance(x, UnmaskedArray):
+                nextinputs.append(x.content)
+                nextparameters.append(x._parameters)
+            elif isinstance(x, Content):
+                nextinputs.append(x)
+                nextparameters.append(x._parameters)
+            else:
+                nextinputs.append(x)
+                nextparameters.append(NO_PARAMETERS)
+
+        outcontent = apply_step(
+            backend,
+            nextinputs,
+            action,
+            depth,
+            copy.copy(depth_context),
+            lateral_context,
+            options,
+        )
+        assert isinstance(outcontent, tuple)
+        parameters = parameters_factory(nextparameters, len(outcontent))
+
+        return tuple(
+            UnmaskedArray(x, parameters=p) for x, p in zip(outcontent, parameters)
+        )
+
     def broadcast_any_option():
         mask = None
         for x in contents:
@@ -712,7 +742,7 @@ def broadcast_any_option():
                     mask = backend.index_nplike.logical_or(mask, m, maybe_out=mask)
 
         nextmask = Index8(mask.view(np.int8))
-        index = backend.index_nplike.full(mask.shape[0], -1, dtype=np.int64)
+        index = backend.index_nplike.full(mask.shape[0], np.int64(-1), dtype=np.int64)
         index[~mask] = backend.index_nplike.arange(
             backend.index_nplike.shape_item_as_index(mask.shape[0])
             - backend.index_nplike.count_nonzero(mask),
@@ -1045,7 +1075,9 @@ def continuation():
 
         # Any option-types?
         elif any(x.is_option for x in contents):
-            if options["function_name"] == "ak.where":
+            if all(not x.is_option or isinstance(x, UnmaskedArray) for x in contents):
+                return broadcast_any_option_all_UnmaskedArray()
+            elif options["function_name"] == "ak.where":
                 return broadcast_any_option_akwhere()
             else:
                 return broadcast_any_option()

diff --git a/src/awkward/_connect/cuda/cuda_kernels/awkward_NumpyArray_subrange_equal_bool.cu b/src/awkward/_connect/cuda/cuda_kernels/awkward_NumpyArray_subrange_equal_bool.cu
@@ -4,9 +4,9 @@
 // def f(grid, block, args):
 //     (tmpptr, fromstarts, fromstops, length, toequal, invocation_index, err_code) = args
 //     if length > 1:
-//         scan_in_array = cupy.full((length - 1) * (length - 2), 0, dtype=cupy.int64)
+//         scan_in_array = cupy.full((length - 1) * (length - 2), cupy.array(0), dtype=cupy.int64)
 //     else:
-//         scan_in_array = cupy.full(0, 0, dtype=cupy.int64)
+//         scan_in_array = cupy.full(0, cupy.array(0), dtype=cupy.int64)
 //     cuda_kernel_templates.get_function(fetch_specialization(["awkward_NumpyArray_subrange_equal_bool", bool_, fromstarts.dtype, fromstops.dtype, bool_]))(grid, block, (tmpptr, fromstarts, fromstops, length, toequal, scan_in_array, invocation_index, err_code))
 //     toequal[0] = cupy.any(scan_in_array == True)
 // out["awkward_NumpyArray_subrange_equal_bool", {dtype_specializations}] = None