Repository: google-deepmind/optax

Issue:
Fix softmax_cross_entropy to handle -inf logits correctly when corresponding label is 0.
#896


Discussion:
Would perhaps using logsumexp (https://jax.readthedocs.io/en/latest/_autosummary/jax.scipy.special.logsumexp.html) be a solution here?

Also, please add tests :-)
@fabianp Added tests.

`log_softmax` uses `logsumexp` internally: https://github.com/google-deepmind/optax/issues/896#issuecomment-2032688959.

Resolving patch:
diff --git a/optax/losses/_classification.py b/optax/losses/_classification.py
--- a/optax/losses/_classification.py
+++ b/optax/losses/_classification.py
@@ -134,6 +134,52 @@ def binary_sparsemax_loss(logits, labels):
   return sparsemax_loss(logits, labels)
 
 
+@jax.custom_jvp
+def weighted_logsoftmax(x: chex.Array, weights: chex.Array) -> chex.Array:
+  r"""Weighted logsoftmax.
+
+  Computes
+  .. math::
+    (w_i \log(\exp x_i /(\sum_i \exp x_i )) )_{i=1}^n
+
+  for :math:`x` the input ``x``, :math:`w` the ``weights``.
+  For :math:`w_i = 0`, :math:`x_i=-\infty`, this implementation ensures that the
+  output is 0 and not nan at the ith entry following the convention that
+  :math:`0 \log 0 = 0`.
+
+  Args:
+    x: input array.
+    weights: weights.
+
+  Returns:
+    logsoftmax of x multiplied elementwise by weights
+  """
+  logsoftmax_x = jax.nn.log_softmax(x, axis=-1)
+  return jnp.where(
+      weights != 0.0, weights * logsoftmax_x, jnp.zeros_like(logsoftmax_x)
+  )
+
+
+def _weighted_logsoftmax_jvp(primals, tangents):
+  """Custom JVP of weighted logsoftmax."""
+  (x, weights) = primals
+  (x_dot, weights_dot) = tangents
+  logsoftmax_x = jax.nn.log_softmax(x, axis=-1)
+  result = jnp.where(
+      weights != 0.0, weights * logsoftmax_x, jnp.zeros_like(logsoftmax_x)
+  )
+  out_tangents = (
+      weights * x_dot
+      - weights
+      * jnp.sum(x_dot * jax.nn.softmax(x, axis=-1), axis=-1, keepdims=True)
+      + weights_dot * logsoftmax_x
+  )
+  return result, out_tangents
+
+
+weighted_logsoftmax.defjvp(_weighted_logsoftmax_jvp)
+
+
 def softmax_cross_entropy(
     logits: chex.Array,
     labels: chex.Array,
@@ -159,7 +205,7 @@ def softmax_cross_entropy(
     distributions, with shape `[...]`.
   """
   chex.assert_type([logits], float)
-  return -jnp.sum(labels * jax.nn.log_softmax(logits, axis=-1), axis=-1)
+  return -jnp.sum(weighted_logsoftmax(logits, labels), axis=-1)
 
 
 def softmax_cross_entropy_with_integer_labels(


Test patch:
diff --git a/optax/losses/_classification_test.py b/optax/losses/_classification_test.py
--- a/optax/losses/_classification_test.py
+++ b/optax/losses/_classification_test.py
@@ -14,14 +14,15 @@
 # ==============================================================================
 """Tests for optax.losses._classification."""
 
+import functools
+
 from absl.testing import absltest
 from absl.testing import parameterized
-
 import chex
 import jax
 import jax.numpy as jnp
+import jax.test_util as jaxtest
 import numpy as np
-
 from optax.losses import _classification
 
 
@@ -29,92 +30,185 @@ class SoftmaxCrossEntropyTest(parameterized.TestCase):
 
   def setUp(self):
     super().setUp()
-    self.ys = np.array([[10., 1., -2.], [1., 4., 0.2]], dtype=np.float32)
-    self.ts = np.array([[0., 1., 0.], [1., 0., 0.]], dtype=np.float32)
-    # taken expected outputs from rlax.
-    self.exp = np.array([9.00013, 3.0696733], dtype=np.float32)
+    self.ys = np.array(
+        [
+            [10.0, 1.0, -2.0],
+            [1.0, 4.0, 0.2],
+            [-np.inf, 0.0, 0.0],
+            [-np.inf, 0.0, 0.0],
+            [-np.inf, 0.0, -np.inf],
+        ],
+        dtype=np.float32,
+    )
+    self.ts = np.array(
+        [
+            [0.0, 1.0, 0.0],
+            [1.0, 0.0, 0.0],
+            [0.0, 0.5, 0.5],
+            [0.4, 0.3, 0.3],
+            [0.0, 1.0, 0.0],
+        ],
+        dtype=np.float32,
+    )
+    self.exp = np.array(
+        [
+            9.00013,
+            3.0696733,
+            0.693147,
+            np.inf,
+            0.0,
+        ],
+        dtype=np.float32,
+    )
 
   @chex.all_variants
   def test_scalar(self):
     """Tests for a full batch."""
     np.testing.assert_allclose(
-        self.variant(
-            _classification.softmax_cross_entropy)(self.ys[0], self.ts[0]),
-        self.exp[0], atol=1e-4)
+        self.variant(_classification.softmax_cross_entropy)(
+            self.ys[0], self.ts[0]
+        ),
+        self.exp[0],
+        atol=1e-4,
+    )
 
   @chex.all_variants
   def test_batched(self):
     """Tests for a full batch."""
     np.testing.assert_allclose(
         self.variant(_classification.softmax_cross_entropy)(self.ys, self.ts),
-        self.exp, atol=1e-4)
+        self.exp,
+        atol=1e-4,
+    )
+
+  def test_gradient(self):
+    """Tests gradient ok."""
+    jaxtest.check_grads(
+        _classification.softmax_cross_entropy,
+        (self.ys[:2], self.ts[:2]),
+        order=1,
+    )
+
+  def test_against_plain_implementation(self):
+    """Tests against plain implementation which does not handle -inf."""
+
+    def _softmax_cross_entropy(logits, labels):
+      return -jnp.sum(labels * jax.nn.log_softmax(logits, axis=-1), axis=-1)
+
+    plain_val_and_grad = jax.value_and_grad(_softmax_cross_entropy)(
+        self.ys[0], self.ts[0]
+    )
+    val_and_grad = jax.value_and_grad(_classification.softmax_cross_entropy)(
+        self.ys[0], self.ts[0]
+    )
+    chex.assert_trees_all_close(plain_val_and_grad, val_and_grad, atol=1e-4)
 
 
 class SoftmaxCrossEntropyWithIntegerLabelsTest(parameterized.TestCase):
 
   def setUp(self):
     super().setUp()
-    self.ys = np.array([[10., 1., -2.], [1., 4., 0.2]], dtype=np.float32)
+    self.ys = np.array([[10.0, 1.0, -2.0], [1.0, 4.0, 0.2]], dtype=np.float32)
     self.ts = np.array([1, 0], dtype=np.int32)
 
   @chex.all_variants
   def test_consistent_with_softmax_cross_entropy_scalar(self):
     """Tests for a scalar."""
     exp = _classification.softmax_cross_entropy(
-        self.ys[0], jax.nn.one_hot(self.ts[0], 3))
+        self.ys[0], jax.nn.one_hot(self.ts[0], 3)
+    )
     np.testing.assert_allclose(
         self.variant(_classification.softmax_cross_entropy_with_integer_labels)(
-            self.ys[0], self.ts[0]),
-        exp, rtol=1e-6)
+            self.ys[0], self.ts[0]
+        ),
+        exp,
+        rtol=1e-6,
+    )
 
   @chex.all_variants
   def test_consistent_with_softmax_cross_entropy_batched(self):
     """Tests for a full batch."""
     exp = _classification.softmax_cross_entropy(
-        self.ys, jax.nn.one_hot(self.ts, 3))
+        self.ys, jax.nn.one_hot(self.ts, 3)
+    )
     np.testing.assert_allclose(
         self.variant(_classification.softmax_cross_entropy_with_integer_labels)(
-            self.ys, self.ts),
-        exp, rtol=1e-6)
+            self.ys, self.ts
+        ),
+        exp,
+        rtol=1e-6,
+    )
+
+  def test_gradient(self):
+    """Tests gradient ok."""
+    jaxtest.check_grads(
+        functools.partial(
+            _classification.softmax_cross_entropy_with_integer_labels,
+            labels=self.ts,
+        ),
+        (self.ys,),
+        order=1,
+    )
 
 
 class SigmoidCrossEntropyTest(parameterized.TestCase):
 
   @parameterized.parameters(
-      dict(preds=np.array([-1e+09, -1e-09]),
-           labels=np.array([1., 0.]),
-           expected=5e+08),
-      dict(preds=np.array([-1e+09, -1e-09]),
-           labels=np.array([0., 1.]),
-           expected=0.3465736),
-      dict(preds=np.array([1e+09, 1e-09]),
-           labels=np.array([1., 0.]),
-           expected=0.3465736),
-      dict(preds=np.array([1e+09, 1e-09]),
-           labels=np.array([0., 1.]),
-           expected=5e+08),
-      dict(preds=np.array([-1e+09, 1e-09]),
-           labels=np.array([1., 0.]),
-           expected=5e+08),
-      dict(preds=np.array([-1e+09, 1e-09]),
-           labels=np.array([0., 1.]),
-           expected=0.3465736),
-      dict(preds=np.array([1e+09, -1e-09]),
-           labels=np.array([1., 0.]),
-           expected=0.3465736),
-      dict(preds=np.array([1e+09, -1e-09]),
-           labels=np.array([0., 1.]),
-           expected=5e+08),
-      dict(preds=np.array([0., 0.]),
-           labels=np.array([1., 0.]),
-           expected=0.6931472),
-      dict(preds=np.array([0., 0.]),
-           labels=np.array([0., 1.]),
-           expected=0.6931472),
+      dict(
+          preds=np.array([-1e09, -1e-09]),
+          labels=np.array([1.0, 0.0]),
+          expected=5e08,
+      ),
+      dict(
+          preds=np.array([-1e09, -1e-09]),
+          labels=np.array([0.0, 1.0]),
+          expected=0.3465736,
+      ),
+      dict(
+          preds=np.array([1e09, 1e-09]),
+          labels=np.array([1.0, 0.0]),
+          expected=0.3465736,
+      ),
+      dict(
+          preds=np.array([1e09, 1e-09]),
+          labels=np.array([0.0, 1.0]),
+          expected=5e08,
+      ),
+      dict(
+          preds=np.array([-1e09, 1e-09]),
+          labels=np.array([1.0, 0.0]),
+          expected=5e08,
+      ),
+      dict(
+          preds=np.array([-1e09, 1e-09]),
+          labels=np.array([0.0, 1.0]),
+          expected=0.3465736,
+      ),
+      dict(
+          preds=np.array([1e09, -1e-09]),
+          labels=np.array([1.0, 0.0]),
+          expected=0.3465736,
+      ),
+      dict(
+          preds=np.array([1e09, -1e-09]),
+          labels=np.array([0.0, 1.0]),
+          expected=5e08,
+      ),
+      dict(
+          preds=np.array([0.0, 0.0]),
+          labels=np.array([1.0, 0.0]),
+          expected=0.6931472,
+      ),
+      dict(
+          preds=np.array([0.0, 0.0]),
+          labels=np.array([0.0, 1.0]),
+          expected=0.6931472,
+      ),
   )
   def testSigmoidCrossEntropy(self, preds, labels, expected):
     tested = jnp.mean(
-        _classification.sigmoid_binary_cross_entropy(preds, labels))
+        _classification.sigmoid_binary_cross_entropy(preds, labels)
+    )
     np.testing.assert_allclose(tested, expected, rtol=1e-6, atol=1e-6)
 
 
@@ -197,9 +291,9 @@ def test_batched(self, eps, expected):
   def test_equals_to_cross_entropy_when_eps0(self, logits, labels):
     np.testing.assert_allclose(
         self.variant(_classification.poly_loss_cross_entropy)(
-            logits, labels, epsilon=0.0),
-        self.variant(_classification.softmax_cross_entropy)(
-            logits, labels),
+            logits, labels, epsilon=0.0
+        ),
+        self.variant(_classification.softmax_cross_entropy)(logits, labels),
         atol=1e-4,
     )
 
@@ -209,9 +303,11 @@ class HingeTest(parameterized.TestCase):
   def test_binary(self):
     label = jnp.array(1)
     signed_label = jnp.array(2.0 * label - 1.0)
-    score = jnp.array(10.)
+    score = jnp.array(10.0)
+
     def reference_impl(label, logit):
       return jax.nn.relu(1 - logit * (2.0 * label - 1.0))
+
     expected = reference_impl(label, score)
     result = _classification.hinge_loss(score, signed_label)
     np.testing.assert_allclose(result, expected, atol=1e-4)
@@ -219,9 +315,11 @@ def reference_impl(label, logit):
   def test_batched_binary(self):
     labels = jnp.array([1, 0])
     signed_labels = jnp.array(2.0 * labels - 1.0)
-    scores = jnp.array([10., 20.])
+    scores = jnp.array([10.0, 20.0])
+
     def reference_impl(label, logit):
       return jax.nn.relu(1 - logit * (2.0 * label - 1.0))
+
     expected = jax.vmap(reference_impl)(labels, scores)
     # no need to vmap the optax loss. leading dimensions automatically handled.
     result = _classification.hinge_loss(scores, signed_labels)
@@ -229,20 +327,24 @@ def reference_impl(label, logit):
 
   def test_multi_class(self):
     label = jnp.array(1)
-    scores = jnp.array([10., 3.])
+    scores = jnp.array([10.0, 3.0])
+
     def reference_impl(label, scores):
       one_hot_label = jax.nn.one_hot(label, scores.shape[-1])
       return jnp.max(scores + 1.0 - one_hot_label) - scores[label]
+
     expected = reference_impl(label, scores)
     result = _classification.multiclass_hinge_loss(scores, label)
     np.testing.assert_allclose(result, expected, atol=1e-4)
 
   def test_batched_multi_class(self):
     label = jnp.array([1, 0])
-    scores = jnp.array([[10., 3.], [11., -2.]])
+    scores = jnp.array([[10.0, 3.0], [11.0, -2.0]])
+
     def reference_impl(label, scores):
       one_hot_label = jax.nn.one_hot(label, scores.shape[-1])
       return jnp.max(scores + 1.0 - one_hot_label) - scores[label]
+
     expected = jax.vmap(reference_impl)(label, scores)
     # no need to vmap the optax loss. leading dimensions automatically handled.
     result = _classification.multiclass_hinge_loss(scores, label)
@@ -253,34 +355,40 @@ class SparsemaxTest(parameterized.TestCase):
 
   def test_binary(self):
     label = 1
-    score = 10.
+    score = 10.0
+
     def reference_impl(label, logit):
-      scores = -(2*label-1)*logit
+      scores = -(2 * label - 1) * logit
       if scores <= -1.0:
         return 0.0
       elif scores >= 1.0:
         return scores
       else:
         return (scores + 1.0) ** 2 / 4
+
     expected = reference_impl(label, score)
     result = _classification.sparsemax_loss(
-        jnp.asarray(score), jnp.asarray(label))
+        jnp.asarray(score), jnp.asarray(label)
+    )
     np.testing.assert_allclose(result, expected, atol=1e-4)
 
   def test_batched_binary(self):
     labels = jnp.array([1, 0])
-    scores = jnp.array([10., 20.])
+    scores = jnp.array([10.0, 20.0])
+
     def reference_impl(label, logit):
-      scores = -(2*label-1)*logit
+      scores = -(2 * label - 1) * logit
       if scores <= -1.0:
         return 0.0
       elif scores >= 1.0:
         return scores
       else:
         return (scores + 1.0) ** 2 / 4
+
     expected = jnp.asarray([
         reference_impl(labels[0], scores[0]),
-        reference_impl(labels[1], scores[1])])
+        reference_impl(labels[1], scores[1]),
+    ])
     # in the optax loss the leading dimensions are automatically handled.
     result = _classification.sparsemax_loss(scores, labels)
     np.testing.assert_allclose(result, expected, atol=1e-4)
@@ -305,7 +413,8 @@ def setUp(self):
   def test_scalar(self):
     np.testing.assert_allclose(
         self.variant(_classification.convex_kl_divergence)(
-            self.log_ps[0], self.qs[0]),
+            self.log_ps[0], self.qs[0]
+        ),
         self.exp[0],
         atol=1e-4,
     )
@@ -314,7 +423,8 @@ def test_scalar(self):
   def test_batched(self):
     np.testing.assert_allclose(
         self.variant(_classification.convex_kl_divergence)(
-            self.log_ps, self.qs),
+            self.log_ps, self.qs
+        ),
         self.exp,
         atol=1e-4,
     )
@@ -325,9 +435,11 @@ class PerceptronTest(parameterized.TestCase):
   def test_binary(self):
     label = jnp.array(1)
     signed_label = jnp.array(2.0 * label - 1.0)
-    score = jnp.array(10.)
+    score = jnp.array(10.0)
+
     def reference_impl(label, logit) -> float:
-      return jax.nn.relu(- logit * (2.0 * label - 1.0))
+      return jax.nn.relu(-logit * (2.0 * label - 1.0))
+
     expected = reference_impl(label, score)
     result = _classification.perceptron_loss(score, signed_label)
     np.testing.assert_allclose(result, expected, atol=1e-4)
@@ -335,9 +447,11 @@ def reference_impl(label, logit) -> float:
   def test_batched_binary(self):
     labels = jnp.array([1, 0])
     signed_labels = jnp.array(2.0 * labels - 1.0)
-    scores = jnp.array([10., 20.])
+    scores = jnp.array([10.0, 20.0])
+
     def reference_impl(label, logit) -> float:
-      return jax.nn.relu(- logit * (2.0 * label - 1.0))
+      return jax.nn.relu(-logit * (2.0 * label - 1.0))
+
     expected = jax.vmap(reference_impl)(labels, scores)
     # no need to vmap the optax loss. leading dimensions automatically handled.
     result = _classification.perceptron_loss(scores, signed_labels)
@@ -345,18 +459,22 @@ def reference_impl(label, logit) -> float:
 
   def test_multi_class(self):
     label = jnp.array(1)
-    scores = jnp.array([10., 3.])
+    scores = jnp.array([10.0, 3.0])
+
     def reference_impl(label, scores):
       return jnp.max(scores) - scores[label]
+
     expected = reference_impl(label, scores)
     result = _classification.multiclass_perceptron_loss(scores, label)
     np.testing.assert_allclose(result, expected, atol=1e-4)
 
   def test_batched_multi_class(self):
     label = jnp.array([1, 0])
-    scores = jnp.array([[10., 3.], [11., -2.]])
+    scores = jnp.array([[10.0, 3.0], [11.0, -2.0]])
+
     def reference_impl(label, scores):
       return jnp.max(scores) - scores[label]
+
     expected = jax.vmap(reference_impl)(label, scores)
     # no need to vmap the optax loss. leading dimensions automatically handled.
     result = _classification.multiclass_perceptron_loss(scores, label)
@@ -367,11 +485,13 @@ class KLDivergenceTest(parameterized.TestCase):
 
   def setUp(self):
     super().setUp()
-    self.log_ps = np.array(
-        [[-2.9957, -3.5066, -3.9120, -1.2040, -0.6931, -2.3026],
-         [-1.6094, -1.6094, -1.6094, -2.3026, -1.8971, -1.8971]])
-    self.qs = np.array([[0.2, 0.2, 0.2, 0.1, 0.15, 0.15],
-                        [0.05, 0.03, 0.02, 0.3, 0.5, 0.]])
+    self.log_ps = np.array([
+        [-2.9957, -3.5066, -3.9120, -1.2040, -0.6931, -2.3026],
+        [-1.6094, -1.6094, -1.6094, -2.3026, -1.8971, -1.8971],
+    ])
+    self.qs = np.array(
+        [[0.2, 0.2, 0.2, 0.1, 0.15, 0.15], [0.05, 0.03, 0.02, 0.3, 0.5, 0.0]]
+    )
     # Computed kullback-leibler divergence of P from Q.
     self.exp = np.array([0.8875577, 0.7592807])
 
@@ -380,25 +500,30 @@ def test_scalar(self):
     np.testing.assert_allclose(
         self.variant(_classification.kl_divergence)(self.log_ps[0], self.qs[0]),
         self.exp[0],
-        atol=1e-4)
+        atol=1e-4,
+    )
 
   @chex.all_variants
   def test_batched(self):
     np.testing.assert_allclose(
         self.variant(_classification.kl_divergence)(self.log_ps, self.qs),
         self.exp,
-        atol=1e-4)
+        atol=1e-4,
+    )
 
 
 class KLDivergenceWithLogTargetsTest(parameterized.TestCase):
 
   def setUp(self):
     super().setUp()
-    self.log_ps = np.array(
-        [[-2.9957, -3.5066, -3.9120, -1.2040, -0.6931, -2.3026],
-         [-1.6094, -1.6094, -1.6094, -2.3026, -1.8971, -1.8971]])
-    self.qs = np.array([[-1.6094, -1.6094, -1.6094, -2.3026, -1.8971, -1.8971],
-                        [-2.9957, -3.5066, -3.9120, -1.2040, -0.6931, -2.3026]])
+    self.log_ps = np.array([
+        [-2.9957, -3.5066, -3.9120, -1.2040, -0.6931, -2.3026],
+        [-1.6094, -1.6094, -1.6094, -2.3026, -1.8971, -1.8971],
+    ])
+    self.qs = np.array([
+        [-1.6094, -1.6094, -1.6094, -2.3026, -1.8971, -1.8971],
+        [-2.9957, -3.5066, -3.9120, -1.2040, -0.6931, -2.3026],
+    ])
     # Computed kullback-leibler divergence of P from Q.
     self.exp = np.array([0.8875625, 0.7187435584901326])
 
@@ -406,17 +531,21 @@ def setUp(self):
   def test_scalar(self):
     np.testing.assert_allclose(
         self.variant(_classification.kl_divergence_with_log_targets)(
-            self.log_ps[0], self.qs[0]),
+            self.log_ps[0], self.qs[0]
+        ),
         self.exp[0],
-        atol=1e-4)
+        atol=1e-4,
+    )
 
   @chex.all_variants
   def test_batched(self):
     np.testing.assert_allclose(
         self.variant(_classification.kl_divergence_with_log_targets)(
-            self.log_ps, self.qs),
+            self.log_ps, self.qs
+        ),
         self.exp,
-        atol=1e-4)
+        atol=1e-4,
+    )
 
 
 def _lengths_to_paddings(lengths: chex.Array, maxlength: int) -> chex.Array:
@@ -427,13 +556,16 @@ def _lengths_to_paddings(lengths: chex.Array, maxlength: int) -> chex.Array:
 
 
 def _average_ctc_loss(
-    logprobs: chex.Array, logprob_paddings: chex.Array,
+    logprobs: chex.Array,
+    logprob_paddings: chex.Array,
     labels: chex.Array,
-    label_paddings: chex.Array
+    label_paddings: chex.Array,
 ) -> chex.Array:
   return jnp.average(
       _classification.ctc_loss(
-          logprobs, logprob_paddings, labels, label_paddings))
+          logprobs, logprob_paddings, labels, label_paddings
+      )
+  )
 
 
 class CTCTest(parameterized.TestCase):
@@ -450,8 +582,9 @@ def test_with_one_to_one_alignment(self):
     steps = 50
     nclasses = 40
     logits = np.random.randn(batchsize, steps, nclasses)
-    labels = np.random.uniform(
-        1, nclasses, size=(batchsize, steps)).astype(np.int32)
+    labels = np.random.uniform(1, nclasses, size=(batchsize, steps)).astype(
+        np.int32
+    )
 
     # This function only covers the cases without same-label repetition.
     # `test_repeat_with_one_to_one_alignment` below complements those cases.
@@ -462,8 +595,8 @@ def test_with_one_to_one_alignment(self):
           labels[n, t] = np.random.uniform(1, nclasses)
 
     results = self.variant(_classification.ctc_loss_with_forward_probs)(
-        logits, np.zeros(logits.shape[:2]),
-        labels, np.zeros(labels.shape))
+        logits, np.zeros(logits.shape[:2]), labels, np.zeros(labels.shape)
+    )
     (per_seq_loss, logalpha_blank, logalpha_emit) = results
 
     logprobs = jax.nn.log_softmax(logits)
@@ -471,28 +604,27 @@ def test_with_one_to_one_alignment(self):
       p = 0.0
       for t in range(steps):
         p += logprobs[b, t, labels[b, t]]
-      np.testing.assert_allclose(
-          np.array(-p), per_seq_loss[b], rtol=self._rtol)
+      np.testing.assert_allclose(np.array(-p), per_seq_loss[b], rtol=self._rtol)
 
       # Check forward-probabilities.
       # 1. All-phi path: logalpha_blank[-1, b, 0] must be a probability of
       #   the path that outputs blank symbols for all the frames.
-      np.testing.assert_allclose(logalpha_blank[-1, b, 0],
-                                 np.sum(logprobs[b, :, 0]),
-                                 rtol=self._rtol)
+      np.testing.assert_allclose(
+          logalpha_blank[-1, b, 0], np.sum(logprobs[b, :, 0]), rtol=self._rtol
+      )
 
       # 2. After emitting all the labels
       #   the negated loss must be identical with the forward probability of
       #   paths after consuming all the labels (because one-to-one alignment
       #   doesn't allow extra blank symbols)
-      np.testing.assert_allclose(logalpha_emit[-1, b, steps - 1],
-                                 -per_seq_loss[b],
-                                 rtol=self._rtol)
+      np.testing.assert_allclose(
+          logalpha_emit[-1, b, steps - 1], -per_seq_loss[b], rtol=self._rtol
+      )
       #   and, this forward probability must be copied to the blank forward
       #   probability of the next step.
-      np.testing.assert_allclose(logalpha_blank[-1, b, steps],
-                                 -per_seq_loss[b],
-                                 rtol=self._rtol)
+      np.testing.assert_allclose(
+          logalpha_blank[-1, b, steps], -per_seq_loss[b], rtol=self._rtol
+      )
 
   @chex.all_variants
   def test_with_one_to_one_alignment_and_paddings(self):
@@ -513,24 +645,28 @@ def test_with_one_to_one_alignment_and_paddings(self):
     paddings = _lengths_to_paddings(lengths, steps)
 
     actual_loss = self.variant(_classification.ctc_loss)(
-        logits, paddings, labels, paddings)
+        logits, paddings, labels, paddings
+    )
 
     value_and_grad = self.variant(jax.value_and_grad(_average_ctc_loss))
     unused_avg_loss, actual_gradients = value_and_grad(
-        logits, paddings, labels, paddings)
+        logits, paddings, labels, paddings
+    )
 
     for n in range(batch_size):
-      expected_loss = -sum(logprobs[n, t, k]
-                           for t, k in enumerate(labels[n, :lengths[n]]))
+      expected_loss = -sum(
+          logprobs[n, t, k] for t, k in enumerate(labels[n, : lengths[n]])
+      )
       np.testing.assert_allclose(expected_loss, actual_loss[n], rtol=self._rtol)
 
       expected_gradients = np.array(jax.nn.softmax(logits[n]))
-      expected_gradients[lengths[n]:] = 0.0
-      for t, k in enumerate(labels[n, :lengths[n]]):
+      expected_gradients[lengths[n] :] = 0.0
+      for t, k in enumerate(labels[n, : lengths[n]]):
         expected_gradients[t, k] -= 1.0
       expected_gradients /= batch_size
       np.testing.assert_allclose(
-          expected_gradients, actual_gradients[n], rtol=self._rtol)
+          expected_gradients, actual_gradients[n], rtol=self._rtol
+      )
 
   @chex.all_variants
   def test_repeat_with_one_to_one_alignment(self):
@@ -563,14 +699,17 @@ def test_repeat_with_one_to_one_alignment(self):
     logit_paddings = _lengths_to_paddings(logit_lengths, logit_steps)
 
     per_seq_loss = self.variant(_classification.ctc_loss)(
-        logits, logit_paddings, labels, label_paddings)
+        logits, logit_paddings, labels, label_paddings
+    )
 
     logprobs = jax.nn.log_softmax(logits)
     for n in range(batch_size):
-      expected_loss = -sum(logprobs[n, t, k]
-                           for t, k in enumerate(expected_alignment[n]))
+      expected_loss = -sum(
+          logprobs[n, t, k] for t, k in enumerate(expected_alignment[n])
+      )
       np.testing.assert_allclose(
-          jnp.array(expected_loss), per_seq_loss[n], rtol=self._rtol)
+          jnp.array(expected_loss), per_seq_loss[n], rtol=self._rtol
+      )
 
 
 class SigmoidFocalLossTest(parameterized.TestCase):