Fix imports

Author: ziatdinovmax

gpax/utils/fn.py

@@ -10,9 +10,10 @@
 import inspect
 import re
 
-from typing import List, Callable, Optional
+from typing import List, Callable, Optional, Dict
 
 import jax
+import jax.numpy as jnp
 
 from ..kernels.kernels import square_scaled_distance, add_jitter, _sqrt