Matthew Johnson 7c2f842353 shard_map and other fixes to direct-linearize ...

Co-authored-by: Dougal Maclaurin <dougalm@google.com>

2025-03-07 21:02:40 +00:00

..

__init__.py

Cleanup: fix unused imports & mark exported names

2024-10-16 17:42:41 -07:00

discharge.py

Removed unused `# type: ignore` comments

2025-02-13 21:12:27 +00:00

indexing.py

Migrated to mypy 1.14.1 with --allow_redefinition

2025-02-13 15:38:28 +00:00

primitives.py

shard_map and other fixes to direct-linearize

2025-03-07 21:02:40 +00:00

types.py

Add swap as method to TransformedRef

2025-02-26 19:19:10 -08:00

utils.py

[better_errors] Continue adding debug info to Jaxprs (step 6)

2025-02-11 11:28:58 +01:00