braintools.trainer.broadcast# braintools.trainer.broadcast(tensor, src=0, axis_name='batch', num_devices=None)[source]# Broadcast tensor from source device to all devices. Parameters: tensor (Any) – Tensor to broadcast. src (int) – Source device index. axis_name (str) – Name of the axis for broadcast (used with pmap). num_devices (int | None) – Number of devices. Returns: Broadcasted tensor. Return type: Any
