diff --git a/recml/core/ops/binary_cross_entropy_ops.py b/recml/core/ops/binary_cross_entropy_ops.py index 070c63f..eb12af0 100644 --- a/recml/core/ops/binary_cross_entropy_ops.py +++ b/recml/core/ops/binary_cross_entropy_ops.py @@ -249,6 +249,9 @@ class BCEConfig: ``embeddings.shape[0]``. optimize_large_l: Whether to use the target-free vocabulary pass + sparse positive-target correction. Defaults to False. + force_bwd_block_v: Whether to force the backward pass to use `block_v` + directly (aligned to lane width and capped by vocabulary size) instead of + clamping to `_max_safe_block_v`. Defaults to False. """ block_v: int @@ -257,6 +260,7 @@ class BCEConfig: use_pallas: bool = True global_vocab: int | None = None optimize_large_l: bool = False + force_bwd_block_v: bool = False # Sharding specs for VJP backward pass optimization mesh: jax.sharding.Mesh | None = None @@ -595,6 +599,7 @@ def cut_binary_cross_entropy( *, return_per_target_losses: bool = False, return_metrics: bool = False, + return_per_target_metrics: bool = False, block_v: int | None = None, block_n: int | None = None, mesh: jax.sharding.Mesh | None = None, @@ -602,23 +607,24 @@ def cut_binary_cross_entropy( emb_spec: jax.sharding.PartitionSpec | None = None, use_pallas: bool = True, optimize_large_l: bool = False, + force_bwd_block_v: bool = False, ) -> ( 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, ""], + jt.Float[jt.Array, "... B N"] | jt.Float[jt.Array, ""], + jt.Float[jt.Array, "... B N"] | jt.Float[jt.Array, ""], + jt.Float[jt.Array, "... B N"] | jt.Float[jt.Array, ""], + jt.Float[jt.Array, "... B N"] | 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, ""], + jt.Float[jt.Array, "... B N"] | jt.Float[jt.Array, ""], + jt.Float[jt.Array, "... B N"] | jt.Float[jt.Array, ""], + jt.Float[jt.Array, "... B N"] | jt.Float[jt.Array, ""], + jt.Float[jt.Array, "... B N"] | jt.Float[jt.Array, ""], ] ): """Computes binary cross entropy loss over unmaterialized logits. @@ -643,6 +649,8 @@ def cut_binary_cross_entropy( position's weight. 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. + return_per_target_metrics: If True (with ``return_metrics=True``), return + per-position TP, FP, FN, TN tensors instead of summing them to scalars. block_v: Vocab-axis block size. Auto-picked if omitted. block_n: Sequence-axis block size for backward kernel. mesh: Optional mesh to use for sharding. @@ -652,9 +660,21 @@ def cut_binary_cross_entropy( implementation. optimize_large_l: If True, use the target-free vocabulary pass + sparse positive-target correction. Defaults to False. + force_bwd_block_v: If True, force the backward pass to use `block_v` + directly instead of clamping to `_max_safe_block_v`. Defaults to False. Returns: - Scalar loss, optionally paired with per-target losses and/or metrics. + - ``loss`` if ``return_per_target_losses=False`` and + ``return_metrics=False``. + - ``(loss, losses)`` if ``return_per_target_losses=True`` and + ``return_metrics=False``. + - ``(loss, tp, fp, fn, tn)`` if ``return_per_target_losses=False`` and + ``return_metrics=True``. + - ``(loss, losses, tp, fp, fn, tn)`` if ``return_per_target_losses=True`` + and ``return_metrics=True``. + Here ``loss`` is a scalar ``[]``, ``losses`` has shape ``[..., B, N]``, and + ``(tp, fp, fn, tn)`` each have shape ``[..., B, N]`` when + ``return_per_target_metrics=True`` or scalar shape ``[]`` when ``False``. Raises: ValueError: If ``activations`` has rank 5 or above, or if ``weights`` does @@ -715,6 +735,7 @@ def cut_binary_cross_entropy( use_pallas=use_pallas, global_vocab=vocab_size, optimize_large_l=optimize_large_l, + force_bwd_block_v=force_bwd_block_v, ) res = _cut_binary_cross_entropy( @@ -730,22 +751,24 @@ def cut_binary_cross_entropy( 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) + tp = tp * weights + fp = fp * weights + fn = fn * weights + tn = tn * weights else: weight_sum = np.prod(losses_shape) - 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 not return_per_target_metrics: + tp = jnp.sum(tp) + fp = jnp.sum(fp) + fn = jnp.sum(fn) + tn = jnp.sum(tn) + 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 + return loss, losses, tp, fp, fn, tn + return loss, tp, fp, fn, tn else: losses = res @@ -871,6 +894,9 @@ def _effective_block_v( """ lane = _pallas_lane() padded_d = pl.align_to(hidden, lane) + if config.force_bwd_block_v: + block_v = pl.align_to(min(config.block_v, vocab), lane) + return padded_d, min(block_v, pl.align_to(vocab, lane)) max_safe = _max_safe_block_v(_pallas_vmem_budget(), padded_d, config.block_n) block_v = pl.align_to(min(config.block_v, vocab, max_safe), lane) return padded_d, min(block_v, pl.align_to(vocab, lane), max_safe) diff --git a/recml/core/ops/binary_cross_entropy_ops_benchmark_test.py b/recml/core/ops/binary_cross_entropy_ops_benchmark_test.py index d580936..f90ed5e 100644 --- a/recml/core/ops/binary_cross_entropy_ops_benchmark_test.py +++ b/recml/core/ops/binary_cross_entropy_ops_benchmark_test.py @@ -96,7 +96,7 @@ def test_mini_benchmark_cut_bce(self): def run_cut(act, emb, bv=_BLOCK_V): return binary_cross_entropy_ops.cut_binary_cross_entropy( - act, emb, targets, block_v=bv + act, emb, targets, block_v=bv, optimize_large_l=True ) grad_cut = jax.jit(jax.grad(run_cut, argnums=(0, 1))) diff --git a/recml/core/ops/binary_cross_entropy_ops_test.py b/recml/core/ops/binary_cross_entropy_ops_test.py index 122bad4..8e735aa 100644 --- a/recml/core/ops/binary_cross_entropy_ops_test.py +++ b/recml/core/ops/binary_cross_entropy_ops_test.py @@ -854,6 +854,72 @@ def test_cut_bce_metrics(self): self.assertIsNotNone(fn) self.assertIsNotNone(tn) + @parameterized.named_parameters( + ('unweighted_no_per_target_losses', False, False), + ('weighted_no_per_target_losses', True, False), + ('weighted_with_per_target_losses', True, True), + ) + def test_cut_bce_per_target_metrics( + self, use_weights: bool, return_per_target_losses: bool + ): + batch, seq_len, hidden_dim, vocab_size, num_labels = 2, 16, 32, 128, 2 + block_v = 64 + key_act, key_emb, key_tgt, key_w = jax.random.split( + jax.random.PRNGKey(17), 4 + ) + 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 + ) + weights = None + if use_weights: + weights = jax.random.bernoulli(key_w, 0.7, (batch, seq_len)).astype( + jnp.float32 + ) + + res = binary_cross_entropy_ops.cut_binary_cross_entropy( + activations, + embeddings, + targets, + weights, + block_v=block_v, + return_metrics=True, + return_per_target_metrics=True, + return_per_target_losses=return_per_target_losses, + use_pallas=False, + ) + if return_per_target_losses: + loss, losses, tp, fp, fn, tn = res + else: + loss, tp, fp, fn, tn = res + losses = None + + loss_ref, losses_ref = _keras_bce( + activations, embeddings, targets, weights=weights + ) + np.testing.assert_allclose(loss, loss_ref, atol=1e-5, rtol=1e-5) + if losses is not None: + np.testing.assert_allclose(losses, losses_ref, atol=1e-5, rtol=1e-5) + + logits = jnp.matmul(activations, embeddings.T) + multi_hot = jnp.max(jax.nn.one_hot(targets, vocab_size), axis=-2) > 0 + preds = logits > 0.0 + tp_ref = jnp.sum(multi_hot & preds, axis=-1, dtype=jnp.float32) + fp_ref = jnp.sum(~multi_hot & preds, axis=-1, dtype=jnp.float32) + fn_ref = jnp.sum(multi_hot & ~preds, axis=-1, dtype=jnp.float32) + tn_ref = jnp.sum(~multi_hot & ~preds, axis=-1, dtype=jnp.float32) + if weights is not None: + tp_ref = tp_ref * weights + fp_ref = fp_ref * weights + fn_ref = fn_ref * weights + tn_ref = tn_ref * weights + + np.testing.assert_allclose(tp, tp_ref, rtol=1e-6) + np.testing.assert_allclose(fp, fp_ref, rtol=1e-6) + np.testing.assert_allclose(fn, fn_ref, rtol=1e-6) + np.testing.assert_allclose(tn, tn_ref, rtol=1e-6) + def test_cut_bce_block_v_capped_at_vocab(self): activations = jnp.ones((2, 64, 32)) embeddings = jnp.ones((200, 32)) @@ -1088,6 +1154,62 @@ def loss_fn(act, emb, opt): np.testing.assert_allclose(g_act_opt, g_act_ref, atol=1e-5, rtol=1e-5) np.testing.assert_allclose(g_emb_opt, g_emb_ref, atol=1e-5, rtol=1e-5) + def test_force_bwd_block_v(self): + config_clamped = binary_cross_entropy_ops.BCEConfig( + block_v=11008, block_n=256, force_bwd_block_v=False + ) + config_forced = binary_cross_entropy_ops.BCEConfig( + block_v=11008, block_n=256, force_bwd_block_v=True + ) + _, bwd_v_clamped = binary_cross_entropy_ops._effective_block_v( + config_clamped, hidden=256, vocab=16384 + ) + _, bwd_v_forced = binary_cross_entropy_ops._effective_block_v( + config_forced, hidden=256, vocab=16384 + ) + self.assertLess(bwd_v_clamped, 11008) + self.assertEqual(bwd_v_forced, 11008) + + batch, seq_len, hidden_dim, vocab_size, num_labels = 2, 16, 128, 2048, 4 + key_act, key_emb, key_tgt = jax.random.split(jax.random.PRNGKey(99), 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 + ) + + def loss_fn(act, emb, force_bwd): + return binary_cross_entropy_ops.cut_binary_cross_entropy( + act, + emb, + targets, + block_v=1024, + block_n=128, + force_bwd_block_v=force_bwd, + ) + + with mock.patch.object( + binary_cross_entropy_ops, '_max_safe_block_v', return_value=512 + ): + loss_clamped = loss_fn(activations, embeddings, False) + loss_forced = loss_fn(activations, embeddings, True) + np.testing.assert_allclose( + loss_forced, loss_clamped, atol=1e-6, rtol=1e-6 + ) + + g_act_clamped, g_emb_clamped = jax.jit( + jax.grad(lambda a, e: loss_fn(a, e, False), argnums=(0, 1)) + )(activations, embeddings) + g_act_forced, g_emb_forced = jax.jit( + jax.grad(lambda a, e: loss_fn(a, e, True), argnums=(0, 1)) + )(activations, embeddings) + np.testing.assert_allclose( + g_act_forced, g_act_clamped, atol=1e-5, rtol=1e-5 + ) + np.testing.assert_allclose( + g_emb_forced, g_emb_clamped, atol=1e-5, rtol=1e-5 + ) + if __name__ == '__main__': absltest.main() diff --git a/recml/core/ops/binary_focal_cross_entropy_ops.py b/recml/core/ops/binary_focal_cross_entropy_ops.py index c0c211a..311af8e 100644 --- a/recml/core/ops/binary_focal_cross_entropy_ops.py +++ b/recml/core/ops/binary_focal_cross_entropy_ops.py @@ -67,7 +67,8 @@ class FocalBCEConfig(bce_ops.BCEConfig): """Configuration for the binary focal cross-entropy loss. - Inherits common fields from `bce_ops.BCEConfig`. + Inherits common fields (including `optimize_large_l` and `force_bwd_block_v`) + from `bce_ops.BCEConfig`. Attributes: gamma: Focusing parameter of the focal loss. The per-target BCE loss is @@ -412,6 +413,7 @@ def cut_binary_focal_cross_entropy( apply_class_balancing: bool = False, return_per_target_losses: bool = False, return_metrics: bool = False, + return_per_target_metrics: bool = False, block_v: int | None = None, block_n: int | None = None, mesh: jax.sharding.Mesh | None = None, @@ -419,23 +421,24 @@ def cut_binary_focal_cross_entropy( emb_spec: jax.sharding.PartitionSpec | None = None, use_pallas: bool = True, optimize_large_l: bool = False, + force_bwd_block_v: bool = False, ) -> ( 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, ""], + jt.Float[jt.Array, "... B N"] | jt.Float[jt.Array, ""], + jt.Float[jt.Array, "... B N"] | jt.Float[jt.Array, ""], + jt.Float[jt.Array, "... B N"] | jt.Float[jt.Array, ""], + jt.Float[jt.Array, "... B N"] | 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, ""], + jt.Float[jt.Array, "... B N"] | jt.Float[jt.Array, ""], + jt.Float[jt.Array, "... B N"] | jt.Float[jt.Array, ""], + jt.Float[jt.Array, "... B N"] | jt.Float[jt.Array, ""], + jt.Float[jt.Array, "... B N"] | jt.Float[jt.Array, ""], ] ): """Computes binary focal cross entropy loss over unmaterialized logits. @@ -463,6 +466,8 @@ def cut_binary_focal_cross_entropy( apply_class_balancing: If True, apply class balancing. 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. + return_per_target_metrics: If True (with ``return_metrics=True``), return + per-position TP, FP, FN, TN tensors instead of summing them to scalars. block_v: Vocab-axis block size. Auto-picked if omitted. block_n: Sequence-axis block size for backward kernel. mesh: Optional mesh to use for sharding. @@ -472,9 +477,21 @@ def cut_binary_focal_cross_entropy( implementation. optimize_large_l: If True, use the target-free vocabulary pass + sparse positive-target correction. Defaults to False. + force_bwd_block_v: If True, force the backward pass to use `block_v` + directly instead of clamping to `_max_safe_block_v`. Defaults to False. Returns: - Scalar loss, optionally paired with per-target losses and/or metrics. + - ``loss`` if ``return_per_target_losses=False`` and + ``return_metrics=False``. + - ``(loss, losses)`` if ``return_per_target_losses=True`` and + ``return_metrics=False``. + - ``(loss, tp, fp, fn, tn)`` if ``return_per_target_losses=False`` and + ``return_metrics=True``. + - ``(loss, losses, tp, fp, fn, tn)`` if ``return_per_target_losses=True`` + and ``return_metrics=True``. + Here ``loss`` is a scalar ``[]``, ``losses`` has shape ``[..., B, N]``, and + ``(tp, fp, fn, tn)`` each have shape ``[..., B, N]`` when + ``return_per_target_metrics=True`` or scalar shape ``[]`` when ``False``. Raises: ValueError: If ``activations`` has rank 5 or above, or if ``weights`` does @@ -537,6 +554,7 @@ def cut_binary_focal_cross_entropy( use_pallas=use_pallas, global_vocab=vocab_size, optimize_large_l=optimize_large_l, + force_bwd_block_v=force_bwd_block_v, ) res = _cut_binary_focal_cross_entropy( @@ -552,22 +570,24 @@ def cut_binary_focal_cross_entropy( 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) + tp = tp * weights + fp = fp * weights + fn = fn * weights + tn = tn * weights else: weight_sum = np.prod(losses_shape) - 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 not return_per_target_metrics: + tp = jnp.sum(tp) + fp = jnp.sum(fp) + fn = jnp.sum(fn) + tn = jnp.sum(tn) + 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 + return loss, losses, tp, fp, fn, tn + return loss, tp, fp, fn, tn else: losses = res diff --git a/recml/core/ops/binary_focal_cross_entropy_ops_test.py b/recml/core/ops/binary_focal_cross_entropy_ops_test.py index 6f1c2d0..b04a9a8 100644 --- a/recml/core/ops/binary_focal_cross_entropy_ops_test.py +++ b/recml/core/ops/binary_focal_cross_entropy_ops_test.py @@ -787,6 +787,72 @@ def run_naive(act, emb): tn, jnp.sum((~multi_hot & ~preds) * w), rtol=1e-6 ) + @parameterized.named_parameters( + ('unweighted_no_per_target_losses', False, False), + ('weighted_no_per_target_losses', True, False), + ('weighted_with_per_target_losses', True, True), + ) + def test_cut_focal_bce_per_target_metrics( + self, use_weights: bool, return_per_target_losses: bool + ): + batch, seq_len, hidden_dim, vocab_size, num_labels = 2, 16, 32, 128, 2 + block_v = 64 + key_act, key_emb, key_tgt, key_w = jax.random.split( + jax.random.PRNGKey(19), 4 + ) + 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 + ) + weights = None + if use_weights: + weights = jax.random.bernoulli(key_w, 0.7, (batch, seq_len)).astype( + jnp.float32 + ) + + res = binary_focal_cross_entropy_ops.cut_binary_focal_cross_entropy( + activations, + embeddings, + targets, + weights, + block_v=block_v, + return_metrics=True, + return_per_target_metrics=True, + return_per_target_losses=return_per_target_losses, + use_pallas=False, + ) + if return_per_target_losses: + loss, losses, tp, fp, fn, tn = res + else: + loss, tp, fp, fn, tn = res + losses = None + + loss_ref, losses_ref = _naive_focal_bce( + activations, embeddings, targets, weights=weights + ) + np.testing.assert_allclose(loss, loss_ref, atol=1e-5, rtol=1e-5) + if losses is not None: + np.testing.assert_allclose(losses, losses_ref, atol=1e-5, rtol=1e-5) + + logits = jnp.matmul(activations, embeddings.T) + multi_hot = jnp.max(jax.nn.one_hot(targets, vocab_size), axis=-2) > 0 + preds = logits > 0.0 + tp_ref = jnp.sum(multi_hot & preds, axis=-1, dtype=jnp.float32) + fp_ref = jnp.sum(~multi_hot & preds, axis=-1, dtype=jnp.float32) + fn_ref = jnp.sum(multi_hot & ~preds, axis=-1, dtype=jnp.float32) + tn_ref = jnp.sum(~multi_hot & ~preds, axis=-1, dtype=jnp.float32) + if weights is not None: + tp_ref = tp_ref * weights + fp_ref = fp_ref * weights + fn_ref = fn_ref * weights + tn_ref = tn_ref * weights + + np.testing.assert_allclose(tp, tp_ref, rtol=1e-6) + np.testing.assert_allclose(fp, fp_ref, rtol=1e-6) + np.testing.assert_allclose(fn, fn_ref, rtol=1e-6) + np.testing.assert_allclose(tn, tn_ref, rtol=1e-6) + @parameterized.named_parameters( ('4d_tgt', True), ('3d_shared_tgt', False), @@ -1036,6 +1102,64 @@ def loss_fn(act, emb, opt): np.testing.assert_allclose(g_act_opt, g_act_ref, atol=1e-4, rtol=1e-4) np.testing.assert_allclose(g_emb_opt, g_emb_ref, atol=1e-4, rtol=1e-4) + def test_force_bwd_block_v(self): + config_clamped = binary_focal_cross_entropy_ops.FocalBCEConfig( + block_v=11008, block_n=256, force_bwd_block_v=False + ) + config_forced = binary_focal_cross_entropy_ops.FocalBCEConfig( + block_v=11008, block_n=256, force_bwd_block_v=True + ) + _, bwd_v_clamped = binary_focal_cross_entropy_ops._effective_block_v( + config_clamped, hidden=256, vocab=16384 + ) + _, bwd_v_forced = binary_focal_cross_entropy_ops._effective_block_v( + config_forced, hidden=256, vocab=16384 + ) + self.assertLess(bwd_v_clamped, 11008) + self.assertEqual(bwd_v_forced, 11008) + + batch, seq_len, hidden_dim, vocab_size, num_labels = 2, 16, 128, 2048, 4 + key_act, key_emb, key_tgt = jax.random.split(jax.random.PRNGKey(99), 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 + ) + + def loss_fn(act, emb, force_bwd): + return binary_focal_cross_entropy_ops.cut_binary_focal_cross_entropy( + act, + emb, + targets, + block_v=1024, + block_n=128, + force_bwd_block_v=force_bwd, + ) + + with mock.patch.object( + binary_focal_cross_entropy_ops.bce_ops, + '_max_safe_block_v', + return_value=512, + ): + loss_clamped = loss_fn(activations, embeddings, False) + loss_forced = loss_fn(activations, embeddings, True) + np.testing.assert_allclose( + loss_forced, loss_clamped, atol=1e-6, rtol=1e-6 + ) + + g_act_clamped, g_emb_clamped = jax.jit( + jax.grad(lambda a, e: loss_fn(a, e, False), argnums=(0, 1)) + )(activations, embeddings) + g_act_forced, g_emb_forced = jax.jit( + jax.grad(lambda a, e: loss_fn(a, e, True), argnums=(0, 1)) + )(activations, embeddings) + np.testing.assert_allclose( + g_act_forced, g_act_clamped, atol=1e-4, rtol=1e-4 + ) + np.testing.assert_allclose( + g_emb_forced, g_emb_clamped, atol=1e-4, rtol=1e-4 + ) + if __name__ == '__main__': absltest.main()