wehub-resource-sync
44 lines
1.1 KiB
Python
Executable File
44 lines
1.1 KiB
Python
Executable File
import numpy as np
|
|
from prml.nn.tensor.constant import Constant
|
|
from prml.nn.tensor.tensor import Tensor
|
|
from prml.nn.function import Function
|
|
|
|
|
|
class Solve(Function):
|
|
|
|
def forward(self, a, b):
|
|
a = self._convert2tensor(a)
|
|
b = self._convert2tensor(b)
|
|
self._equal_ndim(a, 2)
|
|
self._equal_ndim(b, 2)
|
|
self.a = a
|
|
self.b = b
|
|
self.output = np.linalg.solve(a.value, b.value)
|
|
if isinstance(self.a, Constant) and isinstance(self.b, Constant):
|
|
return Constant(self.output)
|
|
return Tensor(self.output, function=self)
|
|
|
|
def backward(self, delta):
|
|
db = np.linalg.solve(self.a.value.T, delta)
|
|
da = -db @ self.output.T
|
|
self.a.backward(da)
|
|
self.b.backward(db)
|
|
|
|
|
|
def solve(a, b):
|
|
"""
|
|
solve a linear matrix equation
|
|
ax = b
|
|
Parameters
|
|
----------
|
|
a : (d, d) tensor_like
|
|
coefficient matrix
|
|
b : (d, k) tensor_like
|
|
dependent variable
|
|
Returns
|
|
-------
|
|
output : (d, k) tensor_like
|
|
solution of the equation
|
|
"""
|
|
return Solve().forward(a, b)
|