brainpy.math.parallels.pmap#

brainpy.math.parallels.pmap(func, dyn_vars=None, axis_name=None, in_axes=0, out_axes=0, static_broadcasted_argnums=(), devices=None, backend=None, axis_size=None, donate_argnums=(), global_arg_shapes=None, reduce_func=None)[source]#

Parallel compilation for class objects.

Parallel compile a function or a module to run on multiple devices in parallel.

Parameters
  • func

  • axis_name

  • in_axes

  • out_axes

  • static_broadcasted_argnums

  • devices

  • backend

  • axis_size

  • donate_argnums

  • global_arg_shapes

Examples