Skip to content

logic

JAX visitors for logic expressions.

Visitors: All, Any, Cond

Lowers the logic AST nodes to traceable JAX functions: All/Any reduce predicate residuals (satisfied when <= 0) with jnp.all/jnp.any, and Cond wraps jax.lax.cond so a predicate selects between two lowered branches at runtime, optionally gated to a set of node ranges. These branching/boolean nodes have no disciplined-convex form, so they are lowered only by this JAX backend, never by CVXPy.

Both cond branches are cast to a shared default float dtype (:func:set_default_float_dtype, configured by Problem.__init__) to avoid dtype-mismatch errors when JAX traces the two branches separately.

get_default_float_dtype()

Get the current default float dtype for conditional branches.

Source code in openscvx/symbolic/lowerers/jax/logic.py 
46
47
48
def get_default_float_dtype():
    """Get the current default float dtype for conditional branches."""
    return _DEFAULT_FLOAT_DTYPE

set_default_float_dtype(dtype: str) -> None

Set the default float dtype for conditional branches.

This is called by Problem.init to configure the dtype used in jax.lax.cond branches.

Parameters:

Name Type Description Default
dtype str

String like "float32" or "float64"

 required
Source code in openscvx/symbolic/lowerers/jax/logic.py 
30
31
32
33
34
35
36
37
38
39
40
41
42
43
def set_default_float_dtype(dtype: str) -> None:
    """Set the default float dtype for conditional branches.

    This is called by Problem.__init__ to configure the dtype used in jax.lax.cond branches.

    Args:
        dtype: String like "float32" or "float64"
    """
    global _DEFAULT_FLOAT_DTYPE
    dtype_lower = dtype.lower()
    if dtype_lower in ("float64", "f64", "double"):
        _DEFAULT_FLOAT_DTYPE = jnp.float64
    else:
        _DEFAULT_FLOAT_DTYPE = jnp.float32