refactor(jax-utils): ensure vectorize_method is private

src/galax/potential/_potential/base.py

@@ -33,7 +33,7 @@
     Vec6,
 )
 from galax.units import UnitSystem, dimensionless
-from galax.utils import vectorize_method
+from galax.utils._jax import vectorize_method
 from galax.utils._shape import batched_shape, expand_arr_dims, expand_batch_dims
 from galax.utils.dataclasses import ModuleMeta
 

src/galax/potential/_potential/builtin.py

@@ -29,7 +29,7 @@
     FloatScalar,
     Vec3,
 )
-from galax.utils import vectorize_method
+from galax.utils._jax import vectorize_method
 from galax.utils.dataclasses import field
 
 mass = u.get_physical_type("mass")

src/galax/potential/_potential/param/core.py

@@ -19,7 +19,7 @@
     FloatScalar,
     Unit,
 )
-from galax.utils import vectorize_method
+from galax.utils._jax import vectorize_method
 from galax.utils.dataclasses import converter_float_array
 
 

src/galax/utils/_jax.py

@@ -1,7 +1,7 @@
 """galax: Galactic Dynamix in Jax."""
 
 
-__all__ = ["vectorize_method"]
+__all__: list[str] = []
 
 from collections.abc import Callable, Sequence
 from functools import partial

tests/unit/utils/test_jax.py

@@ -4,7 +4,7 @@
 from galax.utils import vectorize_method
 
 
-def test_vectorize_method():
+def test_vectorize_method() -> None:
     """Test the vectorize_method function."""
 
     class A: