jax

b275aec2 - Use Pyrefly instead of mypy

Use Pyrefly instead of mypy A few caveats * We do not type check anything under jax/experimental at the moment. This will be addressed separately in follow up PRs. * We need a nightly jaxlib for accurate type stubs for MLIR dialects.