Abmax: JAX-based ABM Framework & Beyond

• Abmax is a term with multiple meanings, primarily defining a JAX-based agent-based modeling framework that leverages JIT-compatible static-shape updates.
• The framework employs innovative RM and SCI update kernels to efficiently handle dynamic agent mutations while conforming to JAX's immutable array constraints.
• Distinct interpretations of Abmax span maximal monotone operator problems, branching Brownian motion, and maximal reliability in information theory.

Searching arXiv for papers on “Abmax” and related ABMax usage. Abmax is a term with multiple technical uses. In current computational-modeling usage, it denotes “an ABM framework based on JAX” that implements “multiple just-in-time (JIT) compilable algorithms” for updating “a dynamically selected number of agents by applying distinct changes to them during a simulation”; in other literatures, the term appears in the “Abmax problem” for monotone operators, in summaries of maximal displacement in branching Brownian motion with absorption, and as $Ab_{\max}$ for the maximal reliability function in multiple-access channels (Chaturvedi et al., 22 Aug 2025, Borwein et al., 2012, Yang, 2024, 0803.3645). The dominant contemporary use in the supplied record is the JAX-based framework, while the broader terminological landscape is relevant because identical or near-identical labels refer to unrelated maximality or maxima problems across analysis, probability, and information theory.

1. Terminological scope

In monotone operator theory, the sum problem is “often called the ‘Abmax problem’,” and asks when the sum $A+B$ of two maximally monotone operators is itself maximally monotone (Borwein et al., 2012). In branching Brownian motion with absorption, the label “Abmax” is used in the supplied summary for the maximal displacement of the absorbed process (Yang, 2024). In coding theory for discrete memoryless multiple-access channels, $Ab_{\max}$ denotes the maximal reliability function associated with maximal error probability (0803.3645). By contrast, “Abmax: A JAX-based Agent-based Modeling Framework” introduces Abmax as software infrastructure for agent-based modeling (Chaturvedi et al., 22 Aug 2025).

This multiplicity of meanings suggests that the term is not a stable discipline-independent noun. A plausible implication is that references to “Abmax” require immediate contextual disambiguation, especially in bibliographic search, citation analysis, and interdisciplinary discussion.

2. JAX-based agent-based modeling framework

Abmax is introduced for the setting in which “Agent-based modeling (ABM) is a principal approach for studying complex systems,” and where “high-performance array computing libraries like JAX can help scale such computational models to a large number of agents by using automatic vectorization and just-in-time (JIT) compilation” (Chaturvedi et al., 22 Aug 2025). The central difficulty is that, in JAX, “the shapes of arrays used in the computational model should remain immutable throughout the simulation.” In ABM, this constrains “agent manipulation operations that require flexible data structures,” especially “the ability to update a dynamically selected number of agents by applying distinct changes to them during a simulation” (Chaturvedi et al., 22 Aug 2025).

The framework addresses that difficulty through JAX-compatible data structures and static-shape collections. All Abmax data structures use flax.struct, and the framework defines helper classes State, Params, and Signal, together with core classes Policy, Agent, and Set (Chaturvedi et al., 22 Aug 2025). State holds “per-step mutating agent quantities,” Params stores “constants/parameters for agents,” and Signal is “used for messages between agents/environment.” Policy encapsulates an agent’s decision rule, Agent bundles State, Params, Policy, “a unique ID, a type tag, and an age counter,” and Set is a “collection (array) of Agent objects with static shape (max capacity), a count of active agents, and copies of State, Params, and Policy for the collective” (Chaturvedi et al., 22 Aug 2025).

The framework enforces static shapes by preallocating agent collections “to a maximum size”; “unused slots are filled with ‘placeholder’ agents,” while only the “active” count changes dynamically (Chaturvedi et al., 22 Aug 2025). This design is JAX-idiomatic: it preserves compatibility with jit, vmap, and related transformations while still supporting nontrivial ABM state evolution.

3. Static-shape update mechanisms

Abmax exposes several functions for agent creation, stepping, and mutation. The supplied summary lists create_agents, step_agents, set_agents_sci, set_agents_rm, set_agents_mask, sort_agents, and select_agents; “all functions except initialization (create_agents) are JIT-compatible and support vectorized operation (via JAX’s vmap/pmap)” (Chaturvedi et al., 22 Aug 2025).

Two update kernels are central: Rank-Match (RM) and Sort-Count-Iterate (SCI). RM is described as a vectorized matching procedure in which a selection mask is converted into ranks and updates are applied when agent and update ranks coincide. In the example given in the paper summary,

$$\mathbf{m_a} \gets \mathbf{a}\bmod 2 == 0,\qquad \mathbf{m_b} \gets \mathbf{b}\bmod 2 == 1,$$

followed by

$$\mathbf{r_a} \gets \operatorname{cumsum}(\mathbf{m_a})\odot\mathbf{m_a},\qquad \mathbf{r_b} \gets \operatorname{cumsum}(\mathbf{m_b})\odot\mathbf{m_b}.$$

The update then applies when $r_a == r_b$ (Chaturvedi et al., 22 Aug 2025). RM is characterized as “very fast,” “JIT-friendly,” and “shape-static,” and the summary emphasizes that it is suitable for “ABM steps where unique distinct updates need to be mapped agent-wise” (Chaturvedi et al., 22 Aug 2025).

SCI takes a more explicit index-based route. It forms masks, computes

$r \gets \min(\sum \mathbf{m}_a,\sum \mathbf{m}_b),$

sorts candidate indices with

$$\mathbf{i_a} \gets \operatorname{argsort}(\mathbf{m_a}),\qquad \mathbf{i_b} \gets \operatorname{argsort}(\mathbf{m_b}),$$

and then assigns the first $r$ updates to the first $r$ selected agents in a shape-static loop (Chaturvedi et al., 22 Aug 2025). The summary states that SCI is “more explicit index handling, more flexible, but slower,” whereas RM is “faster, less explicit control” (Chaturvedi et al., 22 Aug 2025). The framework also includes a mask-based path “for independent, non-conflicting updates” (Chaturvedi et al., 22 Aug 2025).

The design principle is explicit: “always work with maximal-size arrays + ‘active agent’ masks, never true Python lists or append/delete patterns inside JIT” (Chaturvedi et al., 22 Aug 2025). This is less a generic ABM abstraction than a concrete response to JAX’s immutable-shape constraint.

4. Performance profile and example models

On the canonical predation benchmark, the summary reports a comparison among Agents.jl, Abmax RM, Abmax SCI, and Mesa over 100 steps (Chaturvedi et al., 22 Aug 2025). The reported runtimes are as follows.

Environment	Implementation	Runtime (ms)
small (600 sheep, 400 wolves, 100x100 grid)	Agents.jl	14.93
small (600 sheep, 400 wolves, 100x100 grid)	Abmax RM	50.26
small (600 sheep, 400 wolves, 100x100 grid)	Abmax SCI	726.78
small (600 sheep, 400 wolves, 100x100 grid)	Mesa	1333.05
large (6000 sheep, 4000 wolves, 1000x1000 grid)	Agents.jl	685.03
large (6000 sheep, 4000 wolves, 1000x1000 grid)	Abmax RM	3315.88
large (6000 sheep, 4000 wolves, 1000x1000 grid)	Abmax SCI	5455.01
large (6000 sheep, 4000 wolves, 1000x1000 grid)	Mesa	170070.95

The summary’s interpretation is precise: “Abmax RM comes significantly closer to Agents.jl performance,” “Abmax SCI is slower than RM,” and “Abmax’s strong point: Not absolute top speed, but enabling JIT-friendly, highly parallel, and scalable ABM in the Python/JAX ecosystem” (Chaturvedi et al., 22 Aug 2025). The same source states that, with JAX’s vmap, Abmax can simulate “many similar agent-based models in parallel,” and that vectorization is a first-class capability rather than an afterthought (Chaturvedi et al., 22 Aug 2025).

Two demonstration models are highlighted. The traffic-flow model is a cellular-automata-style road system in which “road and cars” are agents, movement and conflict resolution are vectorized across lanes, and “random batch insertion of cars into available entry cells” is supported (Chaturvedi et al., 22 Aug 2025). The financial market model simulates “noisy traders interacting with multiple Limit Order Books (LOBs) in parallel,” with orders “sorted, matched, and cumulative shares managed in a fully vectorized (parallelized) fashion” (Chaturvedi et al., 22 Aug 2025). In both cases, the point is not merely that the framework can host diverse domains, but that it can express domain-specific update logic under JAX’s compilation model.

5. ABMax in predator-prey tuning toward Lotka-Volterra dynamics

A subsequent predator-prey study states that the model “is implemented in ABMax, a JAX-based agent-based modelling framework that enables efficient batched simulation on hardware accelerators” (Mandl et al., 11 Jun 2026). Here the capitalization shifts to ABMax, but the description is continuous with the framework summarized above: “all simulation code is written using JAX,” multiple model instances can be simulated “in parallel on the same accelerator,” and the “rank-match update algorithm” is explicitly named as the mechanism that “maintains different numbers of active agents per instance while keeping computation efficient” (Mandl et al., 11 Jun 2026).

The model itself is a continuous predator-prey system. Sheep and wolves are “active agents with local sensing, internal energy, and recurrent neural network-based controllers” in a “2D continuous Cartesian plane” (Mandl et al., 11 Jun 2026). Birth and death are energy-threshold processes with fixed species buffers $A+B$0 and $A+B$1 (Mandl et al., 11 Jun 2026). Motion obeys

$A+B$2

and

$A+B$3

with speed and angular velocity generated by the controller (Mandl et al., 11 Jun 2026). Sheep gain energy from a grass patch,

$A+B$4

while wolves gain energy by preying on sheep through the nearest-target rule encoded in $A+B$5 (Mandl et al., 11 Jun 2026). The per-step updates are

$A+B$6

for sheep and

$A+B$7

for wolves, with clipping to $A+B$8 (Mandl et al., 11 Jun 2026).

The study’s main question is whether “environmental and demographic parameters can be tuned so that the resulting population dynamics resemble classical Lotka-Volterra cycles” (Mandl et al., 11 Jun 2026). The tuning objective is a feature-based loss combining “phase lag, bounded populations, and long-term persistence,” with normalized populations

$A+B$9

and a correlation term

$Ab_{\max}$0

The total loss is

$Ab_{\max}$1

and optimization proceeds with CMA-ES over ecological parameters such as $Ab_{\max}$2, $Ab_{\max}$3, $Ab_{\max}$4, $Ab_{\max}$5, and species-specific birth and death thresholds and dwell times (Mandl et al., 11 Jun 2026). In this usage, ABMax is not merely a simulation engine; it is the computational substrate for large batched parameter search on GPU hardware.

6. Other research uses of “Abmax”

Several unrelated research programs use the same or a near-identical label.

Usage	Meaning
“Abmax problem”	The sum problem in monotone operator theory (Borwein et al., 2012)
“Abmax” in BBM with absorption	Maximal displacement / maximum of the absorbed process (Yang, 2024)
$Ab_{\max}$6	Maximal reliability function for DM-MAC maximal error (0803.3645)

In monotone operator theory, the core question is: “Given two maximally monotone operators $Ab_{\max}$7 and $Ab_{\max}$8 on a Banach space $Ab_{\max}$9, under what conditions is their sum $$\mathbf{m_a} \gets \mathbf{a}\bmod 2 == 0,\qquad \mathbf{m_b} \gets \mathbf{b}\bmod 2 == 1,$$0 also maximally monotone?” (Borwein et al., 2012). One result establishes that if $$\mathbf{m_a} \gets \mathbf{a}\bmod 2 == 0,\qquad \mathbf{m_b} \gets \mathbf{b}\bmod 2 == 1,$$1 is “a maximally monotone linear relation,” $$\mathbf{m_a} \gets \mathbf{a}\bmod 2 == 0,\qquad \mathbf{m_b} \gets \mathbf{b}\bmod 2 == 1,$$2 is “a maximally monotone operator,” and

$$\mathbf{m_a} \gets \mathbf{a}\bmod 2 == 0,\qquad \mathbf{m_b} \gets \mathbf{b}\bmod 2 == 1,$$3

then $$\mathbf{m_a} \gets \mathbf{a}\bmod 2 == 0,\qquad \mathbf{m_b} \gets \mathbf{b}\bmod 2 == 1,$$4 is maximally monotone, and moreover “$$\mathbf{m_a} \gets \mathbf{a}\bmod 2 == 0,\qquad \mathbf{m_b} \gets \mathbf{b}\bmod 2 == 1,$$5 is of type (FPV)” (Borwein et al., 2012). A later result broadens this direction by proving maximal monotonicity of $$\mathbf{m_a} \gets \mathbf{a}\bmod 2 == 0,\qquad \mathbf{m_b} \gets \mathbf{b}\bmod 2 == 1,$$6 when $$\mathbf{m_a} \gets \mathbf{a}\bmod 2 == 0,\qquad \mathbf{m_b} \gets \mathbf{b}\bmod 2 == 1,$$7 is of type (FPV) under the same Rockafellar constraint qualification, and presents an equivalent formulation of the general sum problem in terms of sums with normal cones (Yao, 2014).

In branching Brownian motion with absorption, the summary labeled “maximal displacement (‘Abmax’)” studies the supercritical regime $$\mathbf{m_a} \gets \mathbf{a}\bmod 2 == 0,\qquad \mathbf{m_b} \gets \mathbf{b}\bmod 2 == 1,$$8, where particles undergo Brownian motion, branch, and are killed at a moving absorbing barrier $$\mathbf{m_a} \gets \mathbf{a}\bmod 2 == 0,\qquad \mathbf{m_b} \gets \mathbf{b}\bmod 2 == 1,$$9 (Yang, 2024). For the centered maximum $$\mathbf{r_a} \gets \operatorname{cumsum}(\mathbf{m_a})\odot\mathbf{m_a},\qquad \mathbf{r_b} \gets \operatorname{cumsum}(\mathbf{m_b})\odot\mathbf{m_b}.$$0, the time-averaged empirical distribution converges almost surely to a randomly shifted Gumbel law: $$\mathbf{r_a} \gets \operatorname{cumsum}(\mathbf{m_a})\odot\mathbf{m_a},\qquad \mathbf{r_b} \gets \operatorname{cumsum}(\mathbf{m_b})\odot\mathbf{m_b}.$$1 with $$\mathbf{r_a} \gets \operatorname{cumsum}(\mathbf{m_a})\odot\mathbf{m_a},\qquad \mathbf{r_b} \gets \operatorname{cumsum}(\mathbf{m_b})\odot\mathbf{m_b}.$$2 the almost sure limit of the derivative martingale (Yang, 2024).

In information theory, $$\mathbf{r_a} \gets \operatorname{cumsum}(\mathbf{m_a})\odot\mathbf{m_a},\qquad \mathbf{r_b} \gets \operatorname{cumsum}(\mathbf{m_b})\odot\mathbf{m_b}.$$3 refers to the maximal reliability function or maximal achievable error exponent under the maximal-error criterion for a discrete memoryless multiple-access channel (0803.3645). The cited result derives a new sphere-packing lower bound for maximal error probability that “explicitly imposes independence of the users’ input distributions,” yielding a tighter exponent than Haroutunian’s bound. The single-letter exponent is

$$\mathbf{r_a} \gets \operatorname{cumsum}(\mathbf{m_a})\odot\mathbf{m_a},\qquad \mathbf{r_b} \gets \operatorname{cumsum}(\mathbf{m_b})\odot\mathbf{m_b}.$$4

where the product form of $$\mathbf{r_a} \gets \operatorname{cumsum}(\mathbf{m_a})\odot\mathbf{m_a},\qquad \mathbf{r_b} \gets \operatorname{cumsum}(\mathbf{m_b})\odot\mathbf{m_b}.$$5 enforces the independence of separate encoders (0803.3645).

Across these domains, “Abmax” is therefore not a unified concept. In one line of work it is software for static-shape, JIT-compilable, vectorized ABM in JAX; in others it is shorthand for maximality of operator sums, maximal displacement of absorbed branching systems, or maximal-error reliability exponents. The common lexical root is “max,” but the underlying mathematics and computational objects are entirely distinct.