diff --git a/pyproject.toml b/pyproject.toml
index 197b89c..035f6f6 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -49,7 +49,6 @@ dev = [
     "pylint>=2.6.0",
     "pyink",
     # Lazy deps
-    "chex",
     "jax[cpu]",
     "jupyter",
     "tf-nightly",
diff --git a/visu3d/plotly/fig_utils_test.py b/visu3d/plotly/fig_utils_test.py
index a35d4e5..77d73eb 100644
--- a/visu3d/plotly/fig_utils_test.py
+++ b/visu3d/plotly/fig_utils_test.py
@@ -12,11 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for plotly."""
-
 from __future__ import annotations
 
-import chex
 import dataclass_array as dca
 from etils import enp
 import numpy as np
@@ -28,7 +25,7 @@
 
 @enp.testing.parametrize_xnp()
 def test_to_xyz_dict(xnp: enp.NpModule):
-  chex.assert_trees_all_close(
+  dca.testing.assert_allclose(
       v3d.plotly.to_xyz_dict([
           [0, 1, 2],
           [0, 10, 20],
@@ -41,7 +38,7 @@ def test_to_xyz_dict(xnp: enp.NpModule):
       },
   )
 
-  chex.assert_trees_all_close(
+  dca.testing.assert_allclose(
       v3d.plotly.to_xyz_dict(
           [
               [0, 1, 2],