Remove use of initial argument to jax.nn.softmax and `jax.nn.log_…

…softmax`

This has been deprecated and has had no effect since JAX v0.4.27, and will result in a TypeError in JAX v0.4.36.

PiperOrigin-RevId: 690740882

rax/_src/losses.py

@@ -311,7 +311,6 @@ def unique_softmax_loss(
       scores_repeated,
       axis=-1,
       where=identity_mask | labels_lt,
-      initial=jnp.min(scores),
   )
   log_softmax = jnp.diagonal(log_softmax, axis1=-2, axis2=-1)
 

rax/_src/segment_utils_test.py

@@ -109,8 +109,8 @@ def test_computes_log_softmax_with_mask(self):
     segments = jnp.array([0, 0, 0, 1, 1])
     mask = jnp.array([1, 1, 0, 1, 1])
 
-    expected_1 = jax.nn.log_softmax(a[0:3], where=mask[0:3], initial=jnp.min(a))
-    expected_2 = jax.nn.log_softmax(a[3:5], where=mask[3:5], initial=jnp.min(a))
+    expected_1 = jax.nn.log_softmax(a[0:3], where=mask[0:3])
+    expected_2 = jax.nn.log_softmax(a[3:5], where=mask[3:5])
     actual = segment_utils.segment_log_softmax(a, segments, where=mask)
 
     np.testing.assert_allclose(actual[0:2], expected_1[0:2])
@@ -146,8 +146,8 @@ def test_computes_softmax_with_mask(self):
     segments = jnp.array([0, 0, 0, 1, 1])
     mask = jnp.array([1, 1, 0, 1, 1])
 
-    expected_1 = jax.nn.softmax(a[0:3], where=mask[0:3], initial=jnp.min(a))
-    expected_2 = jax.nn.softmax(a[3:5], where=mask[3:5], initial=jnp.min(a))
+    expected_1 = jax.nn.softmax(a[0:3], where=mask[0:3])
+    expected_2 = jax.nn.softmax(a[3:5], where=mask[3:5])
     actual = segment_utils.segment_softmax(a, segments, where=mask)
 
     np.testing.assert_allclose(actual[0:2], expected_1[0:2])