brainstate 0.1.10__py2.py3-none-any.whl → 0.2.1__py2.py3-none-any.whl

brainstate/compile/_ad_checkpoint.py

@@ -1,204 +0,0 @@
-# Copyright 2024 BDP 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.
-# ==============================================================================
-
-import functools
-from typing import Callable, Tuple, Union
-
-import jax
-
-from brainstate.typing import Missing
-from ._make_jaxpr import StatefulFunction, _ensure_index_tuple
-from ._util import write_back_state_values
-
-__all__ = [
-    'checkpoint',
-    'remat'
-]
-
-
-def checkpoint(
-    fun: Callable = Missing(),
-    *,
-    prevent_cse: bool = True,
-    policy: Callable[..., bool] | None = None,
-    static_argnums: int | Tuple[int, ...] = (),
-) -> Union[Callable, Callable[[Callable], Callable]]:
-    """Make ``fun`` recompute internal linearization points when differentiated.
-
-    The :func:`jax.checkpoint` decorator, aliased to :func:`jax.remat`, provides a
-    way to trade off computation time and memory cost in the context of automatic
-    differentiation, especially with reverse-mode autodiff like :func:`jax.grad`
-    and :func:`jax.vjp` but also with :func:`jax.linearize`.
-
-    When differentiating a function in reverse-mode, by default all the
-    linearization points (e.g. inputs to elementwise nonlinear primitive
-    operations) are stored when evaluating the forward pass so that they can be
-    reused on the backward pass. This evaluation strategy can lead to a high
-    memory cost, or even to poor performance on hardware accelerators where memory
-    access is much more expensive than FLOPs.
-
-    An alternative evaluation strategy is for some of the linearization points to
-    be recomputed (i.e. rematerialized) rather than stored. This approach can
-    reduce memory usage at the cost of increased computation.
-
-    This function decorator produces a new version of ``fun`` which follows
-    the rematerialization strategy rather than the default store-everything
-    strategy. That is, it returns a new version of ``fun`` which, when
-    differentiated, doesn't store any of its intermediate linearization points.
-    Instead, these linearization points are recomputed from the function's saved
-    inputs.
-
-    See the examples below.
-
-    Args:
-      fun: Function for which the autodiff evaluation strategy is to be changed
-        from the default of storing all intermediate linearization points to
-        recomputing them. Its arguments and return value should be arrays,
-        scalars, or (nested) standard Python containers (tuple/list/dict) thereof.
-      prevent_cse: Optional, boolean keyword-only argument indicating whether to
-        prevent common subexpression elimination (CSE) optimizations in the HLO
-        generated from differentiation. This CSE prevention has costs because it
-        can foil other optimizations, and because it can incur high overheads on
-        some backends, especially GPU. The default is True because otherwise,
-        under a :func:`~jax.jit` or :func:`~jax.pmap`, CSE can defeat the purpose
-        of this decorator.
-        But in some settings, like when used inside a :func:`~jax.lax.scan`, this
-        CSE prevention mechanism is unnecessary, in which case ``prevent_cse`` can
-        be set to False.
-      static_argnums: Optional, int or sequence of ints, a keyword-only argument
-        indicating which argument values on which to specialize for tracing and
-        caching purposes. Specifying arguments as static can avoid
-        ConcretizationTypeErrors when tracing, but at the cost of more retracing
-        overheads. See the example below.
-      policy: Optional, callable keyword-only argument. It should be one of the
-        attributes of ``jax.checkpoint_policies``. The callable takes as input a
-        type-level specification of a first-order primitive application and
-        returns a boolean indicating whether the corresponding output value(s) can
-        be saved as residuals (or instead must be recomputed in the (co)tangent
-        computation if needed).
-
-    Returns:
-      A function (callable) with the same input/output behavior as ``fun`` but
-      which, when differentiated using e.g. :func:`jax.grad`, :func:`jax.vjp`, or
-      :func:`jax.linearize`, recomputes rather than stores intermediate
-      linearization points, thus potentially saving memory at the cost of extra
-      computation.
-
-    Here is a simple example:
-
-    >>> import jax
-    >>> import jax.numpy as jnp
-
-    >>> @jax.checkpoint
-    ... def g(x):
-    ...   y = jnp.sin(x)
-    ...   z = jnp.sin(y)
-    ...   return z
-    ...
-    >>> jax.value_and_grad(g)(2.0)
-    (Array(0.78907233, dtype=float32, weak_type=True), Array(-0.2556391, dtype=float32, weak_type=True))
-
-    Here, the same value is produced whether or not the :func:`jax.checkpoint`
-    decorator is present. When the decorator is not present, the values
-    ``jnp.cos(2.0)`` and ``jnp.cos(jnp.sin(2.0))`` are computed on the forward
-    pass and are stored for use in the backward pass, because they are needed
-    on the backward pass and depend only on the primal inputs. When using
-    :func:`jax.checkpoint`, the forward pass will compute only the primal outputs
-    and only the primal inputs (``2.0``) will be stored for the backward pass.
-    At that time, the value ``jnp.sin(2.0)`` is recomputed, along with the values
-    ``jnp.cos(2.0)`` and ``jnp.cos(jnp.sin(2.0))``.
-
-    While :func:`jax.checkpoint` controls what values are stored from the
-    forward-pass to be used on the backward pass, the total amount of memory
-    required to evaluate a function or its VJP depends on many additional internal
-    details of that function. Those details include which numerical primitives are
-    used, how they're composed, where jit and control flow primitives like scan
-    are used, and other factors.
-
-    The :func:`jax.checkpoint` decorator can be applied recursively to express
-    sophisticated autodiff rematerialization strategies. For example:
-
-    >>> def recursive_checkpoint(funs):
-    ...   if len(funs) == 1:
-    ...     return funs[0]
-    ...   elif len(funs) == 2:
-    ...     f1, f2 = funs
-    ...     return lambda x: f1(f2(x))
-    ...   else:
-    ...     f1 = recursive_checkpoint(funs[:len(funs)//2])
-    ...     f2 = recursive_checkpoint(funs[len(funs)//2:])
-    ...     return lambda x: f1(jax.checkpoint(f2)(x))
-    ...
-
-    If ``fun`` involves Python control flow that depends on argument values,
-    it may be necessary to use the ``static_argnums`` parameter. For example,
-    consider a boolean flag argument::
-
-      from functools import partial
-
-      @partial(jax.checkpoint, static_argnums=(1,))
-      def foo(x, is_training):
-        if is_training:
-          ...
-        else:
-          ...
-
-    Here, the use of ``static_argnums`` allows the ``if`` statement's condition
-    to depends on the value of ``is_training``. The cost to using
-    ``static_argnums`` is that it introduces re-tracing overheads across calls:
-    in the example, ``foo`` is re-traced every time it is called with a new value
-    of ``is_training``. In some situations, ``jax.ensure_compile_time_eval``
-    is needed as well::
-
-      @partial(jax.checkpoint, static_argnums=(1,))
-      def foo(x, y):
-        with jax.ensure_compile_time_eval():
-          y_pos = y > 0
-        if y_pos:
-          ...
-        else:
-          ...
-
-    As an alternative to using ``static_argnums`` (and
-    ``jax.ensure_compile_time_eval``), it may be easier to compute some values
-    outside the :func:`jax.checkpoint`-decorated function and then close over them.
-    """
-    if isinstance(fun, Missing):
-        return lambda f: checkpoint(f, prevent_cse=prevent_cse, policy=policy, static_argnums=static_argnums)
-
-    static_argnums = _ensure_index_tuple(tuple() if static_argnums is None else static_argnums)
-    fun = StatefulFunction(fun, static_argnums=static_argnums, name='checkpoint')
-    checkpointed_fun = jax.checkpoint(
-        fun.jaxpr_call,
-        prevent_cse=prevent_cse,
-        policy=policy,
-        static_argnums=tuple(i + 1 for i in static_argnums)
-    )
-
-    @functools.wraps(fun.fun)
-    def remat_fun(*args, **params):
-        # compile the function and get the state trace
-        state_trace = fun.compile_function_and_get_state_trace(*args, **params, return_only_write=True)
-        read_state_vals = state_trace.get_read_state_values()
-        # call the checkpointed function
-        write_state_vals, outs = checkpointed_fun(state_trace.get_state_values(), *args, **params)
-        # write the state values back to the states
-        write_back_state_values(state_trace, read_state_vals, write_state_vals)
-        return outs
-
-    return remat_fun
-
-
-remat = checkpoint

brainstate/compile/_conditions.py

@@ -1,256 +0,0 @@
-# Copyright 2024 BDP 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 collections.abc import Callable, Sequence
-
-import jax
-import jax.numpy as jnp
-import numpy as np
-
-from brainstate._compatible_import import to_concrete_aval, Tracer
-from brainstate._utils import set_module_as
-from ._error_if import jit_error_if
-from ._make_jaxpr import StatefulFunction
-from ._util import wrap_single_fun_in_multi_branches, write_back_state_values
-
-__all__ = [
-    'cond', 'switch', 'ifelse',
-]
-
-
-@set_module_as('brainstate.compile')
-def cond(pred, true_fun: Callable, false_fun: Callable, *operands):
-    """
-    Conditionally apply ``true_fun`` or ``false_fun``.
-
-    Provided arguments are correctly typed, ``cond()`` has equivalent
-    semantics to this Python implementation, where ``pred`` must be a
-    scalar type::
-
-      def cond(pred, true_fun, false_fun, *operands):
-        if pred:
-          return true_fun(*operands)
-        else:
-          return false_fun(*operands)
-
-
-    In contrast with :func:`jax.lax.select`, using ``cond`` indicates that only one of
-    the two branches is executed (up to compiler rewrites and optimizations).
-    However, when transformed with :func:`~jax.vmap` to operate over a batch of
-    predicates, ``cond`` is converted to :func:`~jax.lax.select`.
-
-    Args:
-      pred: Boolean scalar type, indicating which branch function to apply.
-      true_fun: Function (A -> B), to be applied if ``pred`` is True.
-      false_fun: Function (A -> B), to be applied if ``pred`` is False.
-      operands: Operands (A) input to either branch depending on ``pred``. The
-        type can be a scalar, array, or any pytree (nested Python tuple/list/dict)
-        thereof.
-
-    Returns:
-      Value (B) of either ``true_fun(*operands)`` or ``false_fun(*operands)``,
-      depending on the value of ``pred``. The type can be a scalar, array, or any
-      pytree (nested Python tuple/list/dict) thereof.
-    """
-    if not (callable(true_fun) and callable(false_fun)):
-        raise TypeError("true_fun and false_fun arguments should be callable.")
-
-    if pred is None:
-        raise TypeError("cond predicate is None")
-    if isinstance(pred, Sequence) or np.ndim(pred) != 0:
-        raise TypeError(f"Pred must be a scalar, got {pred} of " +
-                        (f"type {type(pred)}" if isinstance(pred, Sequence)
-                         else f"shape {np.shape(pred)}."))
-
-    # check pred
-    try:
-        pred_dtype = jax.dtypes.result_type(pred)
-    except TypeError as err:
-        raise TypeError("Pred type must be either boolean or number, got {}.".format(pred)) from err
-    if pred_dtype.kind != 'b':
-        if pred_dtype.kind in 'iuf':
-            pred = pred != 0
-        else:
-            raise TypeError("Pred type must be either boolean or number, got {}.".format(pred_dtype))
-
-    # not jit
-    if jax.config.jax_disable_jit and not isinstance(to_concrete_aval(pred), Tracer):
-        if pred:
-            return true_fun(*operands)
-        else:
-            return false_fun(*operands)
-
-    # evaluate jaxpr
-    stateful_true = StatefulFunction(true_fun, name='cond:true').make_jaxpr(*operands)
-    stateful_false = StatefulFunction(false_fun, name='conda:false').make_jaxpr(*operands)
-
-    # state trace and state values
-    state_trace = stateful_true.get_state_trace() + stateful_false.get_state_trace()
-    read_state_vals = state_trace.get_read_state_values(True)
-    write_state_vals = state_trace.get_write_state_values(True)
-
-    # wrap the functions
-    true_fun = wrap_single_fun_in_multi_branches(stateful_true, state_trace, read_state_vals, True)
-    false_fun = wrap_single_fun_in_multi_branches(stateful_false, state_trace, read_state_vals, True)
-
-    # cond
-    write_state_vals, out = jax.lax.cond(pred, true_fun, false_fun, write_state_vals, *operands)
-
-    # assign the written state values and restore the read state values
-    write_back_state_values(state_trace, read_state_vals, write_state_vals)
-    return out
-
-
-@set_module_as('brainstate.compile')
-def switch(index, branches: Sequence[Callable], *operands):
-    """
-    Apply exactly one of ``branches`` given by ``index``.
-
-    If ``index`` is out of bounds, it is clamped to within bounds.
-
-    Has the semantics of the following Python::
-
-      def switch(index, branches, *operands):
-        index = clamp(0, index, len(branches) - 1)
-        return branches[index](*operands)
-
-    Internally this wraps XLA's `Conditional
-    <https://www.tensorflow.org/xla/operation_semantics#conditional>`_
-    operator. However, when transformed with :func:`~jax.vmap` to operate over a
-    batch of predicates, ``cond`` is converted to :func:`~jax.lax.select`.
-
-    Args:
-      index: Integer scalar type, indicating which branch function to apply.
-      branches: Sequence of functions (A -> B) to be applied based on ``index``.
-      operands: Operands (A) input to whichever branch is applied.
-
-    Returns:
-      Value (B) of ``branch(*operands)`` for the branch that was selected based
-      on ``index``.
-    """
-    # check branches
-    if not all(callable(branch) for branch in branches):
-        raise TypeError("branches argument should be a sequence of callables.")
-
-    # check index
-    if len(np.shape(index)) != 0:
-        raise TypeError(f"Branch index must be scalar, got {index} of shape {np.shape(index)}.")
-    try:
-        index_dtype = jax.dtypes.result_type(index)
-    except TypeError as err:
-        msg = f"Index type must be an integer, got {index}."
-        raise TypeError(msg) from err
-    if index_dtype.kind not in 'iu':
-        raise TypeError(f"Index type must be an integer, got {index} as {index_dtype}")
-
-    # format branches
-    branches = tuple(branches)
-    if len(branches) == 0:
-        raise ValueError("Empty branch sequence")
-    elif len(branches) == 1:
-        return branches[0](*operands)
-
-    # format index
-    index = jax.lax.convert_element_type(index, np.int32)
-    lo = np.array(0, np.int32)
-    hi = np.array(len(branches) - 1, np.int32)
-    index = jax.lax.clamp(lo, index, hi)
-
-    # not jit
-    if jax.config.jax_disable_jit and isinstance(jax.core.core.get_aval(index), jax.core.ConcreteArray):
-        return branches[int(index)](*operands)
-
-    # evaluate jaxpr
-    wrapped_branches = [StatefulFunction(branch, name='switch') for branch in branches]
-    for wrapped_branch in wrapped_branches:
-        wrapped_branch.make_jaxpr(*operands)
-
-    # wrap the functions
-    state_trace = wrapped_branches[0].get_state_trace() + wrapped_branches[1].get_state_trace()
-    state_trace.merge(*[wrapped_branch.get_state_trace() for wrapped_branch in wrapped_branches[2:]])
-    read_state_vals = state_trace.get_read_state_values(True)
-    write_state_vals = state_trace.get_write_state_values(True)
-    branches = [
-        wrap_single_fun_in_multi_branches(wrapped_branch, state_trace, read_state_vals, True)
-        for wrapped_branch in wrapped_branches
-    ]
-
-    # switch
-    write_state_vals, out = jax.lax.switch(index, branches, write_state_vals, *operands)
-
-    # write back state values or restore them
-    write_back_state_values(state_trace, read_state_vals, write_state_vals)
-    return out
-
-
-@set_module_as('brainstate.compile')
-def ifelse(conditions, branches, *operands, check_cond: bool = True):
-    """
-    ``If-else`` control flows looks like native Pythonic programming.
-
-    Examples
-    --------
-
-    >>> import brainstate
-    >>> def f(a):
-    >>>    return brainstate.compile.ifelse(conditions=[a > 10, a > 5, a > 2, a > 0],
-    >>>                               branches=[lambda: 1,
-    >>>                                         lambda: 2,
-    >>>                                         lambda: 3,
-    >>>                                         lambda: 4,
-    >>>                                         lambda: 5])
-    >>> f(1)
-    4
-    >>> f(0)
-    5
-
-    Parameters
-    ----------
-    conditions: bool, sequence of bool, Array
-      The boolean conditions.
-    branches: Any
-      The branches, at least has two elements. Elements can be functions,
-      arrays, or numbers. The number of ``branches`` and ``conditions`` has
-      the relationship of `len(branches) == len(conditions) + 1`.
-      Each branch should receive one arguement for ``operands``.
-    *operands: optional, Any
-      The operands for each branch.
-    check_cond: bool
-      Whether to check the conditions. Default is True.
-
-    Returns
-    -------
-    res: Any
-      The results of the control flow.
-    """
-    # check branches
-    if not all(callable(branch) for branch in branches):
-        raise TypeError("branches argument should be a sequence of callables.")
-
-    # format branches
-    branches = tuple(branches)
-    if len(branches) == 0:
-        raise ValueError("Empty branch sequence")
-    elif len(branches) == 1:
-        return branches[0](*operands)
-    if len(conditions) != len(branches):
-        raise ValueError("The number of conditions should be equal to the number of branches.")
-
-    # format index
-    conditions = jnp.asarray(conditions, np.int32)
-    if check_cond:
-        jit_error_if(jnp.sum(conditions) != 1, "Only one condition can be True. But got {}.", err_arg=conditions)
-    index = jnp.where(conditions, size=1, fill_value=len(conditions) - 1)[0][0]
-    return switch(index, branches, *operands)