jax
7b7a5d8c - Add pjit_p to extend.core.primitives

Commit

1 year ago

Add pjit_p to extend.core.primitives

References

#29456 - Add `pjit_p` to `jax.extend.core.primitives`

Author

j-towns

j-towns

Parents

Loading