Serialization support (#2)

* prepare support for key paths

* add flax serialization support

* update poetry

* run pre-commit

Author: cgarciae

.github/workflows/run-tests.yml

@@ -31,7 +31,7 @@ jobs:
       - name: Install Poetry
         uses: snok/[email protected]
         with:
-          version: 1.2.1
+          version: 1.4.0
 
       - name: Setup Poetry
         run: |
@@ -81,7 +81,7 @@ jobs:
       - name: Install Poetry
         uses: snok/[email protected]
         with:
-          version: 1.2.1
+          version: 1.4.0
 
       - name: Setup Poetry
         run: |
@@ -98,11 +98,11 @@ jobs:
         if: steps.cache.outputs.cache-hit != 'true'
         run: |
           poetry run pip install -U certifi
-          poetry install --no-dev
+          poetry install --only main
 
       - name: Install Package
         run: |
-          poetry install --no-dev
+          poetry install --only main
         
       - name: Test Import
         run: |