### Install MPAX

Source: https://github.com/mit-lu-lab/mpax/blob/main/README.md

Install the latest released version of MPAX from PyPI or the latest development version from GitHub.

```bash
pip install mpax
```

```bash
pip install git+https://github.com/MIT-Lu-Lab/mpax.git
```

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

### Device Parallelism with JAX Sharding

Source: https://github.com/mit-lu-lab/mpax/blob/main/README.md

Distribute computations across devices using JAX's sharding capabilities for data parallelism. This example demonstrates sharding the constraint matrix A.

```python
import jax
from mpax import create_lp

# Data sharding
mesh = jax.make_mesh((2,), ('x',))
sharding = jax.sharding.NamedSharding(mesh, P('x',))

A_sharded = jax.device_put(A, sharding)
lp_sharded = create_lp(c, A_sharded, b, G, h, l, u)

solver = r2HPDHG(eps_abs=1e-4, eps_rel=1e-4, verbose=True)
jit_optimize = jax.jit(solver.optimize)
result = jit_optimize(lp_sharded)
```

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

### Solve Single LP/QP Problem

Source: https://github.com/mit-lu-lab/mpax/blob/main/README.md

Create and solve a single LP or QP problem using either sparse or dense matrix formats. Choose between r2HPDHG or raPDHG solvers.

```python
from mpax import create_lp, create_qp, raPDHG, r2HPDHG

# Create LP using sparse matrix format (default)
lp = create_lp(c, A, b, G, h, l, u) # use_sparse_matrix=True by default
# Create LP using dense matrix format
lp = create_lp(c, A, b, G, h, l, u, use_sparse_matrix=False)
solver = r2HPDHG(eps_abs=1e-4, eps_rel=1e-4, verbose=True)
result = solver.optimize(lp)

# Or create QP
qp = create_qp(Q, c, A, b, G, h, l, u)
qp = create_qp(Q, c, A, b, G, h, l, u, use_sparse_matrix=False)
solver = raPDHG(eps_abs=1e-4, eps_rel=1e-4, verbose=True)
result = solver.optimize(lp)
```

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

### Batch Solving LP Problems with JAX vmap

Source: https://github.com/mit-lu-lab/mpax/blob/main/README.md

Solve multiple LP problems of the same shape simultaneously by vectorizing a single optimization function using jax.vmap.

```python
import jax.numpy as jnp
from mpax import create_lp, r2HPDHG

def single_optimize(c_vector):
    lp = create_lp(c_vector, A, b, G, h, l, u)
    solver = r2HPDHG(eps_abs=1e-4, eps_rel=1e-4, verbose=True)
    result = solver.optimize(lp)
    obj = jnp.dot(c_vector, result.primal_solution)
    return result.primal_solution, obj

batch_size = 100
batch_c = jnp.tile(c, (batch_size, 1))
batch_optimize = jax.vmap(single_optimize)

result = batch_optimize(batch_c)
```

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

### Custom VJP for Smart Predict-then-Optimize+

Source: https://github.com/mit-lu-lab/mpax/blob/main/README.md

Defines a custom vector-Jacobian product for the Smart Predict-then-Optimize+ loss function using jax.custom_vjp. This allows for efficient gradient computation in optimization problems.

```python
import jax
import jax.numpy as jnp

@jax.custom_vjp
def pso_fun(pred_cost, true_cost, true_sol, true_obj):
    sol, obj = batch_optimize(2*pred_cost - true_cost)
    loss = -obj + 2 * jnp.sum(pred_cost * true_sol, axis=1) - true_obj
    loss = jnp.mean(loss)
    return loss, sol

def spo_fwd(pred_cost, true_cost, true_sol, true_obj):
    loss, sol = pso_fun(pred_cost, true_cost, true_sol, true_obj)
    return loss, (sol, true_sol)

def spo_bwd(res, g):
    sol, true_sol = res
    grad = 2 * (true_sol - sol)
    # No gradients needed for true_cost, true_sol, or true_obj
    return grad * g, None, None, None

pso_fun.defvjp(spo_fwd, spo_bwd)
```

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

### Ensure Deterministic GPU Operations in JAX

Source: https://github.com/mit-lu-lab/mpax/blob/main/README.md

Sets an environment variable to enforce deterministic operations on GPUs. Note: This may degrade performance, especially for batch solving.

```python
import os
os.environ["XLA_FLAGS"] = "--xla_gpu_deterministic_ops=true"
```

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

### Enable Double Precision in JAX

Source: https://github.com/mit-lu-lab/mpax/blob/main/README.md

Enables 64-bit precision for JAX computations. This is useful for applications requiring higher numerical accuracy.

```python
jax.config.update("jax_enable_x64", True)
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.