rocm_jax

mirror of https://github.com/ROCm/jax.git synced 2025-04-24 23:56:06 +00:00

History

Peter Hawkins 8ca6622c0b Change lax.select_p to be an n-ary predicate, 'lax.select_n_p'. Change lax.select() to be a thin shim around the new n-ary version.

Note that one key difference between `lax.select_p` and `lax.select_n_p` is that the order of the cases is reversed for boolean predicates. This merited a new name to minimize confusion.

Use lax.select_n() in conditional batching. This means that we only produce one `select_n()` primitive for each conditional output, rather than a tree. While this has no effect on the number of HLO operators we generate, it can reduces the number of jaxpr equations significantly.

PiperOrigin-RevId: 427517899

2022-02-09 11:03:09 -08:00

__init__.py

2021-02-22 21:11:41 -08:00

control_flow.py

Change lax.select_p to be an n-ary predicate, 'lax.select_n_p'. Change lax.select() to be a thin shim around the new n-ary version.

2022-02-09 11:03:09 -08:00

convolution.py

[MLIR] Change signature of lowering rules.

2021-12-15 19:06:58 -08:00

fft.py

Support string fft_type values in lax.fft.

2022-02-03 08:52:38 -05:00