[jax] Relax einsum checks for >=0.8.2 and update docs (#55)

* Remove non-deterministic assertion in _einops.py and update README for ecosystem terminology

* Update version to 0.1.3 and adjust version_info to reflect new format

* Update changelog for version 0.1.3 with compatibility note for jax

Author: chaoming0625

brainunit/README.md

@@ -48,7 +48,7 @@ The official documentation of BrainUnit is hosted on Read the Docs: [https://bra
 
 ## Features
 
-`brainunit` can be seamlessly integrated into every aspect of our [brain dynamics programming ecosystem](https://ecosystem-for-brain-dynamics.readthedocs.io/), such as, the checkpointing of [braintools](https://github.com/chaobrain/braintools), the event-driven operators in [brainevent](https://github.com/chaobrain/brainevent), the state-based JIT compilation in [brainstate](https://github.com/chaobrain/brainstate), online learning rules in [brainscale](https://github.com/chaobrain/brainscale), or event more. 
+`brainunit` can be seamlessly integrated into every aspect of our [brain modeling ecosystem](https://brainmodeling.readthedocs.io/), such as, the checkpointing of [braintools](https://github.com/chaobrain/braintools), the event-driven operators in [brainevent](https://github.com/chaobrain/brainevent), the state-based JIT compilation in [brainstate](https://github.com/chaobrain/brainstate), online learning rules in [brainscale](https://github.com/chaobrain/brainscale), or event more. 
 
 A quick example for this kind of integration:
 
@@ -128,6 +128,6 @@ If you use `brainunit` in your research, please consider citing the following pa
 
 
 
-## See also the BDP ecosystem
+## See also the ecosystem
 
-We are building the [brain dynamics programming (BDP) ecosystem](https://ecosystem-for-brain-dynamics.readthedocs.io/). [brainunit](https://github.com/chaobrain/brainunit) has been deeply integrated into our BDP ecosystem.
+``brainunit`` is one part of our [brain modeling ecosystem](https://brainmodeling.readthedocs.io/). 

brainunit/brainunit/__init__.py

@@ -17,7 +17,7 @@
 import saiunit
 
 __version__ = saiunit.__version__
-__version_info__ = saiunit.__version_info__
+__version_info__ = tuple(map(int, __version__.split(".")))
 
 from . import autograd
 from . import constants

changelog.md

@@ -1,5 +1,10 @@
 # Release Notes
 
+## Version 0.1.3
+
+- Compatible with `jax>=0.8.2`
+
+
 ## Version 0.1.2
 
 - Renamed ``CustomArray.value`` to ``CustomArray.data`` for API consistency

saiunit/__init__.py

@@ -13,8 +13,8 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "0.1.2"
-__version_info__ = (0, 1, 2)
+__version__ = "0.1.3"
+__version_info__ = tuple(map(int, __version__.split(".")))
 
 from . import _matplotlib_compat
 from . import autograd

saiunit/math/_einops.py

@@ -938,16 +938,6 @@ def _einsum(
 
             lhs_batch, rhs_batch = unzip2((lhs_names.find(n), rhs_names.find(n)) for n in batch_names)
 
-            # NOTE: this can fail non-deterministically in python3, maybe
-            # due to opt_einsum
-            assert jax.config.jax_dynamic_shapes or all(
-                name in lhs_names and name in rhs_names and
-                lhs.shape[lhs_names.index(name)] == rhs.shape[rhs_names.index(name)]
-                for name in contracted_names), \
-                ("Incompatible reduction dimensions: "
-                 f"lhs.shape={lhs.shape} lhs_names={lhs_names} "
-                 f"rhs.shape={rhs.shape} rhs_names={rhs_names}")
-
             # contract using dot_general
             batch_names_str = ''.join(batch_names)
             lhs_cont, rhs_cont = unzip2((lhs_names.index(n), rhs_names.index(n)) for n in contracted_names)