Roy Frostig b79fe34b94 re-export RNG primitives from jax.extend.core.primitives

2024-07-17 12:24:20 -07:00

..

core

re-export RNG primitives from jax.extend.core.primitives

2024-07-17 12:24:20 -07:00

mlir

Deprecate support for the mhlo dialect.

2024-01-15 02:13:40 -08:00

__init__.py

Move jax.ffi submodule to jax.extend.ffi

2024-05-31 12:34:59 -04:00

backend.py

jax.clear_backends() is not doing what it is intended to do, users should try to avoid using it.

2024-03-14 16:11:31 -07:00

BUILD

Move jax.ffi submodule to jax.extend.ffi

2024-05-31 12:34:59 -04:00

ffi.py

Add ffi_call function with a similar signature to pure_callback.

2024-07-01 09:40:31 -04:00

linear_util.py

JEX: move jax.linear_util to jax.extend.linear_util

2023-08-30 18:32:12 -07:00

random.py

[jex] replace extend.random.PRNGImpl with extend.random.define_prng_impl

2023-10-31 17:21:54 -07:00

source_info_util.py

export jax.extend.source_info_util.current

2023-10-13 12:31:11 -07:00