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 29, 2024 · 16fa34a · 16fa34a
1 parent 45099b8
commit 16fa34a
Show file tree

Hide file tree

Showing 2 changed files with 173 additions and 154 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",