ri-kernels 0.2.0__tar.gz → 0.2.2__tar.gz

{ri_kernels-0.2.0 → ri_kernels-0.2.2}/.github/workflows/wheels.yml

@@ -81,9 +81,10 @@ jobs:
         # (plus 7.5, the CUDA 13 floor): cubins are binary compatible upward
         # within a major, so 8.0 covers 8.6/8.7/8.9, 10.0 covers 10.3 (B300)
         # and 12.0 covers 12.1 (GB10). 10.x and 12.x are separate families and
-        # need separate cubins. This is the set the jax CUDA plugin itself
-        # ships, minus 7.0, which only CUDA 12 can still build. The trailing
-        # `-virtual` embeds PTX once, as the JIT fallback for future GPUs.
+        # need separate cubins. CUDA 12 starts at JAX's 5.2 support floor and
+        # CUDA 13 starts at 7.5 because that toolkit dropped older targets. The
+        # trailing `-virtual` embeds PTX once, as the JIT fallback for future
+        # GPUs.
         cuda:
           - major: "12"
             pkg: "12-9"
@@ -95,8 +96,11 @@ jobs:
             # CUDA 13 dropped Maxwell/Pascal/Volta; Turing is the floor.
             archs: "75-real;80-real;90-real;100-real;120-real;120-virtual"
         platform:
-          - { os: ubuntu-latest, arch: x86_64, repo_arch: x86_64 }
-          - { os: ubuntu-24.04-arm, arch: aarch64, repo_arch: sbsa }
+          # Maxwell has no Linux SBSA target, so SM 5.2 belongs only in the
+          # x86_64 CUDA 12 wheel. Both platforms otherwise share the list from
+          # the CUDA entry above.
+          - { os: ubuntu-latest, arch: x86_64, repo_arch: x86_64, cuda12_prefix: "52-real;" }
+          - { os: ubuntu-24.04-arm, arch: aarch64, repo_arch: sbsa, cuda12_prefix: "" }
     steps:
       - uses: actions/checkout@v4
 
@@ -127,7 +131,7 @@ jobs:
           CIBW_ENVIRONMENT_LINUX: >
             CUDACXX=${{ matrix.cuda.home }}/bin/nvcc
             PATH=${{ matrix.cuda.home }}/bin:$PATH
-            CMAKE_ARGS="-DCMAKE_CUDA_ARCHITECTURES=${{ matrix.cuda.archs }}"
+            CMAKE_ARGS="-DCMAKE_CUDA_ARCHITECTURES=${{ matrix.cuda.major == '12' && matrix.platform.cuda12_prefix || '' }}${{ matrix.cuda.archs }}"
           # No test here: installing the wheel would pull ri_kernels==<version>
           # from the index, which does not exist yet for an unreleased version. The
           # library is smoke-tested below instead.

ri_kernels-0.2.2/.gitignore

@@ -0,0 +1,4 @@
+**/__pycache__
+
+# Local FFI builds are platform-specific artifacts, never source changes.
+/ri_kernels/*.so

{ri_kernels-0.2.0 → ri_kernels-0.2.2}/CMakeLists.txt

@@ -57,6 +57,7 @@ endif()
 option(RI_KERNELS_CPU "Build CPU extension" ON)
 option(RI_KERNELS_CUDA "Build CUDA extension" OFF)
 option(RI_KERNELS_ROCM "Build ROCm extension" OFF)
+option(RI_KERNELS_STATIC_CUDART "Link the CUDA runtime statically" ON)
 
 option(RI_KERNELS_BUNDLED_LIBS "Use all bundled libraries" ON)
 cmake_dependent_option(RI_KERNELS_BUNDLED_HIGHWAY "Use bundled highway lib" ON "RI_KERNELS_BUNDLED_LIBS" OFF)
@@ -95,12 +96,14 @@ set(RI_KERNELS_CPU_SOURCES
   ./src/rfi_kernel.cpp
   ./src/rfi_jvp_kernel.cpp
   ./src/rfi_transpose_kernel.cpp
+  ./src/rfi_analytic_kernel.cpp
 )
 
 set(RI_KERNELS_GPU_SOURCES
   ./src/rfi_kernel_gpu.cu
   ./src/rfi_jvp_kernel_gpu.cu
   ./src/rfi_transpose_kernel_gpu.cu
+  ./src/rfi_analytic_kernel_gpu.cu
   ./src/util_gpu.cu
 )
 
@@ -150,7 +153,13 @@ if(RI_KERNELS_CUDA)
   add_library(ri_kernels_cuda SHARED ${RI_KERNELS_GPU_SOURCES})
   # set_target_properties(ri_kernels_cuda PROPERTIES OUTPUT_NAME "ri_kernels_cuda_${CUDAToolkit_VERSION_MAJOR}")
   target_include_directories(ri_kernels_cuda PRIVATE ${JAXLIB_HOME}/include ./src)
-  target_link_libraries(ri_kernels_cuda PRIVATE CUDA::cudart_static)
+  if(RI_KERNELS_STATIC_CUDART)
+    set_property(TARGET ri_kernels_cuda PROPERTY CUDA_RUNTIME_LIBRARY Static)
+    target_link_libraries(ri_kernels_cuda PRIVATE CUDA::cudart_static)
+  else()
+    set_property(TARGET ri_kernels_cuda PROPERTY CUDA_RUNTIME_LIBRARY Shared)
+    target_link_libraries(ri_kernels_cuda PRIVATE CUDA::cudart)
+  endif()
   target_compile_options(ri_kernels_cuda PRIVATE
 	$<$<COMPILE_LANGUAGE:CUDA>:--diag-suppress=940>
 	$<$<COMPILE_LANGUAGE:CUDA>:--diag-suppress=2473>

{ri_kernels-0.2.0 → ri_kernels-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: ri_kernels
-Version: 0.2.0
+Version: 0.2.2
 Summary: CPU and GPU kernels for radio interferometry, exposed to JAX as FFI custom calls
 Author-Email: Chris Finlay <christopher.finlay@epfl.ch>, Simon Frasch <simon.frasch@cscs.ch>
 Maintainer-Email: Chris Finlay <christopher.finlay@epfl.ch>, Simon Frasch <simon.frasch@cscs.ch>
@@ -679,7 +679,7 @@ License:                     GNU GENERAL PUBLIC LICENSE
          Public License instead of this License.  But first, please read
          <https://www.gnu.org/licenses/why-not-lgpl.html>.
          
-Project-URL: Repository, https://github.com/epfl-radio-astro/ri_kernels.git
+Project-URL: Repository, https://github.com/epfl-radio-astro/ri-kernels
 Requires-Python: >=3.10
 Requires-Dist: jax>=0.6.0
 Provides-Extra: test
@@ -687,10 +687,10 @@ Requires-Dist: pytest; extra == "test"
 Requires-Dist: pytest-cov; extra == "test"
 Provides-Extra: cuda12
 Requires-Dist: jax[cuda12]; extra == "cuda12"
-Requires-Dist: ri_kernels_cuda12==0.2.0; extra == "cuda12"
+Requires-Dist: ri_kernels_cuda12==0.2.2; extra == "cuda12"
 Provides-Extra: cuda13
 Requires-Dist: jax[cuda13]; extra == "cuda13"
-Requires-Dist: ri_kernels_cuda13==0.2.0; extra == "cuda13"
+Requires-Dist: ri_kernels_cuda13==0.2.2; extra == "cuda13"
 Description-Content-Type: text/markdown
 
 # RI Kernels
@@ -729,7 +729,9 @@ needs a published release. To build the GPU kernels from the checkout, use
 Other CMake options of note: `RI_KERNELS_CPU` (default `ON`),
 `RI_KERNELS_MULTI_ARCH` (dynamic SIMD dispatch, default `ON` — turn it off and
 set arch flags via `CMAKE_CXX_FLAGS` for a single-target build),
-`RI_KERNELS_BUNDLED_HIGHWAY`, and `CMAKE_CUDA_ARCHITECTURES`.
+`RI_KERNELS_BUNDLED_HIGHWAY`, `CMAKE_CUDA_ARCHITECTURES`, and
+`RI_KERNELS_STATIC_CUDART` (default `ON`; set it to `OFF` when a package
+manager should provide the shared CUDA runtime).
 
 ## Tests
 
@@ -738,3 +740,138 @@ Tests can be run with
 ```bash
 python -m pytest tests
 ```
+
+## RFI visibilities
+
+`RFIVisOp` computes the per-baseline RFI visibility
+
+```
+vis[bl, f, t] = mean over (n_int_freq, n_int_time) of
+                sum over n_rfi of A[a1] conj(A[a2]) exp(i (φ[a1] - φ[a2]))
+```
+
+for amplitudes `A` shaped `(n_ant, n_freq, n_time, n_rfi, n_int_freq,
+n_int_time)` and phases `φ` in radians of the same shape, giving an output of
+shape `(n_baselines, n_freq, n_time)`. It is constructed from the baseline
+layout and evaluated through `eval`:
+
+```python
+from ri_kernels.jax_api import RFIVisOp
+
+vis = RFIVisOp(n_ant, a1, a2).eval(rfi_amp_fine, rfi_phase)
+```
+
+The operator has native primal, JVP, and transpose kernels for CPU, CUDA, and
+ROCm, so both forward- and reverse-mode differentiation stay inside the
+kernels. Precision has to match across the inputs: complex64 with float32, or
+complex128 with float64. Both amplitudes and phases are differentiated.
+
+## Analytic RFI visibilities
+
+`RFIAnalyticVisOp` integrates the time polynomial against phase moments, with
+frequency quadrature unchanged. It stages amplitude coefficients per antenna
+and cell, then contracts pairs of antenna tiles. The default time scratch axis
+has three coefficients, independent of fringe winding.
+
+```python
+from ri_kernels.jax_api import RFIAnalyticVisOp
+
+vis = RFIAnalyticVisOp(n_ant, a1, a2).eval(
+    amp, phase, delay_us, w_freq, start_freq, g_time, start_time,
+    dnu_mhz, int_time, freq_mhz,
+    segments=2, terms=6, cubic_terms=3,
+)
+```
+
+- `amp`, complex `(n_ant, n_rfi, n_freq, n_time)`: the signal on the data
+  grid.
+- `phase`, real `(n_ant, n_rfi, n_freq, n_time)`: the phase at the channel and
+  cell centre, reduced to one turn (in float64, before casting).
+- `delay_us`, real `(n_ant, n_rfi, n_time, n_path)`: the geometric delay in
+  microseconds, with the sign that makes the phase `2π f τ`, and its first
+  `n_path - 1` time derivatives at the cell centre, relative to the array mean
+  (a common term cancels in every baseline; the full delay's change across a
+  cell is ~1e4 wavelengths, beyond float32).
+- `w_freq`, `start_freq`: per channel, the weights that turn its stencil of
+  `n_sf` neighbouring channels into its fine channels, and the first channel
+  of the stencil.
+- `g_time`, `start_time`: per cell, monomial coefficients in
+  `x = 2*tau/int_time`, as `tabascal.poly_interp.monomial_tables` produces
+  them, including the shifted edge stencils, and the first cell of the
+  stencil. Each stencil must lie inside its axis and contain its cell.
+- `dnu_mhz`: the fine offsets from the channel centre; `freq_mhz`: the channel
+  centres. Both are in **MHz**; divide `config.freqs` in Hz by `1e6`. MHz × μs
+  is cycles.
+- `int_time`: a positive scalar in seconds.
+
+Every real operand shares the precision of `phase`; the stencil starts are
+int32. The constructor builds, per pair of 32-antenna tiles, the list of the
+pairs the baselines cover there (`tile_pair_list`), which the staged GPU
+kernels work through a tile pair per block; each `(a1, a2)` baseline may
+appear at most once.
+
+`amp` and `phase` are differentiated; `delay_us`, `int_time` and the tables
+are constants whose gradients `eval` stops. Two kernel pairs carry the
+derivatives: a JVP and transpose for the signal alone,
+and a full pair that takes the phase tangent as well and returns its
+cotangent. The phase enters every cell's weights as the one factor
+`exp(i (phase[a1] - phase[a2]))`, so the phase tangent turns the pair's
+product by `i` (`dV = i (dphase[a1] - dphase[a2]) V` per source and fine
+channel) and the cotangent is `-Im(g V)` on `a1` and `+Im(g V)` on `a2`,
+summed over the baseline's sources and fine channels; both cost one extra
+contraction against weights the kernel has already formed. The JVP rule binds
+the signal-only pair whenever the phase tangent is a symbolic zero, so a
+fixed orbit computes no phase derivative. The compiled JVPs and transposes
+follow JAX's complex cotangent convention.
+
+The moment branches, segmentation and cubic translation follow
+`tabascal.coarse_rfi_vis.analytic_rfi_vis`. `segments`, `terms`, and
+`cubic_terms` are static options: close over them or mark them static when
+jitting. The defaults use two pieces, six curvature terms and three cubic
+terms. `segments=4, terms=16` selects the reference's conservative settings.
+Zero cubic terms deliberately omits cubic phase; delay derivatives above
+order three are always omitted. Coefficient counts 1–9, segments 1–1024,
+curvature terms 1–32 and cubic terms 0–8 are supported. These limits bound
+local storage; convergence still depends on the supplied delay and interval.
+
+Every angle -- the centre phase, the winding and the segment translations --
+is formed and reduced in double for both input precisions: a single-precision
+angle has no fractional turn left at the magnitudes the delay reaches. The
+moment recurrences that follow only multiply the coefficient buffers, so the
+CPU kernels keep them in double, where it is free, while the GPU kernels run
+them at the operator's own precision: a consumer card retires one
+double-precision instruction per sixty-four single-precision ones, and a
+complex64 weight costs well over an order of magnitude more in its double form
+there. Fresnel seeds stay in double either way, and use a
+small-argument series and rational auxiliary functions: a two-term asymptotic
+near 2.5 is insufficient for the higher moments. The common three-coefficient,
+six-term, three-cubic-term path has unrolled moment orders so its CUDA
+recurrence can stay in registers. Wider configurations use a bounded general
+implementation.
+
+GPU scratch is bounded by the `scratch_mb` argument of `eval` (256 MiB by
+default), with time chunks and, when necessary, frequency chunks. A smaller
+budget means more chunks, and each chunk of a transpose passes over the whole
+signal cotangent. The transpose reduces coefficient cotangents within each tile
+in shared memory, writes separate partials for partner tiles, and gathers
+through the interpolation stencils. Shared-memory atomics mean the last bits of
+the GPU transpose can vary between runs.
+
+After building on the GPU host, run:
+
+```bash
+python -m pytest tests/test_rfi_analytic_vis_op.py
+python tests/benchmark_rfi_analytic_vis_op.py --antennas 256 512 --iterations 100
+```
+
+The analytic tests compare both precisions and the amplitude and phase
+derivatives against a frozen float64 JAX reference. They cover zero winding, both recurrence
+branches, the Fresnel boundary, cubic translation, wider coefficient counts,
+sparse/reversed/autocorrelation baselines, partial antenna tiles and forced
+scratch splitting. The benchmark reports synchronised forward, JVP and VJP
+times, signal-only and with the phase, after compilation and warmup. It uses
+synthetic inputs; repeat the scientific accuracy check and the 100-iteration
+optimisation with the production SKA-Low data. Inspect ptxas register/spill
+reports and Nsight local-memory traffic for the default specialisation before
+interpreting its speedup. CUDA compilation, device tests and performance need
+the GPU host; CPU tests do not establish those results.

ri_kernels-0.2.2/README.md

@@ -0,0 +1,182 @@
+# RI Kernels
+
+CPU and GPU kernels for radio interferometry, exposed to JAX as FFI custom
+calls with full support for automatic differentiation.
+
+## Installation
+
+The CPU kernels and all of the Python code live in `ri_kernels`. The GPU kernels
+ship as add-on packages, one per CUDA major version, selected with an extra:
+
+```bash
+pip install ri_kernels               # CPU kernels only
+pip install "ri_kernels[cuda12]"     # + CUDA 12 kernels and jax[cuda12]
+pip install "ri_kernels[cuda13]"     # + CUDA 13 kernels and jax[cuda13]
+```
+
+## Building from source
+
+Requires CMake >= 3.20, a C++20 compiler, and network access at configure time
+(Google Highway is fetched by CMake).
+
+```bash
+pip install .                      # CPU only
+RI_KERNELS_CUDA=1 pip install .    # CPU + CUDA, needs nvcc
+RI_KERNELS_ROCM=1 pip install .    # CPU + ROCm, needs hipcc
+```
+
+Both libraries land in the `ri_kernels` package directory in that case, which is
+also a location the loader searches. Note that `pip install .[cuda12]` is a
+different thing: the extra resolves `ri_kernels_cuda12` from the index, so it
+needs a published release. To build the GPU kernels from the checkout, use
+`RI_KERNELS_CUDA=1`.
+
+Other CMake options of note: `RI_KERNELS_CPU` (default `ON`),
+`RI_KERNELS_MULTI_ARCH` (dynamic SIMD dispatch, default `ON` — turn it off and
+set arch flags via `CMAKE_CXX_FLAGS` for a single-target build),
+`RI_KERNELS_BUNDLED_HIGHWAY`, `CMAKE_CUDA_ARCHITECTURES`, and
+`RI_KERNELS_STATIC_CUDART` (default `ON`; set it to `OFF` when a package
+manager should provide the shared CUDA runtime).
+
+## Tests
+
+Tests can be run with
+
+```bash
+python -m pytest tests
+```
+
+## RFI visibilities
+
+`RFIVisOp` computes the per-baseline RFI visibility
+
+```
+vis[bl, f, t] = mean over (n_int_freq, n_int_time) of
+                sum over n_rfi of A[a1] conj(A[a2]) exp(i (φ[a1] - φ[a2]))
+```
+
+for amplitudes `A` shaped `(n_ant, n_freq, n_time, n_rfi, n_int_freq,
+n_int_time)` and phases `φ` in radians of the same shape, giving an output of
+shape `(n_baselines, n_freq, n_time)`. It is constructed from the baseline
+layout and evaluated through `eval`:
+
+```python
+from ri_kernels.jax_api import RFIVisOp
+
+vis = RFIVisOp(n_ant, a1, a2).eval(rfi_amp_fine, rfi_phase)
+```
+
+The operator has native primal, JVP, and transpose kernels for CPU, CUDA, and
+ROCm, so both forward- and reverse-mode differentiation stay inside the
+kernels. Precision has to match across the inputs: complex64 with float32, or
+complex128 with float64. Both amplitudes and phases are differentiated.
+
+## Analytic RFI visibilities
+
+`RFIAnalyticVisOp` integrates the time polynomial against phase moments, with
+frequency quadrature unchanged. It stages amplitude coefficients per antenna
+and cell, then contracts pairs of antenna tiles. The default time scratch axis
+has three coefficients, independent of fringe winding.
+
+```python
+from ri_kernels.jax_api import RFIAnalyticVisOp
+
+vis = RFIAnalyticVisOp(n_ant, a1, a2).eval(
+    amp, phase, delay_us, w_freq, start_freq, g_time, start_time,
+    dnu_mhz, int_time, freq_mhz,
+    segments=2, terms=6, cubic_terms=3,
+)
+```
+
+- `amp`, complex `(n_ant, n_rfi, n_freq, n_time)`: the signal on the data
+  grid.
+- `phase`, real `(n_ant, n_rfi, n_freq, n_time)`: the phase at the channel and
+  cell centre, reduced to one turn (in float64, before casting).
+- `delay_us`, real `(n_ant, n_rfi, n_time, n_path)`: the geometric delay in
+  microseconds, with the sign that makes the phase `2π f τ`, and its first
+  `n_path - 1` time derivatives at the cell centre, relative to the array mean
+  (a common term cancels in every baseline; the full delay's change across a
+  cell is ~1e4 wavelengths, beyond float32).
+- `w_freq`, `start_freq`: per channel, the weights that turn its stencil of
+  `n_sf` neighbouring channels into its fine channels, and the first channel
+  of the stencil.
+- `g_time`, `start_time`: per cell, monomial coefficients in
+  `x = 2*tau/int_time`, as `tabascal.poly_interp.monomial_tables` produces
+  them, including the shifted edge stencils, and the first cell of the
+  stencil. Each stencil must lie inside its axis and contain its cell.
+- `dnu_mhz`: the fine offsets from the channel centre; `freq_mhz`: the channel
+  centres. Both are in **MHz**; divide `config.freqs` in Hz by `1e6`. MHz × μs
+  is cycles.
+- `int_time`: a positive scalar in seconds.
+
+Every real operand shares the precision of `phase`; the stencil starts are
+int32. The constructor builds, per pair of 32-antenna tiles, the list of the
+pairs the baselines cover there (`tile_pair_list`), which the staged GPU
+kernels work through a tile pair per block; each `(a1, a2)` baseline may
+appear at most once.
+
+`amp` and `phase` are differentiated; `delay_us`, `int_time` and the tables
+are constants whose gradients `eval` stops. Two kernel pairs carry the
+derivatives: a JVP and transpose for the signal alone,
+and a full pair that takes the phase tangent as well and returns its
+cotangent. The phase enters every cell's weights as the one factor
+`exp(i (phase[a1] - phase[a2]))`, so the phase tangent turns the pair's
+product by `i` (`dV = i (dphase[a1] - dphase[a2]) V` per source and fine
+channel) and the cotangent is `-Im(g V)` on `a1` and `+Im(g V)` on `a2`,
+summed over the baseline's sources and fine channels; both cost one extra
+contraction against weights the kernel has already formed. The JVP rule binds
+the signal-only pair whenever the phase tangent is a symbolic zero, so a
+fixed orbit computes no phase derivative. The compiled JVPs and transposes
+follow JAX's complex cotangent convention.
+
+The moment branches, segmentation and cubic translation follow
+`tabascal.coarse_rfi_vis.analytic_rfi_vis`. `segments`, `terms`, and
+`cubic_terms` are static options: close over them or mark them static when
+jitting. The defaults use two pieces, six curvature terms and three cubic
+terms. `segments=4, terms=16` selects the reference's conservative settings.
+Zero cubic terms deliberately omits cubic phase; delay derivatives above
+order three are always omitted. Coefficient counts 1–9, segments 1–1024,
+curvature terms 1–32 and cubic terms 0–8 are supported. These limits bound
+local storage; convergence still depends on the supplied delay and interval.
+
+Every angle -- the centre phase, the winding and the segment translations --
+is formed and reduced in double for both input precisions: a single-precision
+angle has no fractional turn left at the magnitudes the delay reaches. The
+moment recurrences that follow only multiply the coefficient buffers, so the
+CPU kernels keep them in double, where it is free, while the GPU kernels run
+them at the operator's own precision: a consumer card retires one
+double-precision instruction per sixty-four single-precision ones, and a
+complex64 weight costs well over an order of magnitude more in its double form
+there. Fresnel seeds stay in double either way, and use a
+small-argument series and rational auxiliary functions: a two-term asymptotic
+near 2.5 is insufficient for the higher moments. The common three-coefficient,
+six-term, three-cubic-term path has unrolled moment orders so its CUDA
+recurrence can stay in registers. Wider configurations use a bounded general
+implementation.
+
+GPU scratch is bounded by the `scratch_mb` argument of `eval` (256 MiB by
+default), with time chunks and, when necessary, frequency chunks. A smaller
+budget means more chunks, and each chunk of a transpose passes over the whole
+signal cotangent. The transpose reduces coefficient cotangents within each tile
+in shared memory, writes separate partials for partner tiles, and gathers
+through the interpolation stencils. Shared-memory atomics mean the last bits of
+the GPU transpose can vary between runs.
+
+After building on the GPU host, run:
+
+```bash
+python -m pytest tests/test_rfi_analytic_vis_op.py
+python tests/benchmark_rfi_analytic_vis_op.py --antennas 256 512 --iterations 100
+```
+
+The analytic tests compare both precisions and the amplitude and phase
+derivatives against a frozen float64 JAX reference. They cover zero winding, both recurrence
+branches, the Fresnel boundary, cubic translation, wider coefficient counts,
+sparse/reversed/autocorrelation baselines, partial antenna tiles and forced
+scratch splitting. The benchmark reports synchronised forward, JVP and VJP
+times, signal-only and with the phase, after compilation and warmup. It uses
+synthetic inputs; repeat the scientific accuracy check and the 100-iteration
+optimisation with the production SKA-Low data. Inspect ptxas register/spill
+reports and Nsight local-memory traffic for the default specialisation before
+interpreting its speedup. CUDA compilation, device tests and performance need
+the GPU host; CPU tests do not establish those results.

{ri_kernels-0.2.0 → ri_kernels-0.2.2}/pyproject.toml

@@ -24,7 +24,7 @@ cmake.define.RI_KERNELS_ROCM = "ON"
 
 [project]
 name = "ri_kernels"
-version = "0.2.0"
+version = "0.2.2"
 license = {file = "LICENSE"}
 readme = "README.md"
 authors = [
@@ -57,11 +57,11 @@ test  = ["pytest", "pytest-cov"]
 #
 # The pins here must track [project].version; ci/make_variant.py fails the build
 # if they drift.
-cuda12 = ["jax[cuda12]", "ri_kernels_cuda12==0.2.0"]
-cuda13 = ["jax[cuda13]", "ri_kernels_cuda13==0.2.0"]
+cuda12 = ["jax[cuda12]", "ri_kernels_cuda12==0.2.2"]
+cuda13 = ["jax[cuda13]", "ri_kernels_cuda13==0.2.2"]
 
 [project.urls]
-Repository = "https://github.com/epfl-radio-astro/ri_kernels.git"
+Repository = "https://github.com/epfl-radio-astro/ri-kernels"
 
 
 # ── Wheel building ─────────────────────────────────────────────────────────────

ri_kernels-0.2.2/ri_kernels/jax_api/__init__.py

@@ -0,0 +1,6 @@
+"""JAX bindings for the RI kernels."""
+
+from .rfi_analytic_vis_op import RFIAnalyticVisOp
+from .rfi_vis_op import RFIVisOp, prepare_indices
+
+__all__ = ["RFIAnalyticVisOp", "RFIVisOp", "prepare_indices"]