brainstate 0.0.1.1.post20240804__py2.py3-none-any.whl → 0.0.2.post20240814__py2.py3-none-any.whl

brainstate/__init__.py

@@ -17,7 +17,7 @@
 A ``State``-based Transformation System for Brain Dynamics Programming
 """
 
-__version__ = "0.0.1.1"
+__version__ = "0.0.2"
 
 from . import environ
 from . import functional

brainstate/_random_for_unit.py

@@ -30,13 +30,14 @@ def uniform_for_unit(
     maxval: ArrayLike = 1.
 ) -> jax.Array | bu.Quantity:
   if isinstance(minval, bu.Quantity) and isinstance(maxval, bu.Quantity):
-    return bu.Quantity(jr.uniform(key, shape, dtype, minval.value, maxval.value), dim=minval.dim)
+    maxval = maxval.in_unit(minval.unit)
+    return bu.Quantity(jr.uniform(key, shape, dtype, minval.mantissa, maxval.mantissa), unit=minval.unit)
   elif isinstance(minval, bu.Quantity):
     assert minval.is_unitless, f'minval must be unitless when maxval is not a Quantity, got {minval}'
-    minval = minval.value
+    minval = minval.mantissa
   elif isinstance(maxval, bu.Quantity):
     assert maxval.is_unitless, f'maxval must be unitless when minval is not a Quantity, got {maxval}'
-    maxval = maxval.value
+    maxval = maxval.mantissa
   return jr.uniform(key, shape, dtype, minval, maxval)
 
 
@@ -47,5 +48,5 @@ def permutation_for_unit(
     independent: bool = False
 ) -> jax.Array | bu.Quantity:
   if isinstance(x, bu.Quantity):
-    return bu.Quantity(jr.permutation(key, x.value, axis, independent=independent), dim=x.dim)
+    return bu.Quantity(jr.permutation(key, x.mantissa, axis, independent=independent), unit=x.unit)
   return jr.permutation(key, x, axis, independent=independent)

brainstate/init/_random_inits.py

@@ -289,8 +289,8 @@ class VarianceScaling(Initializer):
       denominator = (fan_in + fan_out) / 2
     else:
       raise ValueError("invalid mode for variance scaling initializer: {}".format(self.mode))
-    scale = self.scale.value if isinstance(self.scale, bu.Quantity) else self.scale
-    dim = self.scale.dim if isinstance(self.scale, bu.Quantity) else bu.DIMENSIONLESS
+    scale = self.scale.mantissa if isinstance(self.scale, bu.Quantity) else self.scale
+    unit = bu.get_unit(self.scale)
     variance = (scale / denominator).astype(self.dtype)
     if self.distribution == "truncated_normal":
       stddev = (jnp.sqrt(variance) / .87962566103423978).astype(self.dtype)
@@ -302,7 +302,7 @@ class VarianceScaling(Initializer):
              jnp.sqrt(3 * variance).astype(self.dtype))
     else:
       raise ValueError("invalid distribution for variance scaling initializer")
-    return res if dim == bu.DIMENSIONLESS else res * dim
+    return res if unit.is_unitless else bu.Quantity(res, unit=unit)
 
   def __repr__(self):
     name = self.__class__.__name__
@@ -445,8 +445,8 @@ class Orthogonal(Initializer):
     matrix_shape = (n_rows, n_cols) if n_rows > n_cols else (n_cols, n_rows)
     norm_dst = random.normal(size=matrix_shape, dtype=self.dtype)
 
-    scale = self.scale.value if isinstance(self.scale, bu.Quantity) else self.scale
-    dim = self.scale.dim if isinstance(self.scale, bu.Quantity) else bu.DIMENSIONLESS
+    scale = self.scale.mantissa if isinstance(self.scale, bu.Quantity) else self.scale
+    unit = bu.get_unit(self.scale)
     q_mat, r_mat = jnp.linalg.qr(norm_dst)
     # Enforce Q is uniformly distributed
     q_mat *= jnp.sign(jnp.diag(r_mat))
@@ -455,7 +455,7 @@ class Orthogonal(Initializer):
     q_mat = jnp.reshape(q_mat, (n_rows,) + tuple(np.delete(shape, self.axis)))
     q_mat = jnp.moveaxis(q_mat, 0, self.axis)
     r = jnp.asarray(scale, dtype=self.dtype) * q_mat
-    return r if dim == bu.DIMENSIONLESS else r * dim
+    return r if unit.is_unitless else bu.Quantity(r, unit=unit)
 
   def __repr__(self):
     return f'{self.__class__.__name__}(scale={self.scale}, axis={self.axis}, dtype={self.dtype})'
@@ -480,8 +480,8 @@ class DeltaOrthogonal(Initializer):
       raise ValueError("Delta orthogonal initializer requires a 3D, 4D or 5D shape.")
     if shape[-1] < shape[-2]:
       raise ValueError("`fan_in` must be less or equal than `fan_out`. ")
-    scale = self.scale.value if isinstance(self.scale, bu.Quantity) else self.scale
-    dim = self.scale.dim if isinstance(self.scale, bu.Quantity) else bu.DIMENSIONLESS
+    scale = self.scale.mantissa if isinstance(self.scale, bu.Quantity) else self.scale
+    unit = bu.get_unit(self.scale)
     ortho_matrix = Orthogonal(scale=scale, axis=self.axis, dtype=self.dtype)(*shape[-2:])
     W = jnp.zeros(shape, dtype=self.dtype)
     if len(shape) == 3:
@@ -493,7 +493,7 @@ class DeltaOrthogonal(Initializer):
     else:
       k1, k2, k3 = shape[:3]
       W = W.at[(k1 - 1) // 2, (k2 - 1) // 2, (k3 - 1) // 2].set(ortho_matrix)
-    return W if dim == bu.DIMENSIONLESS else W * dim
+    return W if unit.is_unitless else bu.Quantity(W, unit=unit)
 
   def __repr__(self):
     return f'{self.__class__.__name__}(scale={self.scale}, axis={self.axis}, dtype={self.dtype})'

brainstate/random.py

@@ -490,14 +490,13 @@ class RandomState(State):
     upper = bu.math.asarray(upper, dtype=dtype)
     loc = bu.math.asarray(loc, dtype=dtype)
     scale = bu.math.asarray(scale, dtype=dtype)
-    bu.fail_for_dimension_mismatch(lower, upper)
-    bu.fail_for_dimension_mismatch(lower, loc)
-    bu.fail_for_dimension_mismatch(lower, scale)
-    dim = lower.dim if isinstance(lower, bu.Quantity) else bu.DIMENSIONLESS
-    lower = lower.value if isinstance(lower, bu.Quantity) else lower
-    upper = upper.value if isinstance(upper, bu.Quantity) else upper
-    loc = loc.value if isinstance(loc, bu.Quantity) else loc
-    scale = scale.value if isinstance(scale, bu.Quantity) else scale
+    unit = bu.get_unit(lower)
+    lower, upper, loc, scale = (
+      lower.mantissa if isinstance(lower, bu.Quantity) else lower,
+      bu.Quantity(upper).in_unit(unit).mantissa,
+      bu.Quantity(loc).in_unit(unit).mantissa,
+      bu.Quantity(scale).in_unit(unit).mantissa
+    )
 
     jit_error(
       bu.math.any(bu.math.logical_or(loc < lower - 2 * scale, loc > upper + 2 * scale)),
@@ -535,10 +534,12 @@ class RandomState(State):
     out = out * scale * sqrt2 + loc
 
     # Clamp to ensure it's in the proper range
-    out = jnp.clip(out,
-                   lax.nextafter(lax.stop_gradient(lower), np.array(np.inf, dtype=dtype)),
-                   lax.nextafter(lax.stop_gradient(upper), np.array(-np.inf, dtype=dtype)))
-    return out if dim == bu.DIMENSIONLESS else bu.Quantity(out, dim=dim)
+    out = jnp.clip(
+      out,
+      lax.nextafter(lax.stop_gradient(lower), np.array(np.inf, dtype=dtype)),
+      lax.nextafter(lax.stop_gradient(upper), np.array(-np.inf, dtype=dtype))
+    )
+    return out if unit.is_unitless else bu.Quantity(out, unit=unit)
 
   def _check_p(self, p):
     raise ValueError(f'Parameter p should be within [0, 1], but we got {p}')
@@ -555,30 +556,33 @@ class RandomState(State):
     r = jr.bernoulli(key, p=p, shape=_size2shape(size))
     return r
 
-  def lognormal(self,
-                mean=None,
-                sigma=None,
-                size: Optional[Size] = None,
-                key: Optional[SeedOrKey] = None,
-                dtype: DTypeLike = None):
+  def lognormal(
+      self,
+      mean=None,
+      sigma=None,
+      size: Optional[Size] = None,
+      key: Optional[SeedOrKey] = None,
+      dtype: DTypeLike = None
+  ):
     mean = _check_py_seq(mean)
     sigma = _check_py_seq(sigma)
     mean = bu.math.asarray(mean, dtype=dtype)
     sigma = bu.math.asarray(sigma, dtype=dtype)
-    bu.fail_for_dimension_mismatch(mean, sigma)
-    dim = mean.dim if isinstance(mean, bu.Quantity) else bu.DIMENSIONLESS
-    mean = mean.value if isinstance(mean, bu.Quantity) else mean
-    sigma = sigma.value if isinstance(sigma, bu.Quantity) else sigma
+    unit = mean.unit if isinstance(mean, bu.Quantity) else bu.Unit()
+    mean = mean.mantissa if isinstance(mean, bu.Quantity) else mean
+    sigma = sigma.in_unit(unit).mantissa if isinstance(sigma, bu.Quantity) else sigma
 
     if size is None:
-      size = jnp.broadcast_shapes(jnp.shape(mean),
-                                  jnp.shape(sigma))
+      size = jnp.broadcast_shapes(
+        jnp.shape(mean),
+        jnp.shape(sigma)
+      )
     key = self.split_key() if key is None else _formalize_key(key)
     dtype = dtype or environ.dftype()
     samples = jr.normal(key, shape=_size2shape(size), dtype=dtype)
     samples = _loc_scale(mean, sigma, samples)
     samples = jnp.exp(samples)
-    return samples if dim == bu.DIMENSIONLESS else bu.Quantity(samples, dim=dim)
+    return samples if unit.is_unitless else bu.Quantity(samples, unit=unit)
 
   def binomial(self,
                n,
@@ -678,10 +682,10 @@ class RandomState(State):
     cov = bu.math.asarray(_check_py_seq(cov), dtype=dtype)
     if isinstance(mean, bu.Quantity):
       assert isinstance(cov, bu.Quantity)
-      assert mean.dim ** 2 == cov.dim
-    mean = mean.value if isinstance(mean, bu.Quantity) else mean
-    cov = cov.value if isinstance(cov, bu.Quantity) else cov
-    dim = mean.dim if isinstance(mean, bu.Quantity) else bu.DIMENSIONLESS
+      assert mean.unit ** 2 == cov.unit
+    mean = mean.mantissa if isinstance(mean, bu.Quantity) else mean
+    cov = cov.mantissa if isinstance(cov, bu.Quantity) else cov
+    unit = mean.unit if isinstance(mean, bu.Quantity) else bu.Unit()
 
     key = self.split_key() if key is None else _formalize_key(key)
     if not jnp.ndim(mean) >= 1:
@@ -708,7 +712,7 @@ class RandomState(State):
       factor = jnp.linalg.cholesky(cov)
     normal_samples = jr.normal(key, size + mean.shape[-1:], dtype=dtype)
     r = mean + jnp.einsum('...ij,...j->...i', factor, normal_samples)
-    return r if dim == bu.DIMENSIONLESS else bu.Quantity(r, dim=dim)
+    return r if unit.is_unitless else bu.Quantity(r, unit=unit)
 
   def rayleigh(self,
                scale=1.0,

{brainstate-0.0.1.1.post20240804.dist-info → brainstate-0.0.2.post20240814.dist-info}/METADATA

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

{brainstate-0.0.1.1.post20240804.dist-info → brainstate-0.0.2.post20240814.dist-info}/RECORD

@@ -1,14 +1,14 @@
-brainstate/__init__.py,sha256=oxslZrm6wxtBQDqwJFb2BaAKZFmnp4d_esDkaeuGMWE,1410
+brainstate/__init__.py,sha256=zipNSih9Tyvi4-5cXqNPGsDF7VeestkLp-lcjJ4-dA0,1408
 brainstate/_module.py,sha256=YJDp9aD38wBa_lY6BojWjWV9LJ2aFMAMYh-KZe5a4eM,52443
 brainstate/_module_test.py,sha256=oQaoaZBTo1o3wHrMEJTInQCc7RdcVs1gcfQGvdSb1SI,7843
-brainstate/_random_for_unit.py,sha256=eW4NJkX27VCCNWUwAlyt2otkeEthGKOpUoX6XJ6i95Y,1946
+brainstate/_random_for_unit.py,sha256=1rHr7gfH_bYrJfpxbDhQUk_j00Yosx-GzyZCXrLxsd0,2007
 brainstate/_state.py,sha256=C0widCOj_ca6zfqh95jzFXf_G5vi0hJyuQ5GIqEqOUs,12102
 brainstate/_state_test.py,sha256=HDdipndRLhEHWEdTmyT1ayEBkbv6qJKykfCWKI6yJ_E,1253
 brainstate/_utils.py,sha256=RLorgGJkt2BhbX4C-ygd-PPG0wfcGCghjSP93sRvzqM,833
 brainstate/environ.py,sha256=LwRwnFaTbv8l7nHRIbSV46WzcN7pGLQFhT_xDUox2yA,10240
 brainstate/mixin.py,sha256=OumTTSVyYSbtudjfS_MRThsBaeVJ_0JggeMClY7xtBA,10758
 brainstate/mixin_test.py,sha256=-Ej9oUOu8O1M4oy37SVMj7xNRYhHHyAHwrjS_aISayo,2923
-brainstate/random.py,sha256=rqwSsiUoeZwxhk9ot8NnOJA8iWMdZB0HaHOVuweJdZQ,188387
+brainstate/random.py,sha256=BqEBYVD9TGe8dSzp8U0suK0O4r6Ox59GCq0mwfUndVQ,188073
 brainstate/random_test.py,sha256=cCeuYvlZkCS2_RgG0vipZFNSHG8b-uJ7SXM9SZDCYQM,17866
 brainstate/surrogate.py,sha256=1kgbn82GSlpReIytIVl29yozk75gkdZv0gTBlixQ4C4,43798
 brainstate/typing.py,sha256=6BlkLSN5TiaNO49q8b0OYyzcuSxmdoG3noIJTbyhE3s,7895
@@ -21,7 +21,7 @@ brainstate/functional/_spikes.py,sha256=70qGvo4B--QtxfJMjLwGmk9pVsf2x2YNEEgjT-il
 brainstate/init/__init__.py,sha256=R1dHgub47o-WJM9QkFLc7x_Q7GsyaKKDtrRHTFPpC5g,1097
 brainstate/init/_base.py,sha256=jRTmfoUsH_315vW9YMZzyIn2KDAAsv56SplBnvOyBW0,1148
 brainstate/init/_generic.py,sha256=LB7IQfswOG6X-q0QX5N8T5vZmxdygetsSBQ6iXlZ0oU,7324
-brainstate/init/_random_inits.py,sha256=LsfvKSX4wsR7Kh5jgKgdyXTCEEa5Nn_iYcp_2GgLQKY,16030
+brainstate/init/_random_inits.py,sha256=vNUVDdUOCXTx2i3i1enzxgg1USCzugYd56r0-2lBL-0,15919
 brainstate/init/_regular_inits.py,sha256=u77aSM0BkK9VULFJQZ1lIEYA_sJJzEZBTEttBSJ79RI,3090
 brainstate/nn/__init__.py,sha256=YJHoI8cXKVRS8f2vUl3Zegp5wm0svMz3qo9JmQJiMQk,2162
 brainstate/nn/_base.py,sha256=lzbZpku3Q2arH6ZaAwjs6bhbV0RcFChxo2UcpnX5t84,8481
@@ -59,8 +59,8 @@ brainstate/transform/_jit_test.py,sha256=5ltT7izh_OS9dcHnRymmVhq01QomjwZGdA8XzwJ
 brainstate/transform/_make_jaxpr.py,sha256=ZkrOZu4_0xcILuPUA3RFEkorJ-xbDuDtXorJI_qVThE,30450
 brainstate/transform/_make_jaxpr_test.py,sha256=K3vRUBroDTCCx0lnmhgHtgrlWvWglJO2f1K2phTvU70,3819
 brainstate/transform/_progress_bar.py,sha256=VGoRZPRBmB8ELNwLc6c7S8QhUUTvn0FY46IbBm9cuYM,3502
-brainstate-0.0.1.1.post20240804.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
-brainstate-0.0.1.1.post20240804.dist-info/METADATA,sha256=RTuqQrR0-syn5SyxoyEbfbdAUpXBRxNMzpaqnVM2cqQ,3807
-brainstate-0.0.1.1.post20240804.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-brainstate-0.0.1.1.post20240804.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
-brainstate-0.0.1.1.post20240804.dist-info/RECORD,,
+brainstate-0.0.2.post20240814.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
+brainstate-0.0.2.post20240814.dist-info/METADATA,sha256=skMWlfxiGaJHxzQS7dY95V91umhXVjN_HuhGO0xHP1M,3805
+brainstate-0.0.2.post20240814.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+brainstate-0.0.2.post20240814.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
+brainstate-0.0.2.post20240814.dist-info/RECORD,,