brainstate 0.0.2.post20241009__py2.py3-none-any.whl → 0.0.2.post20241010__py2.py3-none-any.whl

brainstate/random/_rand_state.py

@@ -15,6 +15,8 @@
 
 # -*- coding: utf-8 -*-
 
+from collections import namedtuple
+from functools import partial
 from operator import index
 from typing import Optional
 
@@ -23,7 +25,8 @@ import jax
 import jax.numpy as jnp
 import jax.random as jr
 import numpy as np
-from jax import lax, core
+from jax import jit, vmap
+from jax import lax, core, dtypes
 
 from brainstate import environ
 from brainstate._state import State
@@ -75,6 +78,10 @@ class RandomState(State):
   def clone(self):
     return type(self)(self.split_key())
 
+  def set_key(self, key: SeedOrKey):
+    self.value = key
+
+
   def seed(self, seed_or_key: Optional[SeedOrKey] = None):
     """Sets a new random seed.
 
@@ -492,7 +499,7 @@ class RandomState(State):
     # Get upper and lower cdf values
     sqrt2 = np.array(np.sqrt(2), dtype=dtype)
     l = self.__norm_cdf((lower - loc) / scale, sqrt2, dtype)
-    u = self.__norm_cdf((upper - loc) / scale, sqrt2, dtype)
+    u_ = self.__norm_cdf((upper - loc) / scale, sqrt2, dtype)
 
     # Uniformly fill tensor with values from [l, u], then translate to
     # [2l-1, 2u-1].
@@ -500,7 +507,7 @@ class RandomState(State):
     out = uniform_for_unit(
       key, size, dtype,
       minval=lax.nextafter(2 * l - 1, np.array(np.inf, dtype=dtype)),
-      maxval=lax.nextafter(2 * u - 1, np.array(-np.inf, dtype=dtype))
+      maxval=lax.nextafter(2 * u_- 1, np.array(-np.inf, dtype=dtype))
     )
 
     # Use inverse cdf transform for normal distribution to get truncated
@@ -617,8 +624,8 @@ class RandomState(State):
       size = jnp.shape(p)
     key = self.split_key() if key is None else _formalize_key(key)
     dtype = dtype or environ.dftype()
-    u = uniform_for_unit(key, size, dtype=dtype)
-    r = jnp.floor(jnp.log1p(-u) / jnp.log1p(-p))
+    u_ = uniform_for_unit(key, size, dtype=dtype)
+    r = jnp.floor(jnp.log1p(-u_) / jnp.log1p(-p))
     return r
 
   def _check_p2(self, p):
@@ -680,8 +687,8 @@ class RandomState(State):
       _check_shape("normal", size, mean.shape[:-1], cov.shape[:-2])
 
     if method == 'svd':
-      (u, s, _) = jnp.linalg.svd(cov)
-      factor = u * jnp.sqrt(s[..., None, :])
+      (u_, s, _) = jnp.linalg.svd(cov)
+      factor = u_ * jnp.sqrt(s[..., None, :])
     elif method == 'eigh':
       (w, v) = jnp.linalg.eigh(cov)
       factor = v * jnp.sqrt(w[..., None, :])

{brainstate-0.0.2.post20241009.dist-info → brainstate-0.0.2.post20241010.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: brainstate
-Version: 0.0.2.post20241009
+Version: 0.0.2.post20241010
 Summary: A State-based Transformation System for Brain Dynamics Programming.
 Home-page: https://github.com/brainpy/brainstate
 Author: BDP

{brainstate-0.0.2.post20241009.dist-info → brainstate-0.0.2.post20241010.dist-info}/RECORD

@@ -61,7 +61,7 @@ brainstate/optim/_sgd_optimizer.py,sha256=JiK_AVGregL0wn8uHhRQvK9Qq7Qja7dEyLW6Aa
 brainstate/random/__init__.py,sha256=c5q-RC3grRIjx-HBb2IhKZpi_xzbFmUUxzRAzqfREic,1045
 brainstate/random/_rand_funs.py,sha256=mIoENR3iEVeVR-qCQ2UQVP0SEosWry4xhzhYr0UXPAI,132072
 brainstate/random/_rand_seed.py,sha256=jRXP4zsQde1XyiOdx4aWpHXDvz5PAB1pogqQw0ywYnk,4633
-brainstate/random/_rand_state.py,sha256=D3-cvu-0mZZjnB2Wy4xmeqfm0oTvDHHPvV6SqXb7Of8,53716
+brainstate/random/_rand_state.py,sha256=LDt3p7JCEidg7tBQ8GNUJ8aAIlj-PWUh21T5-yutmlI,53887
 brainstate/random/_random_for_unit.py,sha256=Nm02GmMFiPx5LdfEDmbdK9hqLsebPdVd0bAQduGhASI,2017
 brainstate/random/random_test.py,sha256=vw45gTn-39tu2WUuamR6rr0gb2h1eeClvcaJDj16Vz8,18402
 brainstate/transform/__init__.py,sha256=hqef3a4sLQ_Oihuqs8E5IghSLr9o2bS7CWmwRL8jX6E,1887
@@ -80,8 +80,8 @@ brainstate/transform/_make_jaxpr_test.py,sha256=K3vRUBroDTCCx0lnmhgHtgrlWvWglJO2
 brainstate/transform/_mapping.py,sha256=G9XUsD1xKLCprwwE0wv0gSXS0NYZ-ZIsv-PKKRlOoTA,3821
 brainstate/transform/_progress_bar.py,sha256=VGoRZPRBmB8ELNwLc6c7S8QhUUTvn0FY46IbBm9cuYM,3502
 brainstate/transform/_unvmap.py,sha256=8Se_23QrwDdcJpFcUnnMgD6EP-4XylbhP9K5TDhW358,3311
-brainstate-0.0.2.post20241009.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
-brainstate-0.0.2.post20241009.dist-info/METADATA,sha256=2RH74ehNlwwLBu-0oRYZdpmpEVVsGKZcrAv3BHXVPcg,3311
-brainstate-0.0.2.post20241009.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-brainstate-0.0.2.post20241009.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
-brainstate-0.0.2.post20241009.dist-info/RECORD,,
+brainstate-0.0.2.post20241010.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
+brainstate-0.0.2.post20241010.dist-info/METADATA,sha256=yfpoOk0ZcT8Pd3_kC6AGm4CfGb6IrMbIgZKBgAxgBsQ,3311
+brainstate-0.0.2.post20241010.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+brainstate-0.0.2.post20241010.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
+brainstate-0.0.2.post20241010.dist-info/RECORD,,