flâneur

jax.Array — JAX documentation

jax.readthedocs.io · 810 words · saved by 1 readers

jax.Array is the public interface for instance checks and type annotation of JAX arrays and tracers. Its main applications are in instance checks and type annotations; for example: jax.Array should not be used directly for creation of arrays; instead you should use array creation routines offered in jax.numpy, such as jax.numpy.array(), jax.numpy.zeros(), jax.numpy.ones(), jax.numpy.full(), jax.numpy.arange(), etc. Methods __init__() addressable_data(index) Return an array of the addressable data at a particular index. all([axis, out, keepdims, where]) Test whether all array elements along a given axis evaluate to True. any([axis, out, keepdims, where]) Test whether any of the array elements along a given axis evaluate to True. argmax([axis, out, keepdims]) Returns the indices of the maximum values along an axis. argmin([axis, out, keepdims]) Returns the indices of the minimum values along an axis. argpartition(kth[, axis]) Returns indices that partially sort an array. argsort([axis, k

jax.Array # class jax. Array # Array base class for JAX jax.Array is the public interface for instance checks and type annotation of JAX arrays and tracers. Its main applications are in instance checks and type annotations; for example: x = jnp . arange ( 5 ) isinstance ( x , jax . Array ) # returns True both inside and outside traced functions. def f ( x : Array ) -> Array : # type annotations are valid for traced and non-traced types. return x jax.Array should not be used directly for creation of arrays; instead you should use array creation routines offered in jax.numpy , such as jax.numpy.

related reading