ksuit.distributed.gather¶

Functions¶

`get_device_and_bfloat16supported`()
`get_bool_gather_supported`()
`all_gather_grad`(x[, batch_dim])
`all_gather_grad_autograd`(x)
`all_gather_nograd`(x)
`all_gather_nograd_clipped`(x, max_length)
`all_reduce_sum_nograd`(x)
`all_reduce_sum_grad`(x)
`all_reduce_mean_grad`(x)
`all_reduce_mean_nograd`(x)

Module Contents¶

ksuit.distributed.gather.get_device_and_bfloat16supported()¶

ksuit.distributed.gather.get_bool_gather_supported()¶

ksuit.distributed.gather.all_gather_grad(x, batch_dim=0)¶

ksuit.distributed.gather.all_gather_grad_autograd(x)¶

ksuit.distributed.gather.all_gather_nograd(x)¶

ksuit.distributed.gather.all_gather_nograd_clipped(x, max_length)¶

ksuit.distributed.gather.all_reduce_sum_nograd(x)¶

ksuit.distributed.gather.all_reduce_sum_grad(x)¶

ksuit.distributed.gather.all_reduce_mean_grad(x)¶

ksuit.distributed.gather.all_reduce_mean_nograd(x)¶