pylops.JaxOperator

class pylops.JaxOperator(Op)

Enable JAX backend for PyLops operator.

This class can be used to wrap a pylops operator to enable the JAX backend. Doing so, users can run all of the methods of a pylops operator with JAX arrays. Moreover, the forward and adjoint are internally just-in-time compiled, and other JAX functionalities such as automatic differentiation and automatic vectorization are enabled.

Parameters:

Oppylops.LinearOperator

PyLops operator

dimstuple

Shape of the array after the adjoint, but before flattening.

For example, x_reshaped = (Op.H * y.ravel()).reshape(Op.dims).

dimsdtuple

Shape of the array after the forward, but before flattening.

For example, y_reshaped = (Op * x.ravel()).reshape(Op.dimsd).

shapetuple

Operator shape.

Methods

__init__(Op)
adjoint()
apply_columns(cols)	Apply subset of columns of operator
cond([uselobpcg])	Condition number of linear operator.
conj()	Complex conjugate operator
div(y[, niter, densesolver])	Solve the linear problem \(\mathbf{y}=\mathbf{A}\mathbf{x}\).
dot(x)	Matrix-matrix or matrix-vector multiplication.
eigs([neigs, symmetric, niter, uselobpcg])	Most significant eigenvalues of linear operator.
matmat(X)	Matrix-matrix multiplication.
matvec(x)	Matrix-vector multiplication.
reset_count()	Reset counters
rmatmat(X)	Matrix-matrix multiplication.
rmatvec(x)	Adjoint matrix-vector multiplication.
rmatvecad(x, y)	Vector-Jacobian product
todense([backend])	Return dense matrix.
toimag([forw, adj])	Imag operator
toreal([forw, adj])	Real operator
tosparse()	Return sparse matrix.
trace([neval, method, backend])	Trace of linear operator.
transpose()

Examples using pylops.JaxOperator

21. JAX Operator

21. JAX Operator