ejkernel 0.0.78__tar.gz → 0.0.79__tar.gz

{ejkernel-0.0.78 → ejkernel-0.0.79}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: ejkernel
-Version: 0.0.78
+Version: 0.0.79
 Summary: Accelerate, Optimize performance with streamlined training and serving options with JAX.
 Keywords: Deep Learning,Machine Learning,JAX,CUDA,XLA,Triton,Pallas
 Author: Erfan Zare Chavoshi
@@ -16,21 +16,20 @@ Classifier: Programming Language :: Python :: 3.13
 Requires-Dist: beartype>=0.22.2
 Requires-Dist: chex>=0.1.91
 Requires-Dist: einops>=0.8.1
-Requires-Dist: jax~=0.9.0
-Requires-Dist: jaxlib~=0.9.0
+Requires-Dist: jax~=0.10.0
+Requires-Dist: jaxlib~=0.10.0
 Requires-Dist: jaxtyping>=0.3.2
 Requires-Dist: pydantic>=2.11.10
 Requires-Dist: tqdm>=4.67.1
-Requires-Dist: jax[cuda13]>=0.9.0 ; extra == 'cuda'
+Requires-Dist: jax[cuda13]>=0.10.0 ; extra == 'cuda'
 Requires-Dist: triton==3.6.0 ; extra == 'cuda'
 Requires-Dist: nvidia-cutlass-dsl[cu13]>=4.4.0.dev1 ; extra == 'cuda'
 Requires-Dist: jax-tvm-ffi==0.1.2 ; extra == 'cuda'
 Requires-Dist: eformer ; extra == 'dev'
-Requires-Dist: jax-metallib==0.9.4.5 ; extra == 'mps'
 Requires-Dist: xprof>=2.20.6 ; extra == 'profile'
 Requires-Dist: tb-nightly>=2.21.0a20250820 ; extra == 'profile'
 Requires-Dist: xprof-nightly>=2.21.6a20250820 ; extra == 'profile'
-Requires-Dist: jax[tpu]>=0.9.0 ; extra == 'tpu'
+Requires-Dist: jax[tpu]>=0.10.0 ; extra == 'tpu'
 Requires-Dist: triton==3.6.0 ; extra == 'triton'
 Requires-Python: >=3.11, <3.14
 Project-URL: Documentation, https://ejkernel.readthedocs.io/en/latest/
@@ -38,7 +37,6 @@ Project-URL: Homepage, https://github.com/erfanzar/ejkernel
 Project-URL: Repository, https://github.com/erfanzar/ejkernel
 Provides-Extra: cuda
 Provides-Extra: dev
-Provides-Extra: mps
 Provides-Extra: profile
 Provides-Extra: tpu
 Provides-Extra: triton

{ejkernel-0.0.78 → ejkernel-0.0.79}/ejkernel/__init__.py

@@ -43,7 +43,7 @@ import os as _os
 _os.environ.setdefault("TF_GPU_ALLOCATOR", "cuda_malloc_async")
 _os.environ.setdefault("CUTE_DSL_ENABLE_TVM_FFI", "1")
 
-__version__ = "0.0.78"
+__version__ = "0.0.79"
 
 from . import errors, kernels, modules, types, utils, xla_utils
 from .errors import EjkernelRuntimeError

{ejkernel-0.0.78 → ejkernel-0.0.79}/ejkernel/kernels/_pallas/tpu/reduce_scatter_matmul/_pallas_impl.py

@@ -673,13 +673,13 @@ def reduce_scatter_matmul(
         grid_spec=pltpu.PrefetchScalarGridSpec(
             num_scalar_prefetch=0,
             in_specs=[
-                pl.BlockSpec(memory_space=pltpu.ANY),
-                pl.BlockSpec(memory_space=pltpu.ANY),
+                pl.BlockSpec(memory_space=pl.ANY),
+                pl.BlockSpec(memory_space=pl.ANY),
             ],
             out_specs=[
-                pl.BlockSpec(memory_space=pltpu.ANY),
-                pl.BlockSpec(memory_space=pltpu.ANY),
-                pl.BlockSpec(memory_space=pltpu.ANY),
+                pl.BlockSpec(memory_space=pl.ANY),
+                pl.BlockSpec(memory_space=pl.ANY),
+                pl.BlockSpec(memory_space=pl.ANY),
             ],
             scratch_shapes=[
                 x_vmem_shape,

{ejkernel-0.0.78 → ejkernel-0.0.79}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "uv_build"
 
 [project]
 name = "ejkernel"
-version = "0.0.78"
+version = "0.0.79"
 authors = [{ name = "Erfan Zare Chavoshi", email = "Erfanzare810@gmail.com" }]
 description = "Accelerate, Optimize performance with streamlined training and serving options with JAX."
 readme = "README.md"
@@ -32,8 +32,8 @@ dependencies = [
     "beartype>=0.22.2",
     "chex>=0.1.91",
     "einops>=0.8.1",
-    "jax~=0.9.0",
-    "jaxlib~=0.9.0",
+    "jax~=0.10.0",
+    "jaxlib~=0.10.0",
     "jaxtyping>=0.3.2",
     "pydantic>=2.11.10",
     "tqdm>=4.67.1",
@@ -62,14 +62,14 @@ where = ["."]
 [project.optional-dependencies]
 
 cuda = [
-    "jax[cuda13]>=0.9.0",
+    "jax[cuda13]>=0.10.0",
     "triton==3.6.0", 
     "nvidia-cutlass-dsl[cu13]>=4.4.0.dev1",
     "jax-tvm-ffi==0.1.2",
 ]
-tpu = ["jax[tpu]>=0.9.0"]
+tpu = ["jax[tpu]>=0.10.0"]
 triton = ["triton==3.6.0"]
-mps = ["jax-metallib==0.9.4.5"]
+
 profile = [
     "xprof>=2.20.6",
     "tb-nightly>=2.21.0a20250820",