Fwd-port jnp.isscalar

src/ott/geometry/geometry.py

@@ -220,7 +220,7 @@ def inv_scale_cost(self) -> jnp.ndarray:
       return 1.0 / jnp.mean(self._cost_matrix)
     if self._scale_cost == "median":
       return 1.0 / jnp.median(self._cost_matrix)
-    if jnp.isscalar(self._scale_cost):
+    if utils.is_scalar(self._scale_cost):
       return 1.0 / self._scale_cost
     raise ValueError(f"Scaling {self._scale_cost} not implemented.")
 

src/ott/geometry/low_rank.py

@@ -112,7 +112,7 @@ def inv_scale_cost(self) -> jnp.ndarray:  # noqa: D102
       return 1.0 / (mean + self._bias)
     if self._scale_cost == "max_cost":
       return 1.0 / self._max_cost_matrix
-    if jnp.isscalar(self._scale_cost):
+    if utils.is_scalar(self._scale_cost):
       return 1.0 / self._scale_cost
     raise ValueError(f"Scaling {self._scale_cost} not implemented.")
 

src/ott/geometry/pointcloud.py

@@ -358,7 +358,7 @@ def inv_scale_cost(self) -> jnp.ndarray:  # noqa: D102
           "the cost matrix when the cost is not squared euclidean "
           "is not implemented."
       )
-    if jnp.isscalar(self._scale_cost):
+    if utils.is_scalar(self._scale_cost):
       return 1.0 / self._scale_cost
     raise ValueError(f"Scaling {self._scale_cost} not implemented.")
 

src/ott/utils.py

@@ -49,6 +49,7 @@
     "default_progress_fn",
     "tqdm_progress_fn",
     "batched_vmap",
+    "is_scalar",
 ]
 
 IOStatus = Tuple[np.ndarray, np.ndarray, np.ndarray, NamedTuple]
@@ -422,3 +423,13 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
   batched_fun = _apply_scan(vmapped_fun, in_axes=in_axes)
 
   return wrapper
+
+
+# TODO(michalk8): remove when `jax>=0.4.31`
+def is_scalar(x: Any) -> bool:  # noqa: D103
+  if (
+      isinstance(x, (np.ndarray, jax.Array)) or hasattr(x, "__jax_array__") or
+      np.isscalar(x)
+  ):
+    return jnp.asarray(x).ndim == 0
+  return False