mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-17 04:16:07 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/examples
History
Dan Foreman-Mackey 28afd25259 Add FFI example demonstrating the use of XLA's FFI state. 
Support for this was added in JAX v0.5.0.

PiperOrigin-RevId: 722597704
2025-02-03 04:06:10 -08:00
..
ffi
Add FFI example demonstrating the use of XLA's FFI state.
2025-02-03 04:06:10 -08:00
jax_cpp
[pjrt] Use the PjRtMemorySpace version of BufferFromHostLiteral
2025-01-31 12:22:23 -08:00
__init__.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
advi.py
update top-level examples to use new-style typed keys
2024-03-07 12:40:10 -08:00
datasets.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
differentially_private_sgd.py
update top-level examples to use new-style typed keys
2024-03-07 12:40:10 -08:00
examples_test.py
Force float32 matmuls in examples_test.
2024-05-10 19:30:02 +00:00
gaussian_process_regression.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
kernel_lsq.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
mnist_classifier_fromscratch.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
mnist_classifier.py
update top-level examples to use new-style typed keys
2024-03-07 12:40:10 -08:00
mnist_vae.py
update top-level examples to use new-style typed keys
2024-03-07 12:40:10 -08:00
onnx2xla.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
spmd_mnist_classifier_fromscratch.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00