Merge pull request #5 from shyamsn97/notebook-fix

fixing notebooks

hypernn/jax/dynamic_hypernet.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import Any, Dict, List, Optional, Tuple
 
 import flax.linen as nn
@@ -91,7 +93,7 @@ def from_target(
         weight_chunk_dim: Optional[int] = None,
         *args,
         **kwargs,
-    ) -> JaxDynamicEmbeddingModule:
+    ) -> JaxDynamicHyperNetwork:
         num_target_parameters, variables = cls.count_params(
             target_network, target_input_shape, inputs=inputs, return_variables=True
         )

hypernn/jax/hypernet.py

@@ -125,7 +125,10 @@ def count_params(
         return_variables: bool = False,
     ):
         return count_jax_params(
-            target, target_input_shape, inputs=inputs, return_variables=return_variables
+            target,
+            target_input_shape,
+            inputs=inputs,
+            return_variables=return_variables,
         )
 
     @classmethod
@@ -139,8 +142,12 @@ def from_target(
         **kwargs,
     ) -> JaxHyperNetwork:
         num_target_parameters, variables = cls.count_params(
-            target_network, target_input_shape, inputs=inputs, return_variables=True
+            target_network,
+            target_input_shape,
+            inputs=inputs,
+            return_variables=True,
         )
+
         _value_flat, target_treedef = jax.tree_util.tree_flatten(variables)
         target_weight_shapes = [v.shape for v in _value_flat]
 

hypernn/jax/utils.py

@@ -22,12 +22,13 @@ def count_jax_params(
     input_shape: Optional[Tuple[int, ...]] = None,
     inputs: Optional[List[jnp.array]] = None,
     return_variables: bool = False,
+    **kwargs
 ) -> int:
     if input_shape is None and inputs is None:
         raise ValueError("Input shape or inputs must be specified")
     if inputs is None:
         inputs = [jnp.zeros(shape) for shape in input_shape]
-    variables = model.init(jax.random.PRNGKey(0), *inputs)
+    variables = model.init(jax.random.PRNGKey(0), *inputs, **kwargs)
 
     def count_recursive(d):
         s = 0

notebooks/Intro.ipynb

@@ -185,7 +185,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -251,7 +251,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -295,7 +295,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [
     {
@@ -368,10 +368,13 @@
   }
  ],
  "metadata": {
+  "interpreter": {
+   "hash": "7e5f2d8038e6c8941d283a9a145e7dfd2f60905a23729a9c846dae091a9571f4"
+  },
   "kernelspec": {
    "display_name": "Python [conda env:py39] *",
    "language": "python",
-   "name": "conda-env-py39-py"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {