def pad_input(x):
pad_widths = [(0, 0)] * len(x.shape)
pad_widths[-2] = (0, pad_len) # Padding on axis=-2
return np.pad(x,
pad_widths,
mode='constant',
constant_values=x.dtype.type(0))
def forward(self, x, weights):
assert self._padding == 'VALID'
# Left pad with 0s. Applying an unmasked valid convolution on top of this
# yields a causal convolution.
# TODO(ddohan): Support strided and dilated convolutions.
rate = 1
effective_kernel_size = int((self._kernel_size[0] - 1) * rate + 1)
pad = effective_kernel_size - 1
x_leftpad = np.pad(x,
pad_width=[[0, 0], [pad, 0], [0, 0]],
mode='constant')
return super(CausalConv, self).forward(x_leftpad, weights)
def ShiftRight(x, n_shifts=1, mode='train', **unused_kwargs):
"""Layer to shift the tensor to the right by padding on axis 1."""
if mode == 'predict':
# Do nothing in predict mode, as then the sequence length is 1.
return x
pad_widths = [(0, 0)] * len(x.shape)
pad_widths[1] = (n_shifts, 0) # Padding on axis=1
padded = np.pad(x,
pad_widths,
mode='constant',
constant_values=x.dtype.type(0))
return padded[:, :-n_shifts]
def DiagonalGate(x, **kwargs):
"""Split channels in 3 parts. Shifts 1st and 3rd sections to left/right."""
del kwargs
# x : [batch, 1, length, depth]
x = np.pad(x, [(0, 0), (0, 0), (1, 1), (0, 0)],
mode='constant',
constant_values=0.0)
depth = x.shape[-1] // 3
assert 3 * depth == x.shape[-1], ('Depth must be divisible by 3', depth,
x.shape)
xs = [
x[:, :, :-2, :depth], x[:, :, 1:-1, depth:2 * depth],
x[:, :, 2:, 2 * depth:3 * depth]
]
return np.concatenate(xs, axis=3)
