push.particle

class push.particle.Particle(node_event_loop, device: int, pid: int, module: Module, state: Dict[str, any])

Bases: Waitable

User-facing particle interface.

Implements concurrent particles.

forward(x: Tensor, *args: any) → PFuture

Performs a forward pass.

Parameters:

x (torch.Tensor) – Input to the particle.

Returns:

Future that eventually contains the value of the forward pass.

Return type:

PFuture

get(pid: int) → PFuture

Obtains particle pid’s parameters.

Parameters:

pid (int) – Particle identifier to obtain parameters of.

Returns:

Handle to future result.

Return type:

PFuture

other_particles() → List[int]

Returns all particles except current particle.

Returns:

List of all particle identifiers visible to current particle except current particle.

Return type:

List[int]

particle_ids() → List[int]

Returns all particles.

Returns:

List of all particle identifiers visible to current particle.

Return type:

List[int]

register_receive(msg: str, fn: Callable, state: dict[str, any]) → None

Register receive functionality for current particle.

Parameters:

• msg (str) – Message for current particle to respond to.

• fn (Callable) – Function to execute on msg.

• state (dict[str, any]) – User state.

Returns:

None.

scheduler_step(*args) → PFuture

Performs a forward and backward pass using the registered optimizer.

Parameters:

• loss_fn (Callable) – Loss function to take a step with respect to.

• data (torch.Tensor) – Data.

• label (torch.Tensor) – label.

Returns:

Future that eventually contains the loss of the step.

Return type:

PFuture

send(pid: int, msg: str, *args: any) → PFuture

Send a msg to pid with arguments *args from current particle.

Parameters:

• pid (int) – Particle identifier of receiver.

• msg (str) – Message.

• *args (any) – Any additional arguments to give to receiver.

Returns:

Handle to future result.

Return type:

PFuture

step(loss_fn: Callable, data: Tensor, label: Tensor, *args) → PFuture

Performs a forward and backward pass using the registered optimizer.

Parameters:

• loss_fn (Callable) – Loss function to take a step with respect to.

• data (torch.Tensor) – Data.

• label (torch.Tensor) – label.

Returns:

Future that eventually contains the loss of the step.

Return type:

PFuture

wait(pfutures: List[PFuture]) → List[any]

Pause execution of current particle until all pfutures have been resolved.

Parameters:

pfutures (List[PFuture]) – Future values to wait on.

Returns:

List of resolved future values.

Return type:

List[any]

zero_grad() → PFuture

Zeros gradients on current particle.

Returns:

Wait ensures that the gradients have been zeroed.

Return type:

PFuture

class push.particle.ParticleView(view_cache, pid: int)

Bases: object

User-facing particle view interface.

Enables you to view another particle’s parameters.

view() → Module

Call to access it’s latest state.

Returns:

The other particle.

Return type:

nn.Module