jax.transfer_guard#

jax.transfer_guard(new_val)[source]#

A contextmanager to control the transfer guard level for all transfers.

For more information, see https://jax.readthedocs.io/en/latest/transfer_guard.html

Parameters

new_val (str) – The new thread-local transfer guard level for all transfers.

Yields

None.

Return type

Iterator[None]