lamm-mit/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

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:

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

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

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

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

sj.greater(x, y)
sj.equal(x, y)
sj.less(x, y)
sj.isclose(x, y)

Logical operators

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

sj.where(condition, x, y)

Straight-through estimators (hard forward, soft backward)

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)

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

# 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

python3 scripts/softjax_client.py sj.sort(x, method='neuralsort', softness=0.1, mode='soft')

Example Output

array([-0.8792, -0.1641, 0.2767, 0.8738])