2019-01-15 20:14:19 -05:00
|
|
|
|
|
|
|
|
jax.numpy package
|
|
|
|
|
=================
|
|
|
|
|
|
|
|
|
|
.. currentmodule:: jax.numpy
|
|
|
|
|
|
|
|
|
|
.. automodule:: jax.numpy
|
|
|
|
|
|
2019-07-20 14:40:31 +01:00
|
|
|
Implements the NumPy API, using the primitives in :mod:`jax.lax`.
|
|
|
|
|
|
|
|
|
|
While JAX tries to follow the NumPy API as closely as possible, sometimes JAX
|
|
|
|
|
cannot follow NumPy exactly.
|
|
|
|
|
|
|
|
|
|
* Notably, since JAX arrays are immutable, NumPy APIs that mutate arrays
|
|
|
|
|
in-place cannot be implemented in JAX. However, often JAX is able to provide a
|
|
|
|
|
alternative API that is purely functional. For example, instead of in-place
|
|
|
|
|
array updates (:code:`x[i] = y`), JAX provides an alternative pure indexed
|
|
|
|
|
update function :func:`jax.ops.index_update`.
|
|
|
|
|
|
|
|
|
|
* NumPy is very aggressive at promoting values to :code:`float64` type. JAX
|
|
|
|
|
sometimes is less aggressive about type promotion.
|
|
|
|
|
|
|
|
|
|
A small number of NumPy operations that have data-dependent output shapes are
|
|
|
|
|
incompatible with :func:`jax.jit` compilation. The XLA compiler requires that
|
|
|
|
|
shapes of arrays be known at compile time. While it would be possible to provide
|
|
|
|
|
a JAX implementation of an API such as :func:`numpy.nonzero`, we would be unable
|
|
|
|
|
to JIT-compile it because the shape of its output depends on the contents of the
|
|
|
|
|
input data.
|
|
|
|
|
|
|
|
|
|
Not every function in NumPy is implemented; contributions are welcome!
|
|
|
|
|
|
2019-01-15 20:14:19 -05:00
|
|
|
.. autosummary::
|
|
|
|
|
:toctree: _autosummary
|
|
|
|
|
|
|
|
|
|
abs
|
|
|
|
|
absolute
|
|
|
|
|
add
|
|
|
|
|
all
|
|
|
|
|
allclose
|
|
|
|
|
alltrue
|
|
|
|
|
amax
|
|
|
|
|
amin
|
|
|
|
|
angle
|
|
|
|
|
any
|
|
|
|
|
append
|
|
|
|
|
arange
|
|
|
|
|
arccos
|
|
|
|
|
arccosh
|
|
|
|
|
arcsin
|
|
|
|
|
arcsinh
|
|
|
|
|
arctan
|
|
|
|
|
arctan2
|
|
|
|
|
arctanh
|
|
|
|
|
argmax
|
|
|
|
|
argmin
|
|
|
|
|
argsort
|
|
|
|
|
around
|
|
|
|
|
array
|
2019-02-04 10:45:50 -05:00
|
|
|
array_repr
|
|
|
|
|
array_str
|
2019-01-15 20:14:19 -05:00
|
|
|
asarray
|
|
|
|
|
atleast_1d
|
|
|
|
|
atleast_2d
|
2019-02-06 08:40:43 -05:00
|
|
|
atleast_3d
|
2019-02-04 21:26:58 -05:00
|
|
|
bartlett
|
2019-01-15 20:14:19 -05:00
|
|
|
bitwise_and
|
|
|
|
|
bitwise_not
|
|
|
|
|
bitwise_or
|
|
|
|
|
bitwise_xor
|
2019-02-04 21:26:58 -05:00
|
|
|
blackman
|
2020-01-29 11:55:53 -05:00
|
|
|
block
|
2019-01-15 20:14:19 -05:00
|
|
|
broadcast_arrays
|
|
|
|
|
broadcast_to
|
2019-10-29 20:53:20 -04:00
|
|
|
can_cast
|
2019-01-15 20:14:19 -05:00
|
|
|
ceil
|
|
|
|
|
clip
|
|
|
|
|
column_stack
|
|
|
|
|
concatenate
|
|
|
|
|
conj
|
|
|
|
|
conjugate
|
2019-07-30 21:23:44 -04:00
|
|
|
corrcoef
|
2019-01-15 20:14:19 -05:00
|
|
|
cos
|
|
|
|
|
cosh
|
|
|
|
|
count_nonzero
|
2019-07-08 13:37:32 -04:00
|
|
|
cov
|
2019-03-25 17:42:08 -05:00
|
|
|
cross
|
2019-02-01 08:57:22 -05:00
|
|
|
cumsum
|
|
|
|
|
cumprod
|
|
|
|
|
cumproduct
|
2019-02-04 10:45:50 -05:00
|
|
|
deg2rad
|
|
|
|
|
degrees
|
2019-01-15 20:14:19 -05:00
|
|
|
diag
|
|
|
|
|
diag_indices
|
|
|
|
|
diagonal
|
|
|
|
|
divide
|
|
|
|
|
divmod
|
|
|
|
|
dot
|
2019-02-04 10:45:50 -05:00
|
|
|
dsplit
|
2019-02-06 08:40:43 -05:00
|
|
|
dstack
|
2019-01-15 20:14:19 -05:00
|
|
|
einsum
|
|
|
|
|
equal
|
2019-02-05 08:58:57 -05:00
|
|
|
empty
|
|
|
|
|
empty_like
|
2019-01-15 20:14:19 -05:00
|
|
|
exp
|
|
|
|
|
exp2
|
|
|
|
|
expand_dims
|
|
|
|
|
expm1
|
|
|
|
|
eye
|
|
|
|
|
fabs
|
2019-04-30 22:33:25 +05:30
|
|
|
fix
|
2019-01-15 20:14:19 -05:00
|
|
|
flip
|
2019-02-01 08:57:22 -05:00
|
|
|
fliplr
|
|
|
|
|
flipud
|
2019-02-22 08:39:18 -05:00
|
|
|
float_power
|
2019-01-15 20:14:19 -05:00
|
|
|
floor
|
|
|
|
|
floor_divide
|
|
|
|
|
fmod
|
2020-04-06 12:39:28 -04:00
|
|
|
frexp
|
2019-01-15 20:14:19 -05:00
|
|
|
full
|
|
|
|
|
full_like
|
2019-02-19 15:57:22 -05:00
|
|
|
gcd
|
2019-01-15 20:14:19 -05:00
|
|
|
geomspace
|
|
|
|
|
greater
|
|
|
|
|
greater_equal
|
2019-02-04 21:26:58 -05:00
|
|
|
hamming
|
|
|
|
|
hanning
|
2019-02-06 09:05:53 -05:00
|
|
|
heaviside
|
2019-02-04 10:45:50 -05:00
|
|
|
hsplit
|
2019-01-15 20:14:19 -05:00
|
|
|
hstack
|
2019-02-04 10:45:50 -05:00
|
|
|
hypot
|
2019-01-15 20:14:19 -05:00
|
|
|
identity
|
|
|
|
|
imag
|
|
|
|
|
inner
|
|
|
|
|
isclose
|
2019-02-05 08:58:57 -05:00
|
|
|
iscomplex
|
2019-01-15 20:14:19 -05:00
|
|
|
isfinite
|
|
|
|
|
isinf
|
|
|
|
|
isnan
|
|
|
|
|
isneginf
|
|
|
|
|
isposinf
|
2019-02-05 08:58:57 -05:00
|
|
|
isreal
|
2019-10-29 20:53:20 -04:00
|
|
|
isscalar
|
2019-02-04 10:45:50 -05:00
|
|
|
issubdtype
|
|
|
|
|
issubsctype
|
2019-06-17 17:19:21 -04:00
|
|
|
ix_
|
2019-02-04 21:26:58 -05:00
|
|
|
kaiser
|
2019-01-15 20:14:19 -05:00
|
|
|
kron
|
2019-02-19 15:57:22 -05:00
|
|
|
lcm
|
2020-04-06 12:39:28 -04:00
|
|
|
ldexp
|
2019-01-15 20:14:19 -05:00
|
|
|
left_shift
|
|
|
|
|
less
|
|
|
|
|
less_equal
|
|
|
|
|
linspace
|
|
|
|
|
log
|
|
|
|
|
log10
|
|
|
|
|
log1p
|
|
|
|
|
log2
|
|
|
|
|
logaddexp
|
|
|
|
|
logaddexp2
|
|
|
|
|
logical_and
|
|
|
|
|
logical_not
|
|
|
|
|
logical_or
|
|
|
|
|
logical_xor
|
|
|
|
|
logspace
|
|
|
|
|
matmul
|
|
|
|
|
max
|
|
|
|
|
maximum
|
|
|
|
|
mean
|
2019-08-02 08:55:22 -04:00
|
|
|
median
|
2019-01-15 20:14:19 -05:00
|
|
|
meshgrid
|
|
|
|
|
min
|
|
|
|
|
minimum
|
|
|
|
|
mod
|
|
|
|
|
moveaxis
|
2020-03-09 17:07:12 +03:00
|
|
|
msort
|
2019-01-15 20:14:19 -05:00
|
|
|
multiply
|
|
|
|
|
nan_to_num
|
2019-02-01 08:57:22 -05:00
|
|
|
nancumprod
|
|
|
|
|
nancumsum
|
2019-01-15 20:14:19 -05:00
|
|
|
nanmax
|
|
|
|
|
nanmin
|
|
|
|
|
nanprod
|
|
|
|
|
nansum
|
|
|
|
|
negative
|
2019-12-11 16:41:24 -05:00
|
|
|
nextafter
|
2019-12-20 18:42:33 -05:00
|
|
|
nonzero
|
2019-01-15 20:14:19 -05:00
|
|
|
not_equal
|
|
|
|
|
ones
|
|
|
|
|
ones_like
|
|
|
|
|
outer
|
|
|
|
|
pad
|
2019-07-08 13:37:32 -04:00
|
|
|
percentile
|
2019-01-15 20:14:19 -05:00
|
|
|
polyval
|
|
|
|
|
power
|
2019-02-05 08:58:57 -05:00
|
|
|
positive
|
2019-01-15 20:14:19 -05:00
|
|
|
prod
|
2019-02-04 10:45:50 -05:00
|
|
|
product
|
2019-10-29 20:53:20 -04:00
|
|
|
promote_types
|
2019-02-05 08:58:57 -05:00
|
|
|
ptp
|
2019-07-08 13:37:32 -04:00
|
|
|
quantile
|
2019-02-04 10:45:50 -05:00
|
|
|
rad2deg
|
|
|
|
|
radians
|
2019-01-15 20:14:19 -05:00
|
|
|
ravel
|
|
|
|
|
real
|
2019-02-04 10:45:50 -05:00
|
|
|
reciprocal
|
2019-01-15 20:14:19 -05:00
|
|
|
remainder
|
|
|
|
|
repeat
|
|
|
|
|
reshape
|
2019-10-29 20:53:20 -04:00
|
|
|
result_type
|
2019-01-15 20:14:19 -05:00
|
|
|
right_shift
|
2019-02-19 07:29:30 -05:00
|
|
|
roll
|
2019-01-15 20:14:19 -05:00
|
|
|
rot90
|
|
|
|
|
round
|
|
|
|
|
row_stack
|
2019-06-24 10:51:21 -04:00
|
|
|
select
|
2019-01-15 20:14:19 -05:00
|
|
|
sign
|
2019-11-21 09:51:26 -05:00
|
|
|
signbit
|
2019-01-15 20:14:19 -05:00
|
|
|
sin
|
2019-02-05 08:58:57 -05:00
|
|
|
sinc
|
2019-01-15 20:14:19 -05:00
|
|
|
sinh
|
|
|
|
|
sometrue
|
|
|
|
|
sort
|
|
|
|
|
split
|
|
|
|
|
sqrt
|
|
|
|
|
square
|
|
|
|
|
squeeze
|
|
|
|
|
stack
|
|
|
|
|
std
|
|
|
|
|
subtract
|
|
|
|
|
sum
|
|
|
|
|
swapaxes
|
2019-02-02 09:47:55 -05:00
|
|
|
take
|
2019-01-15 20:14:19 -05:00
|
|
|
take_along_axis
|
|
|
|
|
tan
|
|
|
|
|
tanh
|
|
|
|
|
tensordot
|
2019-04-30 13:31:34 -07:00
|
|
|
tile
|
2019-01-15 20:14:19 -05:00
|
|
|
trace
|
|
|
|
|
transpose
|
|
|
|
|
tri
|
|
|
|
|
tril
|
2019-08-02 08:55:22 -04:00
|
|
|
tril_indices
|
2019-01-15 20:14:19 -05:00
|
|
|
triu
|
2019-08-02 08:55:22 -04:00
|
|
|
triu_indices
|
2019-01-15 20:14:19 -05:00
|
|
|
true_divide
|
2019-02-05 08:58:57 -05:00
|
|
|
vander
|
2019-01-15 20:14:19 -05:00
|
|
|
var
|
|
|
|
|
vdot
|
2019-02-04 10:45:50 -05:00
|
|
|
vsplit
|
2019-01-15 20:14:19 -05:00
|
|
|
vstack
|
|
|
|
|
where
|
|
|
|
|
zeros
|
|
|
|
|
zeros_like
|
|
|
|
|
|
2019-11-01 09:04:44 -07:00
|
|
|
jax.numpy.fft
|
|
|
|
|
-------------
|
|
|
|
|
|
|
|
|
|
.. automodule:: jax.numpy.fft
|
|
|
|
|
|
|
|
|
|
.. autosummary::
|
|
|
|
|
:toctree: _autosummary
|
|
|
|
|
|
2020-01-02 17:35:22 -08:00
|
|
|
fft
|
|
|
|
|
ifft
|
2020-01-04 18:21:30 -08:00
|
|
|
fft2
|
|
|
|
|
ifft2
|
2020-01-13 14:59:00 -08:00
|
|
|
fftn
|
|
|
|
|
ifftn
|
|
|
|
|
rfft
|
|
|
|
|
irfft
|
|
|
|
|
rfft2
|
|
|
|
|
irfft2
|
|
|
|
|
rfftn
|
|
|
|
|
irfftn
|
2020-01-06 12:38:12 -08:00
|
|
|
fftfreq
|
2020-01-10 16:31:47 -08:00
|
|
|
rfftfreq
|
2020-03-12 15:05:59 -04:00
|
|
|
fftshift
|
|
|
|
|
ifftshift
|
2019-11-01 09:04:44 -07:00
|
|
|
|
2019-01-15 20:14:19 -05:00
|
|
|
jax.numpy.linalg
|
|
|
|
|
----------------
|
|
|
|
|
|
|
|
|
|
.. automodule:: jax.numpy.linalg
|
|
|
|
|
|
|
|
|
|
.. autosummary::
|
|
|
|
|
:toctree: _autosummary
|
|
|
|
|
|
|
|
|
|
cholesky
|
|
|
|
|
det
|
2019-05-13 16:02:02 -04:00
|
|
|
eig
|
2019-01-15 20:14:19 -05:00
|
|
|
eigh
|
2020-03-12 15:05:59 -04:00
|
|
|
eigvals
|
|
|
|
|
eigvalsh
|
2019-01-15 20:14:19 -05:00
|
|
|
inv
|
2020-01-24 16:52:40 -05:00
|
|
|
matrix_power
|
2020-01-26 14:29:33 -05:00
|
|
|
matrix_rank
|
2019-02-07 11:02:56 -05:00
|
|
|
norm
|
2020-03-12 15:05:59 -04:00
|
|
|
pinv
|
2019-01-15 20:14:19 -05:00
|
|
|
qr
|
|
|
|
|
slogdet
|
|
|
|
|
solve
|
|
|
|
|
svd
|
