# -*- coding: utf-8 -*-
# (C) Copyright 2020, 2021, 2022, 2023, 2024 IBM. All Rights Reserved.
#
# Licensed under the MIT license. See LICENSE file in the project root for details.
""" Defines an analog matrix
"""
from typing import Any, Union, Tuple
from scipy.sparse.linalg import LinearOperator
from numpy import ndarray, dtype
# Imports from PyTorch.
from torch import Tensor, from_numpy
from torch.autograd import no_grad
# Imports from aihwkit.
from aihwkit.simulator.parameters.base import RPUConfigBase
[docs]class AnalogMatrix(LinearOperator):
"""Defines an analog matrix for linear algebra.
Args:
matrix: matrix to store into the analog crossbar array.
rpu_config: RPU Config specifying the analog hardware properties.
realistic: Whether to use realistic writing while storing the
matrix elements. Otherwise, the elements will be stored
exactly without write noise(fake write)
to_kwargs: arguments for the torch .to call such as `device`
"""
def __init__(
self,
matrix: Union[Tensor, ndarray],
rpu_config: RPUConfigBase,
realistic: bool = False,
**to_kwargs: Any,
):
# pylint: disable=super-init-not-called
out_features, in_features = matrix.shape
tile_module_class = rpu_config.get_default_tile_module_class(out_features, in_features)
if isinstance(matrix, ndarray):
matrix = from_numpy(matrix)
self._tile = tile_module_class(out_features, in_features, rpu_config, False)
self._tile = self._tile.to(**to_kwargs)
self._tile.set_weights(matrix, realistic=realistic)
self.realistic = realistic
[docs] def cuda(self, *args: Any, **kwargs: Any) -> "AnalogMatrix":
"""Move to GPU."""
self._tile = self._tile.cuda(*args, **kwargs)
return self
[docs] def to(self, *args: Any, **kwargs: Any) -> "AnalogMatrix":
"""Move to device, datatype, or rpu_config."""
# pylint: disable=invalid-name
self._tile = self._tile.to(*args, **kwargs)
return self
def __matmul__(self, matrix: Union[Tensor, ndarray]) -> Union[Tensor, ndarray]:
output_numpy = False
if isinstance(matrix, ndarray):
output_numpy = True
matrix = from_numpy(matrix).to(self._tile.device)
with no_grad():
output = self._tile.forward(matrix)
if output_numpy:
return output.cpu().numpy()
return output
def __rmatmul__(self, matrix: Union[Tensor, ndarray]) -> Union[Tensor, ndarray]:
output_numpy = False
if isinstance(matrix, ndarray):
output_numpy = True
matrix = from_numpy(matrix).to(self._tile.device)
with no_grad():
output = self._tile.backward(matrix)
if output_numpy:
return output.cpu().numpy()
return output
def _matmat(self, X: Union[Tensor, ndarray]) -> Union[Tensor, ndarray]:
return self @ X
def _matvec(self, x: Union[Tensor, ndarray]) -> Union[Tensor, ndarray]:
return self @ x
def _rmatvec(self, x: Union[Tensor, ndarray]) -> Union[Tensor, ndarray]:
return x @ self
def _rmatmat(self, X: Union[Tensor, ndarray]) -> Union[Tensor, ndarray]:
return X @ self
@property
def shape(self) -> Tuple[int, int]:
"""Returns shape of analog matrix."""
return (self._tile.out_size, self._tile.in_size)
@property
def dtype(self) -> dtype:
"""Return data type."""
return self._tile.rpu_config.get_data_type().as_torch()
[docs] def ger(
self, x_matrix: Union[Tensor, ndarray], d_matrix: Union[Tensor, ndarray], alpha: float = 1.0
) -> None:
"""GER (rank update) function performed on the analog matrix.
Args:
x_matrix: left input matrix
d_matrix: right input matrix
alpha: scale factor
"""
if isinstance(x_matrix, ndarray):
x_matrix = from_numpy(x_matrix)
if isinstance(d_matrix, ndarray):
d_matrix = from_numpy(d_matrix)
self._tile.set_learning_rate(alpha)
self._tile.update(x_matrix, -d_matrix)