Replace usage of jax.experimental.host_callback.call with jax.experim…

…ental.io_callback.

The jax.experimental.host_callback module is deprecated and will be removed.

See jax-ml/jax#20385.

PiperOrigin-RevId: 683196778

t5x/decoding_test.py

@@ -21,7 +21,7 @@
 from absl.testing import absltest
 from absl.testing import parameterized
 import jax
-from jax.experimental import host_callback as hcb
+from jax.experimental import io_callback
 import jax.numpy as jnp
 import numpy as np
 from t5x import decoding
@@ -156,12 +156,10 @@ def callback_fn(current_index_and_sequences):
             sequences[i, current_index[i] + 1] = EOS_ID
         return sequences
 
-      sequences = hcb.call(
+      sequences = io_callback(
           callback_fn,
+          jax.ShapeDtypeStruct(state.sequences.shape, state.sequences.dtype),
           (state.cur_index, state.sequences),
-          result_shape=jax.ShapeDtypeStruct(
-              state.sequences.shape, state.sequences.dtype
-          ),
       )
       return state.replace(sequences=sequences)