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

brainstate/transform/_jit.py

@@ -1,265 +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 __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 ._make_jaxpr import StatefulFunction, _ensure_index_tuple, _assign_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
-
-  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)
-    states = fun.compile_and_get_states_by_static_args(*args, **params)
-    state_vals, outs = jit_fun([st.value for st in states], *args, **params)
-    _assign_state_values(states, state_vals)
-    return outs
-
-  def clear_cache():
-    # clear the cache of the stateful function
-    fun.clear_cache()
-    # clear the cache of the jitted function
-    jit_fun.clear_cache()
-
-  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
-
-  return jitted_fun
-
-
-@set_module_as('brainstate.transform')
-def jit(
-    fun: Callable = None,
-    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 fun is None:
-    def wrapper(fun_again: Callable) -> JittedFunction:
-      return _get_jitted_fun(fun_again,
-                             in_shardings,
-                             out_shardings,
-                             static_argnums,
-                             donate_argnums,
-                             donate_argnames,
-                             keep_unused,
-                             device,
-                             backend,
-                             inline,
-                             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/transform/_jit_test.py

@@ -1,118 +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 unittest
-
-import jax.numpy as jnp
-
-import brainstate as bc
-
-
-class TestJIT(unittest.TestCase):
-  def test_inner_state_are_not_catched(self):
-    a = bc.State(bc.random.randn(10))
-
-    @bc.transform.jit
-    def fun1(inp):
-      a.value += inp
-
-      b = bc.State(bc.random.randn(1))
-
-      def inner_fun(x):
-        b.value += x
-
-      bc.transform.for_loop(inner_fun, bc.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 = bc.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 = bc.State(bc.random.randn(10))
-
-    @bc.transform.jit
-    def fun1(inp):
-      a.value += inp
-
-      b = bc.State(bc.random.randn(1))
-
-      def inner_fun(x):
-        b.value += x
-
-      bc.transform.for_loop(inner_fun, bc.random.randn(100))
-
-      return a.value + b.value
-
-    # test kwargs
-    print(fun1(inp=bc.random.randn(10)))
-
-  def test_jit_compile_sensitive_to_input_shape(self):
-    global_data = [0]
-
-    @bc.transform.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(bc.random.randn(10)))
-    self.assertTrue(global_data[0] == 2)
-
-    print(fun1(bc.random.randn(10, 10)))
-    self.assertTrue(global_data[0] == 3)
-
-  def test_jit_clear_cache(self):
-    a = bc.State(bc.random.randn(1))
-    compiling = []
-
-    @bc.transform.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 = bc.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 = bc.transform.jit(fun1)
-    self.assertTrue(jitted_fun.origin_fun is fun1)
-    self.assertTrue(isinstance(jitted_fun.stateful_fun, bc.transform.StatefulFunction))
-    self.assertTrue(callable(jitted_fun.jitted_fun))
-    self.assertTrue(callable(jitted_fun.clear_cache))