Fix jax breakage (assert_trees_all_close fail for Jax tensor), remo…

…ve chex dep

PiperOrigin-RevId: 691413236

pyproject.toml

@@ -49,7 +49,6 @@ dev = [
     "pylint>=2.6.0",
     "pyink",
     # Lazy deps
-    "chex",
     "jax[cpu]",
     "jupyter",
     "tf-nightly",

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],