jax
fc74a0aa - [pallas] Relax check for block spec equality in fuser.

Commit
1 year ago
[pallas] Relax check for block spec equality in fuser. Most functions in Python use identity equality, so the block spec equality will fail even when the index map functions are equivalent (e.g. from closures or `partial`). PiperOrigin-RevId: 776397038
Author
Parents
Loading