brainstate 0.0.1__py2.py3-none-any.whl

brainstate/transform/_make_jaxpr_test.py

@@ -0,0 +1,133 @@
+# 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
+import jax.numpy as jnp
+import pytest
+
+import brainstate as bc
+
+
+class TestMakeJaxpr(unittest.TestCase):
+
+  def test_compar_jax_make_jaxpr(self):
+    def func4(arg):  # Arg is a pair
+      temp = arg[0] + jnp.sin(arg[1]) * 3.
+      c = bc.random.rand_like(arg[0])
+      return jnp.sum(temp + c)
+
+    key = bc.random.DEFAULT.value
+    jaxpr = jax.make_jaxpr(func4)((jnp.zeros(8), jnp.ones(8)))
+    print(jaxpr)
+    self.assertTrue(len(jaxpr.in_avals) == 2)
+    self.assertTrue(len(jaxpr.consts) == 1)
+    self.assertTrue(len(jaxpr.out_avals) == 1)
+    self.assertTrue(jnp.allclose(jaxpr.consts[0], key))
+
+    jaxpr2, states = bc.transform.make_jaxpr(func4)((jnp.zeros(8), jnp.ones(8)))
+    print(jaxpr2)
+    self.assertTrue(len(jaxpr2.in_avals) == 3)
+    self.assertTrue(len(jaxpr2.out_avals) == 2)
+    self.assertTrue(len(jaxpr2.consts) == 0)
+
+  def test_StatefulFunction_1(self):
+    def func4(arg):  # Arg is a pair
+      temp = arg[0] + jnp.sin(arg[1]) * 3.
+      c = bc.random.rand_like(arg[0])
+      return jnp.sum(temp + c)
+
+    fun = bc.transform.StatefulFunction(func4).make_jaxpr((jnp.zeros(8), jnp.ones(8)))
+    print(fun.get_states())
+    print(fun.get_jaxpr())
+
+  def test_StatefulFunction_2(self):
+    st1 = bc.State(jnp.ones(10))
+
+    def f1(x):
+      st1.value = x + st1.value
+
+    def f2(x):
+      jaxpr = bc.transform.make_jaxpr(f1)(x)
+      c = 1. + x
+      return c
+
+    def f3(x):
+      jaxpr = bc.transform.make_jaxpr(f1)(x)
+      c = 1.
+      return c
+
+    print()
+    jaxpr = bc.transform.make_jaxpr(f1)(jnp.zeros(1))
+    print(jaxpr)
+    jaxpr = jax.make_jaxpr(f2)(jnp.zeros(1))
+    print(jaxpr)
+    jaxpr = jax.make_jaxpr(f3)(jnp.zeros(1))
+    print(jaxpr)
+    jaxpr, _ = bc.transform.make_jaxpr(f3)(jnp.zeros(1))
+    print(jaxpr)
+    self.assertTrue(jnp.allclose(jax.core.jaxpr_as_fun(jaxpr)(jnp.zeros(1), st1.value)[0],
+                                 f3(jnp.zeros(1))))
+
+  def test_compar_jax_make_jaxpr2(self):
+    st1 = bc.State(jnp.ones(10))
+
+    def fa(x):
+      st1.value = x + st1.value
+
+    def ffa(x):
+      jaxpr, states = bc.transform.make_jaxpr(fa)(x)
+      c = 1. + x
+      return c
+
+    jaxpr, states = bc.transform.make_jaxpr(ffa)(jnp.zeros(1))
+    print()
+    print(jaxpr)
+    print(states)
+    print(jax.core.jaxpr_as_fun(jaxpr)(jnp.zeros(1), st1.value))
+    jaxpr = jax.make_jaxpr(ffa)(jnp.zeros(1))
+    print(jaxpr)
+    print(jax.core.jaxpr_as_fun(jaxpr)(jnp.zeros(1)))
+
+  def test_compar_jax_make_jaxpr3(self):
+    def fa(x):
+      return 1.
+
+    jaxpr, states = bc.transform.make_jaxpr(fa)(jnp.zeros(1))
+    print()
+    print(jaxpr)
+    print(states)
+    # print(jax.core.jaxpr_as_fun(jaxpr)(jnp.zeros(1)))
+    jaxpr = jax.make_jaxpr(fa)(jnp.zeros(1))
+    print(jaxpr)
+    # print(jax.core.jaxpr_as_fun(jaxpr)(jnp.zeros(1)))
+
+
+def test_return_states():
+  import jax.numpy
+
+  import brainstate as bc
+
+  a = bc.State(jax.numpy.ones(3))
+
+  @bc.transform.jit
+  def f():
+    return a
+
+  with pytest.raises(ValueError):
+    f()
+
+

brainstate/transform/_progress_bar.py

@@ -0,0 +1,113 @@
+# 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 copy
+from typing import Optional
+
+import jax
+
+from brainstate import environ
+
+try:
+  from tqdm.auto import tqdm
+except (ImportError, ModuleNotFoundError):
+  tqdm = None
+
+__all__ = [
+  'ProgressBar',
+]
+
+
+class ProgressBar(object):
+  __module__ = "brainstate.transform"
+
+  def __init__(self, freq: Optional[int] = None, count: Optional[int] = None, **kwargs):
+    self.print_freq = freq
+    self.print_count = count
+    if self.print_freq is not None and self.print_count is not None:
+      raise ValueError("Cannot specify both count and freq.")
+    for kwarg in ("total", "mininterval", "maxinterval", "miniters"):
+      kwargs.pop(kwarg, None)
+    self.kwargs = kwargs
+    if tqdm is None:
+      raise ImportError("tqdm is not installed.")
+
+  def init(self, n: int):
+    kwargs = copy.copy(self.kwargs)
+    freq = self.print_freq
+    count = self.print_count
+    if count is not None:
+      freq, remainder = divmod(n, count)
+      if freq == 0:
+        raise ValueError(f"Count {count} is too large for n {n}.")
+    elif freq is None:
+      if n > 20:
+        freq = int(n / 20)
+      else:
+        freq = 1
+      remainder = n % freq
+    else:
+      if freq < 1:
+        raise ValueError(f"Print rate should be > 0 got {freq}")
+      elif freq > n:
+        raise ValueError("Print rate should be less than the "
+                         f"number of steps {n}, got {freq}")
+      remainder = n % freq
+    desc = kwargs.pop("desc", f"Running for {n:,} iterations")
+    message = kwargs.pop("message", desc)
+    return ProgressBarRunner(n, message, freq, remainder, **kwargs)
+
+
+class ProgressBarRunner(object):
+  __module__ = "brainstate.transform"
+
+  def __init__(self, n: int, message, print_freq: int, remainder: int, **kwargs):
+    self.tqdm_bars = {}
+    self.kwargs = kwargs
+    self.n = n
+    self.print_freq = print_freq
+    self.remainder = remainder
+    self.message = message
+
+  def _define_tqdm(self):
+    self.tqdm_bars[0] = tqdm(range(self.n), **self.kwargs)
+    self.tqdm_bars[0].set_description(self.message, refresh=False)
+
+  def _update_tqdm(self):
+    self.tqdm_bars[0].update(self.print_freq)
+
+  def _close_tqdm(self):
+    if self.remainder > 0:
+      self.tqdm_bars[0].update(self.remainder)
+    self.tqdm_bars[0].close()
+
+  def __call__(self, iter_num, *args, **kwargs):
+
+    _ = jax.lax.cond(
+      iter_num == 0,
+      lambda: jax.debug.callback(self._define_tqdm),
+      lambda: None,
+    )
+    _ = jax.lax.cond(
+      (iter_num + 1) % self.print_freq == 0,
+      lambda: jax.debug.callback(self._update_tqdm),
+      lambda: None,
+    )
+    _ = jax.lax.cond(
+      iter_num == self.n - 1,
+      lambda: jax.debug.callback(self._close_tqdm),
+      lambda: None,
+    )

brainstate/typing.py

@@ -0,0 +1,69 @@
+# 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 typing import Any, Sequence, Protocol, Union
+
+import jax
+import numpy as np
+
+__all__ = [
+  'Size',
+  'Axes',
+  'SeedOrKey',
+  'ArrayLike',
+  'DType',
+  'DTypeLike',
+]
+
+Size = Union[int, Sequence[int]]
+Axes = Union[int, Sequence[int]]
+SeedOrKey = Union[int, jax.Array, np.ndarray]
+
+# --- Array --- #
+
+# ArrayLike is a Union of all objects that can be implicitly converted to a
+# standard JAX array (i.e. not including future non-standard array types like
+# KeyArray and BInt). It's different than np.typing.ArrayLike in that it doesn't
+# accept arbitrary sequences, nor does it accept string data.
+ArrayLike = Union[
+  jax.Array,  # JAX array type
+  np.ndarray,  # NumPy array type
+  np.bool_, np.number,  # NumPy scalar types
+  bool, int, float, complex,  # Python scalar types
+]
+
+# --- Dtype --- #
+
+
+DType = np.dtype
+
+
+class SupportsDType(Protocol):
+  @property
+  def dtype(self) -> DType: ...
+
+
+# DTypeLike is meant to annotate inputs to np.dtype that return
+# a valid JAX dtype. It's different than numpy.typing.DTypeLike
+# because JAX doesn't support objects or structured dtypes.
+# Unlike np.typing.DTypeLike, we exclude None, and instead require
+# explicit annotations when None is acceptable.
+DTypeLike = Union[
+  str,  # like 'float32', 'int32'
+  type[Any],  # like np.float32, np.int32, float, int
+  np.dtype,  # like np.dtype('float32'), np.dtype('int32')
+  SupportsDType,  # like jnp.float32, jnp.int32
+]