Dan Foreman-Mackey 86f48a85b4 Add support for the DeviceLocalLayout API when lowering FFI calls. ...

This PR updates the FFI lowering rule to support a DeviceLoweringLayout
object as input when specifying the input and output layouts. For now,
this just converts the DLL object to its appropriate list of
minor-to-major integers because that's what the custom call op expects.

2024-09-05 14:30:06 -04:00

..

__init__.py

add jax.extend.random.wrap_key_data

2023-08-26 11:39:25 -07:00

ffi.py

Add support for the DeviceLocalLayout API when lowering FFI calls.

2024-09-05 14:30:06 -04:00

random.py

Run pyupgrade --py310-plus.

2024-06-26 16:10:18 -04:00