aihwkit.optim.context module

Parameter context for analog tiles.

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

Bases: Parameter

Context for analog optimizer.

Parameters:
Return type:

AnalogContext

cpu()[source]

Move the context to CPU.

Note

This is a no-op for CPU context.

Returns:

self

Return type:

AnalogContext

cuda(device=None)[source]

Move the context to a cuda device.

Parameters:

device (str | device | int | None) – the desired device of the tile.

Returns:

This context in the specified device.

Return type:

AnalogContext

get_data()[source]

Get the data value of the underlying Tensor.

Return type:

Tensor

has_gradient()[source]

Return whether a gradient trace was stored.

Return type:

bool

reset(analog_tile=None)[source]

Reset the gradient trace and optionally sets the tile pointer.

Parameters:

analog_tile (SimulatorTileWrapper | None) –

Return type:

None

set_data(data)[source]

Set the data value of the Tensor.

Parameters:

data (Tensor) –

Return type:

None

set_indexed(value=True)[source]

Set the context to forward_indexed.

Parameters:

value (bool) –

Return type:

None

to(*args, **kwargs)[source]

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:

AnalogContext