noether.core.distributed.gather

Functions

get_device_and_bfloat16supported()
get_bool_gather_supported()
all_gather_grad(x[, batch_dim])
all_gather_nograd(x[, batch_dim])
all_gather_nograd_clipped(x[, max_length, batch_dim])

Module Contents

noether.core.distributed.gather.get_device_and_bfloat16supported()

noether.core.distributed.gather.get_bool_gather_supported()

noether.core.distributed.gather.all_gather_grad(x, batch_dim=0)

Parameters:

x (torch.Tensor)

Return type:

torch.Tensor

noether.core.distributed.gather.all_gather_nograd(x, batch_dim=0)

Parameters:

x (torch.Tensor)

Return type:

torch.Tensor

noether.core.distributed.gather.all_gather_nograd_clipped(x, max_length=None, batch_dim=0)

Parameters:

• x (torch.Tensor)

• max_length (int | None)

Return type:

torch.Tensor