MODULE.bazel

module(name = "tpu_raiden")

bazel_dep(name = "protobuf", version = "32.1", repo_name = "com_google_protobuf")
bazel_dep(name = "rules_cc", version = "0.2.9")
bazel_dep(name = "platforms", version = "1.0.0")
bazel_dep(name = "aspect_bazel_lib", version = "2.22.5")
bazel_dep(name = "aspect_rules_esbuild", version = "0.25.1")
bazel_dep(name = "rules_foreign_cc", version = "0.15.1")
bazel_dep(name = "rules_go", version = "0.60.0", repo_name = "io_bazel_rules_go")
bazel_dep(name = "rules_ml_toolchain")
archive_override(
    module_name = "rules_ml_toolchain",
    integrity = "sha256-AD+fcA+AoQ9wvsdh0AUpI30WpuIMBjEkyf/6mKjArzU=",
    patch_strip = 1,
    patches = [
        # Keep tpu-raiden on the host C++ toolchain: rules_ml_toolchain's
        # MODULE.bazel registers its own hermetic clang/sysroot toolchains,
        # which would otherwise win toolchain resolution for the extension.
        "//:third_party/rules_ml_toolchain/no_register_toolchains.patch",
    ],
    strip_prefix = "rules_ml_toolchain-c0eb2743b7b12b2bbcf0e1888e26d36ba6b093de",
    urls = ["https://github.com/google-ml-infra/rules_ml_toolchain/archive/c0eb2743b7b12b2bbcf0e1888e26d36ba6b093de.tar.gz"],
)

bazel_dep(name = "jax")
git_override(
    module_name = "jax",
    commit = "32544801e26115ac1794926d027148abf3baf009",
    patch_strip = 1,
    patches = [
        # This patch makes three jax MODULE.bazel fixups: (1) drop
        # `python.override(add_runtime_manifest_files=...)` which pinned
        # rules_python (1.8.5) doesn't support; (2) drop `dev_dependency = True`
        # on the pip use_extension so the `jax_pypi` hub is created when jax is
        # a dependency (tpu-raiden imports it via use_repo); (3) remove the
        # local_wheels attribute that rules_python (1.8.5) doesn't support, so
        # wheels resolve from the index.
        "//:third_party/py/jax_remove_local_wheels.patch",
        # Backport of jax-ml/jax@53bb8f2a5d (`jaxlib/mosaic/gpu/serde.cc`) so
        # jax-v0.11.2 compiles against the newer LLVM pulled in by the `xla`
        # pin below. Remove once `jax` is bumped past 53bb8f2a5d.
        "//:third_party/py/jax_mosaic_gpu_mbarrier_layout.patch",
    ],
    remote = "https://github.com/jax-ml/jax.git",
)

bazel_dep(name = "xla")
git_override(
    module_name = "xla",
    commit = "f60be94c9c5d1e1340a6259433ead57824c37293",
    patch_strip = 1,
    patches = [
        "//:third_party/xla/future_sfinae.patch",
        "//:third_party/xla/attribute_map_static_assert.patch",
        "//:third_party/xla/attribute_map_cc_static_assert.patch",
    ],
    remote = "https://github.com/openxla/xla.git",
)

bazel_dep(name = "rules_python", version = "1.8.5")

bazel_dep(name = "torch_tpu", version = "0.1.1", dev_dependency = True)

single_version_override(
    module_name = "rules_python",
    patch_strip = 1,
    patches = [
        "//:third_party/py/rules_python_missing_version_safe.patch",
    ],
    version = "1.8.5",
)

bazel_dep(name = "grpc", version = "1.81.0", repo_name = "com_github_grpc_grpc")
single_version_override(
    module_name = "grpc",
    patch_strip = 1,
    patches = [
        "//:third_party/grpc/grpc.patch",
    ],
    version = "1.81.0",
)

single_version_override(
    module_name = "protobuf",
    version = "32.1",
)

third_party = use_extension("@xla//third_party/extensions:third_party.bzl", "third_party_ext")
use_repo(
    third_party,
    "nanobind",
    "stablehlo",
)

bazel_dep(name = "abseil-cpp", version = "20260526.0", repo_name = "com_google_absl")

# The torch extension hands absl objects (Status, StatusOr, the PJRT device
# attribute map) across the torch_tpu boundary, so both binaries must be built
# from the same abseil release: the one XLA pins. peregrine declares a newer
# release, which bzlmod would otherwise select for the whole build.
single_version_override(
    module_name = "abseil-cpp",
    version = "20260526.0",
)

# OSS crypto for global_registry_client.cc SHA256.
bazel_dep(name = "boringssl", version = "0.20240913.0")

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

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

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

bazel_dep(name = "pybind11_bazel", version = "2.13.6")
bazel_dep(name = "nanobind_bazel", version = "2.12.0")
bazel_dep(name = "nanobind_abseil")
bazel_dep(name = "abseil-py", version = "2.1.0", repo_name = "com_google_absl_py")
bazel_dep(name = "googletest", version = "1.17.0.bcr.2", repo_name = "com_google_googletest")
bazel_dep(name = "flatbuffers", version = "25.12.19", repo_name = "com_github_google_flatbuffers")
bazel_dep(name = "highway", version = "1.2.0")
bazel_dep(name = "peregrine", version = "0.1.6")
bazel_dep(name = "libzmq", version = "4.3.5.bcr.5")
single_version_override(
    module_name = "libzmq",
    patch_strip = 1,
    patches = [
        "//:third_party/libzmq/libzmq_oss.patch",
    ],
    version = "4.3.5.bcr.5",
)

bazel_dep(name = "prometheus-cpp", version = "1.3.0", repo_name = "com_github_jupp0r_prometheus_cpp")
bazel_dep(name = "perfetto")

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

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

python = use_extension("@rules_python//python/extensions:python.bzl", "python")
python.defaults(
    python_version = "3.12",
    python_version_env = "HERMETIC_PYTHON_VERSION",
)
python.toolchain(python_version = "3.11")
python.toolchain(python_version = "3.12")
python.toolchain(python_version = "3.13")
python.toolchain(python_version = "3.14")

pip = use_extension("@rules_python//python/extensions:pip.bzl", "pip")
use_repo(pip, pypi = "jax_pypi")

override_repo(
    pip,
    xla_pypi = "pypi",
)

git_override(
    module_name = "nanobind_abseil",
    commit = "7cd4a4749edbeb77da49f209826bf892756e57a2",
    patch_strip = 1,
    patches = [
        "//:third_party/nanobind_abseil/inject_module_block.patch",
    ],
    remote = "https://github.com/google/nanobind_abseil.git",
)

git_override(
    module_name = "perfetto",
    commit = "d0d8c6fa5f7411b59fccacf99233e79bcf94e3fc",
    patch_strip = 1,
    patches = [
        "//:third_party/perfetto/no_android_deps.patch",
    ],
    remote = "https://github.com/google/perfetto.git",
)

# Single-sources the wheel version from pyproject.toml [project].version and
# appends WHEEL_VERSION_EXTRAS (e.g. ".dev<timestamp>"). Consumed by
# //ci/wheel:raiden_wheel via @raiden_version//:version.bzl.
raiden_version_repo = use_repo_rule("//bazel:wheel_version.bzl", "raiden_version_repo")

raiden_version_repo(
    name = "raiden_version",
    pyproject_toml = "//:pyproject.toml",
)