Metadata-Version: 2.4
Name: julax
Version: 0.0.5
Summary: Just Layers over JAX
Author-email: Jun Tian <tianjun.cpp@gmail.com>
License-File: LICENSE
Requires-Python: >=3.12
Requires-Dist: einops>=0.8.1
Requires-Dist: grain>=0.2.12
Requires-Dist: humanize>=4.13.0
Requires-Dist: jax>=0.7.2
Requires-Dist: optax>=0.2.6
Requires-Dist: orbax-checkpoint>=0.11.25
Requires-Dist: plum-dispatch>=2.5.8
Requires-Dist: pydantic>=2.12.0
Provides-Extra: tpu
Requires-Dist: jax[tpu]>=0.7.2; extra == 'tpu'
Description-Content-Type: text/markdown

<div align="center">
  <p>
  <img src="https://github.com/oolong-dev/julax/raw/main/docs/logo.svg?sanitize=true" width="320px">
  </p>
  <p>
  <h1>JULAX: Just Layers over JAX</h1>
  </p>
</div>