statedict2pytree 0.3.0__py3-none-any.whl → 0.4.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (4) hide show
  1. statedict2pytree/statedict2pytree.py +4 -1
  2. {statedict2pytree-0.3.0.dist-info → statedict2pytree-0.4.0.dist-info}/METADATA +1 -1
  3. {statedict2pytree-0.3.0.dist-info → statedict2pytree-0.4.0.dist-info}/RECORD +4 -4
  4. {statedict2pytree-0.3.0.dist-info → statedict2pytree-0.4.0.dist-info}/WHEEL +0 -0
statedict2pytree/statedict2pytree.py CHANGED
@@ -46,7 +46,7 @@ def get_node(
46
46
  return tree
47
47
  else:
48
48
  next_target: str = targets[0]
49
- if bool(re.search(r"\[\d\]", next_target)):
49
+ if bool(re.search(r"\[\d+\]", next_target)):
50
50
  split_index = next_target.rfind("[")
51
51
  name, index = next_target[:split_index], next_target[split_index:]
52
52
  index = index[1:-1]
@@ -157,6 +157,9 @@ def index():
157
157
  def autoconvert(pytree: PyTree, state_dict: dict) -> tuple[PyTree, eqx.nn.State]:
158
158
  jax_fields = pytree_to_fields(pytree)
159
159
  torch_fields = state_dict_to_fields(state_dict)
160
+
161
+ for k, v in state_dict.items():
162
+ state_dict[k] = v.numpy()
160
163
  return convert(jax_fields, torch_fields, pytree, state_dict)
161
164
 
162
165
 
{statedict2pytree-0.3.0.dist-info → statedict2pytree-0.4.0.dist-info}/METADATA RENAMED
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.3
2
2
  Name: statedict2pytree
3
- Version: 0.3.0
3
+ Version: 0.4.0
4
4
  Summary: Converts torch models into PyTrees for Equinox
5
5
  Author-email: "Artur A. Galstyan" <mail@arturgalstyan.dev>
6
6
  Requires-Python: ~=3.10
{statedict2pytree-0.3.0.dist-info → statedict2pytree-0.4.0.dist-info}/RECORD RENAMED
@@ -1,8 +1,8 @@
1
1
  statedict2pytree/__init__.py,sha256=lXxSaFFvkhXweXp5oHSkg_dPjdp49OsF8xoqwX4d_4E,240
2
- statedict2pytree/statedict2pytree.py,sha256=X5Ljf4lYhhH7_V4KgdciChncbTt7YZpIWHcOxcZ3l48,7103
2
+ statedict2pytree/statedict2pytree.py,sha256=u1PddUBY_MHErl3tstSpJt2a6a_H-RizvxP9anPoFpQ,7175
3
3
  statedict2pytree/static/input.css,sha256=zBp60NAZ3bHTLQ7LWIugrCbOQdhiXdbDZjSLJfg6KOw,59
4
4
  statedict2pytree/static/output.css,sha256=KZ9GzeV3q0XKjbEiTdPkC6yV-R6jzXRflRm2S16VkJA,40813
5
5
  statedict2pytree/templates/index.html,sha256=0uG3dB2pAa1f2wcfTpYSO7TBNL77i2ALJP5rIhsbEnk,7506
6
- statedict2pytree-0.3.0.dist-info/METADATA,sha256=YSK4tWzNQemyZ1xKq5BhWiLWWc-RDr4E9q_eV_iOsdw,4232
7
- statedict2pytree-0.3.0.dist-info/WHEEL,sha256=zEMcRr9Kr03x1ozGwg5v9NQBKn3kndp6LSoSlVg-jhU,87
8
- statedict2pytree-0.3.0.dist-info/RECORD,,
6
+ statedict2pytree-0.4.0.dist-info/METADATA,sha256=BbJaUhVv-Qb4-e2DlX0WBNgPOKsTjlKv8ISJZgFNx1o,4232
7
+ statedict2pytree-0.4.0.dist-info/WHEEL,sha256=zEMcRr9Kr03x1ozGwg5v9NQBKn3kndp6LSoSlVg-jhU,87
8
+ statedict2pytree-0.4.0.dist-info/RECORD,,