brainstate 0.0.2__py2.py3-none-any.whl → 0.0.2.post20240824__py2.py3-none-any.whl

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-0.0.2.dist-info → brainstate-0.0.2.post20240824.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: brainstate
-Version: 0.0.2
+Version: 0.0.2.post20240824
 Summary: A State-based Transformation System for Brain Dynamics Programming.
 Home-page: https://github.com/brainpy/brainstate
 Author: BDP
@@ -35,7 +35,7 @@ Requires-Dist: brainunit
 Provides-Extra: cpu
 Requires-Dist: jaxlib ; extra == 'cpu'
 Provides-Extra: cuda12
-Requires-Dist: jaxlib[cuda12_pip] ; extra == 'cuda12'
+Requires-Dist: jaxlib[cuda12] ; extra == 'cuda12'
 Provides-Extra: testing
 Requires-Dist: pytest ; extra == 'testing'
 Provides-Extra: tpu

{brainstate-0.0.2.dist-info → brainstate-0.0.2.post20240824.dist-info}/RECORD

@@ -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.2.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
-brainstate-0.0.2.dist-info/METADATA,sha256=K6yiVOqGj3Qs_vKGgQmFXZtlu8cS4r7EZXl_iyCjwh0,3792
-brainstate-0.0.2.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-brainstate-0.0.2.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
-brainstate-0.0.2.dist-info/RECORD,,
+brainstate-0.0.2.post20240824.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
+brainstate-0.0.2.post20240824.dist-info/METADATA,sha256=XmDSiVoXh250MvzBm2tJGfcNpRQ4FUrQTJsCvphGSSA,3801
+brainstate-0.0.2.post20240824.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+brainstate-0.0.2.post20240824.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
+brainstate-0.0.2.post20240824.dist-info/RECORD,,