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:
analog_tile (SimulatorTileWrapper) –
parameter (Parameter | None) –
- Return type:
- cpu()[source]
Move the context to CPU.
Note
This is a no-op for CPU context.
- Returns:
self
- Return type:
- 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:
- 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: