jaxdecomp 0.2.2__tar.gz → 0.2.4__tar.gz

{jaxdecomp-0.2.2 → jaxdecomp-0.2.4}/.pre-commit-config.yaml

@@ -1,20 +1,16 @@
 repos:
 -   repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v2.3.0
+    rev: 'v5.0.0'
     hooks:
     -   id: check-yaml
     -   id: end-of-file-fixer
     -   id: trailing-whitespace
--   repo: https://github.com/google/yapf
-    rev: v0.40.2
-    hooks:
-    - id: yapf
-      args: ['--parallel', '--in-place']
--   repo: https://github.com/pycqa/isort
-    rev: 5.13.2
-    hooks:
-      - id: isort
-        name: isort (python)
+- repo: https://github.com/astral-sh/ruff-pre-commit
+  rev: v0.8.2
+  hooks:
+    - id: ruff-format
+    - id: ruff
+
 -   repo: https://github.com/pre-commit/mirrors-clang-format
     rev: v18.1.4
     hooks:

{jaxdecomp-0.2.2 → jaxdecomp-0.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxdecomp
-Version: 0.2.2
+Version: 0.2.4
 Summary: JAX bindings for the cuDecomp library
 Author: Wassim Kabalan, Francois Lanusse
 License: MIT License
@@ -34,9 +34,9 @@ Classifier: Intended Audience :: Developers
 Classifier: Intended Audience :: Science/Research
 Project-URL: Homepage, https://github.com/DifferentiableUniverseInitiative/jaxDecomp
 Requires-Dist: jaxtyping>=0.2.0
-Requires-Dist: jax>=0.4.30
+Requires-Dist: jax>=0.4.35
+Requires-Dist: typing-extensions; python_version < "3.11"
 Requires-Dist: pytest>=8.0.0; extra == "test"
-Requires-Dist: jax>=0.4.30; extra == "test"
 Provides-Extra: test
 Description-Content-Type: text/markdown
 

jaxdecomp-0.2.4/examples/lpt_nbody_demo.py

@@ -0,0 +1,292 @@
+import argparse
+import os
+from functools import partial
+from typing import Any, Callable, Hashable
+
+Specs = Any
+AxisName = Hashable
+
+import jax
+
+jax.config.update("jax_enable_x64", False)
+
+import jax.numpy as jnp
+import jax_cosmo as jc
+import numpy as np
+from jax._src import mesh as mesh_lib
+from jax.experimental import mesh_utils
+from jax.experimental.shard_map import shard_map
+from jax.sharding import Mesh
+from jax.sharding import PartitionSpec as P
+from scatter import scatter
+
+import jaxdecomp
+
+
+def shmap(
+    f: Callable,
+    in_specs: Specs,
+    out_specs: Specs,
+    check_rep: bool = True,
+    auto: frozenset[AxisName] = frozenset(),
+):
+    """Helper function to create a shard_map function that extracts the mesh from the
+    context."""
+    mesh = mesh_lib.thread_resources.env.physical_mesh
+    return shard_map(f, mesh, in_specs, out_specs, check_rep, auto)
+
+
+def _global_to_local_size(nc: int):
+    """Helper function to get the local size of a mesh given the global size."""
+    pdims = mesh_lib.thread_resources.env.physical_mesh.devices.shape
+    return [nc // pdims[0], nc // pdims[1], nc]
+
+
+def fttk(nc: int):
+    """
+    Generate Fourier transform wave numbers for a given mesh.
+
+    Args:
+        nc (int): Shape of the mesh grid.
+
+    Returns:
+        list: List of wave number arrays for each dimension in
+        the order [kx, ky, kz].
+    """
+    kd = np.fft.fftfreq(nc) * 2 * np.pi
+
+    @partial(
+        shmap,
+        in_specs=(P("x"), P("y"), P(None)),
+        out_specs=(P("x"), P(None, "y"), P(None)),
+    )
+    def get_kvec(ky, kz, kx):
+        return (ky.reshape([-1, 1, 1]),
+        kz.reshape([1, -1, 1]),
+        kx.reshape([1, 1, -1]))  # yapf: disable
+
+    ky, kz, kx = get_kvec(kd, kd, kd)  # The order of the output
+    # corresponds to the order of dimensions in the transposed FFT
+    # output
+    return kx, ky, kz
+
+
+def gravity_kernel(kx, ky, kz):
+    """Computes a Fourier kernel combining laplace and derivative
+    operators to compute gravitational forces.
+
+    Args:
+        kvec (tuple of float): Wave numbers in Fourier space.
+
+    Returns:
+        tuple of jnp.ndarray: kernels for each dimension.
+    """
+    kk = kx**2 + ky**2 + kz**2
+    laplace_kernel = jnp.where(kk == 0, 1.0, 1.0 / kk)
+
+    grav_kernel = (laplace_kernel * 1j * kx,
+               laplace_kernel * 1j * ky,
+               laplace_kernel * 1j * kz)  # yapf: disable
+    return grav_kernel
+
+
+def gaussian_field_and_forces(key, nc, box_size, power_spectrum):
+    """
+    Generate a Gaussian field with a given power spectrum, along with gravitational forces.
+
+    Args:
+        key (int): Key for the random number generator.
+        nc (int): Number of cells in the mesh.
+        box_size (float): Size of the box.
+        power_spectrum (callable): Power spectrum function.
+
+    Returns:
+        tuple of jnp.ndarray: The generated Gaussian field and the gravitational forces.
+    """
+    local_mesh_shape = _global_to_local_size(nc)
+
+    # Create a distributed field drawn from a Gaussian distribution in real space
+    delta = shmap(
+    partial(jax.random.normal, shape=local_mesh_shape, dtype='float32'),
+    in_specs=P(None),
+    out_specs=P('x', 'y'))(key)  # yapf: disable
+
+    # Compute the Fourier transform of the field
+    delta_k = jaxdecomp.fft.pfft3d(delta.astype(jnp.complex64))
+
+    # Compute the Fourier wavenumbers of the field
+    kx, ky, kz = fttk(nc)
+    kk = jnp.sqrt(kx**2 + ky**2 + kz**2) * (nc / box_size)
+
+    # Apply power spectrum to Fourier modes
+    delta_k *= (power_spectrum(kk) * (nc / box_size) ** 3) ** 0.5
+
+    # Compute inverse Fourier transform to recover the initial conditions in real space
+    delta = jaxdecomp.fft.pifft3d(delta_k).real
+
+    # Compute gravitational forces associated with this field
+    grav_kernel = gravity_kernel(kx, ky, kz)
+    forces_k = [g * delta_k for g in grav_kernel]
+
+    # Retrieve the forces in real space by inverse Fourier transforming
+    forces = jnp.stack([jaxdecomp.fft.pifft3d(f).real for f in forces_k], axis=-1)
+
+    return delta, forces
+
+
+def cic_paint(displacement, halo_size):
+    """Paints particles on a mesh using Cloud-In-Cell interpolation.
+
+    Args:
+        displacement (jnp.ndarray): Displacement of each particle.
+        halo_size (int): Halo size for painting.
+
+    Returns:
+        jnp.ndarray: Density field.
+    """
+    local_mesh_shape = _global_to_local_size(displacement.shape[0])
+    hs = halo_size
+
+    @partial(shmap, in_specs=(P("x", "y"),), out_specs=P("x", "y"))
+    def cic_op(disp):
+        """CiC operation on each local slice of the mesh."""
+        # Create a mesh to paint the particles on for the local slice
+        mesh = jnp.zeros(disp.shape[:-1], dtype="float32")
+
+        # Padding the mesh along the two first dimensions
+        mesh = jnp.pad(mesh, [[hs, hs], [hs, hs], [0, 0]])
+
+        # Compute the position of the particles on a regular grid
+        pos_x, pos_y, pos_z = jnp.meshgrid(
+            jnp.arange(local_mesh_shape[0]),
+            jnp.arange(local_mesh_shape[1]),
+            jnp.arange(local_mesh_shape[2]),
+            indexing="ij",
+        )
+
+        # adding an offset of size halo size
+        pos = jnp.stack([pos_x + hs, pos_y + hs, pos_z], axis=-1)
+
+        # Apply scatter operation to paint the particles on the local mesh
+        field = scatter(pos.reshape([-1, 3]), disp.reshape([-1, 3]), mesh)
+
+        return field
+
+    # Performs painting on a padded mesh, with halos on the two first dimensions
+    field = cic_op(displacement)
+
+    # Run halo exchange to get the correct values at the boundaries
+    field = jaxdecomp.halo_exchange(
+        field, halo_extents=(hs // 2, hs // 2, 0), halo_periods=(True, True, True)
+    )
+
+    @partial(shmap, in_specs=(P("x", "y"),), out_specs=P("x", "y"))
+    def unpad(x):
+        """Removes the padding and reduce the halo regions"""
+        x = x.at[hs : hs + hs // 2].add(x[: hs // 2])
+        x = x.at[-(hs + hs // 2) : -hs].add(x[-hs // 2 :])
+        x = x.at[:, hs : hs + hs // 2].add(x[:, : hs // 2])
+        x = x.at[:, -(hs + hs // 2) : -hs].add(x[:, -hs // 2 :])
+        return x[hs:-hs, hs:-hs, :]
+
+    # Unpad the output array
+    field = unpad(field)
+    return field
+
+
+@partial(jax.jit, static_argnames=("nc", "box_size", "halo_size"))
+def simulation_fn(key, nc, box_size, halo_size, a=1.0):
+    """
+    Run a simulation to generate initial conditions and density field using LPT.
+
+    Args:
+        key (list of int): Jax random key for the random number generator.
+        nc (int): Size of the mesh grid.
+        box_size (float): Size of the box.
+        halo_size (int): Halo size for painting.
+        a (float): Scale factor of final field.
+
+    Returns:
+            tuple of jnp.ndarray: Initial conditions and final density field.
+    """
+    # Build a default cosmology
+    cosmo = jc.Planck15()
+
+    # Create a small function to generate the linear matter power spectrum at arbitrary k
+    k = jnp.logspace(-4, 1, 128)
+    pk = jc.power.linear_matter_power(cosmo, k)
+    pk_fn = jax.jit(
+        lambda x: jc.scipy.interpolate.interp(x.reshape([-1]), k, pk).reshape(x.shape)
+    )
+
+    # Generate a Gaussian field and gravitational forces from a power spectrum
+    intial_conditions, initial_forces = gaussian_field_and_forces(
+        key=key, nc=nc, box_size=box_size, power_spectrum=pk_fn
+    )
+
+    # Compute the LPT displacement that particles initialy placed on a regular grid
+    # would experience at scale factor a, by simple Zeldovich approximation
+    initial_displacement = (
+        jc.background.growth_factor(cosmo, jnp.atleast_1d(a)) * initial_forces
+    )
+
+    # Paints the displaced particles on a mesh to obtain the density field
+    final_field = cic_paint(initial_displacement, halo_size)
+
+    return intial_conditions, final_field
+
+
+def main(args):
+    print(f"Running with arguments {args}")
+
+    # Setting up distributed jax
+    jax.distributed.initialize()
+    rank = jax.process_index()
+    size = jax.process_count()
+
+    # Setting up distributed random numbers
+    master_key = jax.random.PRNGKey(42)
+    key = jax.random.split(master_key, size)[rank]
+
+    # Create computing mesh and sharding information
+    pdims = tuple(map(int, args.pdims.split("x")))
+    devices = mesh_utils.create_device_mesh(pdims)
+    mesh = Mesh(devices.T, axis_names=("x", "y"))
+
+    # Run the simulation on the compute mesh
+    with mesh:
+        initial_conds, final_field = simulation_fn(
+            key=key, nc=args.nc, box_size=args.box_size, halo_size=args.halo_size
+        )
+
+    # Create output directory to save the results
+    output_dir = args.output
+    os.makedirs(output_dir, exist_ok=True)
+    np.save(
+        f"{output_dir}/initial_conditions_{rank}.npy", initial_conds.addressable_data(0)
+    )
+    np.save(f"{output_dir}/field_{rank}.npy", final_field.addressable_data(0))
+    print(f"Finished saved to {output_dir}")
+
+    # Closing distributed jax
+    jax.distributed.shutdown()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser("Distributed LPT N-body simulation.")
+    parser.add_argument(
+        "--pdims", type=str, default="1x1", help="Processor grid dimensions"
+    )
+    parser.add_argument(
+        "--nc", type=int, default=256, help="Number of cells in the mesh"
+    )
+    parser.add_argument(
+        "--box_size", type=float, default=512.0, help="Box size in Mpc/h"
+    )
+    parser.add_argument(
+        "--halo_size", type=int, default=32, help="Halo size for painting"
+    )
+    parser.add_argument("--output", type=str, default="out")
+    args = parser.parse_args()
+
+    main(args)

jaxdecomp-0.2.4/examples/scatter.py

@@ -0,0 +1,163 @@
+# This file is adapted from the scatter implementation of the pmwd library
+# https://github.com/eelregit/pmwd/blob/master/pmwd/scatter.py
+# It provides a simple way to perform a scatter operation by chunks and saves
+# memory compared to a native jax.lax.scatter.
+# Below is the orginal license of the pmwd library:
+###############################################################################
+# BSD 3-Clause License
+#
+# Copyright (c) 2021, the pmwd developers
+# All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+#    list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+#    this list of conditions and the following disclaimer in the documentation
+#    and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+#    contributors may be used to endorse or promote products derived from
+#    this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+import jax.numpy as jnp
+from jax.lax import scan
+
+
+def _chunk_split(ptcl_num, chunk_size, *arrays):
+    """Split and reshape particle arrays into chunks and remainders, with the remainders
+    preceding the chunks. 0D ones are duplicated as full arrays in the chunks."""
+    chunk_size = ptcl_num if chunk_size is None else min(chunk_size, ptcl_num)
+    remainder_size = ptcl_num % chunk_size
+    chunk_num = ptcl_num // chunk_size
+
+    remainder = None
+    chunks = arrays
+    if remainder_size:
+        remainder = [x[:remainder_size] if x.ndim != 0 else x for x in arrays]
+        chunks = [x[remainder_size:] if x.ndim != 0 else x for x in arrays]
+
+    # `scan` triggers errors in scatter and gather without the `full`
+    chunks = [
+        x.reshape(chunk_num, chunk_size, *x.shape[1:])
+        if x.ndim != 0
+        else jnp.full(chunk_num, x)
+        for x in chunks
+    ]
+
+    return remainder, chunks
+
+
+def enmesh(i1, d1, a1, s1, b12, a2, s2):
+    """Multilinear enmeshing."""
+    i1 = jnp.asarray(i1)
+    d1 = jnp.asarray(d1)
+    a1 = jnp.float64(a1) if a2 is not None else jnp.array(a1, dtype=d1.dtype)
+    if s1 is not None:
+        s1 = jnp.array(s1, dtype=i1.dtype)
+    b12 = jnp.float64(b12)
+    if a2 is not None:
+        a2 = jnp.float64(a2)
+    if s2 is not None:
+        s2 = jnp.array(s2, dtype=i1.dtype)
+
+    dim = i1.shape[1]
+    neighbors = (
+        jnp.arange(2**dim, dtype=i1.dtype)[:, jnp.newaxis]
+        >> jnp.arange(dim, dtype=i1.dtype)
+    ) & 1
+
+    if a2 is not None:
+        P = i1 * a1 + d1 - b12
+        P = P[:, jnp.newaxis]  # insert neighbor axis
+        i2 = P + neighbors * a2  # multilinear
+
+        if s1 is not None:
+            L = s1 * a1
+            i2 %= L
+
+        i2 //= a2
+        d2 = P - i2 * a2
+
+        if s1 is not None:
+            d2 -= jnp.rint(d2 / L) * L  # also abs(d2) < a2 is expected
+
+        i2 = i2.astype(i1.dtype)
+        d2 = d2.astype(d1.dtype)
+        a2 = a2.astype(d1.dtype)
+
+        d2 /= a2
+    else:
+        i12, d12 = jnp.divmod(b12, a1)
+        i1 -= i12.astype(i1.dtype)
+        d1 -= d12.astype(d1.dtype)
+
+        # insert neighbor axis
+        i1 = i1[:, jnp.newaxis]
+        d1 = d1[:, jnp.newaxis]
+
+        # multilinear
+        d1 /= a1
+        i2 = jnp.floor(d1).astype(i1.dtype)
+        i2 += neighbors
+        d2 = d1 - i2
+        i2 += i1
+
+        if s1 is not None:
+            i2 %= s1
+
+    f2 = 1 - jnp.abs(d2)
+
+    if s1 is None and s2 is not None:  # all i2 >= 0 if s1 is not None
+        i2 = jnp.where(i2 < 0, s2, i2)
+
+    f2 = f2.prod(axis=-1)
+
+    return i2, f2
+
+
+def scatter(pmid, disp, mesh, chunk_size=2**24, val=1.0, offset=0, cell_size=1.0):
+    ptcl_num, spatial_ndim = pmid.shape
+    val = jnp.asarray(val)
+    mesh = jnp.asarray(mesh)
+
+    remainder, chunks = _chunk_split(ptcl_num, chunk_size, pmid, disp, val)
+
+    carry = mesh, offset, cell_size
+    if remainder is not None:
+        carry = _scatter_chunk(carry, remainder)[0]
+    carry = scan(_scatter_chunk, carry, chunks)[0]
+    mesh = carry[0]
+
+    return mesh
+
+
+def _scatter_chunk(carry, chunk):
+    mesh, offset, cell_size = carry
+    pmid, disp, val = chunk
+    spatial_ndim = pmid.shape[1]
+    spatial_shape = mesh.shape
+
+    # multilinear mesh indices and fractions
+    ind, frac = enmesh(
+        pmid, disp, cell_size, spatial_shape, offset, cell_size, spatial_shape
+    )
+    # scatter
+    ind = tuple(ind[..., i] for i in range(spatial_ndim))
+    mesh = mesh.at[ind].add(val * frac)
+
+    carry = mesh, offset, cell_size
+    return carry, None

{jaxdecomp-0.2.2 → jaxdecomp-0.2.4}/examples/visualizer.ipynb

@@ -70,8 +70,8 @@
     }
    ],
    "source": [
-    "folder = '../out'\n",
-    "pdims=(4,4)\n",
+    "folder = \"../out\"\n",
+    "pdims = (4, 4)\n",
     "\n",
     "init_field_slices = []\n",
     "field_slices = []\n",
@@ -79,11 +79,11 @@
     "for i in range(pdims[0]):\n",
     "    row_init_field = []\n",
     "    row_field = []\n",
-    "    \n",
+    "\n",
     "    for j in range(pdims[1]):\n",
-    "        slice_index = i * pdims[1]  + j \n",
-    "        row_field.append(np.load(f'{folder}/field_{slice_index}.npy'))\n",
-    "        row_init_field.append(np.load(f'{folder}/initial_conditions_{slice_index}.npy'))\n",
+    "        slice_index = i * pdims[1] + j\n",
+    "        row_field.append(np.load(f\"{folder}/field_{slice_index}.npy\"))\n",
+    "        row_init_field.append(np.load(f\"{folder}/initial_conditions_{slice_index}.npy\"))\n",
     "\n",
     "    field_slices.append(np.vstack(row_field))\n",
     "    init_field_slices.append(np.vstack(row_init_field))\n",
@@ -148,23 +148,31 @@
     "    slicing[proj_axis] = slice(None, sum_over)\n",
     "    slicing = tuple(slicing)\n",
     "\n",
-    "\n",
     "    fig, axes = plt.subplots(1, 2, figsize=(12, 6))\n",
     "\n",
     "    # Flatten axes for easy indexing\n",
     "    axes = axes.flatten()\n",
     "\n",
     "    # Plot initial conditions\n",
-    "    axes[0].imshow(initial_conditions[slicing].sum(axis=proj_axis), cmap='magma', extent=[0, box_size + 5, 0, box_size + 5])\n",
-    "    axes[0].set_xlabel('Mpc/h')\n",
-    "    axes[0].set_ylabel('Mpc/h')\n",
-    "    axes[0].set_title('Initial conditions')\n",
+    "    axes[0].imshow(\n",
+    "        initial_conditions[slicing].sum(axis=proj_axis),\n",
+    "        cmap=\"magma\",\n",
+    "        extent=[0, box_size + 5, 0, box_size + 5],\n",
+    "    )\n",
+    "    axes[0].set_xlabel(\"Mpc/h\")\n",
+    "    axes[0].set_ylabel(\"Mpc/h\")\n",
+    "    axes[0].set_title(\"Initial conditions\")\n",
     "\n",
     "    # Plot LPT density field at z=0\n",
-    "    axes[1].imshow(field[slicing].sum(axis=proj_axis), cmap='magma', extent=[0, box_size + 5, 0, box_size + 5])\n",
-    "    axes[1].set_xlabel('Mpc/h')\n",
-    "    axes[1].set_ylabel('Mpc/h')\n",
-    "    axes[1].set_title('LPT density field at z=0')\n",
+    "    axes[1].imshow(\n",
+    "        field[slicing].sum(axis=proj_axis),\n",
+    "        cmap=\"magma\",\n",
+    "        extent=[0, box_size + 5, 0, box_size + 5],\n",
+    "    )\n",
+    "    axes[1].set_xlabel(\"Mpc/h\")\n",
+    "    axes[1].set_ylabel(\"Mpc/h\")\n",
+    "    axes[1].set_title(\"LPT density field at z=0\")\n",
+    "\n",
     "\n",
     "for i in range(3):\n",
     "    plot(i)"
@@ -201,10 +209,14 @@
     "box_size = initial_conditions.shape[0] + 5\n",
     "plt.figure(figsize=(20, 20))\n",
     "# Generate the plot\n",
-    "plt.imshow(np.log10(field[:16].sum(axis=proj_axis) + 1), cmap='magma', extent=[0, box_size, 0, box_size])\n",
-    "plt.xlabel('Mpc/h')\n",
-    "plt.ylabel('Mpc/h')\n",
-    "plt.title('LPT density field at z=0')\n",
+    "plt.imshow(\n",
+    "    np.log10(field[:16].sum(axis=proj_axis) + 1),\n",
+    "    cmap=\"magma\",\n",
+    "    extent=[0, box_size, 0, box_size],\n",
+    ")\n",
+    "plt.xlabel(\"Mpc/h\")\n",
+    "plt.ylabel(\"Mpc/h\")\n",
+    "plt.title(\"LPT density field at z=0\")\n",
     "\n",
     "# Display the plot\n",
     "plt.show()"

{jaxdecomp-0.2.2 → jaxdecomp-0.2.4}/pyproject.toml

@@ -3,7 +3,7 @@ requires = ["scikit-build-core>=0.4.0", "pybind11>=2.9.0"]
 build-backend = "scikit_build_core.build"
 [project]
 name = "jaxdecomp"
-version = "0.2.2"
+version = "0.2.4"
 description = "JAX bindings for the cuDecomp library"
 authors = [
   { name = "Wassim Kabalan" },
@@ -24,11 +24,12 @@ classifiers = [
 ]
 dependencies = [
 "jaxtyping>=0.2.0",
-"jax>=0.4.30",
+"jax>=0.4.35",
+"typing-extensions; python_version < '3.11'",
 ]
 
 [project.optional-dependencies]
-test = ["pytest>=8.0.0" , "jax>=0.4.30"]
+test = ["pytest>=8.0.0"]
 
 [tool.scikit-build]
 minimum-version = "0.8"
@@ -42,3 +43,5 @@ CMAKE_LIBRARY_OUTPUT_DIRECTORY = ""
 [tool.cibuildwheel]
 test-extras = "test"
 test-command = "pytest {project}/tests"
+[tool.ruff]
+extend-ignore = ["E402"]

{jaxdecomp-0.2.2 → jaxdecomp-0.2.4}/scripts/autotune.py

@@ -17,11 +17,12 @@ config.halo_comm_backend = jaxdecomp.config.halo_comm_backend
 config.transpose_comm_backend = jaxdecomp.config.transpose_comm_backend
 
 # Run autotune
-tuned_config = jaxdecomp.get_autotuned_config(config, False, False, True, True,
-                                              (32, 32, 32), (True, True, True))
+tuned_config = jaxdecomp.get_autotuned_config(
+    config, False, False, True, True, (32, 32, 32), (True, True, True)
+)
 
 if rank == 0:
-  print(rank, "*** Results of optimization ***")
-  print(rank, "pdims", tuned_config.pdims)
-  print(rank, "halo_comm_backend", tuned_config.halo_comm_backend)
-  print(rank, "transpose_comm_backend", tuned_config.transpose_comm_backend)
+    print(rank, "*** Results of optimization ***")
+    print(rank, "pdims", tuned_config.pdims)
+    print(rank, "halo_comm_backend", tuned_config.halo_comm_backend)
+    print(rank, "transpose_comm_backend", tuned_config.transpose_comm_backend)