aihwkit.optim.context module¶
Parameter context for analog tiles.
-
class
aihwkit.optim.context.
AnalogContext
(analog_tile, parameter=None)[source]¶ Bases:
torch.nn.parameter.Parameter
Context for analog optimizer.
-
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 (Optional[Union[torch.device, str, int]]) – 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 (Optional[BaseTile]) –
- Return type
None
-
set_data
(data)[source]¶ Set the data value of the Tensor.
- Parameters
data (torch.Tensor) –
- 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
-
-
aihwkit.optim.context.
ones
(*size, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) → Tensor¶ Returns a tensor filled with the scalar value 1, with the shape defined by the variable argument
size
.- Parameters
size (int...) – a sequence of integers defining the shape of the output tensor. Can be a variable number of arguments or a collection like a list or tuple.
- Keyword Arguments
out (Tensor, optional) – the output tensor.
dtype (
torch.dtype
, optional) – the desired data type of returned tensor. Default: ifNone
, uses a global default (seetorch.set_default_tensor_type()
).layout (
torch.layout
, optional) – the desired layout of returned Tensor. Default:torch.strided
.device (
torch.device
, optional) – the desired device of returned tensor. Default: ifNone
, uses the current device for the default tensor type (seetorch.set_default_tensor_type()
).device
will be the CPU for CPU tensor types and the current CUDA device for CUDA tensor types.requires_grad (bool, optional) – If autograd should record operations on the returned tensor. Default:
False
.
Example:
>>> torch.ones(2, 3) tensor([[ 1., 1., 1.], [ 1., 1., 1.]]) >>> torch.ones(5) tensor([ 1., 1., 1., 1., 1.])