aihwkit.nn.functions module

Autograd functions for aihwkit.

class aihwkit.nn.functions.AnalogFunction

Bases: torch.autograd.function.Function

Function that delegates into a RPU unit.

static backward(ctx, grad_output)

Execute the backward pass in the analog tile.

Parameters

• ctx (Any) –

• grad_output (torch.Tensor) –

Return type

Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]

static forward(ctx, analog_tile, input_, weights, _=None, is_test=False)

Execute the forward pass in the analog tile.

Parameters

• ctx (Any) –

• analog_tile (aihwkit.simulator.tiles.floating_point.FloatingPointTile) –

• input_ (torch.Tensor) –

• weights (torch.Tensor) –

• _ (Optional[torch.Tensor]) –

• is_test (bool) –

Return type

torch.Tensor

class aihwkit.nn.functions.AnalogIndexedFunction

Bases: torch.autograd.function.Function

Function that delegates into a RPU unit to use the indexed forward/backward/update.

static backward(ctx, grad_output)

Execute the backward pass in the analog tile.

Parameters

• ctx (Any) –

• grad_output (torch.Tensor) –

Return type

Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]

static forward(ctx, analog_tile, input_, weights, _=None, is_test=False)

Execute the forward pass in the analog tile.

Parameters

• ctx (Any) –

• analog_tile (aihwkit.simulator.tiles.floating_point.FloatingPointTile) –

• input_ (torch.Tensor) –

• weights (torch.Tensor) –

• _ (Optional[torch.Tensor]) –

• is_test (bool) –

Return type

torch.Tensor