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

brainstate/transform/_make_jaxpr_test.py

@@ -1,131 +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
-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/_mapping.py

@@ -1,109 +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 jax
-
-from ._loop_collect_return import scan
-
-__all__ = [
-  'map',
-]
-
-
-def _batch_and_remainder(x, batch_size: int):
-  leaves, treedef = jax.tree.flatten(x)
-
-  scan_leaves = []
-  remainder_leaves = []
-
-  for leaf in leaves:
-    num_batches, _ = divmod(leaf.shape[0], batch_size)
-    total_batch_elems = num_batches * batch_size
-    scan_leaves.append(leaf[:total_batch_elems].reshape(num_batches, batch_size, *leaf.shape[1:]))
-    remainder_leaves.append(leaf[total_batch_elems:])
-
-  scan_tree = treedef.unflatten(scan_leaves)
-  remainder_tree = treedef.unflatten(remainder_leaves)
-  return scan_tree, remainder_tree
-
-
-def map(
-    f,
-    xs,
-    *,
-    batch_size: int | None = None,
-):
-  """Map a function over leading array axes.
-
-  Like Python's builtin map, except inputs and outputs are in the form of
-  stacked arrays. Consider using the :func:`~jax.vmap` transform instead, unless you
-  need to apply a function element by element for reduced memory usage or
-  heterogeneous computation with other control flow primitives.
-
-  When ``xs`` is an array type, the semantics of :func:`~map` are given by this
-  Python implementation::
-
-    def map(f, xs):
-      return np.stack([f(x) for x in xs])
-
-  Like :func:`~scan`, :func:`~map` is implemented in terms of JAX primitives so
-  many of the same advantages over a Python loop apply: ``xs`` may be an
-  arbitrary nested pytree type, and the mapped computation is compiled only
-  once.
-
-  If ``batch_size`` is provided, the computation is executed in batches of that size
-  and parallelized using :func:`~jax.vmap`. This can be used as either a more performant
-  version of ``map`` or as a memory-efficient version of ``vmap``. If the axis is not
-  divisible by the batch size, the remainder is processed in a separate ``vmap`` and
-  concatenated to the result.
-
-    >>> x = jax.numpy.ones((10, 3, 4))
-    >>> def f(x):
-    ...   print('inner shape:', x.shape)
-    ...   return x + 1
-    >>> y = map(f, x, batch_size=3)
-    inner shape: (3, 4)
-    inner shape: (3, 4)
-    >>> y.shape
-    (10, 3, 4)
-
-  In the example above, "inner shape" is printed twice, once while tracing the batched
-  computation and once while tracing the remainder computation.
-
-  Args:
-    f: a Python function to apply element-wise over the first axis or axes of
-      ``xs``.
-    xs: values over which to map along the leading axis.
-    batch_size: (optional) integer specifying the size of the batch for each step to execute
-      in parallel.
-
-  Returns:
-    Mapped values.
-  """
-  if batch_size is not None:
-    scan_xs, remainder_xs = _batch_and_remainder(xs, batch_size)
-    g = lambda _, x: ((), jax.vmap(f)(x))
-    _, scan_ys = scan(g, (), scan_xs)
-    remainder_ys = jax.vmap(f)(remainder_xs)
-    flatten = lambda x: x.reshape(-1, *x.shape[2:])
-    ys = jax.tree.map(
-      lambda x, y: jax.numpy.concatenate([flatten(x), y], axis=0), scan_ys, remainder_ys,
-    )
-  else:
-    g = lambda _, x: ((), f(x))
-    _, ys = scan(g, (), xs)
-  return ys

brainstate/transform/_progress_bar.py

@@ -1,111 +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 copy
-from typing import Optional
-
-import jax
-
-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/transform/_unvmap.py

@@ -1,143 +0,0 @@
-import jax
-import jax.core
-import jax.interpreters.batching as batching
-import jax.interpreters.mlir as mlir
-import jax.numpy as jnp
-
-from brainstate._utils import set_module_as
-
-__all__ = [
-  "unvmap",
-]
-
-
-@set_module_as('brainstate.transform')
-def unvmap(x, op: str = 'any'):
-  if op == 'all':
-    return unvmap_all(x)
-  elif op == 'any':
-    return unvmap_any(x)
-  elif op == 'none':
-    return _without_vmap(x)
-  elif op == 'max':
-    return unvmap_max(x)
-  else:
-    raise ValueError(f'Do not support type: {op}')
-
-
-# unvmap_all
-
-unvmap_all_p = jax.core.Primitive("unvmap_all")
-
-
-def unvmap_all(x):
-  """As `jnp.all`, but ignores batch dimensions."""
-  return unvmap_all_p.bind(x)
-
-
-def _unvmap_all_impl(x):
-  return jnp.all(x)
-
-
-def _unvmap_all_abstract_eval(x):
-  return jax.core.ShapedArray(shape=(), dtype=jax.numpy.bool_.dtype)  # pyright: ignore
-
-
-def _unvmap_all_batch(x, batch_axes):
-  (x,) = x
-  return unvmap_all(x), batching.not_mapped
-
-
-unvmap_all_p.def_impl(_unvmap_all_impl)
-unvmap_all_p.def_abstract_eval(_unvmap_all_abstract_eval)
-batching.primitive_batchers[unvmap_all_p] = _unvmap_all_batch  # pyright: ignore
-mlir.register_lowering(
-  unvmap_all_p,
-  mlir.lower_fun(_unvmap_all_impl, multiple_results=False),
-)
-
-# unvmap_any
-
-unvmap_any_p = jax.core.Primitive("unvmap_any")
-
-
-def unvmap_any(x):
-  """As `jnp.any`, but ignores batch dimensions."""
-  return unvmap_any_p.bind(x)
-
-
-def _unvmap_any_impl(x):
-  return jnp.any(x)
-
-
-def _unvmap_any_abstract_eval(x):
-  return jax.core.ShapedArray(shape=(), dtype=jax.numpy.bool_.dtype)  # pyright: ignore
-
-
-def _unvmap_any_batch(x, batch_axes):
-  (x,) = x
-  return unvmap_any(x), batching.not_mapped
-
-
-unvmap_any_p.def_impl(_unvmap_any_impl)
-unvmap_any_p.def_abstract_eval(_unvmap_any_abstract_eval)
-batching.primitive_batchers[unvmap_any_p] = _unvmap_any_batch  # pyright: ignore
-mlir.register_lowering(
-  unvmap_any_p,
-  mlir.lower_fun(_unvmap_any_impl, multiple_results=False),
-)
-
-# unvmap_max
-
-unvmap_max_p = jax.core.Primitive("unvmap_max")
-
-
-def unvmap_max(x):
-  """As `jnp.max`, but ignores batch dimensions."""
-  return unvmap_max_p.bind(x)
-
-
-def _unvmap_max_impl(x):
-  return jnp.max(x)
-
-
-def _unvmap_max_abstract_eval(x):
-  return jax.core.ShapedArray(shape=(), dtype=x.dtype)
-
-
-def _unvmap_max_batch(x, batch_axes):
-  (x,) = x
-  return unvmap_max(x), batching.not_mapped
-
-
-unvmap_max_p.def_impl(_unvmap_max_impl)
-unvmap_max_p.def_abstract_eval(_unvmap_max_abstract_eval)
-batching.primitive_batchers[unvmap_max_p] = _unvmap_max_batch  # pyright: ignore
-mlir.register_lowering(
-  unvmap_max_p,
-  mlir.lower_fun(_unvmap_max_impl, multiple_results=False),
-)
-
-
-def _without_vmap(x):
-  return _no_vmap_prim.bind(x)
-
-
-def _without_vmap_imp(x):
-  return x
-
-
-def _without_vmap_abs(x):
-  return x
-
-
-def _without_vmap_batch(x, batch_axes):
-  (x,) = x
-  return _without_vmap(x), batching.not_mapped
-
-
-_no_vmap_prim = jax.core.Primitive('no_vmap')
-_no_vmap_prim.def_impl(_without_vmap_imp)
-_no_vmap_prim.def_abstract_eval(_without_vmap_abs)
-batching.primitive_batchers[_no_vmap_prim] = _without_vmap_batch
-mlir.register_lowering(_no_vmap_prim, mlir.lower_fun(_without_vmap_imp, multiple_results=False))