call super on __init_subclas__ (#5)

.gitignore

@@ -127,3 +127,7 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+
+# project specific
+.vscode
+/tmp

simple_pytree/pytree.py

@@ -75,6 +75,7 @@ class Pytree(metaclass=PytreeMeta):
     _pytree__class_is_mutable: bool
 
     def __init_subclass__(cls, mutable: bool = False):
+        super().__init_subclass__()
         # init class variables
         cls._pytree__initialized = False  # initialize mutable
         cls._pytree__class_is_mutable = mutable

tests/test_pytree.py

@@ -1,4 +1,5 @@
 import dataclasses
+from typing import Generic, TypeVar
 
 import jax
 import pytest
@@ -114,6 +115,15 @@ class Foo(Pytree):
         with pytest.raises(ValueError, match="Unknown field"):
             serialization.from_state_dict(foo, state_dict)
 
+    def test_generics(self):
+        T = TypeVar("T")
+
+        class MyClass(Pytree, Generic[T]):
+            def __init__(self, x: T):
+                self.x = x
+
+        MyClass[int]
+
 
 class TestMutablePytree:
     def test_pytree(self):