Getting Started
jax.debug.print
jax.debug.breakpoint
checkify
Reference Documentation
Advanced JAX Tutorials
Notes
Developer documentation
API documentation
from_dlpack(dlpack)
from_dlpack
Returns a DeviceArray representation of a DLPack tensor.
DeviceArray
to_dlpack(x[, take_ownership])
to_dlpack
Returns a DLPack tensor that encapsulates a DeviceArray x.