Update OpenXLA-pin to Nov24

.github/workflows/trusted_partners.js

@@ -53,7 +53,7 @@ const get_email_domain = async ({github, username}) => {
 const filter_action = async ({github, context, domain}) => {
   const labels = ['kokoro:force-run'];
 
-  let assignees = ['radhakrishnaba', 'xla-rotation'];
+  let assignees = ['kamaljeeti', 'xla-rotation'];
   const title =
       context.payload.pull_request && context.payload.pull_request.title;
   const lowercased_title = (title || '').toLowerCase();

.kokoro/jax/build.sh

@@ -0,0 +1,108 @@
+#!/bin/bash
+# Copyright 2022 Google LLC All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+
+# -e: abort script if one command fails
+# -u: error if undefined variable used
+# -o pipefail: entire command fails if pipe fails. watch out for yes | ...
+# -o history: record shell history
+set -euox pipefail -o history
+
+# Builds + tests jaxlib against CL/PR version of XLA + JAX main.
+
+source "${KOKORO_GFILE_DIR}/utils.sh"
+
+function is_linux_gpu_job() {
+  [[ "$KOKORO_JOB_NAME" =~ tensorflow/xla/jax/.*gpu.* ]]
+}
+
+clone_main_jax() {
+  git clone https://github.com/google/jax.git
+}
+
+prelude() {
+  export JAX_ENABLE_X64=0
+
+  if is_linux_gpu_job ; then
+    export JAX_CUDA_VERSION=12
+    export JAX_CUDNN_VERSION=8.9
+    nvidia-smi
+    setup_env_vars_py39
+  else
+    setup_env_vars_py312
+  fi
+
+  cd "${KOKORO_ARTIFACTS_DIR}"
+
+  use_local_or_install_python
+  install_packages "$NUMPY_VERSION" "$SCIPY_VERSION"
+  clone_main_jax
+  # Install bazel
+  update_bazel_linux
+
+  cd jax
+
+}
+
+build_and_test_on_rbe_cpu() {
+  # Run the tests.
+  bazel \
+      test \
+      --verbose_failures=true \
+      --override_repository=xla="${KOKORO_ARTIFACTS_DIR}"/github/xla \
+      --config=avx_posix \
+      --config=mkl_open_source_only \
+      --config="rbe_cpu_linux_py312" \
+      --config=tensorflow_testing_rbe_linux \
+      --test_env=JAX_NUM_GENERATED_CASES=25 \
+      --test_output=errors \
+      -- //tests:cpu_tests //tests:backend_independent_tests
+}
+
+build_and_test_on_rbe_gpu() {
+  # Runs non-multiaccelerator tests with one GPU apiece.
+  # It appears --run_under needs an absolute path.
+
+  # TODO(ddunleavy): reenable `LaxTest.testBitcastConvertType`
+  bazel \
+    test \
+    --verbose_failures=true \
+    --override_repository=xla="${KOKORO_ARTIFACTS_DIR}"/github/xla \
+    --config=avx_posix \
+    --config=mkl_open_source_only \
+    --config="rbe_linux_cuda12.2_nvcc_py3.9" \
+    --config=tensorflow_testing_rbe_linux \
+    --test_env=XLA_PYTHON_CLIENT_ALLOCATOR=platform \
+    --test_output=errors \
+    --test_env=JAX_SKIP_SLOW_TESTS=1 \
+    --test_env=TF_CPP_MIN_LOG_LEVEL=0 \
+    --test_env=JAX_EXCLUDE_TEST_TARGETS="PmapTest.testSizeOverflow|LaxTest.testBitcastConvertType" \
+    --test_tag_filters=-multiaccelerator \
+    -- //tests:gpu_tests //tests:backend_independent_tests
+}
+
+# Generate a templated results file to make output accessible to everyone
+"$KOKORO_ARTIFACTS_DIR"/github/xla/.kokoro/generate_index_html.sh "$KOKORO_ARTIFACTS_DIR"/index.html
+
+prelude
+
+if is_linux_gpu_job ; then
+  build_and_test_on_rbe_gpu
+else
+  build_and_test_on_rbe_cpu
+fi
+
+echo "bazel-testlogs (test results) location:"
+find "$KOKORO_ARTIFACTS_DIR" \
+  -type l,d -name bazel-testlogs || echo "bazel-testlogs not found"

.kokoro/linux/build.sh

@@ -26,6 +26,10 @@ function is_linux_gpu_job() {
   [[ "$KOKORO_JOB_NAME" =~ tensorflow/xla/linux/.*gpu.* ]]
 }
 
+function is_use_nvcc() {
+  [[ "${USE_NVCC:-}" == "true" ]]
+}
+
 # Pull the container (in case it was updated since the instance started) and
 # store its SHA in the Sponge log.
 docker pull "$DOCKER_IMAGE"
@@ -44,16 +48,23 @@ RC_FILE="/usertools/cpu.bazelrc"
 TARGET_FILTER=""
 TAGS_FILTER="-no_oss,-oss_excluded,-oss_serial"
 ADDITIONAL_FLAGS=""
+RBE_CONFIG=""
 
 if is_linux_gpu_job ; then
     TAGS_FILTER="$TAGS_FILTER,gpu,requires-gpu-nvidia,-no_gpu"
     ADDITIONAL_FLAGS="$ADDITIONAL_FLAGS --run_under=//tools/ci_build/gpu_build:parallel_gpu_execute"
     RC_FILE="/usertools/gpu.bazelrc"
+    if is_use_nvcc ; then
+      RBE_CONFIG="rbe_linux_cuda_nvcc"
+    else
+      RBE_CONFIG="rbe_linux_cuda"
+    fi
     echo "***NOTE: nvidia-smi lists the highest CUDA version the driver supports, which may be different than the version of CUDA actually used!!***"
     nvidia-smi
 else
     TAGS_FILTER="$TAGS_FILTER,-gpu,-requires-gpu-nvidia"
     ADDITIONAL_FLAGS="$ADDITIONAL_FLAGS --config=nonccl"
+    RBE_CONFIG="rbe_linux_cpu"
 fi
 
 # Build & test XLA
@@ -65,7 +76,7 @@ docker exec xla bazel --bazelrc=$RC_FILE \
         --features=layering_check \
         --profile=/tf/pkg/profile.json.gz \
         --flaky_test_attempts=3 \
-        --config=rbe \
+        --config=$RBE_CONFIG \
         --jobs=150 \
         --nobuild_tests_only \
         $ADDITIONAL_FLAGS \
@@ -74,9 +85,7 @@ docker exec xla bazel --bazelrc=$RC_FILE \
 
 # Print build time statistics, including critical path.
 docker exec xla bazel analyze-profile "/tf/pkg/profile.json.gz"
-# TODO(ddunleavy): enable once container has clang-tidy
-# docker exec xla git config --global --add safe.directory /tf/xla
-# docker exec xla bash -c "bazel aquery --output=jsonproto \"mnemonic(CppCompile, //xla/...)\" | PYTHONPATH=.. python3 build_tools/lint/clang_tidy.py --changed_lines_only"
+
 # Stop container
 docker stop xla
 

build_tools/lint/BUILD

@@ -13,24 +13,14 @@
 # limitations under the License.
 # ============================================================================
 
-load(
-    "//xla:pytype.default.bzl",
-    "pytype_strict_binary",
-    "pytype_strict_library",
-)
+load("//xla:pytype.default.bzl", "pytype_strict_binary", "pytype_strict_library")
 # Placeholder: load py_test
 
 package(
     # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
     licenses = ["notice"],
 )
 
-pytype_strict_binary(
-    name = "clang_tidy",
-    srcs = ["clang_tidy.py"],
-    deps = [":diff_parser"],
-)
-
 pytype_strict_library(
     name = "check_contents",
     srcs = ["check_contents.py"],
@@ -43,6 +33,11 @@ pytype_strict_library(
     visibility = ["//visibility:public"],
 )
 
+pytype_strict_binary(
+    name = "generate_compile_commands",
+    srcs = ["generate_compile_commands.py"],
+)
+
 py_test(
     name = "check_contents_test",
     srcs = ["check_contents_test.py"],