pylops.JaxOperator#

class pylops.JaxOperator(Op)[source]#

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

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