[XLA] separate out an Execute from ExecutePerReplica
[JAX] reduce the creation of XLA Shape protos on every call (which is slow) PiperOrigin-RevId: 223915944
Loading
Please sign in to comment
[JAX] reduce the creation of XLA Shape protos on every call (which is slow) PiperOrigin-RevId: 223915944