Metadata-Version: 2.3
Name: statedict2pytree
Version: 0.1.2
Summary: Converts torch models into PyTrees for Equinox
Author-email: "Artur A. Galstyan" <mail@arturgalstyan.dev>
Requires-Python: ~=3.10
Requires-Dist: beartype
Requires-Dist: equinox>=0.11.4
Requires-Dist: flask
Requires-Dist: jax
Requires-Dist: jaxlib
Requires-Dist: jaxtyping
Requires-Dist: loguru
Requires-Dist: pydantic
Requires-Dist: torch
Requires-Dist: typing-extensions
Provides-Extra: dev
Requires-Dist: mkdocs; extra == 'dev'
Requires-Dist: nox; extra == 'dev'
Requires-Dist: pre-commit; extra == 'dev'
Requires-Dist: pytest; extra == 'dev'
Description-Content-Type: text/markdown

# statedict2pytree

![statedict2pytree](torch2jax.png "A ResNet demo")

The goal of this package is to simplify the conversion from PyTorch models into JAX PyTrees (which can be used e.g. in Equinox). The way this works is by putting both models side my side and aligning the weights in the right order. Then, all statedict2pytree is doing, is iterating over both lists and matching the weight matrices.

Usually, if you _declared the fields in the same order as in the PyTorch model_, you don't have to rearrange anything -- but the option is there if you need it.

(Theoretically, you can rearrange the model in any way you like - e.g. last layer as the first layer - as long as the shapes match!)

## Get Started

### Installation

Run

```bash
pip install statedict2pytree

```
