aihwkit.optim.context module

Parameter context for analog tiles.

class aihwkit.optim.context.AnalogContext(analog_tile, parameter=None)

Bases: torch.nn.parameter.Parameter

Context for analog optimizer.

Parameters

• analog_tile (BaseTile) –

• parameter (Optional[torch.nn.parameter.Parameter]) –

Return type

AnalogContext

cpu()

Move the context to CPU.

Note

This is a no-op for CPU context.

Returns

self

Return type

aihwkit.optim.context.AnalogContext

cuda(device=None)

Move the context to a cuda device.

Parameters

device (Optional[Union[torch.device, str, int]]) – the desired device of the tile.

Returns

This context in the specified device.

Return type

aihwkit.optim.context.AnalogContext

get_data()

Get the data value of the underlying Tensor.

Return type

torch.Tensor

has_gradient()

Return whether a gradient trace was stored.

Return type

bool

reset(analog_tile=None)

Reset the gradient trace and optionally sets the tile pointer.

Parameters

analog_tile (Optional[BaseTile]) –

Return type

None

set_data(data)

Set the data value of the Tensor.

Parameters

data (torch.Tensor) –

Return type

None

to(*args, **kwargs)

Move analog tiles of the current context to a device.

Note

Please be aware that moving analog tiles from GPU to CPU is currently not supported.

Caution

Other tensor conversions than moving the device to CUDA, such as changing the data type are not supported for analog tiles and will be simply ignored.

Returns

This module in the specified device.

Parameters

• args (Any) –

• kwargs (Any) –

Return type

aihwkit.optim.context.AnalogContext