From 66e9b2c75aaec72d5792d6c73e16fbd4e1963737 Mon Sep 17 00:00:00 2001 From: Maxim Ziatdinov Date: Tue, 19 Mar 2024 20:49:24 +0000 Subject: [PATCH] Fix imports --- gpax/utils/fn.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/gpax/utils/fn.py b/gpax/utils/fn.py index b4f1ca5..3119e09 100644 --- 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