jaxvacua.util.is_static

is_static(arg)

Heuristic test for whether arg should be treated as a JAX static argument: True iff arg is not a jnp.ndarray.

Parameters:

arg (Any) – The argument to test.

Returns:

bool – True if static, False if it should be traced.

Return type:

bool