jax.numpy package¶
Implements the NumPy API, using the primitives in 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 (
x[i] = y
), JAX provides an alternative pure indexed update functionjax.ops.index_update()
.NumPy is very aggressive at promoting values to
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 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 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!
|
Calculate the absolute value element-wise. |
|
Calculate the absolute value element-wise. |
|
Add arguments element-wise. |
|
Test whether all array elements along a given axis evaluate to True. |
|
Returns True if two arrays are element-wise equal within a tolerance. |
|
Test whether all array elements along a given axis evaluate to True. |
|
Return the maximum of an array or maximum along an axis. |
|
Return the minimum of an array or minimum along an axis. |
|
Return the angle of the complex argument. |
|
Test whether any array element along a given axis evaluates to True. |
|
Append values to the end of an array. |
|
Apply a function to 1-D slices along the given axis. |
|
Apply a function repeatedly over multiple axes. |
|
Return evenly spaced values within a given interval. |
|
Trigonometric inverse cosine, element-wise. |
|
Inverse hyperbolic cosine, element-wise. |
|
Inverse sine, element-wise. |
|
Inverse hyperbolic sine element-wise. |
|
Trigonometric inverse tangent, element-wise. |
|
Element-wise arc tangent of |
|
Inverse hyperbolic tangent element-wise. |
|
Returns the indices of the maximum values along an axis. |
|
Returns the indices of the minimum values along an axis. |
|
Returns the indices that would sort an array. |
|
Find the indices of array elements that are non-zero, grouped by element. |
|
Round an array to the given number of decimals. |
|
Create an array. |
|
True if two arrays have the same shape and elements, False otherwise. |
|
Returns True if input arrays are shape consistent and all elements equal. |
|
Return the string representation of an array. |
|
Split an array into multiple sub-arrays. |
|
Return a string representation of the data in an array. |
|
Convert the input to an array. |
|
Convert inputs to arrays with at least one dimension. |
|
View inputs as arrays with at least two dimensions. |
|
View inputs as arrays with at least three dimensions. |
|
Compute the weighted average along the specified axis. |
|
Return the Bartlett window. |
|
Count number of occurrences of each value in array of non-negative ints. |
|
Compute the bit-wise AND of two arrays element-wise. |
|
Compute bit-wise inversion, or bit-wise NOT, element-wise. |
|
Compute the bit-wise OR of two arrays element-wise. |
|
Compute the bit-wise XOR of two arrays element-wise. |
|
Return the Blackman window. |
|
Assemble an nd-array from nested lists of blocks. |
|
Like Numpy’s broadcast_arrays but doesn’t return views. |
|
Broadcast an array to a new shape. |
|
Returns True if cast between data types can occur according to the casting rule. |
|
Return the cube-root of an array, element-wise. |
alias of |
|
|
Return the ceiling of the input, element-wise. |
Abstract base class of all character string scalar types. |
|
|
Construct an array from an index array and a set of arrays to choose from. |
|
Clip (limit) the values in an array. |
|
Stack 1-D arrays as columns into a 2-D array. |
alias of |
|
Abstract base class of all complex number scalar types that are made up of floating-point numbers. |
|
The warning raised when casting a complex dtype to a real dtype. |
|
|
Return selected slices of an array along given axis. |
|
Join a sequence of arrays along an existing axis. |
|
Return the complex conjugate, element-wise. |
|
Return the complex conjugate, element-wise. |
|
Returns the discrete, linear convolution of two one-dimensional sequences. |
|
Change the sign of x1 to that of x2, element-wise. |
|
Return Pearson product-moment correlation coefficients. |
|
Cross-correlation of two 1-dimensional sequences. |
|
Cosine element-wise. |
|
Hyperbolic cosine, element-wise. |
|
Counts the number of non-zero values in the array |
|
Estimate a covariance matrix, given data and weights. |
|
Return the cross product of two (arrays of) vectors. |
alias of |
|
|
Return the cumulative product of elements along a given axis. |
|
Return the cumulative product of elements along a given axis. |
|
Return the cumulative sum of the elements along a given axis. |
|
Convert angles from degrees to radians. |
|
Convert angles from radians to degrees. |
|
Extract a diagonal or construct a diagonal array. |
|
Create a two-dimensional array with the flattened input as a diagonal. |
|
Return the indices to access the main diagonal of an array. |
|
Return the indices to access the main diagonal of an n-dimensional array. |
|
Return specified diagonals. |
|
Calculate the n-th discrete difference along the given axis. |
|
Return the indices of the bins to which each value in input array belongs. |
|
Returns a true division of the inputs, element-wise. |
|
Return element-wise quotient and remainder simultaneously. |
|
Dot product of two arrays. |
alias of |
|
|
Split array into multiple sub-arrays along the 3rd axis (depth). |
|
Stack arrays in sequence depth wise (along third axis). |
|
Create a data type object. |
|
The differences between consecutive elements of an array. |
|
Evaluates the Einstein summation convention on the operands. |
|
Evaluates the lowest cost contraction order for an einsum expression by considering the creation of intermediate arrays. |
|
Return a new array of given shape and type, filled with zeros. |
|
Return an array of zeros with the same shape and type as a given array. |
|
Return (x1 == x2) element-wise. |
|
Calculate the exponential of all elements in the input array. |
|
Calculate 2**p for all p in the input array. |
|
Expand the shape of an array. |
|
Calculate |
|
Return the elements of an array that satisfy some condition. |
|
Return a 2-D array with ones on the diagonal and zeros elsewhere. |
|
Compute the absolute values element-wise. |
|
Machine limits for floating point types. |
|
Round to nearest integer towards zero. |
|
Return indices that are non-zero in the flattened version of a. |
Abstract base class of all scalar types without predefined length. |
|
|
Reverse the order of elements in an array along the given axis. |
|
Flip array in the left/right direction. |
|
Flip array in the up/down direction. |
alias of |
|
Abstract base class of all floating-point scalar types. |
|
|
First array elements raised to powers from second array, element-wise. |
|
Return the floor of the input, element-wise. |
|
Return the largest integer smaller or equal to the division of the inputs. |
|
Element-wise maximum of array elements. |
|
Element-wise minimum of array elements. |
|
Return the element-wise remainder of division. |
|
Decompose the elements of x into mantissa and twos exponent. |
|
Return a new array of given shape and type, filled with fill_value. |
|
Return a full array with the same shape and type as a given array. |
|
Returns the greatest common divisor of |
|
Return numbers spaced evenly on a log scale (a geometric progression). |
|
Return the gradient of an N-dimensional array. |
|
Return the truth value of (x1 > x2) element-wise. |
|
Return the truth value of (x1 >= x2) element-wise. |
|
Return the Hamming window. |
|
Return the Hanning window. |
|
Compute the Heaviside step function. |
|
Compute the histogram of a set of data. |
|
Function to calculate only the edges of the bins used by the histogram function. |
|
Compute the bi-dimensional histogram of two data samples. |
|
Compute the multidimensional histogram of some data. |
|
Split an array into multiple sub-arrays horizontally (column-wise). |
|
Stack arrays in sequence horizontally (column wise). |
|
Given the “legs” of a right triangle, return its hypotenuse. |
|
Modified Bessel function of the first kind, order 0. |
|
Return the identity array. |
|
Machine limits for integer types. |
|
Return the imaginary part of the complex argument. |
|
Test whether each element of a 1-D array is also present in a second array. |
|
Return an array representing the indices of a grid. |
Abstract base class of all numeric scalar types with a (potentially) inexact representation of the values in its range, such as floating-point numbers. |
|
|
Inner product of two arrays. |
alias of |
|
Abstract base class of all integer scalar types. |
|
|
One-dimensional linear interpolation. |
|
Find the intersection of two arrays. |
|
Compute bit-wise inversion, or bit-wise NOT, element-wise. |
|
Returns a boolean array where two arrays are element-wise equal within a tolerance. |
|
Returns a bool array, where True if input element is complex. |
|
Check for a complex type or an array of complex numbers. |
|
Test element-wise for finiteness (not infinity or not Not a Number). |
|
Calculates element in test_elements, broadcasting over element only. |
|
Test element-wise for positive or negative infinity. |
|
Test element-wise for NaN and return result as a boolean array. |
|
Test element-wise for negative infinity, return result as bool array. |
|
Test element-wise for positive infinity, return result as bool array. |
|
Returns a bool array, where True if input element is real. |
|
Return True if x is a not complex type or an array of complex numbers. |
|
Returns True if the type of element is a scalar type. |
|
Returns True if first argument is a typecode lower/equal in type hierarchy. |
|
Determine if the first argument is a subclass of the second argument. |
|
Check whether or not an object can be iterated over. |
|
Construct an open mesh from multiple sequences. |
|
Return the Kaiser window. |
|
Kronecker product of two arrays. |
|
Returns the lowest common multiple of |
|
Returns x1 * 2**x2, element-wise. |
|
Shift the bits of an integer to the left. |
|
Return the truth value of (x1 < x2) element-wise. |
|
Return the truth value of (x1 =< x2) element-wise. |
|
Perform an indirect stable sort using a sequence of keys. |
|
Return evenly spaced numbers over a specified interval. |
|
Load arrays or pickled objects from |
|
Natural logarithm, element-wise. |
|
Return the base 10 logarithm of the input array, element-wise. |
|
Return the natural logarithm of one plus the input array, element-wise. |
|
Base-2 logarithm of x. |
Logarithm of the sum of exponentiations of the inputs. |
|
Logarithm of the sum of exponentiations of the inputs in base-2. |
|
|
Compute the truth value of x1 AND x2 element-wise. |
|
Compute the truth value of NOT x element-wise. |
|
Compute the truth value of x1 OR x2 element-wise. |
|
Compute the truth value of x1 XOR x2, element-wise. |
|
Return numbers spaced evenly on a log scale. |
|
Return the indices to access (n, n) arrays, given a masking function. |
|
Matrix product of two arrays. |
|
Return the maximum of an array or maximum along an axis. |
|
Element-wise maximum of array elements. |
|
Compute the arithmetic mean along the specified axis. |
|
Compute the median along the specified axis. |
|
Return coordinate matrices from coordinate vectors. |
|
Return the minimum of an array or minimum along an axis. |
|
Element-wise minimum of array elements. |
|
Return element-wise remainder of division. |
|
Return the fractional and integral parts of an array, element-wise. |
|
Move axes of an array to new positions. |
|
Return a copy of an array sorted along the first axis. |
|
Multiply arguments element-wise. |
|
Return the indices of the maximum values in the specified axis ignoring NaNs. |
|
Return the indices of the minimum values in the specified axis ignoring NaNs. |
|
Return the cumulative product of array elements over a given axis treating Not a Numbers (NaNs) as one. |
|
Return the cumulative sum of array elements over a given axis treating Not a Numbers (NaNs) as zero. |
|
Return the maximum of an array or maximum along an axis, ignoring any NaNs. |
|
Compute the arithmetic mean along the specified axis, ignoring NaNs. |
|
Compute the median along the specified axis, while ignoring NaNs. |
|
Return minimum of an array or minimum along an axis, ignoring any NaNs. |
|
Compute the qth percentile of the data along the specified axis, while ignoring nan values. |
|
Return the product of array elements over a given axis treating Not a Numbers (NaNs) as ones. |
|
Compute the qth quantile of the data along the specified axis, while ignoring nan values. |
|
Compute the standard deviation along the specified axis, while ignoring NaNs. |
|
Return the sum of array elements over a given axis treating Not a Numbers (NaNs) as zero. |
|
Replace NaN with zero and infinity with large finite numbers (default behaviour) or with the numbers defined by the user using the nan, posinf and/or neginf keywords. |
|
Compute the variance along the specified axis, while ignoring NaNs. |
|
|
|
Return the number of dimensions of an array. |
|
Numerical negative, element-wise. |
|
Return the next floating-point value after x1 towards x2, element-wise. |
|
Return the indices of the elements that are non-zero. |
|
Return (x1 != x2) element-wise. |
Abstract base class of all numeric scalar types. |
|
Any Python object. |
|
|
Return a new array of given shape and type, filled with ones. |
|
Return an array of ones with the same shape and type as a given array. |
|
Compute the outer product of two vectors. |
|
Packs the elements of a binary-valued array into bits in a uint8 array. |
|
Pad an array. |
|
Compute the q-th percentile of the data along the specified axis. |
|
Evaluate a piecewise-defined function. |
|
Find the sum of two polynomials. |
|
Return the derivative of the specified order of a polynomial. |
|
Find the product of two polynomials. |
|
Difference (subtraction) of two polynomials. |
|
Evaluate a polynomial at specific values. |
|
Numerical positive, element-wise. |
|
First array elements raised to powers from second array, element-wise. |
|
Return the product of array elements over a given axis. |
|
Return the product of array elements over a given axis. |
|
Returns the type to which a binary operation should cast its arguments. |
|
Range of values (maximum - minimum) along an axis. |
|
Compute the q-th quantile of the data along the specified axis. |
|
Convert angles from radians to degrees. |
|
Convert angles from degrees to radians. |
|
Return a contiguous flattened array. |
|
Converts a tuple of index arrays into an array of flat indices, applying boundary modes to the multi-index. |
|
Return the real part of the complex argument. |
|
Return the reciprocal of the argument, element-wise. |
|
Return element-wise remainder of division. |
|
Repeat elements of an array. |
|
Gives a new shape to an array without changing its data. |
|
Returns the type that results from applying the NumPy type promotion rules to the arguments. |
|
Shift the bits of an integer to the right. |
|
Round elements of the array to the nearest integer. |
|
Roll array elements along a given axis. |
|
Roll the specified axis backwards, until it lies in a given position. |
|
Return the roots of a polynomial with coefficients given in p. |
|
Rotate an array by 90 degrees in the plane specified by axes. |
|
Round an array to the given number of decimals. |
|
Stack arrays in sequence vertically (row wise). |
|
Save an array to a binary file in NumPy |
|
Save several arrays into a single file in uncompressed |
|
Find indices where elements should be inserted to maintain order. |
|
Return an array drawn from elements in choicelist, depending on conditions. |
|
Set printing options. |
|
Find the set difference of two arrays. |
|
Return the shape of an array. |
|
Returns an element-wise indication of the sign of a number. |
|
Returns element-wise True where signbit is set (less than zero). |
Abstract base class of all signed integer scalar types. |
|
|
Trigonometric sine, element-wise. |
|
Return the sinc function. |
alias of |
|
|
Hyperbolic sine, element-wise. |
|
Return the number of elements along a given axis. |
|
Test whether any array element along a given axis evaluates to True. |
|
Return a sorted copy of an array. |
|
Sort a complex array using the real part first, then the imaginary part. |
|
Split an array into multiple sub-arrays as views into ary. |
|
Return the non-negative square-root of an array, element-wise. |
|
Return the element-wise square of the input. |
|
Remove single-dimensional entries from the shape of an array. |
|
Join a sequence of arrays along a new axis. |
|
Compute the standard deviation along the specified axis. |
|
Subtract arguments, element-wise. |
|
Sum of array elements over a given axis. |
|
Interchange two axes of an array. |
|
Take elements from an array along an axis. |
|
Take values from the input array by matching 1d index and data slices. |
|
Compute tangent element-wise. |
|
Compute hyperbolic tangent element-wise. |
|
Compute tensor dot product along specified axes. |
|
Construct an array by repeating A the number of times given by reps. |
|
Return the sum along diagonals of the array. |
|
Reverse or permute the axes of an array; returns the modified array. |
|
Integrate along the given axis using the composite trapezoidal rule. |
|
An array with ones at and below the given diagonal and zeros elsewhere. |
|
Lower triangle of an array. |
|
Return the indices for the lower-triangle of an (n, m) array. |
|
Return the indices for the lower-triangle of arr. |
|
Trim the leading and/or trailing zeros from a 1-D array or sequence. |
|
Upper triangle of an array. |
|
Return the indices for the upper-triangle of an (n, m) array. |
|
Return the indices for the upper-triangle of arr. |
|
Returns a true division of the inputs, element-wise. |
|
Return the truncated value of the input, element-wise. |
|
Find the unique elements of an array. |
|
Unpacks elements of a uint8 array into a binary-valued output array. |
|
Converts a flat index or array of flat indices into a tuple of coordinate arrays. |
Abstract base class of all unsigned integer scalar types. |
|
|
Unwrap by changing deltas between values to 2*pi complement. |
|
Generate a Vandermonde matrix. |
|
Compute the variance along the specified axis. |
|
Return the dot product of two vectors. |
|
Define a vectorized function with broadcasting. |
|
Split an array into multiple sub-arrays vertically (row-wise). |
|
Stack arrays in sequence vertically (row wise). |
|
Return elements chosen from x or y depending on condition. |
|
Return a new array of given shape and type, filled with zeros. |
|
Return an array of zeros with the same shape and type as a given array. |
jax.numpy.fft¶
|
Compute the one-dimensional discrete Fourier Transform. |
|
Compute the 2-dimensional discrete Fourier Transform |
|
Return the Discrete Fourier Transform sample frequencies. |
|
Compute the N-dimensional discrete Fourier Transform. |
|
Shift the zero-frequency component to the center of the spectrum. |
|
Compute the FFT of a signal that has Hermitian symmetry, i.e., a real spectrum. |
|
Compute the one-dimensional inverse discrete Fourier Transform. |
|
Compute the 2-dimensional inverse discrete Fourier Transform. |
|
Compute the N-dimensional inverse discrete Fourier Transform. |
|
The inverse of fftshift. |
|
Compute the inverse FFT of a signal that has Hermitian symmetry. |
|
Compute the inverse of the n-point DFT for real input. |
|
Compute the 2-dimensional inverse FFT of a real array. |
|
Compute the inverse of the N-dimensional FFT of real input. |
|
Compute the one-dimensional discrete Fourier Transform for real input. |
|
Compute the 2-dimensional FFT of a real array. |
|
Return the Discrete Fourier Transform sample frequencies (for usage with rfft, irfft). |
|
Compute the N-dimensional discrete Fourier Transform for real input. |
jax.numpy.linalg¶
|
Cholesky decomposition. |
|
Compute the condition number of a matrix. |
Compute the determinant of an array. |
|
|
Compute the eigenvalues and right eigenvectors of a square array. |
|
Return the eigenvalues and eigenvectors of a complex Hermitian (conjugate symmetric) or a real symmetric matrix. |
|
Compute the eigenvalues of a general matrix. |
|
Compute the eigenvalues of a complex Hermitian or real symmetric matrix. |
|
Compute the (multiplicative) inverse of a matrix. |
|
Return the least-squares solution to a linear matrix equation. |
|
Raise a square matrix to the (integer) power n. |
|
Return matrix rank of array using SVD method |
|
Compute the dot product of two or more arrays in a single function call, while automatically selecting the fastest evaluation order. |
|
Matrix or vector norm. |
Compute the (Moore-Penrose) pseudo-inverse of a matrix. |
|
|
Compute the qr factorization of a matrix. |
Compute the sign and (natural) logarithm of the determinant of an array. |
|
|
Solve a linear matrix equation, or system of linear scalar equations. |
|
Singular Value Decomposition. |
|
Compute the ‘inverse’ of an N-dimensional array. |
|
Solve the tensor equation |