wehub-resource-sync
49 lines
1.4 KiB
Python
Executable File
49 lines
1.4 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 Nth(Function):
|
|
|
|
def __init__(self, n):
|
|
self.n = n
|
|
|
|
def forward(self, x):
|
|
self.x = x
|
|
if isinstance(self.x, Constant):
|
|
return Constant(x.value)
|
|
return Tensor(x.value, function=self)
|
|
|
|
def backward(self, delta):
|
|
self.x.backward(delta, n=self.n)
|
|
|
|
|
|
class Split(Function):
|
|
|
|
def __init__(self, indices_or_sections, axis=-1):
|
|
self.indices_or_sections = indices_or_sections
|
|
self.axis = axis
|
|
|
|
def forward(self, x):
|
|
x = self._convert2tensor(x)
|
|
self._atleast_ndim(x, 1)
|
|
self.x = x
|
|
output = np.split(x.value, self.indices_or_sections, self.axis)
|
|
if isinstance(self.x, Constant):
|
|
return tuple([Constant(out) for out in output])
|
|
self.n_output = len(output)
|
|
self.delta = [None for _ in output]
|
|
return tuple([Tensor(out, function=self) for out in output])
|
|
|
|
def backward(self, delta, n):
|
|
self.delta[n] = delta
|
|
if all([d is not None for d in self.delta]):
|
|
dx = np.concatenate(self.delta, axis=self.axis)
|
|
self.x.backward(dx)
|
|
|
|
|
|
def split(x, indices_or_sections, axis=-1):
|
|
output = Split(indices_or_sections, axis).forward(x)
|
|
return tuple([Nth(i).forward(out) for i, out in enumerate(output)])
|