Improved extern selection in Pallas GPU

Previously, * weakly typed avals matched the wrong externs; * this was addressed by #23193, which disallowed weakly typed avals entirely. Here we check if a weakly typed aval can be casted to the extern input dtype when selecting an extern. PiperOrigin-RevId: 669067725
jax-ml · Aug 30, 2024 · 2c11277 · 2c11277
1 parent 4c218fb
commit 2c11277
Show file tree

Hide file tree

Showing 2 changed files with 166 additions and 156 deletions.
diff --git a/jax/_src/pallas/triton/BUILD b/jax/_src/pallas/triton/BUILD
@@ -54,6 +54,7 @@ pytype_strict_library(
         "//jax:api_util",
         "//jax:config",
         "//jax:core",
+        "//jax:dtypes",
         "//jax:mlir",
         "//jax:partial_eval",
         "//jax:source_info_util",