Combining assert_scalar_positive with jax.vmap

[thomaspinder]

Is it possible to apply assert_scalar_positive to a vector? The MWE below shows what I'd like to do; however, I cannot currently see how this is possible with Chex.

import jax.numpy as jnp 
import jax
from chex import assert_scalar_positive

x_scaler = 1.
x_vector = jnp.array([1.,1.]) 

assert_scalar_positive(x_scaler) # Works
jax.vmap(assert_scalar_positive)(x_vector) # What I'd like

[Sikandar1310291]

Hi @thomaspinder!

I'd like to work on this issue. I've analyzed the codebase and have a clear approach to solve this.

Problem Analysis

The current assert_scalar_positive is decorated with @_static_assertion and explicitly checks for scalar types (int/float), making it incompatible with jax.vmap.

Proposed Solution

I propose adding a new function assert_array_positive that:

1. Works with arrays and is vmap-compatible
2. Follows the pattern used by other value assertions in chex (like assert_trees_all_close)
3. Has both static and jittable versions

Implementation Plan

Add two variants:

1. Static version - For eager execution:

def _assert_array_positive_static(x: Array) -> None:
    """Checks that all elements in array are positive."""
    if not jnp.all(x > 0):
        raise AssertionError(f"All elements must be positive.")

2. Jittable version - For JIT/vmap compatibility:

def _assert_array_positive_jittable(x: Array) -> Array:
    """Jittable version for assert_array_positive."""
    predicate = jnp.all(x > 0)
    checkify.check(
        pred=predicate,
        msg="Array contains non-positive values: {x}",
        x=x
    )
    return predicate

3. Combine using @_value_assertion:

assert_array_positive = _value_assertion(
    assert_fn=_assert_array_positive_static,
    jittable_assert_fn=_assert_array_positive_jittable,
    name="assert_array_positive"
)

Usage Example

This will enable your use case:

import jax
import jax.numpy as jnp
from chex import assert_array_positive

x_scalar = 1.
x_vector = jnp.array([1., 1.])

assert_array_positive(x_scalar)  # Works
jax.vmap(assert_array_positive)(x_vector)  # Now works!

Additional Work

• Comprehensive unit tests covering:
  • Positive arrays
  • Arrays with zeros/negative values
  • vmap compatibility
  • JIT compatibility
  • Edge cases (empty arrays, different dtypes)
• Documentation and docstrings
• Update exports in __init__.py

Timeline

I can have a PR ready within a few days. Let me know if this approach looks good or if you'd prefer a different design!

cc: @hbq1 @chexdev