Code update

PiperOrigin-RevId: 705208135

.bazelrc

@@ -0,0 +1,123 @@
+##############################################################################
+# All default build options below. These apply to all build commands.
+##############################################################################
+# Make Bazel print out all options from rc files.
+build --announce_rc
+
+# Required by XLA/TSL python configure for experimental 'remotable' feature.
+common --experimental_repo_remote_exec
+
+# By default, execute all actions locally.
+build --spawn_strategy=local
+
+# Enable host OS specific configs. For instance, "build:linux" will be used
+# automatically when building on Linux.
+build --enable_platform_specific_config
+
+build --experimental_cc_shared_library
+
+build --define=tsl_link_protobuf=true
+
+# Enable optimization.
+build -c opt
+
+# Suppress all warning messages.
+build --output_filter=DONT_MATCH_ANYTHING
+
+build --copt=-DMLIR_PYTHON_PACKAGE_PREFIX=jaxlib.mlir.
+
+##############################################################################
+# Platform Specific configs below. These are automatically picked up by Bazel
+# depending on the platform that is running the build.
+##############################################################################
+build:linux --config=posix
+build:linux --copt=-Wno-unknown-warning-option
+
+# Workaround for gcc 10+ warnings related to upb.
+# See https://github.com/tensorflow/tensorflow/issues/39467
+build:linux --copt=-Wno-stringop-truncation
+build:linux --copt=-Wno-array-parameter
+
+build:macos --config=posix
+build:macos --apple_platform_type=macos
+
+# Windows has a relatively short command line limit, which JAX has begun to hit.
+# See https://docs.bazel.build/versions/main/windows.html
+build:windows --features=compiler_param_file
+build:windows --features=archive_param_file
+
+# XLA uses M_* math constants that only get defined by MSVC headers if
+# _USE_MATH_DEFINES is defined.
+build:windows --copt=/D_USE_MATH_DEFINES
+build:windows --host_copt=/D_USE_MATH_DEFINES
+# Make sure to include as little of windows.h as possible
+build:windows --copt=-DWIN32_LEAN_AND_MEAN
+build:windows --host_copt=-DWIN32_LEAN_AND_MEAN
+build:windows --copt=-DNOGDI
+build:windows --host_copt=-DNOGDI
+# https://devblogs.microsoft.com/cppblog/announcing-full-support-for-a-c-c-conformant-preprocessor-in-msvc/
+# otherwise, there will be some compiling error due to preprocessing.
+build:windows --copt=/Zc:preprocessor
+build:windows --cxxopt=/std:c++20
+build:windows --host_cxxopt=/std:c++20
+# Generate PDB files, to generate useful PDBs, in opt compilation_mode
+# --copt /Z7 is needed.
+build:windows --linkopt=/DEBUG
+build:windows --host_linkopt=/DEBUG
+build:windows --linkopt=/OPT:REF
+build:windows --host_linkopt=/OPT:REF
+build:windows --linkopt=/OPT:ICF
+build:windows --host_linkopt=/OPT:ICF
+build:windows --incompatible_strict_action_env=true
+
+##############################################################################
+# Feature-specific configurations.
+##############################################################################
+
+build:posix --copt=-Wno-sign-compare
+build:posix --cxxopt=-std=c++20
+build:posix --host_cxxopt=-std=c++20
+
+build:avx_posix --copt=-mavx
+build:avx_posix --host_copt=-mavx
+
+build:native_arch_posix --copt=-march=native
+build:native_arch_posix --host_copt=-march=native
+
+build:avx_linux --copt=-mavx
+build:avx_linux --host_copt=-mavx
+
+build:avx_windows --copt=/arch:AVX
+
+# Disable clang extention that rejects type definitions within offsetof.
+# This was added in clang-16 by https://reviews.llvm.org/D133574.
+# Can be removed once upb is updated, since a type definition is used within
+# offset of in the current version of ubp.
+# See https://github.com/protocolbuffers/upb/blob/9effcbcb27f0a665f9f345030188c0b291e32482/upb/upb.c#L183.
+build:clang --copt=-Wno-gnu-offsetof-extensions
+# Disable clang extention that rejects unknown arguments.
+build:clang --copt=-Qunused-arguments
+
+#############################################################################
+# Some configs to make getting some forms of debug builds. In general, the
+# codebase is only regularly built with optimizations. Use 'debug_symbols' to
+# just get symbols for the parts of XLA/PJRT that jaxlib uses.
+# Or try 'debug' to get a build with assertions enabled and minimal
+# optimizations.
+# Include these in a local .bazelrc.user file as:
+#   build --config=debug_symbols
+# Or:
+#   build --config=debug
+#
+# Additional files can be opted in for debug symbols by adding patterns
+# to a per_file_copt similar to below.
+#############################################################################
+
+build:debug_symbols --strip=never --per_file_copt="xla/pjrt|xla/python@-g3"
+build:debug --config debug_symbols -c fastbuild
+
+# Load `.jax_configure.bazelrc` file written by build.py
+try-import %workspace%/.jax_configure.bazelrc
+
+# Load rc file with user-specific options.
+try-import %workspace%/.bazelrc.user

.bazelversion

@@ -0,0 +1 @@
+6.5.0

.gitignore

@@ -5,5 +5,20 @@
 _pycache__/
 .cache/
 
+# Poetry, setuptools, PyPI distribution artifacts.
 /*.egg-info
 .eggs/
+dist/
+poetry.lock
+
+# Tests
+.pytest_cache/
+
+# Type checking
+.pytype/
+
+# Other
+*.DS_Store
+
+# PyCharm
+.idea