# Copyright 2025 The JAX Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # https://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. load("@bazel_skylib//lib:selects.bzl", "selects") load("@rules_python//python:py_binary.bzl", "py_binary") load("@rules_python//python:py_library.bzl", "py_library") load( "@xla//third_party/py:py_import.bzl", "py_import", ) load( "//jaxlib:jax.bzl", "jax_source_package", "jax_wheel", "pytype_test", "wheel_sources", ) wheel_sources( name = "jax_sources", data_srcs = ["//jax"], py_srcs = [ "//jax", "//jax/example_libraries:example_libraries", "//jax/example_libraries:optimizers", "//jax/example_libraries:stax", "//jax/experimental:buffer_callback", "//jax/experimental:checkify", "//jax/experimental:colocated_python", "//jax/experimental:compilation_cache", "//jax/experimental:compute_on", "//jax/experimental:overlap", "//jax/experimental:custom_partitioning", "//jax/experimental:fused", "//jax/experimental:hijax", "//jax/experimental:jet", "//jax/experimental:layout", "//jax/experimental:mesh_utils", "//jax/experimental:multihost_utils", "//jax/experimental:ode", "//jax/experimental:pallas_experimental_gpu_ops", "//jax/experimental:pallas_fuser", "//jax/experimental:pallas_gpu_ops", "//jax/experimental:pallas_mosaic_gpu", "//jax/experimental:pallas_tpu_ops", "//jax/experimental:pallas_triton", "//jax/experimental:pjit", "//jax/experimental:profiler", "//jax/experimental:random", "//jax/experimental:rnn", "//jax/experimental:scheduling_groups", "//jax/experimental:serialize_executable", "//jax/experimental:shard_map", "//jax/experimental:source_mapper", "//jax/experimental:sparse_test_util", "//jax/experimental:sparse", "//jax/experimental:topologies", "//jax/experimental:transfer", "//jax/experimental:xla_metadata", "//jax/experimental", "//jax/_src:lax_reference", "//jax/_src:internal_export_back_compat_test_util", "//jax/_src:internal_export_back_compat_test_data", "//jax/_src:internal_test_harnesses", "//jax/_src:internal_test_util", "//jax/_src:test_multiprocess", "//jax/_src:test_util", "//jax/_src/lib", "//jax/_src/pallas/fuser", "//jax/_src/pallas/mosaic_gpu", "//jax/_src/pallas/mosaic_gpu:layout_rules", "//jax/_src/pallas/mosaic_gpu:pallas_call", "//jax/_src/pallas/mosaic_gpu/interpret:interpret_pallas_call", "//jax/experimental/array_serialization:serialization", "//jax/experimental/array_serialization:pytree_serialization", "//jax/experimental/jax2tf", "//jax/experimental/mosaic/gpu/examples:flash_attention", "//jax/experimental/mosaic/gpu/examples:matmul", "//jax/extend", "//jax/extend:ifrt_programs", "//jax/extend/mlir", "//jax/extend/mlir/dialects", "//jax/tools:jax_to_ir", "//jax/tools:pgo_nsys_converter", ], static_srcs = [ "//jax:py.typed", "AUTHORS", "LICENSE", "README.md", "pyproject.toml", "setup.py", ], ) py_binary( name = "build_wheel", srcs = ["build_wheel.py"], deps = [ "//jaxlib/tools:build_utils", "@pypi//build", "@pypi//setuptools", "@pypi//wheel", ], ) jax_wheel( name = "jax_wheel", platform_independent = True, source_files = [":jax_sources"], wheel_binary = ":build_wheel", wheel_name = "jax", ) jax_wheel( name = "jax_wheel_editable", editable = True, platform_independent = True, source_files = [":jax_sources"], wheel_binary = ":build_wheel", wheel_name = "jax", ) jax_source_package( name = "jax_source_package", source_files = [":jax_sources"], source_package_binary = ":build_wheel", source_package_name = "jax", ) genrule( name = "wheel_additives", testonly = True, srcs = [ "//jax/_src:internal_test_harnesses", "//jax/_src:hypothesis_test_util", "//jax/_src:internal_test_util", "//jax/_src:internal_export_back_compat_test_util", "//jax/_src:internal_export_back_compat_test_data", "//jax/experimental/jax2tf/tests:jax2tf_tests", "//jax/experimental/jax2tf/tests:tf_test_util", "//jax/experimental/mosaic/gpu/examples:flash_attention.py", "//jax/experimental/mosaic/gpu/examples:gpu_examples", "//jax/experimental/mosaic/gpu/examples:matmul.py", "//jax/experimental/mosaic/gpu/examples:matmul_blackwell", "//jax/_src:test_multiprocess", "//jax/_src/pallas:pallas_test_util", "//jax/experimental:mosaic_gpu_test_util", ], outs = ["wheel_additives.zip"], cmd = "$(location @bazel_tools//tools/zip:zipper) c $@ $(SRCS)", tools = ["@bazel_tools//tools/zip:zipper"], visibility = ["//jax:internal"], ) py_import( name = "jax_py_import", testonly = True, wheel = ":jax_wheel", zip_deps = [":wheel_additives"], ) # This target is used to add more sources to the jax wheel. # This is needed for the tests that depend on jax and use modules that are not part of # the jax wheel, but share the same package paths as the modules in the jax wheel. py_import( name = "jax_wheel_with_internal_test_util", testonly = True, wheel = "@pypi//jax:whl", zip_deps = [":wheel_additives"], ) pytype_test( name = "jax_wheel_size_test", srcs = ["//jaxlib/tools:wheel_size_test.py"], args = [ "--wheel-path=$(location :jax_wheel)", "--max-size-mib=5", ], data = [":jax_wheel"], main = "wheel_size_test.py", tags = [ "manual", "notap", ], ) py_library( name = "empty", visibility = ["//:__subpackages__"], ) selects.config_setting_group( name = "is_python_3_14_or_3_15", match_any = [ "@rules_python//python/config_settings:is_python_3.14", "@rules_python//python/config_settings:is_python_3.15", ], ) selects.config_setting_group( name = "is_python_3_14_freethreaded", match_all = [ "@rules_python//python/config_settings:is_python_3.14", "@rules_python//python/config_settings:is_py_freethreaded", ], ) selects.config_setting_group( name = "is_python_3_15_freethreaded", match_all = [ "@rules_python//python/config_settings:is_python_3.15", "@rules_python//python/config_settings:is_py_freethreaded", ], ) # portpicker depends on psutil, which does not currently provide wheels for # free-threaded Python 3.15. alias( name = "pypi_optional_portpicker", actual = select({ ":is_python_3_15_freethreaded": ":empty", "//conditions:default": "@pypi//portpicker", }), visibility = ["//:__subpackages__"], ) # matplotlib does not yet provide wheels for Python 3.15. alias( name = "pypi_optional_matplotlib", actual = select({ "@rules_python//python/config_settings:is_python_3.15": ":empty", "//conditions:default": "@pypi//matplotlib", }), visibility = ["//:__subpackages__"], ) # tensorstore does not yet provide wheels for Python 3.15. alias( name = "pypi_optional_tensorstore", actual = select({ "@rules_python//python/config_settings:is_python_3.15": ":empty", "//conditions:default": "@pypi//tensorstore", }), visibility = ["//:__subpackages__"], ) # tensorflow does not provide wheels for Python 3.14 or 3.15. alias( name = "pypi_optional_tensorflow", actual = select({ ":is_python_3_14_or_3_15": ":empty", "//conditions:default": "@pypi//tensorflow", }), visibility = ["//:__subpackages__"], ) # zstandard is included in the Python standard library in Python 3.14+, so the # external PyPI package is not needed or available on Python 3.14+. alias( name = "pypi_optional_zstandard", actual = select({ ":is_python_3_14_or_3_15": ":empty", "//conditions:default": "@pypi//zstandard", }), visibility = ["//:__subpackages__"], )