def predict(x, weights, state, rng):
    """Predict function JIT-compiled and parallelized as requested."""
    res, state = model_predict(
        reshape_by_device(x, n_devices),
        weights,
        state,
        jnp.stack(fastmath.random.split(rng, n_devices)))
    res = _combine_devices(res)
    if do_mean:
      return fastmath.nested_map(
          lambda y: mean_or_pmean(n_devices, y, axis=0), res), state
    else:
      return res, state
Predict function JIT-compiled and parallelized as requested.
trax/layers/acceleration.py