jax.transfer_guard

jax.transfer_guard(new_val)

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

For more information, see https://docs.jax.dev/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]