jax
a41385c8 - [Pallas/TPU] Allow 1-sized batch dim in vmap of dynamic grid

Commit

1 year ago

[Pallas/TPU] Allow 1-sized batch dim in vmap of dynamic grid PiperOrigin-RevId: 603518847

Author

sharadmv

Committer

a-googler

Parents