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