linox: Linear Operators in JAX
linox is a Python package that provides a collection of linear operators for JAX, enabling efficient and flexible linear algebra operations with lazy evaluation. This package is designed as an JAX alternative to probnum.linops, but it is currently still under development having less and more instable features. It has no dependencies other than JAX and plum for multiple dispatch.
Features
- Lazy Evaluation: All operators support lazy evaluation, allowing for efficient computation of complex linear transformations
- JAX Integration: Built on top of JAX, providing automatic differentiation, parallelization, JIT compilation, and GPU/TPU support
- Composable Operators: Operators can be combined to form complex linear transformations
Linear Operators
Basic Operators
Matrix: General matrix operatorIdentity: Identity matrix operatorDiagonal: Diagonal matrix operatorScalar: Scalar multiple of identityZero: Zero matrix operatorOnes: Matrix of ones operator
Block Operators
BlockMatrix: General block matrix operatorBlockMatrix2x2: 2x2 block matrix operatorBlockDiagonal: Block diagonal matrix operator
Low Rank Operators
LowRank: General low rank operatorSymmetricLowRank: Symmetric low rank operatorIsotropicScalingPlusSymmetricLowRank: Isotropic scaling plus symmetric low rankPositiveDiagonalPlusSymmetricLowRank: Positive diagonal plus symmetric low rank
Special Operators
Kronecker: Kronecker product operatorPermutation: Permutation matrix operatorEigenD: Eigenvalue decomposition operator
Benefits of JAX Integration
- Automatic Differentiation: Compute gradients automatically through operator compositions
- JIT Compilation: Speed up computations with just-in-time compilation
- Vectorization: Efficient batch processing of linear operations via e.g.
jax.vmap - GPU/TPU Support: Run computations on accelerators without code changes
- Functional Programming: Pure functions enable better optimization and parallelization
Quick Example
import jax
import jax.numpy as jnp
from linox import Matrix, Diagonal, BlockMatrix
# Create operators
A = Matrix(jnp.array([[1, 2], [3, 4]], dtype=jnp.float32))
D = Diagonal(jnp.array([1, 2], dtype=jnp.float32))
# Compose operators
B = BlockMatrix([[A, D], [D, A]])
# Apply to vector
x = jnp.ones((4,), dtype=jnp.float32)
y = B @ x # Lazy evaluation
# Parallelize over batch of vectors
x_batched = jnp.ones((10, 4), dtype=jnp.float32)
y_batched = jax.vmap(B)(x_batched)
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
File details
Details for the file linox-0.0.1.tar.gz.
File metadata
- Download URL: linox-0.0.1.tar.gz
- Upload date:
- Size: 36.0 kB
- Tags: Source
- Uploaded using Trusted Publishing? Yes
- Uploaded via:
twine/7.0.0 CPython/3.13.14
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
de873f43764556d8f4d15d9eae8eaf7bcbbbb1a9b1dcfd7721e3dd4837307f2d
|
|
| MD5 |
53c0199dbcd66b722ddc0edbf1052fdb
|
|
| BLAKE2b-256 |
a0ddd3985761cc3d2e32ce711fec989a0c04739b5ddacbfa1f2c334e20f6d187
|
Provenance
The following attestation bundles were made for linox-0.0.1.tar.gz:
Publisher:
publish.yml on lenardrommel/linox
-
Statement:
-
Statement type:
https://in-toto.io/Statement/v1 -
Predicate type:
https://docs.pypi.org/attestations/publish/v1 -
Subject name:
linox-0.0.1.tar.gz -
Subject digest:
de873f43764556d8f4d15d9eae8eaf7bcbbbb1a9b1dcfd7721e3dd4837307f2d - Sigstore transparency entry: 2548057628
- Sigstore integration time:
-
Permalink:
lenardrommel/linox@5036fdd2e7fb625e305c8d85877bdd7444f81349 -
Branch / Tag:
refs/heads/main - Owner: https://github.com/lenardrommel
-
Access:
public
-
Token Issuer:
https://token.actions.githubusercontent.com -
Runner Environment:
github-hosted -
Publication workflow:
publish.yml@5036fdd2e7fb625e305c8d85877bdd7444f81349 -
Trigger Event:
workflow_dispatch
-
Statement type:
File details
Details for the file linox-0.0.1-py3-none-any.whl.
File metadata
- Download URL: linox-0.0.1-py3-none-any.whl
- Upload date:
- Size: 27.3 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? Yes
- Uploaded via:
twine/7.0.0 CPython/3.13.14
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
601c2ef618913174c76a4cac22930d02f8b91a6d8d268aae1152d8fb84c172a5
|
|
| MD5 |
61430c5c5def4aaeb8f3a5ab9bde0f7e
|
|
| BLAKE2b-256 |
9c1c7296fd8513fd2f6a7a69d5b6fb9489c40f3d771d650a8273b38808338f86
|
Provenance
The following attestation bundles were made for linox-0.0.1-py3-none-any.whl:
Publisher:
publish.yml on lenardrommel/linox
-
Statement:
-
Statement type:
https://in-toto.io/Statement/v1 -
Predicate type:
https://docs.pypi.org/attestations/publish/v1 -
Subject name:
linox-0.0.1-py3-none-any.whl -
Subject digest:
601c2ef618913174c76a4cac22930d02f8b91a6d8d268aae1152d8fb84c172a5 - Sigstore transparency entry: 2548057725
- Sigstore integration time:
-
Permalink:
lenardrommel/linox@5036fdd2e7fb625e305c8d85877bdd7444f81349 -
Branch / Tag:
refs/heads/main - Owner: https://github.com/lenardrommel
-
Access:
public
-
Token Issuer:
https://token.actions.githubusercontent.com -
Runner Environment:
github-hosted -
Publication workflow:
publish.yml@5036fdd2e7fb625e305c8d85877bdd7444f81349 -
Trigger Event:
workflow_dispatch
-
Statement type: