CTC API for JAX

keras/backend/jax/nn.py

@@ -558,3 +558,78 @@ def batch_normalization(
         res = res + offset
 
     return x * inv + res
+
+
+def ctc_loss(
+    target,
+    output,
+    target_length,
+    output_length,
+    mask_index=0,
+):
+    batch_size, _, _ = output.shape
+    batch_size, max_target_length = target.shape
+
+    output = output.transpose((1, 0, 2))
+    target = target.transpose((1, 0))
+
+    logits = jnn.log_softmax(output)
+    mgrid_t, mgrid_b = jnp.meshgrid(
+        jnp.arange(max_target_length), jnp.arange(batch_size)
+    )
+    logprobs_emit = logits[mgrid_t, mgrid_b, target[:, :, None]]
+    logprobs_mask = logits[:, :, mask_index]
+
+    logit_paddings = jnp.array(
+        jnp.arange(max_target_length) < output_length[:, None],
+        dtype=jnp.float32,
+    )
+
+    repeat = jnp.array(target[1:] == target[:-1])
+    repeat = jnp.pad(repeat, ((0, 1), (0, 0))).transpose((1, 0))
+
+    _logepsilon = -100000.0
+
+    def _iterate(prev, x):
+        prev_mask, prev_emit = prev
+        logprob_mask, logprob_emit, pad = x
+
+        prev_mask_orig = prev_mask
+        prev_mask = prev_mask.at[:, 1:].set(
+            jnp.logaddexp(prev_mask[:, 1:], prev_emit + _logepsilon * repeat),
+        )
+        emit = jnp.logaddexp(
+            prev_mask[:, :-1] + logprob_emit, prev_emit + logprob_emit
+        )
+
+        mask = prev_mask + logprob_mask[:, None]
+        mask = mask.at[:, 1:].set(
+            jnp.logaddexp(
+                mask[:, 1:],
+                prev_emit + logprob_mask[:, None] + _logepsilon * (1 - repeat),
+            )
+        )
+
+        pad = pad[:, None]
+        emit = emit * pad + prev_emit * (1 - pad)
+        mask = mask * pad + prev_mask_orig * (1 - pad)
+
+        return (mask, emit), (mask, emit)
+
+    mask_init = jnp.full((batch_size, max_target_length + 1), _logepsilon)
+    mask_init = mask_init.at[:, 0].set(0.0)
+    emit_init = jnp.full((batch_size, max_target_length), _logepsilon)
+
+    _, (alphas_mask, alphas_emit) = lax.scan(
+        _iterate,
+        (mask_init, emit_init),
+        (logprobs_mask, logprobs_emit, logit_paddings.transpose()),
+    )
+
+    last_alpha_mask = (
+        alphas_mask[-1]
+        .at[:, 1:]
+        .set(jnp.logaddexp(alphas_mask[-1, :, 1:], alphas_emit[-1]))
+    )
+
+    return -last_alpha_mask[jnp.arange(batch_size), target_length]

keras/backend/tensorflow/nn.py

@@ -838,4 +838,5 @@ def ctc_loss(
         label_length=target_length,
         logit_length=output_length,
         blank_index=mask_index,
+        logits_time_major=False,
     )

keras/backend/torch/nn.py

@@ -759,6 +759,7 @@ def ctc_loss(
     target_length = convert_to_tensor(target_length)
     output_length = convert_to_tensor(output_length)
 
+    output = torch.transpose(output, 1, 0)
     logits = tnn.log_softmax(output, dim=-1)
 
     return tnn.ctc_loss(

keras/ops/nn_test.py

@@ -975,8 +975,8 @@ def test_batch_normalization(self):
         )
 
     @pytest.mark.skipif(
-        backend.backend() not in ["tensorflow", "torch"],
-        reason="Only TF and Torch support CTC loss",
+        backend.backend() == "numpy",
+        reason="Numpy does not support CTC loss",
     )
     def test_ctc_loss(self):
         x = KerasTensor([10, 3, 4])
@@ -1762,16 +1762,15 @@ def test_batch_normalization(self):
         self.assertEqual(tuple(output.shape), (2, 3, 3, 5))
 
     @pytest.mark.skipif(
-        backend.backend() not in ["tensorflow", "torch"],
-        reason="Only TF and Torch support CTC loss",
+        backend.backend() == "numpy",
+        reason="Numpy does not support CTC loss",
     )
     def test_ctc_loss(self):
         labels = np.array([[1, 2, 1], [1, 2, 2]])
         outputs = np.array(
             [
-                [[0.4, 0.8, 0.4], [0.4, 0.8, 0.4]],
-                [[0.2, 0.8, 0.3], [0.2, 0.3, 0.3]],
-                [[0.9, 0.4, 0.5], [0.4, 0.3, 0.2]],
+                [[0.4, 0.8, 0.4], [0.2, 0.8, 0.3], [0.9, 0.4, 0.5]],
+                [[0.4, 0.8, 0.4], [0.2, 0.3, 0.3], [0.4, 0.3, 0.2]],
             ]
         )