# Copyright 2025 BrainX Ecosystem Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
from typing import Sequence, Dict, List, Set
from collections import defaultdict
from jax.extend.core.primitives import dot_general_p, conv_general_dilated_p
from brainstate._compatible_import import is_jit_primitive, JaxprEqn, Jaxpr, ClosedJaxpr, Var, Literal
from brainstate._state import State
from brainstate.transform._ir_utils import IdentitySet, IRValidationError, check_all_vars
__all__ = [
'eqns_to_closed_jaxpr',
'eqns_to_jaxpr',
]
def _resolve_jaxpr_vars(
eqns: List[JaxprEqn],
invars: Sequence[Var] = None,
outvars: Sequence[Var] = None,
constvars: Sequence[Var] = None,
) -> tuple:
"""
Resolve the binder lists for a Jaxpr built from ``eqns``.
Any list left as ``None`` is inferred from the equations; provided lists are
validated instead. Factored out of :func:`eqns_to_jaxpr` so that
:func:`eqns_to_closed_jaxpr` can see the resolved ``constvars`` directly.
Reading them back off the constructed :class:`Jaxpr` is not an option: since
jax 0.11.1 an *open* Jaxpr reports ``constvars`` only for inputs that carry
attached constant *values*, so constvars-without-values read back as ``[]``
(see :func:`eqns_to_jaxpr`).
Returns
-------
tuple
The resolved ``(constvars, invars, outvars)`` lists.
"""
# Insertion-ordered set of variables produced by the equations.
produced_vars = IdentitySet()
for eqn in eqns:
produced_vars.update(eqn.outvars)
# Ordered list of used Vars (first-use order), de-duplicated by identity.
# Using ordered structures (instead of plain ``set`` iteration) makes the
# inferred invars/constvars/outvars deterministic across runs.
used_order = []
used_seen = IdentitySet()
for eqn in eqns:
for var in eqn.invars:
if isinstance(var, Var) and var not in used_seen:
used_seen.add(var)
used_order.append(var)
# Infer invars if not provided (used but not produced), else validate.
if invars is None:
invars = [v for v in used_order if v not in produced_vars]
else:
invars = list(invars)
check_all_vars(invars, 'invars')
invars_set = IdentitySet(invars)
# Infer constvars if not provided (used, not produced, not an invar).
if constvars is None:
constvars = [v for v in used_order
if v not in produced_vars and v not in invars_set]
else:
constvars = list(constvars)
check_all_vars(constvars, 'constvars')
# Infer outvars if not provided: variables produced but never consumed,
# in production order (deterministic). Otherwise validate the provided list.
if outvars is None:
consumed_vars = IdentitySet()
for eqn in eqns:
for var in eqn.invars:
if isinstance(var, Var) and var in produced_vars:
consumed_vars.add(var)
outvars = []
for eqn in eqns:
for v in eqn.outvars:
if v not in consumed_vars:
outvars.append(v)
else:
outvars = list(outvars)
check_all_vars(outvars, 'outvars')
return constvars, invars, outvars
[docs]
def eqns_to_jaxpr(
eqns: Sequence[JaxprEqn],
invars: Sequence[Var] = None,
outvars: Sequence[Var] = None,
constvars: Sequence[Var] = None,
) -> Jaxpr:
"""
Convert a sequence of JaxprEqn into a Jaxpr.
Parameters
----------
eqns
Sequence of Jaxpr equations to convert
invars
Input variables. If None, will be inferred from equations
outvars
Output variables. If None, will be inferred from equations
constvars
Constant variables. If None, will be automatically extracted from equations
Returns
-------
Jaxpr
A Jaxpr object constructed from the equations
Notes
-----
``constvars`` are always placed *before* ``invars`` in the resulting jaxpr's
binder list, but how they read back differs by jax version. jax 0.11.1
dropped ``Jaxpr``'s separate constvar count, so an input now counts as a
constvar only when a constant *value* is attached to it: on jax >= 0.11.1
the ``constvars`` passed here surface as leading ``jaxpr.invars`` and
``jaxpr.constvars`` is empty. On jax < 0.11.1 they surface as
``jaxpr.constvars``.
Either way the binder *order* is preserved, so pairing the result with its
constant values -- which is what :func:`eqns_to_closed_jaxpr` does -- yields
the same ``constvars``/``invars`` split on every supported jax version.
"""
eqns = list(eqns)
constvars, invars, outvars = _resolve_jaxpr_vars(eqns, invars, outvars, constvars)
return Jaxpr(
constvars=constvars,
invars=invars,
outvars=outvars,
eqns=eqns,
)
[docs]
def eqns_to_closed_jaxpr(
eqns: Sequence[JaxprEqn],
invars: Sequence[Var] = None,
outvars: Sequence[Var] = None,
constvars: Sequence[Var] = None,
consts: Sequence = None,
) -> ClosedJaxpr:
"""
Convert a sequence of JaxprEqn into a ClosedJaxpr.
Parameters
----------
eqns
Sequence of Jaxpr equations to convert
invars
Input variables. If None, will be inferred from equations
outvars
Output variables. If None, will be inferred from equations
constvars
Constant variables. If None, will be automatically extracted from equations
consts
Constant values corresponding to constvars. If None, defaults to empty list
Returns
-------
ClosedJaxpr
A ClosedJaxpr object constructed from the equations
Notes
-----
If constvars are automatically extracted from equations but no consts are provided,
the resulting ClosedJaxpr will have empty consts list. This may cause runtime errors
if the equations actually depend on these constants. In such cases, you should
explicitly provide both constvars and consts from the original jaxpr.
"""
# Resolve the binder lists up front. The resolved ``constvars`` must be read
# from here rather than from the open jaxpr below: on jax >= 0.11.1 an open
# Jaxpr reports ``constvars`` only for inputs with attached values, so
# ``jaxpr.constvars`` would always be empty and the length check below would
# silently pass for any ``consts`` (see :func:`eqns_to_jaxpr`).
eqns = list(eqns)
constvars, invars, outvars = _resolve_jaxpr_vars(eqns, invars, outvars, constvars)
# Handle consts
if consts is None:
# If no consts provided, create empty list
# This is safe if there are no constvars, but may cause errors otherwise
consts = []
else:
consts = list(consts)
# Verify consts length matches constvars length
if len(consts) != len(constvars):
raise IRValidationError(
f"consts length ({len(consts)}) does not match constvars length ({len(constvars)})"
)
jaxpr = Jaxpr(constvars=constvars, invars=invars, outvars=outvars, eqns=eqns)
# Pairing the jaxpr with its constant values restores the
# constvars/invars split on jax >= 0.11.1, where the open jaxpr above
# carries the constvars as leading invars.
return ClosedJaxpr(jaxpr, consts)
