jax_like¶
Protocols and base classes for JAX-like array objects.
The concrete methods below are written in terms of the wrappers in
probly.representation.jax_functions, so subclasses only have to implement
__jax_function__.
Classes¶
Protocol for array-like objects that implement JAX-specific APIs. |
|
Protocol for array-like objects that can be converted to JAX arrays. |
|
ABC implementation for array-like objects that behave like JAX arrays. |