add pull_request event

.github/workflows/cloud-tpu-ci-nightly.yml

@@ -0,0 +1,70 @@
+name: Cloud-TPU-CI-Nightly
+
+on:
+  schedule:
+    - cron: "0 12 * * *" # Daily at 12:00 UTC
+  workflow_dispatch: # allows triggering the workflow run manually
+  repository_dispatch: # allows triggering the workflow via HTTP
+  pull_request:
+
+jobs:
+  cloud-tpu-test:
+    runs-on: tpu
+    defaults:
+      run:
+        shell: bash -l {0}
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.10"]  # TODO(jakevdp): update to 3.11 when available.
+    outputs:
+      artifacts_availability: ${{ steps.status.outputs.ARTIFACTS_AVAILABLE }}
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Install JAX test requirements
+        run: |
+          pip install -r build/test-requirements.txt
+          pip install pytest-reportlog
+      - name: Install JAX
+        run: |
+          pip install .[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
+      - name: Run tests
+        if: success()
+        id: status
+        env:
+          JAX_PLATFORMS: tpu,cpu
+        run: |
+          pytest --tb=short \
+              --report-log output-${{ matrix.python-version }}-log.jsonl \
+              tests/compilation_cache_test.py \
+              || (
+                echo '::set-output name=ARTIFACTS_AVAILABLE::true' && false
+              )
+        # run: |
+        #   pytest --tb=short \
+        #       --deselect tests/callback_test.py \
+        #       --deselect tests/checkify_test.py \
+        #       --deselect tests/debugger_test.py \
+        #       --deselect tests/debugging_primitives_test.py \
+        #       --deselect tests/jaxpr_effects_test.py-rf \
+        #       --report-log output-${{ matrix.python-version }}-log.jsonl \
+        #       tests \
+        #       || (
+        #         echo '::set-output name=ARTIFACTS_AVAILABLE::true' && false
+        #       )
+      - name: Upload artifacts
+        # if: |
+        #   failure()
+          # && steps.status.outcome == 'failure'
+          # && github.event_name == 'schedule'
+          # && github.repository == 'google/jax'
+        if: failure()
+        uses: actions/upload-artifact@v3
+        with:
+          name: output-${{ matrix.python-version }}-log.jsonl
+          path: output-${{ matrix.python-version }}-log.jsonl
+          retention-days: 5

.github/workflows/self_hosted_runner_utils/runner.env

@@ -0,0 +1 @@
+ACTIONS_RUNNER_HOOK_JOB_STARTED=~/jax/.github/workflows/self_hosted_runner_utils/validate_job.sh

.github/workflows/self_hosted_runner_utils/start_github_runner.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+
+# Copyright 2022 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.
+
+# More or less copied from
+# https://github.com/iree-org/iree/tree/main/build_tools/github_actions/runner/config
+
+~/actions-runner/run.sh > /tmp/actions-runner.`date +"%Y%m%d-%H%M"`.log

.github/workflows/self_hosted_runner_utils/validate_job.sh

@@ -0,0 +1,47 @@
+#!/bin/bash
+
+# Copyright 2022 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.
+
+# More or less copied from
+# https://github.com/iree-org/iree/tree/main/build_tools/github_actions/runner/config
+
+set -euo pipefail
+
+ALLOWED_EVENTS=(
+  "schedule"
+  "workflow_dispatch"
+)
+
+# Tests if the first argument is contained in the array in the second argument.
+# Usage `is_contained "element" "${array[@]}"`
+is_contained() {
+  local e;
+  local match="$1"
+  shift
+  for e in "$@"; do
+    if [[ "${e}" == "${match}" ]]; then
+      return 0
+    fi
+  done
+  return 1
+}
+
+if ! is_contained "${GITHUB_EVENT_NAME}" "${ALLOWED_EVENTS[@]}"; then
+  echo "Event type '${GITHUB_EVENT_NAME}' is not allowed on this runner. Aborting workflow."
+  # clean up any nefarious stuff we may have fetched in job setup.
+  cd ~/actions-runner/_work
+  rm -rfv _actions/ _temp/
+  exit 1
+fi

build/build_wheel.py

@@ -169,9 +169,7 @@ def prepare_wheel(sources_path):
   copy_to_jaxlib("__main__/jaxlib/init.py", dst_filename="__init__.py")
   copy_to_jaxlib(f"__main__/jaxlib/cpu_feature_guard.{pyext}")
   copy_to_jaxlib("__main__/jaxlib/lapack.py")
-  copy_to_jaxlib(f"__main__/jaxlib/_lapack.{pyext}")
   copy_to_jaxlib("__main__/jaxlib/mhlo_helpers.py")
-  copy_to_jaxlib(f"__main__/jaxlib/_ducc_fft.{pyext}")
   copy_to_jaxlib("__main__/jaxlib/ducc_fft.py")
   copy_to_jaxlib("__main__/jaxlib/gpu_prng.py")
   copy_to_jaxlib("__main__/jaxlib/gpu_linalg.py")
@@ -180,6 +178,10 @@ def prepare_wheel(sources_path):
   copy_to_jaxlib("__main__/jaxlib/version.py")
   copy_to_jaxlib("__main__/jaxlib/xla_client.py")
   copy_to_jaxlib(f"__main__/jaxlib/xla_extension.{pyext}")
+  cpu_dir = os.path.join(jaxlib_dir, "cpu")
+  os.makedirs(cpu_dir)
+  copy_file(f"__main__/jaxlib/cpu/_lapack.{pyext}", dst_dir=cpu_dir)
+  copy_file(f"__main__/jaxlib/cpu/_ducc_fft.{pyext}", dst_dir=cpu_dir)
 
   cuda_dir = os.path.join(jaxlib_dir, "cuda")
   if exists(f"__main__/jaxlib/cuda/_cusolver.{pyext}"):

docs/jax.scipy.rst

@@ -318,6 +318,19 @@ jax.scipy.stats.t
    logpdf
    pdf
 
+jax.scipy.stats.truncnorm
+~~~~~~~~~~~~~~~~~~~~~~~~~
+.. automodule:: jax.scipy.stats.truncnorm
+.. autosummary::
+  :toctree: _autosummary
+
+   cdf
+   logcdf
+   logpdf
+   logsf
+   pdf
+   sf
+
 jax.scipy.stats.uniform
 ~~~~~~~~~~~~~~~~~~~~~~~
 .. automodule:: jax.scipy.stats.uniform

docs/jep/9419-jax-versioning.md

@@ -149,7 +149,7 @@ level.
 [as a Bazel submodule](https://github.com/google/jax/blob/main/WORKSPACE).
 To update the version of XLA used during the build, one must update the pinned
 version in the Bazel `WORKSPACE`. This is done manually on an
-as-needed basis, but can be overriden on a build-by-build basis.
+as-needed basis, but can be overridden on a build-by-build basis.
 
 
 ## How do we make changes across the `jax` and `jaxlib` boundary between releases?

jax/_src/lax/lax.py

@@ -960,13 +960,13 @@ def transpose(operand: ArrayLike, permutation: Sequence[int]) -> Array:
     return transpose_p.bind(operand, permutation=permutation)
 
 def argmin(operand: ArrayLike, axis: int,
-           index_dtype: DTypeLike) -> Tuple[Array, Array]:
+           index_dtype: DTypeLike) -> Array:
   """Computes the index of the minimum element along ``axis``."""
   return argmin_p.bind(operand, axes=(axis,),
                        index_dtype=dtypes.canonicalize_dtype(index_dtype))
 
 def argmax(operand: ArrayLike, axis: int,
-           index_dtype: DTypeLike) -> Tuple[Array, Array]:
+           index_dtype: DTypeLike) -> Array:
   """Computes the index of the maximum element along ``axis``."""
   return argmax_p.bind(operand, axes=(axis,),
                        index_dtype=dtypes.canonicalize_dtype(index_dtype))

jax/_src/lax/linalg.py

@@ -41,6 +41,7 @@
 from jax._src.lax import lax as lax_internal
 from jax._src.lax import svd as lax_svd
 from jax._src.lib import lapack
+from jax._src.lib import mlir_api_version
 
 from jax._src.lib import gpu_linalg
 from jax._src.lib import gpu_solver
@@ -873,10 +874,16 @@ def _triangular_solve_lowering(
     transpose = "NO_TRANSPOSE"
   else:
     transpose = "ADJOINT" if conjugate_a else "TRANSPOSE"
-  return mhlo.TriangularSolveOp(
-      mlir.aval_to_ir_type(out_aval), a, b, ir.BoolAttr.get(left_side),
-      ir.BoolAttr.get(lower), ir.BoolAttr.get(unit_diagonal),
-      mhlo.TransposeAttr.get(transpose)).results
+  if mlir_api_version < 36:
+    return mhlo.TriangularSolveOp(
+        mlir.aval_to_ir_type(out_aval), a, b, ir.BoolAttr.get(left_side),
+        ir.BoolAttr.get(lower), ir.BoolAttr.get(unit_diagonal),
+        mhlo.TransposeAttr.get(transpose)).results
+  else:
+    return mhlo.TriangularSolveOp(
+        a, b, ir.BoolAttr.get(left_side),
+        ir.BoolAttr.get(lower), ir.BoolAttr.get(unit_diagonal),
+        mhlo.TransposeAttr.get(transpose)).results
 
 mlir.register_lowering(triangular_solve_p, _triangular_solve_lowering)
 
@@ -900,10 +907,16 @@ def _triangular_solve_cpu_lower(
       transpose = "ADJOINT" if conjugate_a else "TRANSPOSE"
     else:
       transpose = "NO_TRANSPOSE"
-    return mhlo.TriangularSolveOp(b.type, a, b, ir.BoolAttr.get(left_side),
-                                  ir.BoolAttr.get(lower),
-                                  ir.BoolAttr.get(unit_diagonal),
-                                  mhlo.TransposeAttr.get(transpose)).results
+    if mlir_api_version < 36:
+      return mhlo.TriangularSolveOp(b.type, a, b, ir.BoolAttr.get(left_side),
+                                    ir.BoolAttr.get(lower),
+                                    ir.BoolAttr.get(unit_diagonal),
+                                    mhlo.TransposeAttr.get(transpose)).results
+    else:
+      return mhlo.TriangularSolveOp(a, b, ir.BoolAttr.get(left_side),
+                                    ir.BoolAttr.get(lower),
+                                    ir.BoolAttr.get(unit_diagonal),
+                                    mhlo.TransposeAttr.get(transpose)).results
 
 mlir.register_lowering(triangular_solve_p, _triangular_solve_cpu_lower,
                        platform='cpu')