JAX

https://jax.readthedocs.io/
Admin
JAX is a Python library for accelerator-oriented array computation and program transformation,...
Tokens:4,980
Snippets:45
Trust Score:9.6
Update:5 months ago
Show doc for...
Tokens:
Raw
### jax.experimental.sparse module

Source: https://jax.readthedocs.io

Provides functionality for sparse arrays using the BCOO format.

```APIDOC
## jax.experimental.sparse.BCOO

### Description
Represents a sparse array in the Block Coordinate (BCOO) format.

### Method
`jax.experimental.sparse.BCOO(indices, data, shape, block_size=None)`

### Parameters
- **indices** (ndarray) - Array of shape (N, K) containing the coordinates of the non-zero elements.
- **data** (ndarray) - Array of shape (N,) containing the non-zero values.
- **shape** (tuple) - The shape of the dense array.
- **block_size** (tuple, optional) - The block size for the BCOO format. Default is None.

### Response
#### Success Response (200)
- **result** (BCOO) - The BCOO sparse array.

## jax.experimental.sparse.bcoo_broadcast_in_dim

### Description
Broadcasts a BCOO sparse array into a given dimension.

### Method
`jax.experimental.sparse.bcoo_broadcast_in_dim(x, shape, axis)`

### Parameters
- **x** (BCOO) - Input BCOO sparse array.
- **shape** (tuple) - The target shape.
- **axis** (int or sequence of ints) - The dimension(s) to broadcast over.

### Response
#### Success Response (200)
- **result** (BCOO) - The broadcasted BCOO sparse array.

## jax.experimental.sparse.bcoo_concatenate

### Description
Concatenates a sequence of BCOO sparse arrays along a given axis.

### Method
`jax.experimental.sparse.bcoo_concatenate(arrays, axis)`

### Parameters
- **arrays** (Sequence[BCOO]) - A sequence of BCOO sparse arrays to concatenate.
- **axis** (int) - The axis along which to concatenate.

### Response
#### Success Response (200)
- **result** (BCOO) - The concatenated BCOO sparse array.

## jax.experimental.sparse.bcoo_dot_general

### Description
Performs a dot product (general matrix multiplication) on BCOO sparse arrays.

### Method
`jax.experimental.sparse.bcoo_dot_general(lhs, rhs, dimension_numbers, preferred_element_type=None)`

### Parameters
- **lhs** (BCOO) - The left-hand side BCOO sparse array.
- **rhs** (BCOO) - The right-hand side BCOO sparse array.
- **dimension_numbers** (tuple) - Specifies the contraction and batching dimensions.
- **preferred_element_type** (dtype, optional) - The preferred dtype for the output. Default is None.

### Response
#### Success Response (200)
- **result** (BCOO) - The resulting BCOO sparse array.

```

--------------------------------

### jax.numpy.fft module

Source: https://jax.readthedocs.io

Provides Fast Fourier Transform (FFT) functions, mirroring numpy.fft.

```APIDOC
## jax.numpy.fft.fft

### Description
Computes the one-dimensional discrete Fourier Transform.

### Method
`jax.numpy.fft.fft(a, n=None, axis=-1, norm=None)`

### Parameters
- **a** (array_like) - Input array.
- **n** (int, optional) - Length of the transformed axis of the output. If n is not given, the length of the input along the axis specified by axis is used.
- **axis** (int, optional) - Axis over which to compute the FFT. Default is the last axis.
- **norm** ({"backward", "ortho", "forward"}, optional) - Normalization mode. Default is None.

### Response
#### Success Response (200)
- **result** (ndarray) - The transformed array.

## jax.numpy.fft.fft2

### Description
Computes the two-dimensional discrete Fourier Transform.

### Method
`jax.numpy.fft.fft2(a, s=None, axes=(-2, -1), norm=None)`

### Parameters
- **a** (array_like) - Input array.
- **s** (sequence of ints, optional) - Shape of the output (along the transformed axes). If s is not given, the length of the input along the axes specified by axes is used.
- **axes** (sequence of ints, optional) - Axes for which transformation is applied. Default are the last two axes.
- **norm** ({"backward", "ortho", "forward"}, optional) - Normalization mode. Default is None.

### Response
#### Success Response (200)
- **result** (ndarray) - The transformed array.

## jax.numpy.fft.fftfreq

### Description
Returns the discrete Fourier transform sample frequencies.

### Method
`jax.numpy.fft.fftfreq(n, d=1.0)`

### Parameters
- **n** (int) - Number of sample points.
- **d** (scalar, optional) - Sample spacing. Default is 1.0.

### Response
#### Success Response (200)
- **result** (ndarray) - The transformed array.

## jax.numpy.fft.fftn

### Description
Computes the N-dimensional discrete Fourier Transform.

### Method
`jax.numpy.fft.fftn(a, s=None, axes=None, norm=None)`

### Parameters
- **a** (array_like) - Input array.
- **s** (sequence of ints, optional) - Shape of the output (along the transformed axes). If s is not given, the length of the input along the axes specified by axes is used.
- **axes** (sequence of ints, optional) - Axes for which transformation is applied. Default are all axes.
- **norm** ({"backward", "ortho", "forward"}, optional) - Normalization mode. Default is None.

### Response
#### Success Response (200)
- **result** (ndarray) - The transformed array.

## jax.numpy.fft.fftshift

### Description
Shifts zeros of the zero-frequency component to the center of the spectrum.

### Method
`jax.numpy.fft.fftshift(x, axes=None)`

### Parameters
- **x** (array_like) - Input array.
- **axes** (sequence of ints or None, optional) - Axes over which to shift. Default is all axes.

### Response
#### Success Response (200)
- **result** (ndarray) - The shifted array.

## jax.numpy.fft.hfft

### Description
Computes the Hermitian-symmetric discrete Fourier Transform in one dimension.

### Method
`jax.numpy.fft.hfft(a, n=None, axis=-1, norm=None)`

### Parameters
- **a** (array_like) - Input array.
- **n** (int, optional) - Length of the transformed axis of the output. Default is None.
- **axis** (int, optional) - Axis over which to compute the FFT. Default is the last axis.
- **norm** ({"backward", "ortho", "forward"}, optional) - Normalization mode. Default is None.

### Response
#### Success Response (200)
- **result** (ndarray) - The transformed array.

## jax.numpy.fft.ifft

### Description
Computes the one-dimensional inverse discrete Fourier Transform.

### Method
`jax.numpy.fft.ifft(a, n=None, axis=-1, norm=None)`

### Parameters
- **a** (array_like) - Input array.
- **n** (int, optional) - Length of the transformed axis of the output. Default is None.
- **axis** (int, optional) - Axis over which to compute the IFFT. Default is the last axis.
- **norm** ({"backward", "ortho", "forward"}, optional) - Normalization mode. Default is None.

### Response
#### Success Response (200)
- **result** (ndarray) - The transformed array.

## jax.numpy.fft.ifft2

### Description
Computes the two-dimensional inverse discrete Fourier Transform.

### Method
`jax.numpy.fft.ifft2(a, s=None, axes=(-2, -1), norm=None)`

### Parameters
- **a** (array_like) - Input array.
- **s** (sequence of ints, optional) - Shape of the output (along the transformed axes). Default is None.
- **axes** (sequence of ints, optional) - Axes for which transformation is applied. Default are the last two axes.
- **norm** ({"backward", "ortho", "forward"}, optional) - Normalization mode. Default is None.

### Response
#### Success Response (200)
- **result** (ndarray) - The transformed array.

## jax.numpy.fft.ifftn

### Description
Computes the N-dimensional inverse discrete Fourier Transform.

### Method
`jax.numpy.fft.ifftn(a, s=None, axes=None, norm=None)`

### Parameters
- **a** (array_like) - Input array.
- **s** (sequence of ints, optional) - Shape of the output (along the transformed axes). Default is None.
- **axes** (sequence of ints, optional) - Axes for which transformation is applied. Default are all axes.
- **norm** ({"backward", "ortho", "forward"}, optional) - Normalization mode. Default is None.

### Response
#### Success Response (200)
- **result** (ndarray) - The transformed array.

## jax.numpy.fft.ifftshift

### Description
The inverse of fftshift. Shifts the zero-frequency component back to its original location.

### Method
`jax.numpy.fft.ifftshift(x, axes=None)`

### Parameters
- **x** (array_like) - Input array.
- **axes** (sequence of ints or None, optional) - Axes over which to shift. Default is all axes.

### Response
#### Success Response (200)
- **result** (ndarray) - The shifted array.

## jax.numpy.fft.ihfft

### Description
Computes the inverse of the Hermitian-symmetric discrete Fourier Transform in one dimension.

### Method
`jax.numpy.fft.ihfft(a, n=None, axis=-1, norm=None)`

### Parameters
- **a** (array_like) - Input array.
- **n** (int, optional) - Length of the transformed axis of the output. Default is None.
- **axis** (int, optional) - Axis over which to compute the IFFT. Default is the last axis.
- **norm** ({"backward", "ortho", "forward"}, optional) - Normalization mode. Default is None.

### Response
#### Success Response (200)
- **result** (ndarray) - The transformed array.

## jax.numpy.fft.irfft

### Description
Computes the inverse of the rfft, i.e., the inverse discrete Fourier Transform of a complex array in which the negative-frequency components are not stored.

### Method
`jax.numpy.fft.irfft(a, n=None, axis=-1, norm=None)`

### Parameters
- **a** (array_like) - Input array.
- **n** (int, optional) - Length of the transformed axis of the output. Default is None.
- **axis** (int, optional) - Axis over which to compute the IFFT. Default is the last axis.
- **norm** ({"backward", "ortho", "forward"}, optional) - Normalization mode. Default is None.

### Response
#### Success Response (200)
- **result** (ndarray) - The transformed array.

## jax.numpy.fft.irfft2

### Description
Computes the inverse of the rfft2.

### Method
`jax.numpy.fft.irfft2(a, s=None, axes=(-2, -1), norm=None)`

### Parameters
- **a** (array_like) - Input array.
- **s** (sequence of ints, optional) - Shape of the output (along the transformed axes). Default is None.
- **axes** (sequence of ints, optional) - Axes for which transformation is applied. Default are the last two axes.
- **norm** ({"backward", "ortho", "forward"}, optional) - Normalization mode. Default is None.

### Response
#### Success Response (200)
- **result** (ndarray) - The transformed array.

## jax.numpy.fft.irfftn

### Description
Computes the inverse of the rfftn.

### Method
`jax.numpy.fft.irfftn(a, s=None, axes=None, norm=None)`

### Parameters
- **a** (array_like) - Input array.
- **s** (sequence of ints, optional) - Shape of the output (along the transformed axes). Default is None.
- **axes** (sequence of ints, optional) - Axes for which transformation is applied. Default are all axes.
- **norm** ({"backward", "ortho", "forward"}, optional) - Normalization mode. Default is None.

### Response
#### Success Response (200)
- **result** (ndarray) - The transformed array.

## jax.numpy.fft.rfft

### Description
Computes the one-dimensional discrete Fourier Transform of a real-valued array.

### Method
`jax.numpy.fft.rfft(a, n=None, axis=-1, norm=None)`

### Parameters
- **a** (array_like) - Input array.
- **n** (int, optional) - Length of the transformed axis of the output. Default is None.
- **axis** (int, optional) - Axis over which to compute the FFT. Default is the last axis.
- **norm** ({"backward", "ortho", "forward"}, optional) - Normalization mode. Default is None.

### Response
#### Success Response (200)
- **result** (ndarray) - The transformed array.

## jax.numpy.fft.rfft2

### Description
Computes the two-dimensional discrete Fourier Transform of a real-valued array.

### Method
`jax.numpy.fft.rfft2(a, s=None, axes=(-2, -1), norm=None)`

### Parameters
- **a** (array_like) - Input array.
- **s** (sequence of ints, optional) - Shape of the output (along the transformed axes). Default is None.
- **axes** (sequence of ints, optional) - Axes for which transformation is applied. Default are the last two axes.
- **norm** ({"backward", "ortho", "forward"}, optional) - Normalization mode. Default is None.

### Response
#### Success Response (200)
- **result** (ndarray) - The transformed array.

## jax.numpy.fft.rfftfreq

### Description
Computes the frequencies for a real-valued DFT.

### Method
`jax.numpy.fft.rfftfreq(n, d=1.0)`

### Parameters
- **n** (int) - Number of sample points.
- **d** (scalar, optional) - Sample spacing. Default is 1.0.

### Response
#### Success Response (200)
- **result** (ndarray) - The transformed array.

## jax.numpy.fft.rfftn

### Description
Computes the N-dimensional discrete Fourier Transform of a real-valued array.

### Method
`jax.numpy.fft.rfftn(a, s=None, axes=None, norm=None)`

### Parameters
- **a** (array_like) - Input array.
- **s** (sequence of ints, optional) - Shape of the output (along the transformed axes). Default is None.
- **axes** (sequence of ints, optional) - Axes for which transformation is applied. Default are all axes.
- **norm** ({"backward", "ortho", "forward"}, optional) - Normalization mode. Default is None.

### Response
#### Success Response (200)
- **result** (ndarray) - The transformed array.

```

--------------------------------

### jax.scipy.stats.bernoulli module

Source: https://jax.readthedocs.io

Provides functions for the Bernoulli distribution.

```APIDOC
## jax.scipy.stats.bernoulli.logpmf

### Description
Computes the log probability mass function of the Bernoulli distribution.

### Method
`jax.scipy.stats.bernoulli.logpmf(k, p)`

### Parameters
- **k** (array_like) - Observed values.
- **p** (array_like) - Probability of success.

### Response
#### Success Response (200)
- **result** (ndarray) - Log probability mass function values.

## jax.scipy.stats.bernoulli.pmf

### Description
Computes the probability mass function of the Bernoulli distribution.

### Method
`jax.scipy.stats.bernoulli.pmf(k, p)`

### Parameters
- **k** (array_like) - Observed values.
- **p** (array_like) - Probability of success.

### Response
#### Success Response (200)
- **result** (ndarray) - Probability mass function values.

## jax.scipy.stats.bernoulli.cdf

### Description
Computes the cumulative distribution function of the Bernoulli distribution.

### Method
`jax.scipy.stats.bernoulli.cdf(k, p)`

### Parameters
- **k** (array_like) - Observed values.
- **p** (array_like) - Probability of success.

### Response
#### Success Response (200)
- **result** (ndarray) - Cumulative distribution function values.

## jax.scipy.stats.bernoulli.ppf

### Description
Computes the percent point function (inverse of cdf) of the Bernoulli distribution.

### Method
`jax.scipy.stats.bernoulli.ppf(q, p)`

### Parameters
- **q** (array_like) - Quantiles, with `0 <= q <= 1`.
- **p** (array_like) - Probability of success.

### Response
#### Success Response (200)
- **result** (ndarray) - Percent point function values.

```
Tokens:
Raw