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: 669378582

jax/_src/pallas/triton/lowering.py

@@ -580,6 +580,7 @@ def matches(self, avals: Sequence[jax_core.ShapedArray]) -> bool:
  return False
  return all(
  aval.dtype == jnp.dtype(arg_type)
+ or (aval.weak_type and aval.dtype.kind == jnp.dtype(arg_type).kind)
  for aval, arg_type in zip(avals, self.arg_types)
  )