MODULE.bazel

module(name = "jax")

##############################################################
# Bazel module dependencies
##############################################################

bazel_dep(name = "abseil-cpp", version = "20260526.0", repo_name = "com_google_absl")
bazel_dep(name = "apple_support", version = "2.8.0")
bazel_dep(name = "aspect_bazel_lib", version = "2.22.5")
bazel_dep(name = "bazel_features", version = "1.50.0")
bazel_dep(name = "bazel_skylib", version = "1.9.0")
bazel_dep(name = "flatbuffers", version = "25.2.10", repo_name = "com_github_google_flatbuffers")
bazel_dep(name = "googletest", version = "1.17.0.bcr.2", repo_name = "com_google_googletest")
bazel_dep(name = "grpc", version = "1.83.0", repo_name = "com_github_grpc_grpc")
bazel_dep(name = "platforms", version = "1.1.0")
bazel_dep(name = "protobuf", version = "35.1", repo_name = "com_google_protobuf")
bazel_dep(name = "pybind11_bazel", version = "3.0.0")
bazel_dep(name = "rules_cc", version = "0.2.25")
bazel_dep(name = "rules_python", version = "2.2.0")
bazel_dep(name = "rules_shell", version = "0.8.0")
bazel_dep(name = "zlib", version = "1.3.2")

# TODO: use a released version when available
bazel_dep(name = "rules_ml_toolchain")
archive_override(
    module_name = "rules_ml_toolchain",
    integrity = "sha256-9+4rpwQMv77PxA+AGMWpwcjioakWHi4nJ7UQ0iPHrJo=",
    strip_prefix = "rules_ml_toolchain-3a4df0c8e9e51687ac4c62c793a00a27286e5e83",
    urls = ["https://github.com/google-ml-infra/rules_ml_toolchain/archive/3a4df0c8e9e51687ac4c62c793a00a27286e5e83.tar.gz"],
)

# TODO: use a released version when available
bazel_dep(name = "xla")
archive_override(
    module_name = "xla",
    integrity = "sha256-UKZ/Qqn4U274TL4t5cmI87M9yYRIVzksbmFpCy5/34g=",
    strip_prefix = "xla-d867969800d4a8fd045505eabc65e1e66f30794f",
    urls = ["https://github.com/openxla/xla/archive/d867969800d4a8fd045505eabc65e1e66f30794f.tar.gz"],
)

# flatbuffers transitively depends on aspect_rules_esbuild 0.21.0, which is incompatible with Bazel 9.
# Pin to 0.25.1, the last Bazel 9-compatible release on aspect_rules_js 2.x (required by flatbuffers 25.2.10).
# TODO: Remove this override once flatbuffers is upgraded to >= 25.12.19.
single_version_override(
    module_name = "aspect_rules_esbuild",
    version = "0.25.1",
)

# TODO: upstream, otherwise we have to duplicate the patches in jax
single_version_override(
    module_name = "grpc",
    patch_strip = 1,
    patches = ["//third_party/grpc:grpc.patch"],  # upstreaming in cl/822016400
    version = "1.83.0",
)

# Fix transitive dependency graph error breaking `bazel mod show_repo --all_repos`
bazel_dep(name = "envoy_api", version = "0.0.0-20260130-84e8436")

##############################################################
# Dependencies defined in XLA via module extensions
##############################################################

### TSL
tsl_extension = use_extension("@xla//third_party/extensions:tsl.bzl", "tsl_extension")
use_repo(tsl_extension, "tsl")

### LLVM
llvm = use_extension("@xla//third_party/extensions:llvm.bzl", "llvm_extension")
use_repo(llvm, "llvm-project")

### C++ dependencies
# TODO: convert them to Bazel modules (bazel_dep).
third_party = use_extension("@xla//third_party/extensions:third_party.bzl", "third_party_ext")
use_repo(
    third_party,
    "compute_library",
    "dlpack",
    "eigen_archive",
    "gloo",
    "nanobind",
    "riegeli",
    "shardy",
    "stablehlo",
    "triton",
)

rocm = use_extension("@xla//third_party/extensions:rocm_configure.bzl", "rocm_configure_ext")
use_repo(rocm, "config_rocm_hipcc", "local_config_rocm")

# rules_ml_toolchain's own config_rocm_hipcc is a fixed ROCm tarball it documents as
# test-only, and its ROCm cc toolchain reads it. Point that toolchain at our ROCm.
rocm_hermetic = use_extension("@rules_ml_toolchain//extensions:rocm_hermetic_download.bzl", "rocm_hermetic_download_ext")
override_repo(rocm_hermetic, config_rocm_hipcc = "config_rocm_hipcc")

python_version_ext = use_extension("@rules_ml_toolchain//extensions:python_version.bzl", "python_version_ext")
use_repo(python_version_ext, "python_version_repo")

sycl = use_extension("@rules_ml_toolchain//extensions:sycl_configure.bzl", "sycl_configure_ext")
use_repo(sycl, "level_zero", "local_config_sycl", "oneapi", "zero_loader")

### RBE
rbe_config = use_extension("@xla//third_party/extensions:rbe_config.bzl", "rbe_config_ext")
use_repo(rbe_config, "ml_build_config_platform")

remote_execution_configure = use_extension("@xla//third_party/extensions:remote_execution_configure.bzl", "remote_execution_configure_ext")
use_repo(remote_execution_configure, "local_config_remote_execution")

##############################################################
# Python dependencies
##############################################################

single_version_override(
    module_name = "rules_python",
    patch_strip = 1,
    # Remove the upstreamed patches (if any) when updating to a version that includes their PRs:
    # https://rules-python.readthedocs.io/en/latest/changelog.html
    patches = [
        "//third_party/py:rules_python_whl_abi_tags.patch",
        # Upstreamed at https://github.com/bazel-contrib/rules_python/pull/3768
        "//third_party/py:rules_python_local_wheel.patch",
        # Upstreamed: https://github.com/bazel-contrib/rules_python/pull/4071
        "//third_party/py:rules_python_windows_long_paths.patch",
    ],
    version = "2.2.0",
)

single_version_override(
    module_name = "protobuf",
    patch_strip = 1,
    patches = [
        "//third_party/protobuf:protobuf_arena.patch",
    ],
    version = "35.1",
)

single_version_override(
    module_name = "abseil-cpp",
    patch_strip = 1,
    patches = [
        "//third_party/absl:btree.patch",
        "//third_party/absl:raw_hash_set.patch",
    ],
    version = "20260526.0",
)

### Toolchains

# Keep in sync with jax/version.py (_version) and python.defaults below.
DEFAULT_PYTHON_VERSION = "3.12"

JAXLIB_VERSION = "0.9.2"

python = use_extension("@rules_python//python/extensions:python.bzl", "python")
python.override(
    add_runtime_manifest_files = ["//third_party/py:rules_python_cpython_3_15.txt"],
    minor_mapping = {
        "3.15": "3.15.0rc1",
    },
)
python.defaults(
    # The environment variable takes precedence if set.
    python_version = DEFAULT_PYTHON_VERSION,
    python_version_env = "HERMETIC_PYTHON_VERSION",
)
python.toolchain(python_version = "3.12")
python.toolchain(python_version = "3.13")
python.toolchain(python_version = "3.14")
python.toolchain(python_version = "3.15")

### Pypi
pip = use_extension("@rules_python//python/extensions:pip.bzl", "pip", dev_dependency = True)

[pip.parse(
    download_only = True,
    experimental_index_url = "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple",
    extra_hub_aliases = {
        "numpy": ["numpy_headers"],
    },
    hub_name = "jax_pypi",
    local_wheels = {
        "jax": "dist/jax-*.whl",
        "jaxlib": "dist/jaxlib-*.whl",
        "jax-cuda12-pjrt": "dist/jax_cuda12_pjrt-*.whl",
        "jax-cuda12-plugin": "dist/jax_cuda12_plugin-*.whl",
        "jax-cuda13-pjrt": "dist/jax_cuda13_pjrt-*.whl",
        "jax-cuda13-plugin": "dist/jax_cuda13_plugin-*.whl",
        "libtpu": "dist/libtpu-*.whl",  # On CI, we use the downloaded nightly version of libtpu.
        "ml_dtypes": "dist/ml_dtypes-*.whl",
        "numpy": "dist/numpy-*.whl",  # On TSAN CI, we use a custom TSAN-instrumented NumPy wheel.
        "scipy": "dist/scipy-*.whl",
    },
    python_version = python_version,
    requirements_by_platform = {} if python_version in [
        "3.12",
        "3.13",
    ] else {
        "//build:requirements_lock_{}_ft.txt".format(python_version.replace(".", "_")): "linux_x86_64_freethreaded,windows_x86_64_freethreaded,linux_aarch64_freethreaded,osx_aarch64_freethreaded",
    },
    requirements_lock = "//build:requirements_lock_{}.txt".format(python_version.replace(".", "_")),
    target_platforms = [
        # Explicitly list both x86_64 and aarch64 to ensure cross-compilation works.
        "{os}_x86_64",
        "{os}_aarch64",
        "{os}_x86_64_freethreaded",
        "{os}_aarch64_freethreaded",
    ],
    whl_modifications = {
        "//build:numpy.json": "numpy",  # @pypi_mods is defined in XLA's MODULE.bazel
    },
) for python_version in [
    "3.12",
    "3.13",
    "3.14",
    "3.15",
]]

use_repo(pip, "pypi_mods", pypi = "jax_pypi")

# Override XLA's pypi installation to use the one defined by JAX.
# TODO: have a better way to unify the pypi installation.
override_repo(
    pip,
    xla_pypi = "pypi",
)

### Pybind
pybind11_internal_configure = use_extension(
    "@pybind11_bazel//:internal_configure.bzl",
    "internal_configure_extension",
)
use_repo(pybind11_internal_configure, "pybind11")

##############################################################
# Dependencies defined in rules_ml_toolchain
##############################################################

### CUDA dependencies
cuda_configure = use_extension("@rules_ml_toolchain//extensions:cuda_configure.bzl", "cuda_configure_ext")
use_repo(cuda_configure, "local_config_cuda")

cuda_json_ext = use_extension("@rules_ml_toolchain//extensions:cuda_json.bzl", "cuda_json_ext")
use_repo(cuda_json_ext, "cuda_redist_json")

cuda_redist_init_ext = use_extension("@rules_ml_toolchain//extensions:cuda_redist_init.bzl", "cuda_redist_init_ext")
use_repo(
    cuda_redist_init_ext,
    "cuda_crt",
    "cuda_cublas",
    "cuda_cudart",
    "cuda_cudnn",
    "cuda_cufft",
    "cuda_cupti",
    "cuda_curand",
    "cuda_cusolver",
    "cuda_cusparse",
    "cuda_driver",
    "cuda_nvcc",
    "cuda_nvdisasm",
    "cuda_nvjitlink",
    "cuda_nvml",
    "cuda_nvrtc",
    "cuda_nvtx",
    "cuda_nvvm",
)

nvshmem_redist = use_extension("@rules_ml_toolchain//extensions:nvshmem_redist.bzl", "nvshmem_redist_ext")
use_repo(nvshmem_redist, "nvidia_nvshmem")

# Use the newer version of cuda_cccl from XLA
cccl_extension = use_extension("@xla//third_party/cccl:cccl_extension.bzl", "cccl_extension")
use_repo(cccl_extension, "cuda_cccl")

override_repo(
    cuda_redist_init_ext,
    cuda_cccl = "cuda_cccl",
)

### Toolchains
register_toolchains("@rules_ml_toolchain//cc:linux_x86_64_linux_x86_64")

register_toolchains("@rules_ml_toolchain//cc:linux_x86_64_linux_x86_64_cuda")

register_toolchains("@rules_ml_toolchain//cc:linux_x86_64_linux_x86_64_sycl")

register_toolchains("@rules_ml_toolchain//cc:linux_x86_64_linux_x86_64_rocm")

register_toolchains("@rules_ml_toolchain//cc:linux_x86_64_linux_aarch64")

register_toolchains("@rules_ml_toolchain//cc:linux_aarch64_linux_aarch64")

register_toolchains("@rules_ml_toolchain//cc:linux_aarch64_linux_aarch64_cuda")

apple_cc_configure = use_extension("@apple_support//crosstool:setup.bzl", "apple_cc_configure_extension")
use_repo(
    apple_cc_configure,
    "local_config_apple_cc",
    "local_config_apple_cc_toolchains",
)

register_toolchains("@local_config_apple_cc_toolchains//:all")

##############################################################
# Other dependencies via use_repo_rule
##############################################################

nvidia_wheel_versions_repository = use_repo_rule(
    "@xla//third_party/py:python_wheel.bzl",
    "nvidia_wheel_versions_repository",
)

nvidia_wheel_versions_repository(
    name = "nvidia_wheel_versions",
    versions_source = "//build:nvidia-requirements.txt",
)

jax_python_wheel_repository = use_repo_rule(
    "//jaxlib:jax_python_wheel.bzl",
    "jax_python_wheel_repository",
)

jax_python_wheel_repository(
    name = "jax_wheel",
    version_key = "_version",
    version_source = "//jax:version.py",
)

python_wheel_version_suffix_repository = use_repo_rule(
    "@xla//third_party/py:python_wheel.bzl",
    "python_wheel_version_suffix_repository",
)

python_wheel_version_suffix_repository(
    name = "jax_wheel_version_suffix",
)

test_shard_count_repository = use_repo_rule(
    "//:test_shard_count.bzl",
    "test_shard_count_repository",
)

test_shard_count_repository(
    name = "test_shard_count",
)

external_deps_repository = use_repo_rule(
    "//third_party/external_deps:workspace.bzl",
    "external_deps_repository",
)

rocm_wheels_repository = use_repo_rule(
    "//third_party/rocm_wheels:workspace.bzl",
    "rocm_wheels_repository",
)

# Pre-built ROCm wheels from GitHub (ROCm/rocm-jax).
rocm_wheels_repository(
    name = "rocm_wheels",
    jaxlib_version = JAXLIB_VERSION,
    python_version = DEFAULT_PYTHON_VERSION,
    # rocm_version = "7.2.0",  # Optional: pick a specific ROCm version.
)

# Used for --//jax:build_jaxlib=wheel (locally-built wheels).
external_deps_repository(
    name = "rocm_external_test_deps",
    deps = [
        "//jaxlib/tools:jax_rocm_pjrt_py_import",
        "//jaxlib/tools:jax_rocm_plugin_py_import",
    ],
)

# Used for --//jax:build_jaxlib=false (pre-built wheels from GitHub).
external_deps_repository(
    name = "rocm_prebuilt_test_deps",
    deps = [
        "@rocm_wheels//:rocm_pjrt_py_import",
        "@rocm_wheels//:rocm_plugin_py_import",
    ],
)