brainstate 0.0.2.post20241010__py2.py3-none-any.whl → 0.1.0__py2.py3-none-any.whl

brainstate/compile/_jit.py

@@ -0,0 +1,314 @@
+# 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 __future__ import annotations
+
+import functools
+from collections.abc import Iterable, Sequence
+from typing import (Any, Callable, Union)
+
+import jax
+from jax._src import sharding_impls
+from jax.lib import xla_client as xc
+
+from brainstate._utils import set_module_as
+from brainstate.typing import Missing
+from ._make_jaxpr import StatefulFunction, _ensure_index_tuple
+from ._util import write_back_state_values
+
+__all__ = ['jit']
+
+
+class JittedFunction(Callable):
+    """
+    A wrapped version of ``fun``, set up for just-in-time compilation.
+    """
+    origin_fun: Callable  # the original function
+    stateful_fun: StatefulFunction  # the stateful function for extracting states
+    jitted_fun: jax.stages.Wrapped  # the jitted function
+    clear_cache: Callable  # clear the cache of the jitted function
+    eval_shape: Callable  # evaluate the shape of the jitted function
+    lower: Callable  # lower the jitted function
+    trace: Callable  # trace the jitted
+
+    def __call__(self, *args, **kwargs):
+        pass
+
+
+def _get_jitted_fun(
+    fun: Callable,
+    in_shardings,
+    out_shardings,
+    static_argnums,
+    donate_argnums,
+    donate_argnames,
+    keep_unused,
+    device,
+    backend,
+    inline,
+    abstracted_axes,
+    **kwargs
+) -> JittedFunction:
+    static_argnums = _ensure_index_tuple(tuple() if static_argnums is None else static_argnums)
+    # TODO: add to cache stack for clear_cache
+    fun = StatefulFunction(fun, static_argnums=static_argnums, abstracted_axes=abstracted_axes, cache_type='jit')
+    jit_fun = jax.jit(fun.jaxpr_call,
+                      static_argnums=tuple(i + 1 for i in static_argnums),
+                      donate_argnums=donate_argnums,
+                      donate_argnames=donate_argnames,
+                      keep_unused=keep_unused,
+                      device=device,
+                      backend=backend,
+                      inline=inline,
+                      in_shardings=in_shardings,
+                      out_shardings=out_shardings,
+                      abstracted_axes=abstracted_axes,
+                      **kwargs)
+
+    @functools.wraps(fun.fun)
+    def jitted_fun(*args, **params):
+        if jax.config.jax_disable_jit:
+            return fun.fun(*args, **params)
+
+        # compile the function and get the state trace
+        with jax.ensure_compile_time_eval():
+            state_trace = fun.compile_function_and_get_state_trace(*args, **params, return_only_write=True)
+            read_state_vals = state_trace.get_read_state_values(True)
+        # call the jitted function
+        write_state_vals, outs = jit_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
+
+    def clear_cache():
+        """
+        Clear the cache of the jitted function.
+        """
+        # clear the cache of the stateful function
+        fun.clear_cache()
+        # clear the cache of the jitted function
+        jit_fun.clear_cache()
+
+    def eval_shape():
+        raise NotImplementedError
+
+    def lower():
+        """Lower this function explicitly for the given arguments.
+
+        A lowered function is staged out of Python and translated to a
+        compiler's input language, possibly in a backend-dependent
+        manner. It is ready for compilation but not yet compiled.
+
+        Returns:
+          A ``Lowered`` instance representing the lowering.
+        """
+        raise NotImplementedError
+
+    def trace():
+        """Trace this function explicitly for the given arguments.
+
+        A traced function is staged out of Python and translated to a jaxpr. It is
+        ready for lowering but not yet lowered.
+
+        Returns:
+          A ``Traced`` instance representing the tracing.
+        """
+        raise NotImplementedError
+
+    jitted_fun: JittedFunction
+
+    # the original function
+    jitted_fun.origin_fun = fun.fun
+
+    # the stateful function for extracting states
+    jitted_fun.stateful_fun = fun
+
+    # the jitted function
+    jitted_fun.jitted_fun = jit_fun
+
+    # clear cache
+    jitted_fun.clear_cache = clear_cache
+
+    # evaluate the shape of the jitted function
+    jitted_fun.eval_shape = eval_shape
+
+    # lower the jitted function
+    jitted_fun.lower = lower
+
+    # trace the jitted
+    jitted_fun.trace = trace
+
+    return jitted_fun
+
+
+@set_module_as('brainstate.compile')
+def jit(
+    fun: Callable | Missing = Missing(),
+    in_shardings=sharding_impls.UNSPECIFIED,
+    out_shardings=sharding_impls.UNSPECIFIED,
+    static_argnums: int | Sequence[int] | None = None,
+    donate_argnums: int | Sequence[int] | None = None,
+    donate_argnames: str | Iterable[str] | None = None,
+    keep_unused: bool = False,
+    device: xc.Device | None = None,
+    backend: str | None = None,
+    inline: bool = False,
+    abstracted_axes: Any | None = None,
+    **kwargs
+) -> Union[JittedFunction, Callable[[Callable], JittedFunction]]:
+    """
+    Sets up ``fun`` for just-in-time compilation with XLA.
+
+    Does not support setting ``static_argnames`` as in ``jax.jit()``.
+
+
+    Args:
+      fun: Function to be jitted.
+      in_shardings: Pytree of structure matching that of arguments to ``fun``,
+        with all actual arguments replaced by resource assignment specifications.
+        It is also valid to specify a pytree prefix (e.g. one value in place of a
+        whole subtree), in which case the leaves get broadcast to all values in
+        that subtree.
+
+        The ``in_shardings`` argument is optional. JAX will infer the shardings
+        from the input :py:class:`jax.Array`'s and defaults to replicating the input
+        if the sharding cannot be inferred.
+
+        The valid resource assignment specifications are:
+          - :py:class:`XLACompatibleSharding`, which will decide how the value
+              will be partitioned. With this, using a mesh context manager is not
+              required.
+          - :py:obj:`None`, will give JAX the freedom to choose whatever sharding
+            it wants.
+            For in_shardings, JAX will mark is as replicated but this behavior
+            can change in the future.
+            For out_shardings, we will rely on the XLA GSPMD partitioner to
+            determine the output shardings.
+
+        The size of every dimension has to be a multiple of the total number of
+        resources assigned to it. This is similar to pjit's in_shardings.
+      out_shardings: Like ``in_shardings``, but specifies resource
+        assignment for function outputs. This is similar to pjit's
+        out_shardings.
+
+        The ``out_shardings`` argument is optional. If not specified, :py:func:`jax.jit`
+        will use GSPMD's sharding propagation to figure out what the sharding of the
+        output(s) should be.
+      static_argnums: An optional int or collection of ints that specify which
+        positional arguments to treat as static (compile-time constant).
+        Operations that only depend on static arguments will be constant-folded in
+        Python (during tracing), and so the corresponding argument values can be
+        any Python object.
+
+        Static arguments should be hashable, meaning both ``__hash__`` and
+        ``__eq__`` are implemented, and immutable. Calling the jitted function
+        with different values for these constants will trigger recompilation.
+        Arguments that are not arrays or containers thereof must be marked as
+        static.
+
+        If neither ``static_argnums`` nor ``static_argnames`` is provided, no
+        arguments are treated as static. If ``static_argnums`` is not provided but
+        ``static_argnames`` is, or vice versa, JAX uses
+        :code:`inspect.signature(fun)` to find any positional arguments that
+        correspond to ``static_argnames``
+        (or vice versa). If both ``static_argnums`` and ``static_argnames`` are
+        provided, ``inspect.signature`` is not used, and only actual
+        parameters listed in either ``static_argnums`` or ``static_argnames`` will
+        be treated as static.
+      donate_argnums: Specify which positional argument buffers are "donated" to
+        the computation. It is safe to donate argument buffers if you no longer
+        need them once the computation has finished. In some cases XLA can make
+        use of donated buffers to reduce the amount of memory needed to perform a
+        computation, for example recycling one of your input buffers to store a
+        result. You should not reuse buffers that you donate to a computation, JAX
+        will raise an error if you try to. By default, no argument buffers are
+        donated.
+
+        If neither ``donate_argnums`` nor ``donate_argnames`` is provided, no
+        arguments are donated. If ``donate_argnums`` is not provided but
+        ``donate_argnames`` is, or vice versa, JAX uses
+        :code:`inspect.signature(fun)` to find any positional arguments that
+        correspond to ``donate_argnames``
+        (or vice versa). If both ``donate_argnums`` and ``donate_argnames`` are
+        provided, ``inspect.signature`` is not used, and only actual
+        parameters listed in either ``donate_argnums`` or ``donate_argnames`` will
+        be donated.
+
+        For more details on buffer donation see the
+        `FAQ <https://jax.readthedocs.io/en/latest/faq.html#buffer-donation>`_.
+      donate_argnames: An optional string or collection of strings specifying
+        which named arguments are donated to the computation. See the
+        comment on ``donate_argnums`` for details. If not
+        provided but ``donate_argnums`` is set, the default is based on calling
+        ``inspect.signature(fun)`` to find corresponding named arguments.
+      keep_unused: If `False` (the default), arguments that JAX determines to be
+        unused by `fun` *may* be dropped from resulting compiled XLA executables.
+        Such arguments will not be transferred to the device nor provided to the
+        underlying executable. If `True`, unused arguments will not be pruned.
+      device: This is an experimental feature and the API is likely to change.
+        Optional, the Device the jitted function will run on. (Available devices
+        can be retrieved via :py:func:`jax.devices`.) The default is inherited
+        from XLA's DeviceAssignment logic and is usually to use
+        ``jax.devices()[0]``.
+      backend: This is an experimental feature and the API is likely to change.
+        Optional, a string representing the XLA backend: ``'cpu'``, ``'gpu'``, or
+        ``'tpu'``.
+      inline: Specify whether this function should be inlined into enclosing
+        jaxprs (rather than being represented as an application of the xla_call
+        primitive with its own subjaxpr). Default False.
+      abstracted_axes:
+
+    Returns:
+      A wrapped version of ``fun``, set up for just-in-time compilation.
+      The returned object is a :py:class:`JittedFunction` that can be called with the same arguments
+      and has the following attributes and methods:
+
+      - ``stateful_fun``: the stateful function for extracting states, an instance of :py:class:`StatefulFunction`.
+      - ``origin_fun(*args, **kwargs)``: the original function
+      - ``jitted_fun(*args, **kwargs)``: the jitted function
+      - ``clear_cache(*args, **kwargs)``: clear the cache of the jitted function
+
+    """
+
+    if isinstance(fun, Missing):
+        def wrapper(fun_again: Callable) -> JittedFunction:
+            return _get_jitted_fun(fun_again,
+                                   in_shardings=in_shardings,
+                                   out_shardings=out_shardings,
+                                   static_argnums=static_argnums,
+                                   donate_argnums=donate_argnums,
+                                   donate_argnames=donate_argnames,
+                                   keep_unused=keep_unused,
+                                   device=device,
+                                   backend=backend,
+                                   inline=inline,
+                                   abstracted_axes=abstracted_axes,
+                                   **kwargs)
+
+        return wrapper
+
+    else:
+        return _get_jitted_fun(fun,
+                               in_shardings,
+                               out_shardings,
+                               static_argnums,
+                               donate_argnums,
+                               donate_argnames,
+                               keep_unused,
+                               device,
+                               backend,
+                               inline,
+                               abstracted_axes,
+                               **kwargs)

brainstate/compile/_jit_test.py

@@ -0,0 +1,143 @@
+# 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 __future__ import annotations
+
+import unittest
+
+import jax.numpy as jnp
+
+import brainstate as bst
+
+
+class TestJIT(unittest.TestCase):
+    def test_inner_state_are_not_catched(self):
+        a = bst.State(bst.random.randn(10))
+
+        @bst.compile.jit
+        def fun1(inp):
+            a.value += inp
+
+            b = bst.State(bst.random.randn(1))
+
+            def inner_fun(x):
+                b.value += x
+
+            bst.compile.for_loop(inner_fun, bst.random.randn(100))
+
+            return a.value + b.value
+
+        print(fun1(1.))
+        key = fun1.stateful_fun.get_arg_cache_key(1.)
+        self.assertTrue(len(fun1.stateful_fun.get_states(key)) == 2)
+
+        x = bst.random.randn(10)
+        print(fun1(x))
+        key = fun1.stateful_fun.get_arg_cache_key(x)
+        self.assertTrue(len(fun1.stateful_fun.get_states(key)) == 2)
+
+    def test_kwargs(self):
+        a = bst.State(bst.random.randn(10))
+
+        @bst.compile.jit
+        def fun1(inp):
+            a.value += inp
+
+            b = bst.State(bst.random.randn(1))
+
+            def inner_fun(x):
+                b.value += x
+
+            bst.compile.for_loop(inner_fun, bst.random.randn(100))
+
+            return a.value + b.value
+
+        # test kwargs
+        print(fun1(inp=bst.random.randn(10)))
+
+    def test_jit_compile_sensitive_to_input_shape(self):
+        global_data = [0]
+
+        @bst.compile.jit
+        def fun1(inp):
+            global_data[0] += 1
+            return inp
+
+        print(fun1(1.))
+        self.assertTrue(global_data[0] == 1)
+
+        print(fun1(2.))
+        self.assertTrue(global_data[0] == 1)
+
+        print(fun1(bst.random.randn(10)))
+        self.assertTrue(global_data[0] == 2)
+
+        print(fun1(bst.random.randn(10, 10)))
+        self.assertTrue(global_data[0] == 3)
+
+    def test_jit_clear_cache(self):
+        a = bst.State(bst.random.randn(1))
+        compiling = []
+
+        @bst.compile.jit
+        def log2(x):
+            print('compiling')
+            compiling.append(1)
+            ln_x = jnp.log(x)
+            ln_2 = jnp.log(2.0) + a.value
+            return ln_x / ln_2
+
+        x = bst.random.randn(1)
+        print(log2(x))  # compiling
+        self.assertTrue(len(compiling) == 1)
+        print(log2(x))  # no compiling
+        self.assertTrue(len(compiling) == 1)
+
+        log2.clear_cache()
+        print(log2(x))  # compiling
+        self.assertTrue(len(compiling) == 2)
+
+    def test_jit_attribute_origin_fun(self):
+        def fun1(x):
+            return x
+
+        jitted_fun = bst.compile.jit(fun1)
+        self.assertTrue(jitted_fun.origin_fun is fun1)
+        self.assertTrue(isinstance(jitted_fun.stateful_fun, bst.compile.StatefulFunction))
+        self.assertTrue(callable(jitted_fun.jitted_fun))
+        self.assertTrue(callable(jitted_fun.clear_cache))
+
+    def test_clear_cache(self):
+        a = bst.State(bst.random.randn(1))
+
+        @bst.compile.jit
+        def f_jit(x, y):
+            print('Compiling')
+            a.value = jnp.sin(x) + jnp.cos(y)
+
+        f_jit(0.5, 1.0)
+        f_jit.clear_cache()
+        f_jit(0.5, 1.0)
+
+    def test_cache(self):
+        @bst.compile.jit
+        @bst.compile.jit
+        @bst.compile.jit
+        def f(a):
+            print('Compiling')
+            print(a)
+            return a + 1
+
+        print(f(1.))