"""Experimental Bzlmod support for TensorFlow"""

module(
    name = "tensorflow",
    repo_name = "org_tensorflow",
)

bazel_dep(name = "bazel_features", version = "1.36.0")
bazel_dep(name = "abseil-cpp", version = "20260526.0", repo_name = "com_google_absl")
single_version_override(
    module_name = "abseil-cpp",
    patch_strip = 1,
    patches = [
        "//third_party/absl:build_dll.patch",
        "//third_party/absl:endian.patch",
    ],
)

bazel_dep(name = "nlohmann_json", version = "3.12.0.bcr.1", repo_name = "nlohmann_json_lib")
bazel_dep(name = "abseil-py", version = "2.1.0", repo_name = "absl_py")
bazel_dep(name = "rules_python", version = "1.6.1")
bazel_dep(name = "rules_shell", version = "0.6.1")
bazel_dep(name = "bazel_skylib", version = "1.8.1")
bazel_dep(name = "platforms", version = "1.0.0")
bazel_dep(name = "rules_license", version = "1.0.0")
bazel_dep(name = "rules_pkg", version = "1.0.1")
bazel_dep(name = "rules_java", version = "8.16.1")
bazel_dep(name = "rules_jvm_external", version = "6.8")
bazel_dep(name = "protobuf", version = "31.1", repo_name = "com_google_protobuf")
single_version_override(
    module_name = "protobuf",
    version = "31.1",
)

bazel_dep(name = "google_cloud_cpp", version = "3.0.0-rc1", repo_name = "com_github_googlecloudplatform_google_cloud_cpp")
bazel_dep(name = "crc32c", version = "1.1.0", repo_name = "com_github_google_crc32c")
bazel_dep(name = "brotli", version = "1.1.0", repo_name = "org_brotli")
bazel_dep(name = "rules_cc", version = "0.2.11")
bazel_dep(name = "curl", version = "8.11.0.bcr.5")
bazel_dep(name = "rules_webtesting", version = "0.4.1", repo_name = "io_bazel_rules_webtesting")
bazel_dep(name = "rules_closure", version = "0.15.0", repo_name = "io_bazel_rules_closure")
bazel_dep(name = "grpc", version = "1.78.0", repo_name = "com_github_grpc_grpc")
single_version_override(
    module_name = "grpc",
    patch_strip = 1,
    patches = ["//third_party:grpc.patch"],
    version = "1.78.0",
)

bazel_dep(name = "jsoncpp", version = "1.9.6", repo_name = "jsoncpp_git")
bazel_dep(name = "rules_apple", version = "4.3.3", repo_name = "build_bazel_rules_apple")
bazel_dep(name = "google_benchmark", version = "1.8.5", repo_name = "com_google_benchmark")
bazel_dep(name = "rules_foreign_cc", version = "0.15.1")
bazel_dep(name = "rules_proto", version = "7.1.0")
bazel_dep(name = "zstd", version = "1.5.7", repo_name = "net_zstd")
bazel_dep(name = "pybind11_bazel", version = "2.13.6")
bazel_dep(name = "pybind11_protobuf", version = "0.0.0-20250210-f02a2b7")

# NOTE: This is a newer version compared to what was used in WORKSPACE
# and it breaks wheel tests.
# TODO: Fix the incompatibility with wheel tests.
bazel_dep(name = "pybind11_abseil", version = "202402.0")
bazel_dep(name = "or-tools", version = "9.12", repo_name = "com_google_ortools")
bazel_dep(name = "re2", version = "2024-07-02.bcr.1", repo_name = "com_googlesource_code_re2")
bazel_dep(name = "rules_android", version = "0.6.6", repo_name = "build_bazel_rules_android")
bazel_dep(name = "rules_android_ndk", version = "0.1.3")
bazel_dep(name = "apple_support", version = "1.23.1", repo_name = "build_bazel_apple_support")
bazel_dep(name = "rules_swift", version = "2.4.0")
bazel_dep(name = "googletest", version = "1.17.0.bcr.1", repo_name = "com_google_googletest")
bazel_dep(name = "fuzztest", version = "20250805.0", repo_name = "com_google_fuzztest")
bazel_dep(name = "gflags", version = "2.2.2", repo_name = "com_github_gflags_gflags")
bazel_dep(name = "snappy", version = "1.2.2.bcr.1")
bazel_dep(name = "zlib", version = "1.3.1.bcr.7")
bazel_dep(name = "libpng", version = "1.6.50.bcr.1", repo_name = "png")
bazel_dep(name = "boringssl", version = "0.20250818.0")
bazel_dep(name = "highway", version = "1.3.0", repo_name = "com_google_highway")
single_version_override(
    module_name = "highway",
    patch_strip = 1,
    patches = [
        "//third_party/com_google_highway:build.patch",
    ],
)

bazel_dep(name = "libprotobuf-mutator", version = "1.3", repo_name = "libprotobuf_mutator")
bazel_dep(name = "rules_ml_toolchain")
archive_override(
    module_name = "rules_ml_toolchain",
    integrity = "sha256-C0L2k6YMYFDYfbHgoOrrhKs/VBkfzglNhjNPrtyAfaA=",
    strip_prefix = "rules_ml_toolchain-398d613aea7a4c294da49b79a6d6f3f8732bd84c",
    urls = ["https://github.com/google-ml-infra/rules_ml_toolchain/archive/398d613aea7a4c294da49b79a6d6f3f8732bd84c.tar.gz"],
)

bazel_dep(name = "xla", repo_name = "xla")
local_path_override(
    module_name = "xla",
    path = "third_party/xla",
)

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

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

xla_third_party = use_extension("@xla//third_party/extensions:third_party.bzl", "third_party_ext")
use_repo(
    xla_third_party,
    "FP16",
    "FXdiv",
    "XNNPACK",
    "compute_library",
    "cpuinfo",
    "cudnn_frontend_archive",
    "cutlass_archive",
    "dlpack",
    "ducc",
    "eigen_archive",
    "farmhash_archive",
    "farmhash_gpu_archive",
    "gemmlowp",
    "gloo",
    "highwayhash",
    "hwloc",
    "implib_so",
    "llvm-raw",
    "llvm_openmp",
    "mkl_dnn_acl_compatible",
    "ml_dtypes_py",
    "mpitrampoline",
    "nanobind",
    "nasm",
    "nvshmem",
    "onednn",
    "onednn_async",
    "pthreadpool",
    "raft",
    "riegeli",
    "rocm_device_libs",
    "shardy",
    "slinky",
    "stablehlo",
    "tensorrt_oss_archive",
    "triton",
)

tf_third_party = use_extension("//third_party/extensions:third_party.bzl", "third_party_ext")
use_repo(
    tf_third_party,
    "arm_neon_2_x86_sse",
    "com_github_googlecloudplatform_tensorflow_gcp_tools",
    "com_google_googleapis",
    "com_google_pprof",
    "com_google_testing_compile",
    "com_google_truth",
    "com_squareup_javapoet",
    "coremltools",
    "cython",
    "fft2d",
    "flatbuffers",
    "gif",
    "hexagon_nn",
    "icu",
    "inception_v1",
    "jpegxl",
    "junit",
    "kissfft",
    "libjpeg_turbo",
    "libwebp",
    "linenoise",
    "mobile_multibox",
    "mobile_ssd",
    "nccl_archive",
    "nvtx_archive",
    "opencl_headers",
    "org_checkerframework_qual",
    "org_hamcrest_core",
    "org_sqlite",
    "org_xprof",
    "person_detect_data",
    "ruy",
    "six_archive",
    "skcms",
    "sobol_data",
    "speech_commands",
    "stylize",
    "tf_runtime",
    "tflite_conv_actions_frozen",
    "tflite_mobilenet_float",
    "tflite_mobilenet_quant",
    "tflite_mobilenet_ssd",
    "tflite_mobilenet_ssd_quant",
    "tflite_mobilenet_ssd_quant_protobuf",
    "tflite_ovic_testdata",
    "vulkan_headers",
    "xctestrunner",
)

# TODO: Remove this injection once xprof is updated.
inject_repo(
    tf_third_party,
    com_github_nlohmann_json = "nlohmann_json_lib",
)

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

python = use_extension("@rules_python//python/extensions:python.bzl", "python")
python.defaults(
    # The environment variable takes precedence if set.
    python_version = "3.11",
    python_version_env = "HERMETIC_PYTHON_VERSION",
)
python.toolchain(python_version = "3.10")
python.toolchain(python_version = "3.11")
python.toolchain(python_version = "3.12")
python.toolchain(python_version = "3.13")

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

[pip.parse(
    extra_hub_aliases = {
        "numpy": ["numpy_headers"],
    },
    hub_name = "tf_pypi",
    python_version = python_version,
    requirements_lock = "//:requirements_lock_{}.txt".format(python_version.replace(".", "_")),
    whl_modifications = {
        "@pypi_mods//:numpy.json": "numpy",  # @pypi_mods is defined in XLA's MODULE.bazel
    },
) for python_version in [
    "3.10",
    "3.11",
    "3.12",
    "3.13",
]]

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

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

git_configure = use_repo_rule("//third_party/git:git_configure.bzl", "git_configure")

git_configure(name = "local_config_git")

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

register_execution_platforms("@local_config_platform//:host")

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")

syslibs_configure = use_repo_rule("//third_party/systemlibs:syslibs_configure.bzl", "syslibs_configure")

syslibs_configure(name = "local_config_syslibs")

# register_toolchains("@local_config_python//:py_toolchain")

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

nccl_configure = use_extension("@rules_ml_toolchain//extensions:nccl_configure.bzl", "nccl_configure_ext")
use_repo(nccl_configure, "local_config_nccl")

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

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_cccl",
    "cuda_cublas",
    "cuda_cudart",
    "cuda_cudnn",
    "cuda_cufft",
    "cuda_cupti",
    "cuda_curand",
    "cuda_cusolver",
    "cuda_cusparse",
    "cuda_driver",
    "cuda_nvcc",
    "cuda_nvjitlink",
    "cuda_nvml",
    "cuda_nvrtc",
    "cuda_nvtx",
)

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_aarch64_linux_aarch64")

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

### Other local config repos
rocm_configure = use_extension("@xla//third_party/extensions:rocm_configure.bzl", "rocm_configure_ext")
use_repo(rocm_configure, "local_config_rocm")

tensorrt_configure = use_extension("@xla//third_party/extensions:tensorrt_configure.bzl", "tensorrt_configure_ext")
use_repo(tensorrt_configure, "local_config_tensorrt")

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 = "//ci/official/requirements_updater:nvidia-requirements.txt",
)

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 = "tf_wheel_version_suffix")

def_file_filter_configure = use_repo_rule("@xla//tools/def_file_filter:def_file_filter_configure.bzl", "def_file_filter_configure")

def_file_filter_configure(name = "local_config_def_file_filter")