Fix bug and update documentation

.gitignore

@@ -1,5 +1,6 @@
 __pycache__
 /venv
+/build
 /dist
 *.egg-info
 /docs/_build

README.md

@@ -18,6 +18,8 @@ This library requires at least Python 3.12.
 pip install einshard
 ```
 
+You need to have JAX installed by [choosing the correct installation method](https://jax.readthedocs.io/en/latest/installation.html) before installing Einshard.
+
 ## Usage
 
 For testing purpose, we initialise the JAX CPU backend with 16 devices. This should be run before the actual code (e.g. placed at the top of the script):
@@ -59,11 +61,15 @@ Output:
 
 ## Development
 
+Crente venv:
+
 ```sh
 python3.12 -m venv venv
 . venv/bin/activate
 ```
 
+Install dependencies:
+
 ```sh
 pip install -U pip
 pip install -U wheel

docs/index.rst

@@ -4,7 +4,9 @@
    contain the root `toctree` directive.
 
 Einshard Documentation
-====================================
+======================
+
+**Useful links**: `Source Repository <https://github.com/ayaka14732/einshard>`_ | `Issue Tracker <https://github.com/ayaka14732/einshard/issues>`_ 
 
 API
 ---

src/einshard/__init__.py

@@ -1,3 +1,3 @@
-from .einshard import shard, sharding
+from .shard import shard, sharding
 
-__version__ = '0.1.0'
+__version__ = '0.1.1'

src/einshard/einshard.py → src/einshard/shard.py

@@ -37,7 +37,8 @@ def sharding(expression: str, *, n_dims: int | None = None) -> NamedSharding:
     assert n_left_ellipses == n_right_ellipses and n_left_ellipses <= 1
 
     if n_left_ellipses == 0:
-        assert n_dims == len(elements_left)
+        if n_dims is not None:
+            assert n_dims == len(elements_left)
     else:  # == 1
         assert n_dims is not None
         n_dims_elided = n_dims - len(elements_left) + 1