Merge remote-tracking branch 'origin/issue-1118' into issue-1118

optax/losses/__init__.py

@@ -42,4 +42,5 @@
 from optax.losses._regression import log_cosh
 from optax.losses._regression import squared_error
 from optax.losses._self_supervised import ntxent
+from optax.losses._self_supervised import triplet_margin_loss
 from optax.losses._smoothing import smooth_labels

optax/losses/_self_supervised.py

@@ -111,11 +111,6 @@ def triplet_margin_loss(
 ) -> chex.Array:
   """Triplet margin loss function.
 
-  References:
-    V. Balntas et al. `Learning shallow convolutional feature
-    descriptors with triplet losses 
-    <https://bmva-archive.org.uk/bmvc/2016/papers/paper119/paper119.pdf>`_, 2016
-  
   Measures the relative similarity between an anchor point, 
   a positive point, and a negative point using the distance 
   metric specified by p-norm. The loss encourages
@@ -133,21 +128,28 @@ def triplet_margin_loss(
       swap: Use the distance swap optimization
       reduction: Specifies the reduction to apply to the output:
           'none' | 'mean' | 'sum'
-
+  
   Returns:
       The triplet margin loss value.
       If reduction is 'none': tensor of shape [batch_size]
       If reduction is 'mean' or 'sum': scalar tensor.
   
-  example:
+  Examples:
         >>> import jax.numpy as jnp
+        >>> import optax
         >>> anchor = jnp.array([[1.0, 2.0], [3.0, 4.0]])
         >>> positive = jnp.array([[1.1, 2.1], [3.1, 4.1]])
         >>> negative = jnp.array([[2.0, 3.0], [4.0, 5.0]])
         >>> margin = 1.0
-        >>> loss = triplet_margin_loss(anchor, positive, negative,
+        >>> loss = optax.losses.triplet_margin_loss(anchor, positive, negative,
         >>> margin=margin, reduction='mean')
         >>> print(loss)
+
+  References:
+    V. Balntas et al. `Learning shallow convolutional feature
+    descriptors with triplet losses 
+    <https://bmva-archive.org.uk/bmvc/2016/papers/paper119/paper119.pdf>`_, 2016
+  
   """
   chex.assert_equal_shape([anchor, positive, negative])
 

optax/losses/_self_supervised_test.py

@@ -64,6 +64,20 @@ def setUp(self):
     self.p2 = jnp.ones((2, 2))
     self.n2 = jnp.ones((2, 2))*2
 
+  @chex.all_variants
+  def test_vmap(self):
+    # VMAP applied function result
+    original_loss = _self_supervised.triplet_margin_loss(
+                    self.a1, self.p1, self.n1)
+    self.a1 = self.a1.reshape(1, *self.a1.shape)
+    self.p1 = self.p1.reshape(1, *self.p1.shape)
+    self.n1 = self.n1.reshape(1, *self.n1.shape)
+    vmap_loss = self.variant(jax.vmap(
+                _self_supervised.triplet_margin_loss
+                ))(self.a1, self.p1,
+                self.n1)
+    np.testing.assert_allclose(vmap_loss, original_loss, atol=1e-4)
+
   @chex.all_variants
   def test_batched(self):
     def testing_triplet_loss(a, p, n, margin=1.0, swap=False):
@@ -106,16 +120,6 @@ def testing_triplet_loss(a, p, n, margin=1.0, swap=False):
     np.testing.assert_allclose(jit_loss, original_loss,
                                atol=1e-4)
 
-    # VMAP applied function result
-    self.a1 = self.a1.reshape(1, *self.a1.shape)
-    self.p1 = self.p1.reshape(1, *self.p1.shape)
-    self.n1 = self.n1.reshape(1, *self.n1.shape)
-    vmap_loss = self.variant(jax.vmap(
-                _self_supervised.triplet_margin_loss
-                ))(self.a1, self.p1,
-                self.n1)
-    np.testing.assert_allclose(vmap_loss, original_loss, atol=1e-4)
-
 
 if __name__ == '__main__':
   absltest.main()