softjax

## softjax Soft differentiable drop-in replacements for non-differentiable JAX functions (abs, relu, sort, argmax, comparison, logical operators, etc.) with adjustable softening strength. ### Code repository <https://github.com/a-paulus/softjax> **Use this as the implementation source:** clone the repo and follow its README for install, dependencies, and how to run code or experiments. The generated client prints JSON with a suggested ``git clone`` command. ### Paper (arXiv — explanation) <https://arxiv.org/abs/2603.08824> This is the **paper** reference. The client can optionally fetch live Atom metadata (title, abstract) for agents; it does **not** run training or upstream research code by itself. ### What “running” this client does The `*_client.py` script prints **JSON** that combines a **GitHub repository** (clone URL + suggested ``git clone``) with **optional paper context** from arXiv (live Atom metadata when **reference_url** is arXiv). Run the real code by cloning the repo and following its README — the skill is your agent-facing entrypoint, not a substitute for the repo’s install steps. To call a **REST API** instead, set ``BASE_URL`` in `scripts/softjax_client.py` or wrap the upstream CLI with ``subprocess`` after clone. ### How to run the method (from the source) Extracted for **operators and agents**. Confirm against the upstream repository or paper before relying on it in production. ## Prerequisites - Python 3.11 or higher - JAX library installed ## Installation ```bash pip install softjax ``` ## How to run SoftJAX is a library providing drop-in soft replacements for non-differentiable JAX functions. Use it in Python scripts by importing and calling soft operators: ```python import jax.numpy as jnp import softjax as sj # Example: soft ReLU x = jnp.array([-0.2, -1.0, 0.3, 1.0]) print(sj.relu(x)) # soft mode by default print(sj.relu(x, mode="hard")) # hard (non-differentiable) mode ``` ### Elementwise operators ```python sj.abs(x) sj.relu(x) sj.clip(x, min_val, max_val) sj.sign(x) sj.round(x) sj.heaviside(x) ``` ### Array-valued operators ```python sj.max(x, method="neuralsort", softness=0.1) sj.min(x) sj.sort(x, method="neuralsort", softness=0.1) # supports multiple methods sj.quantile(x, q=0.5) sj.median(x) sj.top_k(x, k=3) sj.rank(x, descending=False) ``` ### Operators returning indices ```python sj.argmax(x) # returns soft distribution over indices sj.argmin(x) sj.argsort(x) sj.top_k(x, k=3) # returns (values, soft_indices) ``` ### Comparison operators ```python sj.greater(x, y) sj.equal(x, y) sj.less(x, y) sj.isclose(x, y) ``` ### Logical operators ```python sj.logical_and(a, b) sj.logical_or(a, b) sj.logical_not(a) sj.logical_xor(a, b) sj.all(a) sj.any(a) ``` ### Selection operators ```python sj.where(condition, x, y) ``` ### Straight-through estimators (hard forward, soft backward) ```python sj.relu_st(x) sj.sort_st(x) sj.top_k_st(x, k=3) sj.greater_st(x, y) ``` ### Autograd-safe operators (stable gradients at boundaries) ```python sj.sqrt(x) sj.arcsin(x) sj.arccos(x) sj.log(x) # safe at 0 sj.div(a, b) # safe at division by zero sj.norm(x) # safe at zero vector ``` ## Configuration Key parameters available across operators: - **`mode`**: Control softening behavior ('hard', 'soft', 'smooth', 'c0', 'c1', 'c2') - **`softness`**: Adjust approximation strength (float, higher = closer to hard function) - **`method`**: For sort-like operators, choose algorithm ('softsort', 'neuralsort', 'fast_soft_sort', 'smooth_sort', 'ot', 'sorting_network') ### Example with parameters ```python # Soft sort with custom softness result = sj.sort(x, method="fast_soft_sort", softness=2.0, mode="c1") # Soft argmax returning distribution soft_argmax = sj.argmax(x, mode="soft") # Straight-through sort (hard forward, soft backward for gradient) result = sj.sort_st(x, method="neuralsort", softness=0.1) ``` See documentation at https://a-paulus.github.io/softjax/ for full API details. *The same text lives in* ``scripts/USAGE.md`` *for tools that prefer reading files under* ``scripts/``*.* ### Parameters --mode (str) [optional, default=soft] Softening mode: 'hard' (non-differentiable), 'soft' (default, smooth approximation), 'smooth', 'c0', 'c1', 'c2' for different continuity guarantees. --softness (float) [optional, default=None] Strength of softening; controls smoothness and boundedness of soft function. Adjustable per operator and method. --method (str) [optional, default=neuralsort] Algorithm for soft operators like sort: 'softsort', 'neuralsort', 'fast_soft_sort', 'smooth_sort', 'ot', 'sorting_network'. --k (int) [optional, default=None] Number of top elements for top_k operations. --q (float) [optional, default=None] Quantile parameter (0–1) for quantile and argquantile operations. --descending (bool) [optional, default=True] Sort order for rank operation: True for descending, False for ascending. ### Usage ```bash python3 scripts/softjax_client.py sj.sort(x, method='neuralsort', softness=0.1, mode='soft') ``` ### Example Output ```json array([-0.8792, -0.1641, 0.2767, 0.8738]) ```