[JAX] Replace uses of jax.xla_computation() with jax.jit().lower().

jax.xla_computation() is deprecated in favor of jax.jit(...).lower(...).

The most common replacements are either jax.jit(...).lower(...).compiler_ir(dialect='hlo') or jax.jit(...).lower(...).cost_analysis().

PiperOrigin-RevId: 509937560

trax/supervised/trainer_lib.py

@@ -451,9 +451,9 @@ def save_computation_graphs(self):
     if self.n_devices > 1:
       batch = _reshape_by_device(batch, self.n_devices)
     weights = self._opt_state.weights[0]
-    forward_computation = jax.xla_computation(self._model_predict_eval)(
+    forward_computation = jax.jit(self._model_predict_eval).lower(
         batch, weights=weights, state=self._model_state[0],
-        rng=self._rngs[0])
+        rng=self._rngs[0]).compiler_ir(dialect='hlo')
     with tf.io.gfile.GFile(os.path.join(output_dir, 'forward.txt'), 'w') as f:
       f.write(forward_computation.as_hlo_text())
     with tf.io.gfile.GFile(os.path.join(output_dir, 'forward.dot'), 'w') as f: