statedict2pytree 0.5.2__tar.gz → 0.5.3__tar.gz

{statedict2pytree-0.5.2 → statedict2pytree-0.5.3}/PKG-INFO

@@ -1,9 +1,10 @@
 Metadata-Version: 2.3
 Name: statedict2pytree
-Version: 0.5.2
+Version: 0.5.3
 Summary: Converts torch models into PyTrees for Equinox
 Author-email: "Artur A. Galstyan" <mail@arturgalstyan.dev>
 Requires-Python: ~=3.10
+Requires-Dist: anthropic
 Requires-Dist: beartype
 Requires-Dist: equinox
 Requires-Dist: flask
@@ -14,6 +15,7 @@ Requires-Dist: jaxtyping
 Requires-Dist: loguru
 Requires-Dist: penzai
 Requires-Dist: pydantic
+Requires-Dist: python-dotenv
 Requires-Dist: torch
 Requires-Dist: torchvision
 Requires-Dist: typing-extensions
@@ -37,6 +39,21 @@ PRs and other contributions are *highly* welcome! :)
 
 ## Info
 
+`statedict2pytree` is a powerful tool for converting PyTorch state dictionaries to JAX pytrees. It provides both programmatic and UI-based methods for mapping between PyTorch and JAX model parameters.
+
+## Features
+
+- Convert PyTorch statedicts to JAX pytrees
+- Handle large models with chunked file conversion
+- Provide an "intuitive-ish" UI for parameter mapping
+- Support both in-memory and file-based conversions
+
+## Installation
+
+```bash
+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.
@@ -49,16 +66,12 @@ Currently, there is no sophisticated shape matching in place. Two matrices are c
 
 (8, 1, 1) and (8, ) match, because (8 _ 1 _ 1 = 8)
 
-## Get Started
-
-### Installation
-
-Run
-
-```bash
-pip install statedict2pytree
-```
 
 ### Docs
 
 Documentation will appear as soon as I have all the necessary features implemented. Until then, check out the "main.py" file for a better example.
+
+### Disclaimer
+
+Some of the docstrings and the docs have been written with the help of
+Claude.

{statedict2pytree-0.5.2 → statedict2pytree-0.5.3}/README.md

@@ -9,6 +9,21 @@ PRs and other contributions are *highly* welcome! :)
 
 ## Info
 
+`statedict2pytree` is a powerful tool for converting PyTorch state dictionaries to JAX pytrees. It provides both programmatic and UI-based methods for mapping between PyTorch and JAX model parameters.
+
+## Features
+
+- Convert PyTorch statedicts to JAX pytrees
+- Handle large models with chunked file conversion
+- Provide an "intuitive-ish" UI for parameter mapping
+- Support both in-memory and file-based conversions
+
+## Installation
+
+```bash
+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.
@@ -21,16 +36,12 @@ Currently, there is no sophisticated shape matching in place. Two matrices are c
 
 (8, 1, 1) and (8, ) match, because (8 _ 1 _ 1 = 8)
 
-## Get Started
-
-### Installation
-
-Run
-
-```bash
-pip install statedict2pytree
-```
 
 ### Docs
 
 Documentation will appear as soon as I have all the necessary features implemented. Until then, check out the "main.py" file for a better example.
+
+### Disclaimer
+
+Some of the docstrings and the docs have been written with the help of
+Claude.