stringjax-tools 0.1.0__py3-none-any.whl

stringjax_tools-0.1.0.dist-info/METADATA

@@ -0,0 +1,215 @@
+Metadata-Version: 2.4
+Name: stringjax-tools
+Version: 0.1.0
+Summary: Reusable JAX transform, vectorisation, cache, and pytree helpers for StringJAX packages.
+Author-email: Andreas Schachner <as3475@cornell.edu>
+License-Expression: GPL-3.0-or-later
+Project-URL: Homepage, https://github.com/StringJAX/stringjax-tools
+Project-URL: Issues, https://github.com/StringJAX/stringjax-tools/issues
+Keywords: jax,vmap,jit,pytrees,automatic-vectorisation,scientific-computing
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: jax>=0.4.26
+Requires-Dist: numpy>=1.24
+Provides-Extra: dev
+Requires-Dist: build; extra == "dev"
+Requires-Dist: nbclient; extra == "dev"
+Requires-Dist: nbformat; extra == "dev"
+Requires-Dist: ipykernel; extra == "dev"
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: twine; extra == "dev"
+Provides-Extra: docs
+Requires-Dist: matplotlib; extra == "docs"
+Requires-Dist: myst-nb; extra == "docs"
+Requires-Dist: roman; extra == "docs"
+Requires-Dist: sphinx; extra == "docs"
+Requires-Dist: sphinx-autodoc-typehints; extra == "docs"
+Requires-Dist: sphinx-book-theme; extra == "docs"
+Requires-Dist: sphinx-copybutton; extra == "docs"
+Requires-Dist: sphinx-design; extra == "docs"
+Requires-Dist: sphinx-togglebutton; extra == "docs"
+Requires-Dist: sphinxcontrib-mermaid; extra == "docs"
+Dynamic: license-file
+
+# StringJAX Tools
+
+`stringjax_tools` is a small utility package for reusable JAX transform
+patterns used across StringJAX projects.  It currently focuses on:
+
+- conservative rank-checked automatic vectorisation;
+- manual cached `jax.jit(jax.vmap(...))` wrappers;
+- small helpers for explicit static-argument JIT compilation;
+- explicit JAX persistent-compilation-cache setup;
+- configurable pytree registration for stateful model classes.
+
+The package keeps reusable transformation and pytree machinery separate from
+domain-specific StringJAX model code.
+
+## Installation
+
+From this directory:
+
+```bash
+pip install -e ".[dev,docs]"
+```
+
+The intended published package name is:
+
+```bash
+pip install stringjax-tools
+```
+
+The Python import name is:
+
+```python
+import stringjax_tools
+```
+
+## Documentation
+
+The Sphinx documentation lives in `documentation/`:
+
+```bash
+make -C documentation html
+```
+
+For strict local validation:
+
+```bash
+make -C documentation html SPHINXOPTS="-W -q"
+```
+
+The feature notebooks are included under `documentation/source/notebooks/` and
+are linked from the Sphinx tutorials page.
+
+## Automatic Vectorisation
+
+Keep the single-sample function explicit, then wrap it at the boundary where
+callers may pass either one sample or a paired leading-axis batch.
+
+```python
+import jax.numpy as jnp
+from stringjax_tools.auto_vectorise import auto_vmap
+
+
+def observable_single(moduli, tau, fluxes, scale=1.0):
+    return scale * (jnp.sum(jnp.abs(moduli) ** 2) + tau + jnp.sum(fluxes**2))
+
+
+observable = auto_vmap(moduli=1, tau=0, fluxes=1)(observable_single)
+```
+
+With these declarations:
+
+- `moduli` has sample rank `1`, so `(h,)` is one sample and `(N, h)` is a batch.
+- `tau` has sample rank `0`, so `()` is one sample and `(N,)` is a batch.
+- `fluxes` has sample rank `1`, so `(n,)` is one sample and `(N, n)` is a batch.
+
+If several checked arguments are batched, all leading batch sizes must agree.
+Checked arguments without a batch axis are broadcast with `in_axes=None`.
+
+## Exact Shape Validation
+
+Ranks decide batching.  Exact trailing dimensions are optional:
+
+```python
+class Model:
+    def __init__(self, h12, n_fluxes):
+        self.h12 = h12
+        self.n_fluxes = n_fluxes
+
+    @auto_vmap(
+        sample_ranks={"moduli": 1, "fluxes": 1},
+        sample_shapes={
+            "moduli": "h12",
+            "fluxes": lambda bound: (2 * bound["self"].n_fluxes,),
+        },
+        validate_shapes=True,
+    )
+    def diagnostic(self, moduli, fluxes):
+        return jnp.sum(jnp.abs(moduli) ** 2) + 0.01 * jnp.sum(fluxes**2)
+```
+
+String dimensions are resolved from a bound object, usually `self`.  Derived
+dimensions should be callables rather than expression strings.
+
+## Manual Transform Helpers
+
+When a function already has a deliberate `in_axes` structure, use:
+
+```python
+from stringjax_tools.vmap import vmapping_func_cached
+
+batched = vmapping_func_cached(single_sample_function, in_axes=(0, None, 0))
+```
+
+Static keyword configuration is snapshotted when the cached wrapper is built.
+Prefer immutable configuration objects, such as tuples, for values that should
+participate in wrapper caching.
+
+For explicit static positional arguments:
+
+```python
+from stringjax_tools.jit import jit_with_static_args
+
+jit_f = jit_with_static_args(f, static_argnums=(2,))
+```
+
+## Pytree Policies
+
+Each StringJAX package should define its own static and ignored attribute policy
+locally.  The tool package supplies only generic registration machinery.
+
+```python
+from stringjax_tools.pytrees import PytreePolicy
+
+PACKAGE_PYTREE_POLICY = PytreePolicy(
+    static_keys=("h11", "h12", "model_ID", "_user_Q"),
+    ignore_defaults={"_sampler": None, "_cache": dict},
+)
+
+PACKAGE_PYTREE_POLICY.register(MyModelClass)
+```
+
+Semantic model state belongs in `static_keys` when it is hashable and immutable,
+or in traced children when it is array-like.  Ignored attributes should be
+limited to recomputable caches, scratch state, and eager-only helpers.  Since
+JAX reconstruction bypasses `__init__`, ignored caches that may be read after a
+round trip should be declared through `ignore_defaults`; callable defaults such
+as `dict` are called for each reconstructed object.  Do not ignore user-supplied
+configuration such as physical bounds or tadpole values.
+Static and ignored attribute names must be disjoint.  Static values are
+validated by default: they must be hashable and have scalar, self-equal
+comparison semantics, so array-like and `NaN`-like values are rejected unless a
+package deliberately opts out with `validate_static=False`.
+Leave `validate_static=True` unless the consuming package deliberately stores
+non-hashable metadata in JAX pytree auxiliary data.
+
+## Compilation Cache
+
+`configure_compilation_cache(...)` is opt-in and has no import-time side
+effects.  Call it before the first JAX compilation in a process:
+
+```python
+from stringjax_tools.cache import configure_compilation_cache
+
+configure_compilation_cache(
+    cache_dir=".jax-cache",
+    max_size_bytes=10_000_000_000,
+)
+```
+
+## License
+
+StringJAX Tools is distributed under the GNU General Public License v3.0 or
+later.

stringjax_tools-0.1.0.dist-info/RECORD

@@ -0,0 +1,14 @@
+stringjax_tools/__init__.py,sha256=hbpBc5H2UdIcaH3t7YvLik4aksJ_itMUiPJ_zS10YMQ,2083
+stringjax_tools/__main__.py,sha256=syzBGU0YdHoJywQKfq7GxSzvZgNAQ6mh-_bpaeeitp4,881
+stringjax_tools/_utils.py,sha256=tZad_PGV7sSVtdgTSvNc3DrLIA6gPNp5w_4EGfvb51I,6580
+stringjax_tools/auto_vectorise.py,sha256=dax2nXmmSIVR5G1tj_KBTCucx-8AYWPC2B86Z14jllg,20604
+stringjax_tools/cache.py,sha256=4KvelK6c7jWzGiSubArIc0IC4P0fqEgy8KNPd72kpSA,3310
+stringjax_tools/jit.py,sha256=u9umUF_31eHwEv7gYSC25pHrwHFBA4O9J4gBJb9vLgw,2613
+stringjax_tools/py.typed,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+stringjax_tools/pytrees.py,sha256=6wJXxGSBRRnHBfIGftK4XIa0CkntkLnukGl-KaWF7mU,11952
+stringjax_tools/vmap.py,sha256=9VhmTaTNI0wRYfHAc53yrXwSnDeu0QqbHog9WVPUTpo,3539
+stringjax_tools-0.1.0.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
+stringjax_tools-0.1.0.dist-info/METADATA,sha256=eUldDTGx1fy7EGJV2j6DeOdapiQHUU0WxG71ifFOMi0,6914
+stringjax_tools-0.1.0.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+stringjax_tools-0.1.0.dist-info/top_level.txt,sha256=i2zKBNLJ0IqvM1p9U77pcITsENj5peoTGYobjhiLLFo,16
+stringjax_tools-0.1.0.dist-info/RECORD,,

stringjax_tools-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+