Use dtypes instead of dtype names in Pallas GPU extern tables · jax-ml/jax@5b1a3b5 · GitHub

Commit

Use dtypes instead of dtype names in Pallas GPU extern tables

Browse files

This allows proper checking/casting of weak dtypes, which will be implemented
in a follow up.

PiperOrigin-RevId: 669304338

Loading branch information

superbobry authored and jax authors committed Aug 30, 2024

1 parent 4c218fb commit 5b1a3b5