From 7587f914daff655e670eb6f79f68ecf1dac35c76 Mon Sep 17 00:00:00 2001
From: Zekun Shi <derek.zekun.shi@gmail.com>
Date: Fri, 12 Jan 2024 16:57:18 +0800
Subject: [PATCH] fix test

---
 autofd/general_array_test.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/autofd/general_array_test.py b/autofd/general_array_test.py
index 465f3ce..eafd16d 100644
--- a/autofd/general_array_test.py
+++ b/autofd/general_array_test.py
@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import optax
 import jax
 import jax.numpy as jnp
 from jaxtyping import Float32, Array