def jit_forward(forward, n_devices, do_mean=True):
  """Returns a JIT-compiled forward function running on ``n_devices``."""
  model_predict = _accelerate(forward, n_devices)
  # n_devices == 0 => CPU
  if n_devices < 2:
    return model_predict

  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

  return predict
Returns a JIT-compiled forward function running on ``n_devices``.
trax/layers/acceleration.py