braintools.trainer.all_reduce

Contents

braintools.trainer.all_reduce#

braintools.trainer.all_reduce(tensor, op='mean', axis_name='batch')[source]#

All-reduce tensor across devices.

Parameters:

  • tensor (Any) – Tensor to reduce.

  • op (str) – Reduction operation (‘mean’, ‘sum’, ‘min’, ‘max’).

  • axis_name (str) – Name of the axis for reduction (used with pmap).

Returns:

Reduced tensor.

Return type:

Any

Examples

>>> @functools.partial(jax.pmap, axis_name='batch')
... def parallel_fn(x):
...     return all_reduce(x, op='mean', axis_name='batch')
Contents