def _accelerate(f, n_devices):
"""JIT-compiled version of `f` running on `n_devices`."""
if n_devices == 1:
return math.jit(f)
return math.pmap(f, axis_name='batch')
def _accelerate(f, n_devices):
"""JITed version of f running on n_devices."""
if n_devices == 1:
return math.jit(f)
return math.pmap(f, axis_name='batch')