aihwkit.simulator.tiles.torch_tile module

Low level implementation of torch-based tile.

class aihwkit.simulator.tiles.torch_tile.TorchSimulatorTile(x_size, d_size, rpu_config, bias=False, analog_mvm=<class 'aihwkit.simulator.tiles.analog_mvm.AnalogMVM'>)

Bases: SimulatorTile, Module

Torch based tile class.

Parameters:

• x_size (int) – input size

• d_size (int) – output size

• rpu_config (TorchInferenceRPUConfig) – resistive processing unit configuration.

• bias (bool) –

• analog_mvm (Type[AnalogMVM]) –

classmethod check_rpu_config_support(rpu_config)

Check the RPUConfig for support with TorchSimulatorTile

Throws an assertion error when there is an incompatibility with the used rpu config and what the current torch tile supports

Parameters:

rpu_config (TorchInferenceRPUConfig) – the rpu config to be checked

Raises:

TorchTileConfigError – in case a feature is not supported

Return type:

None

clip_weights(clip)

Clip the weights.

Parameters:

clip (WeightClipParameter) – parameters specifying the clipping methof and type.

Raises:

• NotImplementedError – For unsupported WeightClipTypes

• ConfigError – If unknown WeightClipType used.

Return type:

None

extra_repr()

Set the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

Return type:

str

forward(x_input, bias=False, in_trans=False, out_trans=False, is_test=False, non_blocking=False)

General simulator tile forward.

Parameters:

• x_input (Tensor) –

• bias (bool) –

• in_trans (bool) –

• out_trans (bool) –

• is_test (bool) –

• non_blocking (bool) –

Return type:

Tensor

get_brief_info()

Returns a brief info

Return type:

str

get_d_size()

Returns output size of tile

Return type:

int

get_weights()

Get the tile weights.

Returns:

a tuple where the first item is the [out_size, in_size] weight matrix; and the second item is either the [out_size] bias vector or None if the tile is set not to use bias.

Return type:

Tensor

get_x_size()

Returns input size of tile

Return type:

int

static modify_weight(inp_weight, modifier, batch_size)

Weight modifier that adds noise to the weights according to rpu config.

Parameters:

• inp_weight (Tensor) – Input weights.

• modifier (WeightModifierParameter) – Noise injection configuration.

• batch_size (int) – Batch size.

Raises:

TorchTileConfigError – Unsupported/unknown weight modifier type.

Returns:

Weights with noise injected.

Return type:

Tensor

register_weight_hook(hook)

Register a hook to the weights.

Parameters:

hook (Callable) –

Return type:

Any

remap_weights(remap, scales)

Remap the weights to the specified range and return new scales.

Parameters:

• remap (WeightRemapParameter) – hyper-parameters defining the remapping

• scales (Tensor) – current scale values.

Raises:

ConfigError – If WeightRemapType is unknown.

Returns:

New scales.

Return type:

Tensor

set_config(rpu_config)

Updated the configuration to allow on-the-fly changes.

Parameters:

rpu_config (TorchInferenceRPUConfig) – configuration to use in the next forward passes.

Return type:

None

set_weights(weight)

Set the tile weights.

Parameters:

weight (Tensor) – [out_size, in_size] weight matrix.

Return type:

None