Dan Foreman-Mackey ac560c0d90 Add helper function for building custom call lowering rules ...

This function provides sensible defaults for custom call lowering rules
with the goal of reducing the amount of boilerplate required for
implementing custom calls.

Co-authored-by: Sergei Lebedev <slebedev@google.com>

2024-06-06 11:34:08 -04:00

..

core

Add more imports to jax extend

2024-04-17 15:13:17 -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 helper function for building custom call lowering rules

2024-06-06 11:34:08 -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