rocm_jax

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

History

Eugene Burmako 1fdc9a2b84 [MHLO] Add CHLO lowering for top_k

This follows the introduction of chlo.top_k in the diffbase.

PiperOrigin-RevId: 447581682

2022-05-09 15:45:04 -07:00

array.filecheck.py

[MHLO] Add CHLO lowering for top_k

2022-05-09 15:45:04 -07:00

jax_filecheck_helpers.py

[JAX] Add a dialect option to jit(...).lower(...).compiler_ir().

2021-12-10 13:02:25 -08:00

math.filecheck.py

[MHLO] Add MHLO lowering for erf and erfc

2022-04-22 07:54:23 -07:00

names.filecheck.py

Remove the jax_enable_mlir flag. MLIR is now the only supported code path.

2022-04-04 08:40:09 -07:00

README.md

Add some initial filecheck tests for JAX->MHLO lowering.

2021-12-10 10:59:24 -08:00

shapes.filecheck.py

Remove the jax_enable_mlir flag. MLIR is now the only supported code path.

2022-04-04 08:40:09 -07:00

subcomputations.filecheck.py

[MHLO] Switch call_tf to use an MHLO lowering (attempt 2).

2022-04-19 14:00:00 -07:00

README.md

This directory contains LLVM FileCheck tests that verify that JAX primitives can be lowered to MHLO.

These tests are intended to be a quick and easy-to-understand way to catch regressions from changes due the MLIR Python bindings and from changes to the various MLIR dialects used by JAX, without needing to run the full JAX test suite.