Fix imports

ziatdinovmax · Mar 19, 2024 · 66e9b2c · 66e9b2c
1 parent 8669f14
commit 66e9b2c
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/gpax/utils/fn.py b/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