mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #23143 from vfdev-5:more-nogil-mlir-dialects

Browse Source 
PiperOrigin-RevId: 665464355
This commit is contained in:
jax authors 2024-08-20 11:58:28 -07:00
commit 607ee3eea1
2 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
jaxlib/mlir/_mlir_libs/tpu_ext.cc
View File

@ -314,7 +314,7 @@ MlirContext getDefaultContext() {
} // namespace
PYBIND11_MODULE(_tpu_ext, m) {
PYBIND11_MODULE(_tpu_ext, m, py::mod_gil_not_used()) {
mlirRegisterTPUPasses(); // Register all passes on load.
py::class_<MlirTpuApplyVectorLayoutContext>(m, "ApplyVectorLayoutCtx",

2
jaxlib/mlir/_mlir_libs/triton_ext.cc
View File

@ -22,7 +22,7 @@ limitations under the License.
namespace py = pybind11;
PYBIND11_MODULE(_triton_ext, m) {
PYBIND11_MODULE(_triton_ext, m, py::mod_gil_not_used()) {
//
// Dialects.
//