remove unused import

Author: eshen1991

axlearn/common/ops/_optimization_barrier.py

@@ -5,8 +5,6 @@
 from typing import Any
 
 import jax
-from jax._src import ad_checkpoint  # pylint: disable=protected-access
-
 
 @jax.custom_jvp
 @jax.custom_batching.custom_vmap  # Must be wrapped in this before custom_jvp.

axlearn/common/utils_test.py

@@ -14,7 +14,6 @@
 
 # pylint: disable=no-self-use
 import jax
-import jaxlib
 import numpy as np
 import pytest
 import tensorflow as tf