Peter Hawkins 74506c7dda Rollback of: Implement the JAX transfer guard API ...

Adds `--jax_transfer_guard` flag and `jax.transfer_guard()` context manager that allows logging or disallowing unintended transfers.

The API distinguishes between two types of transfers:
* explicit transfers: `jax.device_put*()` and `jax.device_get()` calls.
* implicit transfers: Other transfers (e.g., printing a `DeviceArray`).

The transfer guard can take an action based on its guard level:

* "allow": Silently allow all transfers (default; s...

PiperOrigin-RevId: 427576107

2022-02-09 14:44:45 -08:00

..

build_jaxlib_wheels_helpers.sh

Remove -it flag from docker run command in jaxlib build script.

2021-04-14 09:44:42 -07:00

build_jaxlib_wheels_macos.sh

Drop support for NumPy 1.18.

2022-01-04 12:11:38 -08:00

build_jaxlib_wheels.sh

Drop Python 3.6 support.

2021-07-15 14:20:29 -04:00

build_wheel_docker_entrypoint.sh

Allow manylinux_2_12_x86_64 in auditwheel test. This is an alias by PEP600.

2021-05-07 18:22:28 -04:00

build_wheel.py

Rollback of: Implement the JAX transfer guard API

2022-02-09 14:44:45 -08:00

BUILD.bazel

fix mlir capi dll building and linking

2021-11-25 00:07:25 +08:00

build.py

Increase bazel version to 5.0.0 to match TensorFlow

2022-01-28 21:11:02 +00:00

Dockerfile

Drop support for NumPy 1.18.

2022-01-04 12:11:38 -08:00

install_cuda.sh

Update cuDNN versions to highest available.

2021-03-09 19:07:00 +00:00

LICENSE.txt

Refactor JAX CPU kernels to make them usable from C++.

2021-09-03 10:03:54 -07:00

test-requirements.txt

Adds a Wadler-Lindig pretty printer.

2021-09-27 21:09:24 -04:00