From 9ec3d7ae1f5f4da751184327324974811e52166b Mon Sep 17 00:00:00 2001
From: Yaoyiran Li <yaoyiran@google.com>
Date: Thu, 18 Jun 2026 09:15:17 -0700
Subject: [PATCH] Internal change

PiperOrigin-RevId: 934389845
---
 recml/core/ops/binary_cross_entropy_ops.py    | 482 ++++++++++++++++++
 .../core/ops/binary_cross_entropy_ops_test.py | 432 ++++++++++++++++
 2 files changed, 914 insertions(+)
 create mode 100644 recml/core/ops/binary_cross_entropy_ops.py
 create mode 100644 recml/core/ops/binary_cross_entropy_ops_test.py

diff --git a/recml/core/ops/binary_cross_entropy_ops.py b/recml/core/ops/binary_cross_entropy_ops.py
new file mode 100644
index 0000000..1b1cb9d
--- /dev/null
+++ b/recml/core/ops/binary_cross_entropy_ops.py
@@ -0,0 +1,482 @@
+# Copyright 2024 RecML authors <recommendations-ml@google.com>.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Binary cross-entropy loss implementation with optimized memory footprint.
+
+This implementation computes BCE loss without materializing the [B, N, V] logits
+matrix in memory, by chunking the vocabulary dimension.
+
+This work has been published at [Placeholder - Paper Link].
+"""
+
+import dataclasses
+import functools
+import jax
+from jax.experimental.pallas import tpu as pltpu
+import jax.numpy as jnp
+import jaxtyping as jt
+import numpy as np
+
+EPS = 1e-8
+
+
+def _get_mxu_size() -> int:
+  """Returns the MXU tile alignment size based on the active TPU generation."""
+  if not any(d.platform == "tpu" for d in jax.devices()):
+    return 128
+
+  tpu_info = pltpu.get_tpu_info()
+  chip = str(tpu_info.chip_version)
+  match chip:
+    case "v2" | "v3" | "v4" | "v4i" | "v5e" | "v5p":
+      return 128
+    case "v6e" | "7" | "7x" | "8i" | "8t":
+      return 256
+    case _:
+      raise NotImplementedError(
+          f"Unsupported TPU chip version: {chip}. Please explicitly verify MXU "
+          "systolic dimensions and extend _get_mxu_size."
+      )
+
+
+def _auto_block_v(
+    n: int, vocab_size: int, dtype: jnp.dtype = jnp.float32
+) -> int:
+  """Automatically picks block_v for intermediate logits."""
+  # Estimate local n per device to handle sharded runs
+  try:
+    num_devices = jax.device_count()
+  except:  # pylint: disable=bare-except
+    num_devices = 1
+  local_n = max(n // num_devices, 1)
+
+  # Scale target memory based on chip VMEM capacity:
+  # - On <=32 MB chips (TPU v3/v4): target 4 MB to leave ample headroom.
+  # - On 64 MB chips (TPU v5p, TPU 7): target 16 MB.
+  # - On 128 MB chips (TPU v5e, v6e): target 32 MB.
+  if any(d.platform == "tpu" for d in jax.devices()):
+    vmem_bytes = pltpu.get_tpu_info().vmem_capacity_bytes
+    if vmem_bytes <= 32 * 1024 * 1024:
+      target_bytes = 4 * 1024 * 1024
+    elif vmem_bytes <= 64 * 1024 * 1024:
+      target_bytes = 16 * 1024 * 1024
+    else:
+      target_bytes = 32 * 1024 * 1024
+  else:
+    target_bytes = 32 * 1024 * 1024
+
+  bytes_per_element = jnp.dtype(dtype).itemsize
+  target_elements = target_bytes // bytes_per_element
+  block_v = target_elements // local_n
+
+  # Align to MXU systolic dimension (128 for legacy TPUs, 256 for TPU v6e+)
+  mxu_size = _get_mxu_size()
+  block_v = max(mxu_size, (block_v // mxu_size) * mxu_size)
+  # Don't exceed vocab_size
+  block_v = min(vocab_size, block_v)
+  return block_v
+
+
+def _get_sharding(x: jt.ArrayLike) -> jax.sharding.Sharding | None:
+  if hasattr(x, "sharding"):
+    return x.sharding
+  if hasattr(x, "aval") and hasattr(x.aval, "sharding"):
+    return x.aval.sharding
+  return None
+
+
+def _replicate_hidden_dim(
+    x: jt.Float[jt.Array, "... D"],
+) -> jt.Float[jt.Array, "... D"]:
+  """Replicates the hidden dimension of the input tensor if sharded."""
+  sharding = _get_sharding(x)
+  if isinstance(sharding, jax.sharding.NamedSharding):
+    mesh = sharding.mesh
+    if not mesh.empty:
+      spec = sharding.spec
+      new_spec_list = list(spec)
+      if new_spec_list:
+        new_spec_list[-1] = None
+      new_spec = jax.sharding.PartitionSpec(*new_spec_list)
+      return jax.lax.with_sharding_constraint(
+          x, jax.sharding.NamedSharding(mesh, new_spec)
+      )
+  return x
+
+
+@dataclasses.dataclass
+class BCEConfig:
+  """Configuration for the binary cross-entropy loss."""
+
+  block_v: int
+  compute_metrics: bool = False
+
+
+def _bce_fwd_chunk(
+    activations_2d: jt.Float[jt.Array, "N D"],
+    embeddings: jt.Float[jt.Array, "V D"],
+    targets_2d: jt.Int[jt.Array, "N L"],
+    j: jt.Int[jt.Array, ""],
+    block_v: int,
+    vocab: int,
+) -> tuple[
+    jt.Float[jt.Array, "N"],
+    jt.Float[jt.Array, "N block_v"],
+    jt.Bool[jt.Array, "N block_v"],
+    jt.Bool[jt.Array, "block_v"],
+]:
+  """Computes logits and base fused BCE loss for a single vocabulary chunk."""
+  actual_start = jnp.maximum(0, jnp.minimum(j * block_v, vocab - block_v))
+  emb_chunk = jax.lax.dynamic_slice_in_dim(embeddings, actual_start, block_v)
+  logits = jax.lax.dot_general(
+      activations_2d,
+      emb_chunk,
+      (((1,), (1,)), ((), ())),
+      preferred_element_type=jnp.float32,
+      precision=jax.lax.Precision.DEFAULT,
+  )
+
+  chunk_indices = actual_start + jnp.arange(block_v)
+  valid_mask = (chunk_indices >= j * block_v) & (chunk_indices < vocab)
+  # Fused BCE Loss: BCE(x, y) = BCE(x, 0) - y * x
+  loss_zero = jnp.maximum(logits, 0.0) + jnp.log1p(jnp.exp(-jnp.abs(logits)))
+
+  n = activations_2d.shape[0]
+  targets_chunk = jnp.zeros((n, block_v), dtype=jnp.bool_)
+  rel_targets = targets_2d - actual_start
+  chunk_cols = jnp.arange(block_v)[None, :]
+  for l_idx in range(targets_2d.shape[-1]):
+    targets_chunk = targets_chunk | (
+        rel_targets[:, l_idx : l_idx + 1] == chunk_cols
+    )
+
+  loss_chunk = loss_zero - targets_chunk * logits
+  loss_chunk = loss_chunk * valid_mask[None, :]
+  loss_sum = jnp.sum(loss_chunk, axis=-1)
+  return loss_sum, logits, targets_chunk, valid_mask
+
+
+def _bce_fwd_local(
+    config: BCEConfig,
+    activations: jt.Float[jt.Array, "B N D"],
+    embeddings: jt.Float[jt.Array, "V D"],
+    targets: jt.Int[jt.Array, "B N L"],
+) -> (
+    tuple[
+        jt.Float[jt.Array, "B N"],
+        jt.Float[jt.Array, "B N"],
+        jt.Float[jt.Array, "B N"],
+        jt.Float[jt.Array, "B N"],
+        jt.Float[jt.Array, "B N"],
+    ]
+    | jt.Float[jt.Array, "B N"]
+):
+  """Computes the sum of Loss(x_v, target_v) over all V block-wise, and metrics."""
+  block_v = config.block_v
+  batch, seq_len, hidden = activations.shape
+  vocab = embeddings.shape[0]
+
+  n = batch * seq_len
+  # NOMUTANTS -- v_blocks is calculated from block_v and vocab.
+  v_blocks = int(np.ceil(vocab / block_v))
+
+  activations_2d = jnp.reshape(activations, (n, hidden))
+  targets_2d = jnp.reshape(targets, (n, -1))
+
+  if config.compute_metrics:
+
+    def v_body(
+        carry: tuple[jt.Float[jt.Array, "N"], ...],
+        j: jt.Int[jt.Array, ""],
+    ) -> tuple[tuple[jt.Float[jt.Array, "N"], ...], None]:
+      loss_acc, tp_acc, fp_acc, fn_acc, tn_acc = carry
+      loss_sum, logits, targets_chunk, valid_mask = _bce_fwd_chunk(
+          activations_2d, embeddings, targets_2d, j, block_v, vocab
+      )
+
+      predictions_chunk = (logits > 0.0) & valid_mask[None, :]
+      targets_chunk = targets_chunk & valid_mask[None, :]
+
+      tp_chunk = targets_chunk & predictions_chunk
+      fp_chunk = predictions_chunk ^ tp_chunk
+      fn_chunk = targets_chunk ^ tp_chunk
+      tn_chunk = valid_mask[None, :] & (~(targets_chunk | predictions_chunk))
+
+      tp_sum = jnp.sum(tp_chunk, axis=-1).astype(jnp.float32)
+      fp_sum = jnp.sum(fp_chunk, axis=-1).astype(jnp.float32)
+      fn_sum = jnp.sum(fn_chunk, axis=-1).astype(jnp.float32)
+      tn_sum = jnp.sum(tn_chunk, axis=-1).astype(jnp.float32)
+
+      return (
+          loss_acc + loss_sum,
+          tp_acc + tp_sum,
+          fp_acc + fp_sum,
+          fn_acc + fn_sum,
+          tn_acc + tn_sum,
+      ), None
+
+    init = (
+        jnp.zeros((n,), dtype=jnp.float32),
+        jnp.zeros((n,), dtype=jnp.float32),
+        jnp.zeros((n,), dtype=jnp.float32),
+        jnp.zeros((n,), dtype=jnp.float32),
+        jnp.zeros((n,), dtype=jnp.float32),
+    )
+    (loss_final, tp_final, fp_final, fn_final, tn_final), _ = jax.lax.scan(
+        v_body, init, jnp.arange(v_blocks)
+    )
+    return (
+        jnp.reshape(loss_final, (batch, seq_len)),
+        jnp.reshape(tp_final, (batch, seq_len)),
+        jnp.reshape(fp_final, (batch, seq_len)),
+        jnp.reshape(fn_final, (batch, seq_len)),
+        jnp.reshape(tn_final, (batch, seq_len)),
+    )
+  else:
+
+    def v_body_no_metrics(
+        loss_acc: jt.Float[jt.Array, "N"],
+        j: jt.Int[jt.Array, ""],
+    ) -> tuple[jt.Float[jt.Array, "N"], None]:
+      loss_sum, _, _, _ = _bce_fwd_chunk(
+          activations_2d, embeddings, targets_2d, j, block_v, vocab
+      )
+      return loss_acc + loss_sum, None
+
+    init = jnp.zeros((n,), dtype=jnp.float32)
+    loss_final, _ = jax.lax.scan(v_body_no_metrics, init, jnp.arange(v_blocks))
+    return jnp.reshape(loss_final, (batch, seq_len))
+
+
+@functools.partial(jax.custom_vjp, nondiff_argnums=(0,))
+def _cut_binary_cross_entropy(
+    config: BCEConfig,
+    activations: jt.Float[jt.Array, "... B N D"],
+    embeddings: jt.Float[jt.Array, "V D"],
+    targets: jt.Int[jt.Array, "... B N L"],
+) -> (
+    tuple[
+        jt.Float[jt.Array, "... B N"],
+        jt.Float[jt.Array, "... B N"],
+        jt.Float[jt.Array, "... B N"],
+        jt.Float[jt.Array, "... B N"],
+        jt.Float[jt.Array, "... B N"],
+    ]
+    | jt.Float[jt.Array, "... B N"]
+):
+  """Computes the non-differentiable path of cut BCE loss and metrics."""
+  outputs, _ = _cut_binary_cross_entropy_fwd(
+      config, activations, embeddings, targets
+  )
+  return outputs
+
+
+def _cut_binary_cross_entropy_fwd(
+    config: BCEConfig,
+    activations: jt.Float[jt.Array, "B N D"],
+    embeddings: jt.Float[jt.Array, "V D"],
+    targets: jt.Int[jt.Array, "B N L"],
+) -> tuple[
+    tuple[
+        jt.Float[jt.Array, "B N"],
+        jt.Float[jt.Array, "B N"],
+        jt.Float[jt.Array, "B N"],
+        jt.Float[jt.Array, "B N"],
+        jt.Float[jt.Array, "B N"],
+    ]
+    | jt.Float[jt.Array, "B N"],
+    tuple[
+        jt.Float[jt.Array, "B N D"],
+        jt.Float[jt.Array, "V D"],
+        jt.Int[jt.Array, "B N L"],
+    ],
+]:
+  """Computes forward mode of cut BCE loss."""
+  replicated_activations = _replicate_hidden_dim(activations)
+  if activations.ndim == 4:
+    if targets.ndim == 3:
+      targets_in_axis = None
+    else:
+      targets_in_axis = 0
+    fwd_vmap = jax.vmap(
+        functools.partial(_bce_fwd_local, config),
+        in_axes=(0, None, targets_in_axis),
+    )
+    res = fwd_vmap(replicated_activations, embeddings, targets)
+  else:
+    res = _bce_fwd_local(config, replicated_activations, embeddings, targets)
+  vocab_size = embeddings.shape[0]
+  if isinstance(res, tuple):
+    loss_y0, tp, fp, fn, tn = res
+    losses = loss_y0 * (1.0 / vocab_size)
+    return (losses, tp, fp, fn, tn), (
+        activations,
+        embeddings,
+        targets,
+    )
+  else:
+    losses = res * (1.0 / vocab_size)
+    return losses, (
+        activations,
+        embeddings,
+        targets,
+    )
+
+
+def _cut_binary_cross_entropy_bwd(
+    config: BCEConfig,
+    res: tuple[
+        jt.Float[jt.Array, "... B N D"],
+        jt.Float[jt.Array, "V D"],
+        jt.Int[jt.Array, "... B N L"],
+    ],
+    d_outputs: (
+        tuple[
+            jt.Float[jt.Array, "... B N"],
+            jt.Float[jt.Array, "... B N"],
+            jt.Float[jt.Array, "... B N"],
+            jt.Float[jt.Array, "... B N"],
+            jt.Float[jt.Array, "... B N"],
+        ]
+        | jt.Float[jt.Array, "... B N"]
+    ),
+) -> tuple[
+    jt.Float[jt.Array, "... B N D"],
+    jt.Float[jt.Array, "V D"],
+    None,
+]:
+  """Computes the backward mode of cut BCE loss."""
+  del config, res, d_outputs
+  raise NotImplementedError(
+      "The memory-efficient backward pass of cut BCE is not implemented yet; "
+      "only the forward pass is currently supported."
+  )
+
+
+_cut_binary_cross_entropy.defvjp(
+    _cut_binary_cross_entropy_fwd, _cut_binary_cross_entropy_bwd
+)
+
+
+def cut_binary_cross_entropy(
+    activations: jt.Float[jt.Array, "... B N D"],
+    embeddings: jt.Float[jt.Array, "V D"],
+    targets: jt.Int[jt.Array, "... B N L"],
+    weights: jt.Float[jt.Array, "... B N"] | None = None,
+    *,
+    return_per_target_losses: bool = False,
+    return_metrics: bool = False,
+    block_v: int | None = None,
+) -> (
+    jt.Float[jt.Array, ""]
+    | tuple[jt.Float[jt.Array, ""], jt.Float[jt.Array, "... B N"]]
+    | tuple[
+        jt.Float[jt.Array, ""],
+        jt.Float[jt.Array, ""],
+        jt.Float[jt.Array, ""],
+        jt.Float[jt.Array, ""],
+        jt.Float[jt.Array, ""],
+    ]
+    | tuple[
+        jt.Float[jt.Array, ""],
+        jt.Float[jt.Array, "... B N"],
+        jt.Float[jt.Array, ""],
+        jt.Float[jt.Array, ""],
+        jt.Float[jt.Array, ""],
+        jt.Float[jt.Array, ""],
+    ]
+):
+  """Computes binary cross entropy loss over unmaterialized logits.
+
+  Args:
+    activations: Hidden-state outputs of shape ``[B, N, D]``.
+    embeddings: Output embedding / unembedding weights of shape ``[V, D]``.
+    targets: Target token ids of shape ``[B, N, L]``.
+    weights: Per-token loss weights of shape ``[B, N]``.
+    return_per_target_losses: If True, also return the per-target loss tensor.
+    return_metrics: If True, also return TP, FP, FN, TN metric counts.
+    block_v: Vocab-axis block size. Auto-picked if omitted.
+
+  Returns:
+    Scalar loss, optionally paired with per-target losses and/or metrics.
+  """
+  vocab_size = embeddings.shape[0]
+
+  # Prevent collective communications inside the loop by forcing replication of
+  # weights.
+  sharding = _get_sharding(embeddings)
+  if (
+      isinstance(sharding, jax.sharding.NamedSharding)
+      and not sharding.mesh.empty
+  ):
+    replicated_sharding = jax.sharding.NamedSharding(
+        sharding.mesh, jax.sharding.PartitionSpec()
+    )
+    embeddings = jax.lax.with_sharding_constraint(
+        embeddings, replicated_sharding
+    )
+
+  if block_v is None:
+    n = activations.shape[-3] * activations.shape[-2]
+    block_v = _auto_block_v(n, vocab_size)
+  else:
+    block_v = min(block_v, vocab_size)
+
+  config = BCEConfig(
+      block_v=block_v,
+      compute_metrics=return_metrics,
+  )
+
+  res = _cut_binary_cross_entropy(
+      config,
+      activations,
+      embeddings,
+      targets,
+  )
+
+  if return_metrics:
+    losses, tp, fp, fn, tn = res
+
+    if weights is not None:
+      losses = losses * weights
+      weight_sum = jnp.sum(weights)
+      tp_sum = jnp.sum(tp * weights)
+      fp_sum = jnp.sum(fp * weights)
+      fn_sum = jnp.sum(fn * weights)
+      tn_sum = jnp.sum(tn * weights)
+    else:
+      weight_sum = np.prod(targets.shape[:-1])
+      tp_sum = jnp.sum(tp)
+      fp_sum = jnp.sum(fp)
+      fn_sum = jnp.sum(fn)
+      tn_sum = jnp.sum(tn)
+
+    loss = jnp.sum(losses) / (weight_sum + EPS)
+
+    if return_per_target_losses:
+      return loss, losses, tp_sum, fp_sum, fn_sum, tn_sum
+    return loss, tp_sum, fp_sum, fn_sum, tn_sum
+  else:
+    losses = res
+
+    if weights is not None:
+      losses = losses * weights
+      weight_sum = jnp.sum(weights)
+    else:
+      weight_sum = np.prod(targets.shape[:-1])
+
+    loss = jnp.sum(losses) / (weight_sum + EPS)
+
+    if return_per_target_losses:
+      return loss, losses
+    return loss
diff --git a/recml/core/ops/binary_cross_entropy_ops_test.py b/recml/core/ops/binary_cross_entropy_ops_test.py
new file mode 100644
index 0000000..a9d98d7
--- /dev/null
+++ b/recml/core/ops/binary_cross_entropy_ops_test.py
@@ -0,0 +1,432 @@
+# Copyright 2024 RecML authors <recommendations-ml@google.com>.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for binary_cross_entropy_ops."""
+
+import time
+from unittest import mock
+
+from absl import logging
+from absl.testing import absltest
+from absl.testing import parameterized
+import jax
+from jax.experimental.pallas import tpu as pltpu
+import jax.numpy as jnp
+import keras
+import numpy as np
+from recml.core.ops import binary_cross_entropy_ops
+
+
+def _naive_bce(activations, embeddings, targets, weights=None):
+  """Naive implementation that materializes the full logits matrix."""
+  vocab_size = embeddings.shape[0]
+  logits = jnp.matmul(activations, embeddings.T)  # (B, N, V)
+
+  # targets: (B, N, L) -> multi_hot: (B, N, V)
+  one_hot = jax.nn.one_hot(targets, vocab_size, axis=-1)  # (B, N, L, V)
+  multi_hot = jnp.max(one_hot, axis=-2)  # (B, N, V)
+
+  # Compute stable BCE loss per class
+  # Loss = max(x, 0) - x * y + log(1 + exp(-|x|))
+  losses = (
+      jnp.maximum(logits, 0.0)
+      - logits * multi_hot
+      + jnp.log1p(jnp.exp(-jnp.abs(logits)))
+  )
+  loss_per_target = jnp.mean(losses, axis=-1)  # (B, N)
+
+  if weights is not None:
+    loss_per_target = loss_per_target * weights
+    weight_sum = jnp.sum(weights)
+  else:
+    weight_sum = np.prod(targets.shape[:-1])
+
+  loss = jnp.sum(loss_per_target) / (weight_sum + 1e-8)
+  return loss, loss_per_target
+
+
+class BinaryCrossEntropyOpsTest(parameterized.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    if jax.devices()[0].platform == 'tpu':
+
+      vmem = pltpu.get_tpu_info().vmem_capacity_bytes
+      logging.info(
+          'JETS_DEBUG: VMEM capacity: %d bytes (%.2f MB)',
+          vmem,
+          vmem / 1024 / 1024,
+      )
+
+  def test_get_sharding(self):
+    class ObjWithSharding:
+      sharding = 'dummy_sharding_1'
+
+    class ObjWithAvalSharding:
+
+      class Aval:
+        sharding = 'dummy_sharding_2'
+
+      aval = Aval()
+
+    class ObjWithAvalWithoutSharding:
+
+      class Aval:
+        pass
+
+      aval = Aval()
+
+    class ObjWithNoSharding:
+      pass
+
+    self.assertEqual(
+        binary_cross_entropy_ops._get_sharding(ObjWithSharding()),
+        'dummy_sharding_1',
+    )
+    self.assertEqual(
+        binary_cross_entropy_ops._get_sharding(ObjWithAvalSharding()),
+        'dummy_sharding_2',
+    )
+    self.assertIsNone(
+        binary_cross_entropy_ops._get_sharding(ObjWithAvalWithoutSharding())
+    )
+    self.assertIsNone(
+        binary_cross_entropy_ops._get_sharding(ObjWithNoSharding())
+    )
+
+  def test_get_mxu_size(self):
+    """Fails on new TPU generations if _get_mxu_size is not explicitly extended."""
+    if jax.devices()[0].platform != 'tpu':
+      self.assertEqual(binary_cross_entropy_ops._get_mxu_size(), 128)
+    else:
+      mxu_size = binary_cross_entropy_ops._get_mxu_size()
+      self.assertIn(mxu_size, (128, 256))
+
+  @parameterized.named_parameters(
+      ('standard', 2, 256, 128, 1024, 4, 256),
+      ('unaligned_seq_len', 2, 130, 128, 1024, 4, 256),
+      ('unaligned_vocab', 2, 256, 128, 1000, 4, 256),
+      ('single_label', 2, 256, 128, 1024, 1, 256),
+      ('small_block_v', 2, 256, 128, 1024, 4, 128),
+  )
+  def test_cut_bce_correctness(
+      self, batch, seq_len, hidden_dim, vocab_size, num_labels, block_v
+  ):
+    if jax.devices()[0].platform != 'tpu':
+      self.skipTest('Skipping TPU test.')
+
+    key = jax.random.PRNGKey(0)
+    key_act, key_emb, key_tgt = jax.random.split(key, 3)
+
+    activations = jax.random.normal(key_act, (batch, seq_len, hidden_dim))
+    embeddings = jax.random.normal(key_emb, (vocab_size, hidden_dim))
+    targets = jax.random.randint(
+        key_tgt, (batch, seq_len, num_labels), 0, vocab_size
+    )
+
+    # naive BCE
+    def run_naive(act, emb):
+      loss, _ = _naive_bce(act, emb, targets)
+      return loss
+
+    loss_naive = run_naive(activations, embeddings)
+
+    # cut BCE
+    def run_cut(act, emb):
+      return binary_cross_entropy_ops.cut_binary_cross_entropy(
+          act, emb, targets, block_v=block_v
+      )
+
+    loss_cut = run_cut(activations, embeddings)
+
+    # Compare
+    np.testing.assert_allclose(loss_cut, loss_naive, atol=1e-5, rtol=1e-5)
+
+  @parameterized.named_parameters(
+      ('4d_act_4d_tgt', (2, 2, 128, 64), (2, 2, 128, 4)),
+      ('4d_act_3d_tgt', (3, 2, 128, 64), (2, 128, 4)),
+  )
+  def test_cut_bce_4d_correctness(self, act_shape, tgt_shape):
+    vocab_size, block_v = 512, 256
+    hidden_dim = act_shape[-1]
+
+    key = jax.random.PRNGKey(42)
+    key_act, key_emb, key_tgt = jax.random.split(key, 3)
+
+    activations = jax.random.normal(key_act, act_shape)
+    embeddings = jax.random.normal(key_emb, (vocab_size, hidden_dim))
+    targets = jax.random.randint(key_tgt, tgt_shape, 0, vocab_size)
+
+    def run_naive(act, emb):
+      loss, _ = _naive_bce(act, emb, targets)
+      return loss
+
+    def run_cut(act, emb):
+      return binary_cross_entropy_ops.cut_binary_cross_entropy(
+          act, emb, targets, block_v=block_v
+      )
+
+    loss_naive = run_naive(activations, embeddings)
+    loss_cut = run_cut(activations, embeddings)
+    np.testing.assert_allclose(loss_cut, loss_naive, rtol=1e-3, atol=1e-3)
+
+  def test_cut_bce_with_sharded_embeddings(self):
+    if jax.devices()[0].platform != 'tpu':
+      self.skipTest('Skipping TPU test.')
+
+    batch, seq_len, hidden_dim, vocab_size, num_labels = 2, 128, 128, 512, 4
+    key = jax.random.PRNGKey(0)
+    key_act, key_emb, key_tgt = jax.random.split(key, 3)
+
+    activations = jax.random.normal(key_act, (batch, seq_len, hidden_dim))
+    embeddings = jax.random.normal(key_emb, (vocab_size, hidden_dim))
+    targets = jax.random.randint(
+        key_tgt, (batch, seq_len, num_labels), 0, vocab_size
+    )
+
+    devices = jax.devices()
+    mesh = jax.sharding.Mesh(np.array(devices), ('devices',))
+    act_sharding = jax.sharding.NamedSharding(
+        mesh, jax.sharding.PartitionSpec('devices', None, None)
+    )
+    emb_sharding = jax.sharding.NamedSharding(
+        mesh, jax.sharding.PartitionSpec('devices', None)
+    )
+    activations_sharded = jax.device_put(activations, act_sharding)
+    embeddings_sharded = jax.device_put(embeddings, emb_sharding)
+
+    with mock.patch.object(
+        jax.lax,
+        'with_sharding_constraint',
+        wraps=jax.lax.with_sharding_constraint,
+    ) as mock_fwd_constraint:
+      binary_cross_entropy_ops.cut_binary_cross_entropy(
+          activations_sharded,
+          embeddings_sharded,
+          targets,
+          block_v=256,
+      )
+      self.assertEqual(mock_fwd_constraint.call_count, 2)
+
+  def test_cut_bce_correctness_large_sequence(self):
+    if jax.devices()[0].platform != 'tpu':
+      self.skipTest('Skipping TPU test.')
+
+    # Test with sequence length larger than chunk_n to trigger scan loop
+    batch, seq_len, hidden_dim, vocab_size, num_labels = 2, 2048, 128, 512, 4
+    block_v = 256
+
+    key = jax.random.PRNGKey(42)
+    key_act, key_emb, key_tgt = jax.random.split(key, 3)
+
+    activations = jax.random.normal(key_act, (batch, seq_len, hidden_dim))
+    embeddings = jax.random.normal(key_emb, (vocab_size, hidden_dim))
+    targets = jax.random.randint(
+        key_tgt, (batch, seq_len, num_labels), 0, vocab_size
+    )
+
+    # naive BCE
+    def run_naive(act, emb):
+      loss, _ = _naive_bce(act, emb, targets)
+      return loss
+
+    loss_naive = run_naive(activations, embeddings)
+
+    # cut BCE
+    def run_cut(act, emb):
+      return binary_cross_entropy_ops.cut_binary_cross_entropy(
+          act, emb, targets, block_v=block_v
+      )
+
+    loss_cut = run_cut(activations, embeddings)
+
+    # Compare
+    np.testing.assert_allclose(loss_cut, loss_naive, atol=1e-5, rtol=1e-5)
+
+  def test_cut_bce_vs_keras(self):
+    if jax.devices()[0].platform != 'tpu':
+      self.skipTest('Skipping TPU test.')
+
+    batch, seq_len, hidden_dim, vocab_size, num_labels = 2, 128, 64, 512, 4
+    block_v = 256
+
+    key = jax.random.PRNGKey(42)
+    key_act, key_emb, key_tgt = jax.random.split(key, 3)
+
+    activations = jax.random.normal(key_act, (batch, seq_len, hidden_dim))
+    embeddings = jax.random.normal(key_emb, (vocab_size, hidden_dim))
+    targets = jax.random.randint(
+        key_tgt, (batch, seq_len, num_labels), 0, vocab_size
+    )
+
+    # Convert targets to multi-hot for Keras
+    one_hot = jax.nn.one_hot(targets, vocab_size, axis=-1)  # (B, N, L, V)
+    multi_hot = jnp.max(one_hot, axis=-2)  # (B, N, V)
+
+    # Keras BCE version
+    def run_keras(act, emb):
+      logits = jnp.matmul(act, emb.T)  # (B, N, V)
+      loss_per_token = keras.losses.binary_crossentropy(
+          multi_hot, logits, from_logits=True
+      )
+      return jnp.mean(loss_per_token)
+
+    loss_keras = run_keras(activations, embeddings)
+
+    # Cut BCE version
+    def run_cut(act, emb):
+      return binary_cross_entropy_ops.cut_binary_cross_entropy(
+          act, emb, targets, block_v=block_v
+      )
+
+    loss_cut = run_cut(activations, embeddings)
+
+    # Compare
+    np.testing.assert_allclose(loss_cut, loss_keras, atol=1e-5, rtol=1e-5)
+
+  def test_mini_benchmark_keras(self):
+    if jax.devices()[0].platform != 'tpu':
+      self.skipTest('Skipping TPU test.')
+
+    b, m, d, v, l = 32, 32, 128, 100000, 8
+
+    key = jax.random.PRNGKey(42)
+    key_act, key_emb, key_tgt = jax.random.split(key, 3)
+
+    activations = jax.random.normal(key_act, (b, m, d))
+    embeddings = jax.random.normal(key_emb, (v, d))
+    targets = jax.random.randint(key_tgt, (b, m, l), 0, v)
+
+    # Convert targets to multi-hot for Keras
+    one_hot = jax.nn.one_hot(targets, v, axis=-1)  # (B, M, L, V)
+    multi_hot = jnp.max(one_hot, axis=-2)  # (B, M, V)
+
+    # --- Keras version ---
+    def run_keras(act, emb):
+      logits = jnp.matmul(act, emb.T)  # (B, M, V)
+      loss_per_token = keras.losses.binary_crossentropy(
+          multi_hot, logits, from_logits=True
+      )
+      return jnp.mean(loss_per_token)
+
+    fwd_keras = jax.jit(run_keras)
+    logging.info('Compiling Keras...')
+    t0 = time.time()
+    fwd_keras(activations, embeddings).block_until_ready()
+    logging.info('Keras compiled in %.2f s', time.time() - t0)
+
+    num_steps = 200
+    logging.info('Benchmarking Keras (%d steps)...', num_steps)
+    t0 = time.time()
+    for _ in range(num_steps):
+      loss_keras = fwd_keras(activations, embeddings)
+      loss_keras.block_until_ready()
+    t_keras = (time.time() - t0) / num_steps
+    logging.info('Keras step time: %.4f ms', t_keras * 1000)
+
+  def test_mini_benchmark_cut_bce(self):
+    if jax.devices()[0].platform != 'tpu':
+      self.skipTest('Skipping TPU test.')
+
+    b, m, d, v, l = 32, 32, 128, 100000, 8
+
+    key = jax.random.PRNGKey(42)
+    key_act, key_emb, key_tgt = jax.random.split(key, 3)
+
+    activations = jax.random.normal(key_act, (b, m, d))
+    embeddings = jax.random.normal(key_emb, (v, d))
+    targets = jax.random.randint(key_tgt, (b, m, l), 0, v)
+    num_steps = 200
+
+    # --- Cut version ---
+    block_v = 4096
+
+    def run_cut(act, emb, bv=block_v):
+      return binary_cross_entropy_ops.cut_binary_cross_entropy(
+          act, emb, targets, block_v=bv
+      )
+
+    fwd_cut = jax.jit(run_cut)
+    logging.info('Compiling cut (block_v=%d)...', block_v)
+    t0 = time.time()
+    fwd_cut(activations, embeddings).block_until_ready()
+    logging.info(
+        'Cut (block_v=%d) compiled in %.2f s', block_v, time.time() - t0
+    )
+
+    logging.info(
+        'Benchmarking cut (block_v=%d) (%d steps)...', block_v, num_steps
+    )
+    t0 = time.time()
+    for _ in range(num_steps):
+      loss_cut = fwd_cut(activations, embeddings)
+      loss_cut.block_until_ready()
+    t_cut = (time.time() - t0) / num_steps
+    logging.info('Cut (block_v=%d) step time: %.4f ms ', block_v, t_cut * 1000)
+
+  def test_cut_bce_metrics(self):
+    if jax.devices()[0].platform != 'tpu':
+      self.skipTest('Skipping TPU test.')
+
+    batch, seq_len, hidden_dim, vocab_size, num_labels = 2, 64, 32, 128, 2
+    key = jax.random.PRNGKey(1)
+    activations = jax.random.normal(key, (batch, seq_len, hidden_dim))
+    embeddings = jax.random.normal(key, (vocab_size, hidden_dim))
+    targets = jax.random.randint(
+        key, (batch, seq_len, num_labels), 0, vocab_size
+    )
+
+    loss, tp, fp, fn, tn = binary_cross_entropy_ops.cut_binary_cross_entropy(
+        activations,
+        embeddings,
+        targets,
+        return_metrics=True,
+    )
+    self.assertIsNotNone(loss)
+    self.assertIsNotNone(tp)
+    self.assertIsNotNone(fp)
+    self.assertIsNotNone(fn)
+    self.assertIsNotNone(tn)
+
+  def test_cut_bce_block_v_capped_at_vocab(self):
+    activations = jnp.ones((2, 64, 32))
+    embeddings = jnp.ones((200, 32))
+    targets = jnp.zeros((2, 64, 2), dtype=jnp.int32)
+
+    with mock.patch.object(
+        binary_cross_entropy_ops,
+        '_cut_binary_cross_entropy',
+        wraps=binary_cross_entropy_ops._cut_binary_cross_entropy,
+    ) as mock_fn:
+      binary_cross_entropy_ops.cut_binary_cross_entropy(
+          activations, embeddings, targets, block_v=1000
+      )
+      config = mock_fn.call_args[0][0]
+      self.assertEqual(config.block_v, 200)
+
+  def test_cut_bce_backward_not_implemented(self):
+    activations = jnp.ones((2, 64, 32))
+    embeddings = jnp.ones((128, 32))
+    targets = jnp.zeros((2, 64, 2), dtype=jnp.int32)
+
+    def run_cut(act, emb):
+      return binary_cross_entropy_ops.cut_binary_cross_entropy(
+          act, emb, targets, block_v=128
+      )
+
+    with self.assertRaises(NotImplementedError):
+      jax.grad(run_cut, argnums=(0, 1))(activations, embeddings)
+
+
+if __name__ == '__main__':
+  absltest.main()