statedict2pytree
Update:
For examples for statedict2pytree, check out my other repository jaxonmodels.
Docs
Docs can be found here.
Info
statedict2pytree is a powerful tool for converting PyTorch state dictionaries to JAX pytrees, specifically for Equinox
Installation
pip install statedict2pytree
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.
Shape Matching? What's that?
Currently, there is no sophisticated shape matching in place. Two matrices are considered "matching" if the product of their shape match. For example:
(8, 1, 1) and (8, ) match, because (8 _ 1 _ 1 = 8)
Disclaimer
Some of the docstrings and the docs have been written with the help of Claude.
1 maintainer
Metadata
1 maintainer
Release files for statedict2pytree 2.0.2
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | |
|---|---|---|---|
| statedict2pytree-2.0.2.tar.gz | 103.9 kB | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| statedict2pytree-2.0.2-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 109.4 kB
Release files / statedict2pytree-2.0.2.tar.gz
| Download URL | statedict2pytree-2.0.2.tar.gz |
|---|---|
| Size | 103.9 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
6338921db3a90a3913f99c9db2efc05b94311fe047df395ac4738b87d984e166
|
|
BLAKE2b-256 checksum How to use checksums |
a4ffce5a1c865262834cba36bb8a63a96c92eb194b2e3707659cf3587c6901de
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
uv/0.10.5 {"installer":{"name":"uv","version":"0.10.5","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"macOS","version":null,"id":null,"libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":null}
|
Release files / statedict2pytree-2.0.2-py3-none-any.whl
| Download URL | statedict2pytree-2.0.2-py3-none-any.whl |
|---|---|
| Size | 5.5 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
fad816fc00c5e627be2b4f54ff621066fe81fa6ae592c7997bfa0a84d4253c4c
|
|
BLAKE2b-256 checksum How to use checksums |
c88a15877a54cf50e481d6c99dbeaa4ab45608e168edbf5e23766df7b8cb6781
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
uv/0.10.5 {"installer":{"name":"uv","version":"0.10.5","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"macOS","version":null,"id":null,"libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":null}
|
