chore(release): bump version to 0.1.2 and update README

- Update version in pyproject.toml and uv.lock to 0.1.2
- Remove status section from README to reflect current state
- Enhance docstrings across various modules for better clarity and consistency

Author: erfanzar

README.md

@@ -665,14 +665,6 @@ pytest tests/test_conformance.py
 
 ---
 
-## Status
-
-`v0.1.0` — alpha. API may still move; pin the version if you depend
-on behavioural stability. See [CHANGELOG.md](CHANGELOG.md) for the
-release log.
-
----
-
 ## License
 
 This project is licensed under AGPL-3.0-or-later.

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "uv_build"
 
 [project]
 name = "spectrax-lib"
-version = "0.1.1"
+version = "0.1.2"
 authors = [{ name = "Erfan Zare Chavoshi", email = "Erfanzare810@gmail.com" }]
 description = "SpectraX: a JAX-only neural-network library with a True MPMD pipeline parallelism for JAX — with an eager module API."
 readme = "README.md"

spectrax/__init__.py

@@ -72,6 +72,10 @@ def _patched_update(name, value):
         Returns ``None`` (without calling the underlying ``update``)
         when ``name`` is a known-removed flag; otherwise delegates to
         the original ``jax.config.update`` and returns its result.
+
+        Args:
+            name: Name used for lookup, logging, or registration.
+            value: Value consumed by the helper.
         """
         if name in removed_flags:
             return None

spectrax/_internal/logging.py

@@ -206,23 +206,53 @@ def _log_once(self, level: int, message: str, *args: object, **kwargs: object) -
                 self._logger.log(level, message, *args, **kwargs)
 
     def debug_once(self, message: str, *args: object, **kwargs: object) -> None:
-        """Log a ``DEBUG`` message once (deduplicated)."""
+        """Log a ``DEBUG`` message once (deduplicated).
+
+        Args:
+            message: Message value consumed by this operation.
+            *args: Additional positional arguments forwarded to the wrapped callable or backend.
+            **kwargs: Additional keyword arguments forwarded to the wrapped callable or backend.
+        """
         self._log_once(logging.DEBUG, message, *args, **kwargs)
 
     def info_once(self, message: str, *args: object, **kwargs: object) -> None:
-        """Log an ``INFO`` message once (deduplicated)."""
+        """Log an ``INFO`` message once (deduplicated).
+
+        Args:
+            message: Message value consumed by this operation.
+            *args: Additional positional arguments forwarded to the wrapped callable or backend.
+            **kwargs: Additional keyword arguments forwarded to the wrapped callable or backend.
+        """
         self._log_once(logging.INFO, message, *args, **kwargs)
 
     def warn_once(self, message: str, *args: object, **kwargs: object) -> None:
-        """Log a ``WARNING`` message once (deduplicated)."""
+        """Log a ``WARNING`` message once (deduplicated).
+
+        Args:
+            message: Message value consumed by this operation.
+            *args: Additional positional arguments forwarded to the wrapped callable or backend.
+            **kwargs: Additional keyword arguments forwarded to the wrapped callable or backend.
+        """
         self._log_once(logging.WARNING, message, *args, **kwargs)
 
     def warning_once(self, message: str, *args: object, **kwargs: object) -> None:
-        """Alias for :meth:`warn_once`."""
+        """Alias for :meth:`warn_once`.
+
+        Args:
+            message: Message value consumed by this operation.
+            *args: Additional positional arguments forwarded to the wrapped callable or backend.
+            **kwargs: Additional keyword arguments forwarded to the wrapped callable or backend.
+        """
         self._log_once(logging.WARNING, message, *args, **kwargs)
 
     def error_once(self, message: str, *args: object, **kwargs: object) -> None:
-        """Log an ``ERROR`` message once (deduplicated)."""
+        """Log an ``ERROR`` message once (deduplicated).
+
+        Args:
+            message: Message value consumed by this operation.
+            *args: Additional positional arguments forwarded to the wrapped callable or backend.
+            **kwargs: Additional keyword arguments forwarded to the wrapped callable or backend.
+        """
         self._log_once(logging.ERROR, message, *args, **kwargs)
 
     def clear_once_cache(self) -> None:
@@ -251,7 +281,15 @@ def __getattr__(self, name: str) -> Callable[..., object]:
 
             @wraps(getattr(logging.Logger, method_name))
             def wrapped_log_method(*args: object, **kwargs: object) -> object:
-                """Delegate to the underlying logger method."""
+                """Delegate to the underlying logger method.
+
+                Args:
+                    *args: Additional positional arguments forwarded to the wrapped callable or backend.
+                    **kwargs: Additional keyword arguments forwarded to the wrapped callable or backend.
+
+                Returns:
+                    Result described by this helper.
+                """
                 self._ensure_initialized()
                 return getattr(self._logger, method_name)(*args, **kwargs)
 
@@ -264,15 +302,31 @@ def wrapped_log_method(*args: object, **kwargs: object) -> object:
 
                 @wraps(getattr(logging.Logger, method_name))
                 def wrapped_log_method(*args: object, **kwargs: object) -> object:
-                    """Delegate to the named level method of the underlying logger."""
+                    """Delegate to the named level method of the underlying logger.
+
+                    Args:
+                        *args: Additional positional arguments forwarded to the wrapped callable or backend.
+                        **kwargs: Additional keyword arguments forwarded to the wrapped callable or backend.
+
+                    Returns:
+                        Result described by this helper.
+                    """
                     self._ensure_initialized()
                     return getattr(self._logger, method_name)(*args, **kwargs)
 
                 return wrapped_log_method
 
             @wraps(logging.Logger.log)
             def wrapped_log_method(*args: object, **kwargs: object) -> object:
-                """Log at the dynamically-resolved level."""
+                """Log at the dynamically-resolved level.
+
+                Args:
+                    *args: Additional positional arguments forwarded to the wrapped callable or backend.
+                    **kwargs: Additional keyword arguments forwarded to the wrapped callable or backend.
+
+                Returns:
+                    Result described by this helper.
+                """
                 self._ensure_initialized()
                 return self._logger.log(level, *args, **kwargs)
 
@@ -424,7 +478,13 @@ def __enter__(self):
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        """Context-manager exit — auto-completes if no exception was raised."""
+        """Context-manager exit — auto-completes if no exception was raised.
+
+        Args:
+            exc_type: Exc type value consumed by this operation.
+            exc_val: Exc val value consumed by this operation.
+            exc_tb: Exc tb value consumed by this operation.
+        """
         if exc_type is None:
             self.complete()
         return False

spectrax/_version.py

@@ -8,4 +8,4 @@
 ``from ._version import __version__`` and by packaging tooling.
 """
 
-__version__ = "0.1.1"
+__version__ = "0.1.2"

spectrax/api.py

@@ -210,6 +210,11 @@ def _fwd(state, *a, **kw):
 
                 Cached per-``GraphDef`` so the same compiled program is
                 reused across calls.
+
+                Args:
+                    state: SpectraX state tree or transform state passed into the operation.
+                    *a: Additional positional arguments forwarded to the wrapped callable or backend.
+                    **kw: Additional keyword arguments forwarded to the wrapped callable or backend.
                 """
                 return bind(gdef, state)(*a, **kw)
 
@@ -233,6 +238,13 @@ def _step(state, args, kwargs, l_args, l_kwargs):
             Differentiates with respect to the full :class:`State`
             tree; partition out trainable subsets at the call site if
             you want narrower gradients.
+
+            Args:
+                state: SpectraX state tree or transform state passed into the operation.
+                args: Positional arguments forwarded to the wrapped callable.
+                kwargs: Keyword arguments forwarded to the wrapped callable.
+                l_args: L args value consumed by this operation.
+                l_kwargs: L kwargs value consumed by this operation.
             """
 
             def loss(state):
@@ -241,6 +253,9 @@ def loss(state):
                 Captures ``gdef``, the model call args, and the
                 supplied ``loss_fn`` from the enclosing scope so the
                 resulting function depends only on ``state``.
+
+                Args:
+                    state: SpectraX state tree or transform state passed into the operation.
                 """
                 out = bind(gdef, state)(*args, **kwargs)
                 return loss_fn(out, *l_args, **l_kwargs)
@@ -337,6 +352,13 @@ def _wrapped_loss(out: object, *vals: object) -> object:
             the captured ``target_keys`` and calls the original
             ``loss_fn`` with keyword targets, preserving its kwargs
             interface even though the pipeline batch is positional.
+
+            Args:
+                out: Output value from an earlier call or transform.
+                *vals: Additional positional arguments forwarded to the wrapped callable or backend.
+
+            Returns:
+                Result described by this helper.
             """
             return original_loss(out, **dict(zip(target_keys, vals, strict=True)))
 

spectrax/contrib/optimizer.py

@@ -53,7 +53,15 @@ class _OptaxModule(Protocol):
     """Subset of :mod:`optax` used by this wrapper."""
 
     def apply_updates(self, params: State, updates: State) -> State:
-        """Apply optax updates to a :class:`State` tree."""
+        """Apply optax updates to a :class:`State` tree.
+
+        Args:
+            params: Parameter mapping or primitive parameter dictionary.
+            updates: Updates value consumed by this operation.
+
+        Returns:
+            Result described by this helper.
+        """
         ...
 
 

spectrax/core/_typing.py

@@ -84,7 +84,16 @@ class Initializer(Protocol):
     """
 
     def __call__(self, key: PRNGKey, shape: Shape, dtype: DType = jnp.float32) -> Array:
-        """Return a freshly-initialized array of ``shape`` and ``dtype`` from ``key``."""
+        """Return a freshly-initialized array of ``shape`` and ``dtype`` from ``key``.
+
+        Args:
+            key: Logical key, path segment, or PRNG key used by the operation.
+            shape: Array shape requested by the initializer or helper.
+            dtype: Array dtype requested for the produced value.
+
+        Returns:
+            Result of invoking the wrapped callable or module.
+        """
         ...
 
 
@@ -114,7 +123,16 @@ def __call__(
         args: tuple[object, ...],
         kwargs: dict[str, object],
     ) -> tuple[tuple[object, ...], dict[str, object]] | None:
-        """Optionally rewrite ``(args, kwargs)`` before ``forward`` runs."""
+        """Optionally rewrite ``(args, kwargs)`` before ``forward`` runs.
+
+        Args:
+            module: SpectraX module instance operated on by the helper.
+            args: Positional arguments forwarded to the wrapped callable.
+            kwargs: Keyword arguments forwarded to the wrapped callable.
+
+        Returns:
+            Result of invoking the wrapped callable or module.
+        """
         ...
 
 
@@ -133,7 +151,17 @@ def __call__(
         kwargs: dict[str, object],
         output: object,
     ) -> object | None:
-        """Optionally rewrite the forward output. Return ``None`` to keep it unchanged."""
+        """Optionally rewrite the forward output. Return ``None`` to keep it unchanged.
+
+        Args:
+            module: SpectraX module instance operated on by the helper.
+            args: Positional arguments forwarded to the wrapped callable.
+            kwargs: Keyword arguments forwarded to the wrapped callable.
+            output: Output value consumed by this operation.
+
+        Returns:
+            Result of invoking the wrapped callable or module.
+        """
         ...
 
 
@@ -147,5 +175,11 @@ class VariableObserver(Protocol):
     """
 
     def __call__(self, var: object, old: object, new: object) -> None:
-        """React to ``var`` having its value changed from ``old`` to ``new``."""
+        """React to ``var`` having its value changed from ``old`` to ``new``.
+
+        Args:
+            var: Var value consumed by this operation.
+            old: Old value consumed by this operation.
+            new: New value consumed by this operation.
+        """
         ...