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

brainstate/{compile → transform}/_jit.py

@@ -1,4 +1,4 @@
-# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
+# Copyright 2024 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.
@@ -24,7 +24,6 @@ from brainstate._compatible_import import Device
 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']
 
@@ -33,6 +32,8 @@ class JittedFunction(Callable):
     """
     A wrapped version of ``fun``, set up for just-in-time compilation.
     """
+    __module__ = 'brainstate.transform'
+
     origin_fun: Callable  # the original function
     stateful_fun: StatefulFunction  # the stateful function for extracting states
     jitted_fun: jax.stages.Wrapped  # the jitted function
@@ -67,8 +68,8 @@ def _get_jitted_fun(
         static_argnums=static_argnums,
         static_argnames=static_argnames,
         abstracted_axes=abstracted_axes,
-        cache_type='jit',
-        name='jit'
+        name='jit',
+        return_only_write=True
     )
     jit_fun = jax.jit(
         fun.jaxpr_call,
@@ -92,14 +93,14 @@ def _get_jitted_fun(
             return fun.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)
+        state_trace = fun.get_state_trace(*args, **params, compile_if_miss=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)
+        state_trace.assign_state_vals_v2(read_state_vals, write_state_vals)
         return outs
 
     def clear_cache():
@@ -139,7 +140,7 @@ def _get_jitted_fun(
           A ``Lowered`` instance representing the lowering.
         """
         # compile the function and get the state trace
-        state_trace = fun.compile_function_and_get_state_trace(*args, **params, return_only_write=True)
+        state_trace = fun.get_state_trace(*args, **params, compile_if_miss=True)
         read_state_vals = state_trace.get_read_state_values(replace_writen=True)
         write_state_vals = state_trace.get_write_state_values(replace_read=True)
 
@@ -147,7 +148,7 @@ def _get_jitted_fun(
         ret = jit_fun.lower(state_trace.get_state_values(), *args, **params).compile()
 
         # write the state values back to the states
-        write_back_state_values(state_trace, read_state_vals, write_state_vals)
+        state_trace.assign_state_vals_v2(read_state_vals, write_state_vals)
         return ret
 
     jitted_fun: JittedFunction
@@ -176,7 +177,7 @@ def _get_jitted_fun(
     return jitted_fun
 
 
-@set_module_as('brainstate.compile')
+@set_module_as('brainstate.transform')
 def jit(
     fun: Callable | Missing = Missing(),
     in_shardings=sharding_impls.UNSPECIFIED,
@@ -195,9 +196,12 @@ def jit(
     """
     Sets up ``fun`` for just-in-time compilation with XLA.
 
-    Args:
-      fun: Function to be jitted.
-      in_shardings: Pytree of structure matching that of arguments to ``fun``,
+    Parameters
+    ----------
+    fun : callable or Missing, optional
+        Function to be jitted.
+    in_shardings : pytree, optional
+        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
@@ -208,26 +212,29 @@ def jit(
         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.
+
+        - :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
+    out_shardings : pytree, optional
+        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
+    static_argnums : int or sequence of int, optional
+        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
@@ -248,12 +255,8 @@ def jit(
         provided, ``inspect.signature`` is not used, and only actual
         parameters listed in either ``static_argnums`` or ``static_argnames`` will
         be treated as static.
-      static_argnames: An optional string or collection of strings specifying
-        which named arguments are treated 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.
-      donate_argnums: Specify which positional argument buffers are "donated" to
+    donate_argnums : int or sequence of int, optional
+        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
@@ -274,38 +277,88 @@ def jit(
 
         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
+    static_argnames : str or sequence of str, optional
+        An optional string or collection of strings specifying
+        which named arguments are treated 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.
+    donate_argnames : str or iterable of str, optional
+        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
+    keep_unused : bool, default False
+        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.
+    device : Device, optional
+        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.
+    backend : str, optional
+        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
+    inline : bool, default False
+        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
-
+    abstracted_axes : Any, optional
+        Abstracted axes specification.
+    **kwargs
+        Additional keyword arguments passed to the underlying JAX jit function.
+
+    Returns
+    -------
+    JittedFunction or callable
+        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
+
+    Examples
+    --------
+    Basic usage with a simple function:
+
+    .. code-block:: python
+
+        >>> import brainstate
+        >>> import jax.numpy as jnp
+        >>>
+        >>> @brainstate.transform.jit
+        ... def f(x):
+        ...     return x ** 2
+        >>>
+        >>> result = f(jnp.array([1, 2, 3]))
+
+    Using static arguments:
+
+    .. code-block:: python
+
+        >>> @brainstate.transform.jit(static_argnums=(1,))
+        ... def g(x, n):
+        ...     return x ** n
+        >>>
+        >>> result = g(jnp.array([1, 2, 3]), 2)
+
+    Manual jitting:
+
+    .. code-block:: python
+
+        >>> def h(x):
+        ...     return x * 2
+        >>>
+        >>> jitted_h = brainstate.transform.jit(h)
+        >>> result = jitted_h(jnp.array([1, 2, 3]))
     """
 
     if isinstance(fun, Missing):

brainstate/{compile → transform}/_jit_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
+# Copyright 2024 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.
@@ -41,12 +41,12 @@ class TestJIT(unittest.TestCase):
 
         print(fun1(1.))
         key = fun1.stateful_fun.get_arg_cache_key(1.)
-        self.assertTrue(len(fun1.stateful_fun.get_states(key)) == 2)
+        self.assertTrue(len(fun1.stateful_fun.get_states_by_cache(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)
+        self.assertTrue(len(fun1.stateful_fun.get_states_by_cache(key)) == 2)
 
     def test_kwargs(self):
         a = bst.State(bst.random.randn(10))