Minor cleanup for optax.projections.

Author: carlosgmartin

docs/api/projections.rst

@@ -9,7 +9,7 @@ The Euclidean projection onto a set :math:`\mathcal{C}` is:
 .. math::
 
     \text{proj}_{\mathcal{C}}(u) :=
-    \underset{v}{\text{argmin}} ~ ||u - v||^2_2 \textrm{ subject to } v \in \mathcal{C}.
+    \underset{v}{\text{argmin}} ~ \|u - v\|^2_2 \textrm{ subject to } v \in \mathcal{C}.
 
 For instance, here is an example how we can project parameters to the non-negative orthant::
 

optax/projections/_projections.py

@@ -29,7 +29,7 @@ def projection_non_negative(tree: Any) -> Any:
 
   .. math::
 
-    \underset{p}{\text{argmin}} ~ ||x - p||_2^2 \quad
+    \underset{p}{\text{argmin}} ~ \|x - p\|_2^2 \quad
     \textrm{subject to} \quad p \ge 0
 
   where :math:`x` is the input tree.
@@ -43,16 +43,12 @@ def projection_non_negative(tree: Any) -> Any:
   return jax.tree.map(jax.nn.relu, tree)
 
 
-def _clip_safe(leaf, lower, upper):
-  return jnp.clip(jnp.asarray(leaf), lower, upper)
-
-
 def projection_box(tree: Any, lower: Any, upper: Any) -> Any:
   r"""Projection onto box constraints.
 
   .. math::
 
-    \underset{p}{\text{argmin}} ~ ||x - p||_2^2 \quad \textrm{subject to} \quad
+    \underset{p}{\text{argmin}} ~ \|x - p\|_2^2 \quad \textrm{subject to} \quad
     \text{lower} \le p \le \text{upper}
 
   where :math:`x` is the input tree.
@@ -67,38 +63,38 @@ def projection_box(tree: Any, lower: Any, upper: Any) -> Any:
   Returns:
     projected tree, with the same structure as ``tree``.
   """
-  return jax.tree.map(_clip_safe, tree, lower, upper)
+  return jax.tree.map(jnp.clip, tree, lower, upper)
 
 
-def projection_hypercube(tree: Any, scale: Any = 1.0) -> Any:
+def projection_hypercube(tree: Any, scale: Any = 1) -> Any:
   r"""Projection onto the (unit) hypercube.
 
   .. math::
 
-    \underset{p}{\text{argmin}} ~ ||x - p||_2^2 \quad \textrm{subject to} \quad
+    \underset{p}{\text{argmin}} ~ \|x - p\|_2^2 \quad \textrm{subject to} \quad
     0 \le p \le \text{scale}
 
   where :math:`x` is the input tree.
 
-  By default, we project to the unit hypercube (`scale=1.0`).
+  By default, we project to the unit hypercube (`scale=1`).
 
   This is a convenience wrapper around
   :func:`projection_box <optax.projections.projection_box>`.
 
   Args:
     tree: tree to project.
-    scale: scale of the hypercube, a scalar or a tree (default: 1.0).
+    scale: scale of the hypercube, a scalar or a tree (default: 1).
 
   Returns:
     projected tree, with the same structure as ``tree``.
   """
-  return projection_box(tree, lower=0.0, upper=scale)
+  return projection_box(tree, lower=0, upper=scale)
 
 
 @jax.custom_jvp
 def _projection_unit_simplex(values: chex.Array) -> chex.Array:
   """Projection onto the unit simplex."""
-  s = 1.0
+  s = 1
   n_features = values.shape[0]
   u = jnp.sort(values)[::-1]
   cumsum_u = jnp.cumsum(u)
@@ -121,22 +117,22 @@ def _projection_unit_simplex_jvp(
   return primal_out, tangent_out
 
 
-def projection_simplex(tree: Any, scale: chex.Numeric = 1.0) -> Any:
+def projection_simplex(tree: Any, scale: chex.Numeric = 1) -> Any:
   r"""Projection onto a simplex.
 
   This function solves the following constrained optimization problem,
   where ``x`` is the input tree.
 
   .. math::
 
-    \underset{p}{\text{argmin}} ~ ||x - p||_2^2 \quad \textrm{subject to} \quad
+    \underset{p}{\text{argmin}} ~ \|x - p\|_2^2 \quad \textrm{subject to} \quad
     p \ge 0, p^\top 1 = \text{scale}
 
   By default, the projection is onto the probability simplex (unit simplex).
 
   Args:
     tree: tree to project.
-    scale: value the projected tree should sum to (default: 1.0).
+    scale: value the projected tree should sum to (default: 1).
 
   Returns:
     projected tree, a tree with the same structure as ``tree``.
@@ -156,25 +152,21 @@ def projection_simplex(tree: Any, scale: chex.Numeric = 1.0) -> Any:
 
   .. versionadded:: 0.2.3
   """
-  if scale is None:
-    scale = 1.0
-
   values, unravel_fn = flatten_util.ravel_pytree(tree)
   new_values = scale * _projection_unit_simplex(values / scale)
-
   return unravel_fn(new_values)
 
 
-def projection_l1_sphere(tree: Any, scale: float = 1.0) -> Any:
+def projection_l1_sphere(tree: Any, scale: chex.Numeric = 1) -> Any:
   r"""Projection onto the l1 sphere.
 
   This function solves the following constrained optimization problem,
   where ``x`` is the input tree.
 
   .. math::
 
-    \underset{y}{\text{argmin}} ~ ||x - y||_2^2 \quad \textrm{subject to} \quad
-    ||y||_1 = \text{scale}
+    \underset{y}{\text{argmin}} ~ \|x - y\|_2^2 \quad \textrm{subject to} \quad
+    \|y\|_1 = \text{scale}
 
   Args:
     tree: tree to project.
@@ -189,16 +181,16 @@ def projection_l1_sphere(tree: Any, scale: float = 1.0) -> Any:
   return otu.tree_mul(tree_sign, tree_abs_proj)
 
 
-def projection_l1_ball(tree: Any, scale: float = 1.0) -> Any:
+def projection_l1_ball(tree: Any, scale: chex.Numeric = 1) -> Any:
   r"""Projection onto the l1 ball.
 
   This function solves the following constrained optimization problem,
   where ``x`` is the input tree.
 
   .. math::
 
-    \underset{y}{\text{argmin}} ~ ||x - y||_2^2 \quad \textrm{subject to} \quad
-    ||y||_1 \le \text{scale}
+    \underset{y}{\text{argmin}} ~ \|x - y\|_2^2 \quad \textrm{subject to} \quad
+    \|y\|_1 \le \text{scale}
 
   Args:
     tree: tree to project.
@@ -229,16 +221,16 @@ def projection_l1_ball(tree: Any, scale: float = 1.0) -> Any:
   )
 
 
-def projection_l2_sphere(tree: Any, scale: float = 1.0) -> Any:
+def projection_l2_sphere(tree: Any, scale: chex.Numeric = 1) -> Any:
   r"""Projection onto the l2 sphere.
 
   This function solves the following constrained optimization problem,
   where ``x`` is the input tree.
 
   .. math::
 
-    \underset{y}{\text{argmin}} ~ ||x - y||_2^2 \quad \textrm{subject to} \quad
-    ||y||_2 = \text{value}
+    \underset{y}{\text{argmin}} ~ \|x - y\|_2^2 \quad \textrm{subject to} \quad
+    \|y\|_2 = \text{value}
 
   Args:
     tree: tree to project.
@@ -253,16 +245,16 @@ def projection_l2_sphere(tree: Any, scale: float = 1.0) -> Any:
   return otu.tree_scale(factor, tree)
 
 
-def projection_l2_ball(tree: Any, scale: float = 1.0) -> Any:
+def projection_l2_ball(tree: Any, scale: chex.Numeric = 1) -> Any:
   r"""Projection onto the l2 ball.
 
   This function solves the following constrained optimization problem,
   where ``x`` is the input tree.
 
   .. math::
 
-    \underset{y}{\text{argmin}} ~ ||x - y||_2^2 \quad \textrm{subject to} \quad
-    ||y||_2 \le \text{scale}
+    \underset{y}{\text{argmin}} ~ \|x - y\|_2^2 \quad \textrm{subject to} \quad
+    \|y\|_2 \le \text{scale}
 
   Args:
     tree: tree to project.
@@ -273,26 +265,20 @@ def projection_l2_ball(tree: Any, scale: float = 1.0) -> Any:
 
   .. versionadded:: 0.2.4
   """
-  l2_norm = otu.tree_l2_norm(tree)
-  factor = scale / l2_norm
-  return jax.lax.cond(
-      l2_norm <= scale,
-      lambda tree: tree,
-      lambda tree: otu.tree_scale(factor, tree),
-      operand=tree,
-  )
+  factor = scale / otu.tree_l2_norm(tree).clip(min=scale)
+  return otu.tree_scale(factor, tree)
 
 
-def projection_linf_ball(tree: Any, scale: float = 1.0) -> Any:
+def projection_linf_ball(tree: Any, scale: chex.Numeric = 1) -> Any:
   r"""Projection onto the l-infinity ball.
 
   This function solves the following constrained optimization problem,
   where ``x`` is the input tree.
 
   .. math::
 
-    \underset{y}{\text{argmin}} ~ ||x - y||_2^2 \quad \textrm{subject to} \quad
-    ||y||_{\infty} \le \text{scale}
+    \underset{y}{\text{argmin}} ~ \|x - y\|_2^2 \quad \textrm{subject to} \quad
+    \|y\|_{\infty} \le \text{scale}
 
   Args:
     tree: tree to project.
@@ -301,6 +287,6 @@ def projection_linf_ball(tree: Any, scale: float = 1.0) -> Any:
   Returns:
     projected tree, with the same structure as ``tree``.
   """
-  lower_tree = otu.tree_full_like(tree, -scale)
-  upper_tree = otu.tree_full_like(tree, scale)
-  return projection_box(tree, lower=lower_tree, upper=upper_tree)
+  lower = otu.tree_full_like(tree, -scale)
+  upper = otu.tree_full_like(tree, scale)
+  return projection_box(tree, lower, upper)