# flake8: noqa

import torch

import torch.nn as nn

import torch.nn.functional as F

import tensorflow as tf

import numpy as np

#       Gold-standard implementations for testing custom layers       #

#                       (Requires Pytorch)                            #

def torchify(var, requires_grad=True):

return torch.autograd.Variable(torch.FloatTensor(var), requires_grad=requires_grad)

def torch_gradient_generator(fn, **kwargs):

def get_grad(z):

z1 = torch.autograd.Variable(torch.FloatTensor(z), requires_grad=True)

z2 = fn(z1, **kwargs).sum()

grad = z1.grad.numpy()

return grad

return get_grad

def torch_xe_grad(y, z):

z = torch.autograd.Variable(torch.FloatTensor(z), requires_grad=True)

y = torch.LongTensor(y.argmax(axis=1))

loss = F.cross_entropy(z, y, reduction="sum")

loss.backward()

grad = z.grad.numpy()

return grad

def torch_mse_grad(y, z, act_fn):

y = torch.FloatTensor(y)

y_pred = act_fn(z)

loss = F.mse_loss(y_pred, y, reduction="sum")  # size_average=False).sum()

loss.backward()

grad = z.grad.numpy()

return grad

class TorchVAELoss(nn.Module):

def __init__(self):

super(TorchVAELoss, self).__init__()

def extract_grads(self, X, X_recon, t_mean, t_log_var):

eps = np.finfo(float).eps

X = torchify(X, requires_grad=False)

X_recon = torchify(np.clip(X_recon, eps, 1 - eps))

t_mean = torchify(t_mean)

t_log_var = torchify(t_log_var)

BCE = torch.sum(F.binary_cross_entropy(X_recon, X, reduction="none"), dim=1)

# see Appendix B from VAE paper:

# Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014

# https://arxiv.org/abs/1312.6114

# 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)

KLD = -0.5 * torch.sum(1 + t_log_var - t_mean.pow(2) - t_log_var.exp(), dim=1)

loss = torch.mean(BCE + KLD)

loss.backward()

grads = {

"loss": loss.detach().numpy(),

"dX_recon": X_recon.grad.numpy(),

"dt_mean": t_mean.grad.numpy(),

"dt_log_var": t_log_var.grad.numpy(),

return grads

class TorchWGANGPLoss(nn.Module):

def __init__(self, lambda_=10):

self.lambda_ = torchify([lambda_])

super(TorchWGANGPLoss, self).__init__()

def forward(self, Y_real, Y_fake, gradInterp):

GY_fake = Y_fake.copy()

self.Y_real = torchify(Y_real)

self.Y_fake = torchify(Y_fake)

self.GY_fake = torchify(GY_fake)

self.gradInterp = torchify(gradInterp)

# calc grad penalty

norm = self.gradInterp.norm(2, dim=1)

self.norm1 = torch.sqrt(torch.sum(self.gradInterp.pow(2), dim=1))

assert torch.allclose(norm, self.norm1)

self.gpenalty = self.lambda_ * ((self.norm1 - 1).pow(2)).mean()

self.C_loss = self.Y_fake.mean() - self.Y_real.mean() + self.gpenalty

self.G_loss = -self.GY_fake.mean()

def extract_grads(self, Y_real, Y_fake, gradInterp):

self.forward(Y_real, Y_fake, gradInterp)

grads = {

"Y_real": self.Y_real.detach().numpy(),

"Y_fake": self.Y_fake.detach().numpy(),

"gradInterp": self.gradInterp.detach().numpy(),

"GP": self.gpenalty.detach().numpy(),

"C_loss": self.C_loss.detach().numpy(),

"G_loss": self.G_loss.detach().numpy(),

"C_dY_real": self.Y_real.grad.numpy(),

"C_dGradInterp": self.gradInterp.grad.numpy(),

"C_dY_fake": self.Y_fake.grad.numpy(),

"G_dY_fake": self.GY_fake.grad.numpy(),

return grads

class TorchLinearActivation(nn.Module):

def __init__(self):

super(TorchLinearActivation, self).__init__()

pass

@staticmethod

def forward(input):

return input

@staticmethod

def backward(grad_output):

return torch.ones_like(grad_output)

class TorchBatchNormLayer(nn.Module):

def __init__(self, n_in, params, mode, momentum=0.9, epsilon=1e-5):

super(TorchBatchNormLayer, self).__init__()

scaler = params["scaler"]

intercept = params["intercept"]

if mode == "1D":

self.layer1 = nn.BatchNorm1d(

num_features=n_in, momentum=1 - momentum, eps=epsilon, affine=True

elif mode == "2D":

self.layer1 = nn.BatchNorm2d(

self.layer1.weight = nn.Parameter(torch.FloatTensor(scaler))

self.layer1.bias = nn.Parameter(torch.FloatTensor(intercept))

def forward(self, X):

# (N, H, W, C) -> (N, C, H, W)

if X.ndim == 4:

X = np.moveaxis(X, [0, 1, 2, 3], [0, -2, -1, -3])

if not isinstance(X, torch.Tensor):

X = torchify(X)

self.X = X

self.Y = self.layer1(self.X)

def extract_grads(self, X, Y_true=None):

self.forward(X)

if isinstance(Y_true, np.ndarray):

Y_true = np.moveaxis(Y_true, [0, 1, 2, 3], [0, -2, -1, -3])

self.loss1 = (

0.5 * F.mse_loss(self.Y, torchify(Y_true), size_average=False).sum()

else:

self.loss1 = self.Y.sum()

X_np = self.X.detach().numpy()

Y_np = self.Y.detach().numpy()

dX_np = self.X.grad.numpy()

dY_np = self.Y.grad.numpy()

if self.X.dim() == 4:

orig, X_swap = [0, 1, 2, 3], [0, -1, -3, -2]

if isinstance(Y_true, np.ndarray):

Y_true = np.moveaxis(Y_true, orig, X_swap)

X_np = np.moveaxis(X_np, orig, X_swap)

Y_np = np.moveaxis(Y_np, orig, X_swap)

dX_np = np.moveaxis(dX_np, orig, X_swap)

dY_np = np.moveaxis(dY_np, orig, X_swap)

grads = {

"loss": self.loss1.detach().numpy(),

"X": X_np,

"momentum": 1 - self.layer1.momentum,

"epsilon": self.layer1.eps,

"intercept": self.layer1.bias.detach().numpy(),

"scaler": self.layer1.weight.detach().numpy(),

"running_mean": self.layer1.running_mean.detach().numpy(),

"running_var": self.layer1.running_var.detach().numpy(),

"y": Y_np,

"dLdy": dY_np,

"dLdIntercept": self.layer1.bias.grad.numpy(),

"dLdScaler": self.layer1.weight.grad.numpy(),

"dLdX": dX_np,

if isinstance(Y_true, np.ndarray):

grads["Y_true"] = Y_true

return grads

class TorchLayerNormLayer(nn.Module):

def __init__(self, feat_dims, params, mode, epsilon=1e-5):

super(TorchLayerNormLayer, self).__init__()

self.layer1 = nn.LayerNorm(

normalized_shape=feat_dims, eps=epsilon, elementwise_affine=True

scaler = params["scaler"]

intercept = params["intercept"]

if mode == "2D":

scaler = np.moveaxis(scaler, [0, 1, 2], [-2, -1, -3])

intercept = np.moveaxis(intercept, [0, 1, 2], [-2, -1, -3])

assert scaler.shape == self.layer1.weight.shape

assert intercept.shape == self.layer1.bias.shape

self.layer1.weight = nn.Parameter(torch.FloatTensor(scaler))

self.layer1.bias = nn.Parameter(torch.FloatTensor(intercept))

def forward(self, X):

# (N, H, W, C) -> (N, C, H, W)

if X.ndim == 4:

X = np.moveaxis(X, [0, 1, 2, 3], [0, -2, -1, -3])

if not isinstance(X, torch.Tensor):

X = torchify(X)

self.X = X

self.Y = self.layer1(self.X)

def extract_grads(self, X, Y_true=None):

self.forward(X)

if isinstance(Y_true, np.ndarray):

Y_true = np.moveaxis(Y_true, [0, 1, 2, 3], [0, -2, -1, -3])

self.loss1 = (

0.5 * F.mse_loss(self.Y, torchify(Y_true), size_average=False).sum()

else:

self.loss1 = self.Y.sum()

X_np = self.X.detach().numpy()

Y_np = self.Y.detach().numpy()

dX_np = self.X.grad.numpy()

dY_np = self.Y.grad.numpy()

intercept_np = self.layer1.bias.detach().numpy()

scaler_np = self.layer1.weight.detach().numpy()

dIntercept_np = self.layer1.bias.grad.numpy()

dScaler_np = self.layer1.weight.grad.numpy()

if self.X.dim() == 4:

orig, X_swap = [0, 1, 2, 3], [0, -1, -3, -2]

orig_p, p_swap = [0, 1, 2], [-1, -3, -2]

if isinstance(Y_true, np.ndarray):

Y_true = np.moveaxis(Y_true, orig, X_swap)

X_np = np.moveaxis(X_np, orig, X_swap)

Y_np = np.moveaxis(Y_np, orig, X_swap)

dX_np = np.moveaxis(dX_np, orig, X_swap)

dY_np = np.moveaxis(dY_np, orig, X_swap)

scaler_np = np.moveaxis(scaler_np, orig_p, p_swap)

intercept_np = np.moveaxis(intercept_np, orig_p, p_swap)

dScaler_np = np.moveaxis(dScaler_np, orig_p, p_swap)

dIntercept_np = np.moveaxis(dIntercept_np, orig_p, p_swap)

grads = {

"loss": self.loss1.detach().numpy(),

"X": X_np,

"epsilon": self.layer1.eps,

"intercept": intercept_np,

"scaler": scaler_np,

"y": Y_np,

"dLdy": dY_np,

"dLdIntercept": dIntercept_np,

"dLdScaler": dScaler_np,

"dLdX": dX_np,

if isinstance(Y_true, np.ndarray):

grads["Y_true"] = Y_true

return grads

class TorchAddLayer(nn.Module):

def __init__(self, act_fn, **kwargs):

super(TorchAddLayer, self).__init__()

self.act_fn = act_fn

def forward(self, Xs):

self.Xs = []

x = Xs[0].copy()

if not isinstance(x, torch.Tensor):

x = torchify(x)

self.sum = x.clone()

self.Xs.append(x)

for i in range(1, len(Xs)):

x = Xs[i]

if not isinstance(x, torch.Tensor):

x = torchify(x)

self.Xs.append(x)

self.sum += x

self.Y = self.act_fn(self.sum)

return self.Y

def extract_grads(self, X):

self.forward(X)

self.loss = self.Y.sum()

self.loss.backward()

grads = {

"Xs": X,

"Sum": self.sum.detach().numpy(),

"Y": self.Y.detach().numpy(),

"dLdY": self.Y.grad.numpy(),

"dLdSum": self.sum.grad.numpy(),

grads.update(

{"dLdX{}".format(i + 1): xi.grad.numpy() for i, xi in enumerate(self.Xs)}

return grads

class TorchMultiplyLayer(nn.Module):

def __init__(self, act_fn, **kwargs):

super(TorchMultiplyLayer, self).__init__()

self.act_fn = act_fn

def forward(self, Xs):

self.Xs = []

x = Xs[0].copy()

if not isinstance(x, torch.Tensor):

x = torchify(x)

self.prod = x.clone()

self.Xs.append(x)

for i in range(1, len(Xs)):

x = Xs[i]

if not isinstance(x, torch.Tensor):

x = torchify(x)

self.Xs.append(x)

self.prod *= x

self.Y = self.act_fn(self.prod)

return self.Y

def extract_grads(self, X):

self.forward(X)

self.loss = self.Y.sum()

self.loss.backward()

grads = {

"Xs": X,

"Prod": self.prod.detach().numpy(),

"Y": self.Y.detach().numpy(),

"dLdY": self.Y.grad.numpy(),

"dLdProd": self.prod.grad.numpy(),

grads.update(

return grads

class TorchSkipConnectionIdentity(nn.Module):

def __init__(self, act_fn, pad1, pad2, params, hparams, momentum=0.9, epsilon=1e-5):

super(TorchSkipConnectionIdentity, self).__init__()

self.conv1 = nn.Conv2d(

bias=True,

self.act_fn = act_fn

self.batchnorm1 = nn.BatchNorm2d(

momentum=1 - momentum,

eps=epsilon,

affine=True,

self.conv2 = nn.Conv2d(

bias=True,

self.batchnorm2 = nn.BatchNorm2d(

momentum=1 - momentum,

eps=epsilon,

affine=True,

orig, W_swap = [0, 1, 2, 3], [-2, -1, -3, -4]

# (f[0], f[1], n_in, n_out) -> (n_out, n_in, f[0], f[1])

W = params["components"]["conv1"]["W"]

b = params["components"]["conv1"]["b"]

W = np.moveaxis(W, orig, W_swap)

assert self.conv1.weight.shape == W.shape

assert self.conv1.bias.shape == b.flatten().shape

self.conv1.weight = nn.Parameter(torch.FloatTensor(W))

self.conv1.bias = nn.Parameter(torch.FloatTensor(b.flatten()))

scaler = params["components"]["batchnorm1"]["scaler"]

intercept = params["components"]["batchnorm1"]["intercept"]

self.batchnorm1.weight = nn.Parameter(torch.FloatTensor(scaler))

self.batchnorm1.bias = nn.Parameter(torch.FloatTensor(intercept))

W = params["components"]["conv2"]["W"]

b = params["components"]["conv2"]["b"]

W = np.moveaxis(W, orig, W_swap)

assert self.conv2.weight.shape == W.shape

assert self.conv2.bias.shape == b.flatten().shape

self.conv2.weight = nn.Parameter(torch.FloatTensor(W))

self.conv2.bias = nn.Parameter(torch.FloatTensor(b.flatten()))

scaler = params["components"]["batchnorm2"]["scaler"]

intercept = params["components"]["batchnorm2"]["intercept"]

self.batchnorm2.weight = nn.Parameter(torch.FloatTensor(scaler))

self.batchnorm2.bias = nn.Parameter(torch.FloatTensor(intercept))

def forward(self, X):

if not isinstance(X, torch.Tensor):

# (N, H, W, C) -> (N, C, H, W)

X = np.moveaxis(X, [0, 1, 2, 3], [0, -2, -1, -3])

X = torchify(X)

self.X = X

self.conv1_out = self.conv1(self.X)

self.act_fn1_out = self.act_fn(self.conv1_out)

self.batchnorm1_out = self.batchnorm1(self.act_fn1_out)

self.conv2_out = self.conv2(self.batchnorm1_out)

self.batchnorm2_out = self.batchnorm2(self.conv2_out)

self.layer3_in = self.batchnorm2_out + self.X

self.Y = self.act_fn(self.layer3_in)

def extract_grads(self, X):

self.forward(X)

self.loss = self.Y.sum()

self.loss.backward()

orig, X_swap, W_swap = [0, 1, 2, 3], [0, -1, -3, -2], [-1, -2, -4, -3]

grads = {

# layer parameters

"conv1_W": np.moveaxis(self.conv1.weight.detach().numpy(), orig, W_swap),

"conv1_b": self.conv1.bias.detach().numpy().reshape(1, 1, 1, -1),

"bn1_intercept": self.batchnorm1.bias.detach().numpy(),

"bn1_scaler": self.batchnorm1.weight.detach().numpy(),

"bn1_running_mean": self.batchnorm1.running_mean.detach().numpy(),

"bn1_running_var": self.batchnorm1.running_var.detach().numpy(),

"conv2_W": np.moveaxis(self.conv2.weight.detach().numpy(), orig, W_swap),

"conv2_b": self.conv2.bias.detach().numpy().reshape(1, 1, 1, -1),

"bn2_intercept": self.batchnorm2.bias.detach().numpy(),

"bn2_scaler": self.batchnorm2.weight.detach().numpy(),

"bn2_running_mean": self.batchnorm2.running_mean.detach().numpy(),

"bn2_running_var": self.batchnorm2.running_var.detach().numpy(),

# layer inputs/outputs (forward step)

"X": np.moveaxis(self.X.detach().numpy(), orig, X_swap),

"conv1_out": np.moveaxis(self.conv1_out.detach().numpy(), orig, X_swap),

"act1_out": np.moveaxis(self.act_fn1_out.detach().numpy(), orig, X_swap),

"bn1_out": np.moveaxis(self.batchnorm1_out.detach().numpy(), orig, X_swap),

"conv2_out": np.moveaxis(self.conv2_out.detach().numpy(), orig, X_swap),

"bn2_out": np.moveaxis(self.batchnorm2_out.detach().numpy(), orig, X_swap),

"add_out": np.moveaxis(self.layer3_in.detach().numpy(), orig, X_swap),

"Y": np.moveaxis(self.Y.detach().numpy(), orig, X_swap),

# layer gradients (backward step)

"dLdY": np.moveaxis(self.Y.grad.numpy(), orig, X_swap),

"dLdAdd": np.moveaxis(self.layer3_in.grad.numpy(), orig, X_swap),

"dLdBn2_out": np.moveaxis(self.batchnorm2_out.grad.numpy(), orig, X_swap),

"dLdConv2_out": np.moveaxis(self.conv2_out.grad.numpy(), orig, X_swap),

"dLdBn1_out": np.moveaxis(self.batchnorm1_out.grad.numpy(), orig, X_swap),

"dLdActFn1_out": np.moveaxis(self.act_fn1_out.grad.numpy(), orig, X_swap),

"dLdConv1_out": np.moveaxis(self.act_fn1_out.grad.numpy(), orig, X_swap),

"dLdX": np.moveaxis(self.X.grad.numpy(), orig, X_swap),

# layer parameter gradients (backward step)

"dLdBn2_intercept": self.batchnorm2.bias.grad.numpy(),

"dLdBn2_scaler": self.batchnorm2.weight.grad.numpy(),

"dLdConv2_W": np.moveaxis(self.conv2.weight.grad.numpy(), orig, W_swap),

"dLdConv2_b": self.conv2.bias.grad.numpy().reshape(1, 1, 1, -1),

"dLdBn1_intercept": self.batchnorm1.bias.grad.numpy(),

"dLdBn1_scaler": self.batchnorm1.weight.grad.numpy(),

"dLdConv1_W": np.moveaxis(self.conv1.weight.grad.numpy(), orig, W_swap),

"dLdConv1_b": self.conv1.bias.grad.numpy().reshape(1, 1, 1, -1),

return grads

class TorchCausalConv1d(torch.nn.Conv1d):

NB: this is only ensures that the convolution out length is the same as

the input length IFF stride = 1. Otherwise, in/out lengths will differ.

def __init__(

self,

bias=True,

self.__padding = (kernel_size - 1) * dilation

super(TorchCausalConv1d, self).__init__(

stride=stride,

dilation=dilation,

groups=groups,

bias=bias,

def forward(self, input):

result = super(TorchCausalConv1d, self).forward(input)

if self.__padding != 0:

return result[:, :, : -self.__padding]

return result

class TorchWavenetModule(nn.Module):

def __init__(self, params, hparams, conv_1x1_pad):

super(TorchWavenetModule, self).__init__()

self.conv_dilation = TorchCausalConv1d(

in_channels=hparams["components"]["conv_dilation"]["in_ch"],

out_channels=hparams["components"]["conv_dilation"]["out_ch"],

kernel_size=hparams["components"]["conv_dilation"]["kernel_width"],

stride=hparams["components"]["conv_dilation"]["stride"],

dilation=hparams["components"]["conv_dilation"]["dilation"] + 1,

bias=True,

self.conv_1x1 = nn.Conv1d(

in_channels=hparams["components"]["conv_1x1"]["in_ch"],

out_channels=hparams["components"]["conv_1x1"]["out_ch"],

kernel_size=hparams["components"]["conv_1x1"]["kernel_width"],

stride=hparams["components"]["conv_1x1"]["stride"],

dilation=hparams["components"]["conv_1x1"]["dilation"] + 1,

bias=True,

W = params["components"]["conv_dilation"]["W"]

b = params["components"]["conv_dilation"]["b"]

# (f[0], n_in, n_out) -> (n_out, n_in, f[0])

W = np.moveaxis(W, [0, 1, 2], [-1, -2, -3])

self.conv_dilation.weight = nn.Parameter(torch.FloatTensor(W))

self.conv_dilation.bias = nn.Parameter(torch.FloatTensor(b.flatten()))

assert self.conv_dilation.weight.shape == W.shape

assert self.conv_dilation.bias.shape == b.flatten().shape

W = params["components"]["conv_1x1"]["W"]

b = params["components"]["conv_1x1"]["b"]

self.conv_1x1.weight = nn.Parameter(torch.FloatTensor(W))

self.conv_1x1.bias = nn.Parameter(torch.FloatTensor(b.flatten()))

assert self.conv_1x1.weight.shape == W.shape

assert self.conv_1x1.bias.shape == b.flatten().shape

def forward(self, X_main, X_skip):

# (N, W, C) -> (N, C, W)

self.X_main = np.moveaxis(X_main, [0, 1, 2], [0, -1, -2])

self.X_main = torchify(self.X_main)

self.conv_dilation_out = self.conv_dilation(self.X_main)

self.tanh_out = torch.tanh(self.conv_dilation_out)

self.sigm_out = torch.sigmoid(self.conv_dilation_out)

self.multiply_gate_out = self.tanh_out * self.sigm_out

self.conv_1x1_out = self.conv_1x1(self.multiply_gate_out)

self.X_skip = torch.zeros_like(self.conv_1x1_out)

if X_skip is not None:

self.X_skip = torchify(np.moveaxis(X_skip, [0, 1, 2], [0, -1, -2]))

self.Y_skip = self.X_skip + self.conv_1x1_out

self.Y_main = self.X_main + self.conv_1x1_out

def extract_grads(self, X_main, X_skip):

self.forward(X_main, X_skip)

self.loss = (self.Y_skip + self.Y_main).sum()

self.loss.backward()

# W (theirs): (n_out, n_in, f[0]) -> W (mine): (f[0], n_in, n_out)

# X (theirs): (N, C, W)              -> X (mine): (N, W, C)

# Y (theirs): (N, C, W)              -> Y (mine): (N, W, C)

orig, X_swap, W_swap = [0, 1, 2], [0, -1, -2], [-1, -2, -3]

grads = {

"X_main": np.moveaxis(self.X_main.detach().numpy(), orig, X_swap),

"X_skip": np.moveaxis(self.X_skip.detach().numpy(), orig, X_swap),

"conv_dilation_W": np.moveaxis(

self.conv_dilation.weight.detach().numpy(), orig, W_swap

"conv_dilation_b": self.conv_dilation.bias.detach()

.numpy()

.reshape(1, 1, -1),

"conv_1x1_W": np.moveaxis(

self.conv_1x1.weight.detach().numpy(), orig, W_swap

"conv_1x1_b": self.conv_1x1.bias.detach().numpy().reshape(1, 1, -1),

"conv_dilation_out": np.moveaxis(

self.conv_dilation_out.detach().numpy(), orig, X_swap

"tanh_out": np.moveaxis(self.tanh_out.detach().numpy(), orig, X_swap),

"sigm_out": np.moveaxis(self.sigm_out.detach().numpy(), orig, X_swap),

"multiply_gate_out": np.moveaxis(

self.multiply_gate_out.detach().numpy(), orig, X_swap

"conv_1x1_out": np.moveaxis(

self.conv_1x1_out.detach().numpy(), orig, X_swap

"Y_main": np.moveaxis(self.Y_main.detach().numpy(), orig, X_swap),

"Y_skip": np.moveaxis(self.Y_skip.detach().numpy(), orig, X_swap),

"dLdY_skip": np.moveaxis(self.Y_skip.grad.numpy(), orig, X_swap),

"dLdY_main": np.moveaxis(self.Y_main.grad.numpy(), orig, X_swap),

"dLdConv_1x1_out": np.moveaxis(

self.conv_1x1_out.grad.numpy(), orig, X_swap

"dLdConv_1x1_W": np.moveaxis(

self.conv_1x1.weight.grad.numpy(), orig, W_swap

"dLdConv_1x1_b": self.conv_1x1.bias.grad.numpy().reshape(1, 1, -1),

"dLdMultiply_out": np.moveaxis(

self.multiply_gate_out.grad.numpy(), orig, X_swap

"dLdTanh_out": np.moveaxis(self.tanh_out.grad.numpy(), orig, X_swap),

"dLdSigm_out": np.moveaxis(self.sigm_out.grad.numpy(), orig, X_swap),

"dLdConv_dilation_out": np.moveaxis(

self.conv_dilation_out.grad.numpy(), orig, X_swap

"dLdConv_dilation_W": np.moveaxis(

self.conv_dilation.weight.grad.numpy(), orig, W_swap

"dLdConv_dilation_b": self.conv_dilation.bias.grad.numpy().reshape(

1, 1, -1

"dLdX_main": np.moveaxis(self.X_main.grad.numpy(), orig, X_swap),

"dLdX_skip": np.moveaxis(self.X_skip.grad.numpy(), orig, X_swap),

return grads

class TorchSkipConnectionConv(nn.Module):

def __init__(

self, act_fn, pad1, pad2, pad_skip, params, hparams, momentum=0.9, epsilon=1e-5

super(TorchSkipConnectionConv, self).__init__()

self.conv1 = nn.Conv2d(

bias=True,

self.act_fn = act_fn

self.batchnorm1 = nn.BatchNorm2d(

momentum=1 - momentum,

eps=epsilon,

affine=True,

self.conv2 = nn.Conv2d(

bias=True,

self.batchnorm2 = nn.BatchNorm2d(

momentum=1 - momentum,

eps=epsilon,

affine=True,

self.conv_skip = nn.Conv2d(

bias=True,

self.batchnorm_skip = nn.BatchNorm2d(

momentum=1 - momentum,

eps=epsilon,

affine=True,

orig, W_swap = [0, 1, 2, 3], [-2, -1, -3, -4]

# (f[0], f[1], n_in, n_out) -> (n_out, n_in, f[0], f[1])

W = params["components"]["conv1"]["W"]

b = params["components"]["conv1"]["b"]

W = np.moveaxis(W, orig, W_swap)

assert self.conv1.weight.shape == W.shape

assert self.conv1.bias.shape == b.flatten().shape

self.conv1.weight = nn.Parameter(torch.FloatTensor(W))

self.conv1.bias = nn.Parameter(torch.FloatTensor(b.flatten()))

scaler = params["components"]["batchnorm1"]["scaler"]

intercept = params["components"]["batchnorm1"]["intercept"]

self.batchnorm1.weight = nn.Parameter(torch.FloatTensor(scaler))

self.batchnorm1.bias = nn.Parameter(torch.FloatTensor(intercept))

W = params["components"]["conv2"]["W"]

b = params["components"]["conv2"]["b"]

W = np.moveaxis(W, orig, W_swap)

assert self.conv2.weight.shape == W.shape

assert self.conv2.bias.shape == b.flatten().shape

self.conv2.weight = nn.Parameter(torch.FloatTensor(W))

self.conv2.bias = nn.Parameter(torch.FloatTensor(b.flatten()))

scaler = params["components"]["batchnorm2"]["scaler"]

intercept = params["components"]["batchnorm2"]["intercept"]

self.batchnorm2.weight = nn.Parameter(torch.FloatTensor(scaler))

self.batchnorm2.bias = nn.Parameter(torch.FloatTensor(intercept))

W = params["components"]["conv_skip"]["W"]

b = params["components"]["conv_skip"]["b"]

W = np.moveaxis(W, orig, W_swap)

assert self.conv_skip.weight.shape == W.shape

assert self.conv_skip.bias.shape == b.flatten().shape

self.conv_skip.weight = nn.Parameter(torch.FloatTensor(W))

self.conv_skip.bias = nn.Parameter(torch.FloatTensor(b.flatten()))

scaler = params["components"]["batchnorm_skip"]["scaler"]

intercept = params["components"]["batchnorm_skip"]["intercept"]

self.batchnorm_skip.weight = nn.Parameter(torch.FloatTensor(scaler))

self.batchnorm_skip.bias = nn.Parameter(torch.FloatTensor(intercept))

def forward(self, X):

if not isinstance(X, torch.Tensor):

# (N, H, W, C) -> (N, C, H, W)

X = np.moveaxis(X, [0, 1, 2, 3], [0, -2, -1, -3])

X = torchify(X)

self.X = X

self.conv1_out = self.conv1(self.X)

self.act_fn1_out = self.act_fn(self.conv1_out)

self.batchnorm1_out = self.batchnorm1(self.act_fn1_out)

self.conv2_out = self.conv2(self.batchnorm1_out)

self.batchnorm2_out = self.batchnorm2(self.conv2_out)

self.c_skip_out = self.conv_skip(self.X)

self.bn_skip_out = self.batchnorm_skip(self.c_skip_out)

self.layer3_in = self.batchnorm2_out + self.bn_skip_out

self.Y = self.act_fn(self.layer3_in)

def extract_grads(self, X):

self.forward(X)

self.loss = self.Y.sum()

self.loss.backward()

orig, X_swap, W_swap = [0, 1, 2, 3], [0, -1, -3, -2], [-1, -2, -4, -3]

grads = {

# layer parameters

"conv1_b": self.conv1.bias.detach().numpy().reshape(1, 1, 1, -1),

"bn1_intercept": self.batchnorm1.bias.detach().numpy(),

"bn1_scaler": self.batchnorm1.weight.detach().numpy(),

"bn1_running_mean": self.batchnorm1.running_mean.detach().numpy(),

"bn1_running_var": self.batchnorm1.running_var.detach().numpy(),

"conv2_b": self.conv2.bias.detach().numpy().reshape(1, 1, 1, -1),

"bn2_intercept": self.batchnorm2.bias.detach().numpy(),

"bn2_scaler": self.batchnorm2.weight.detach().numpy(),

"bn2_running_mean": self.batchnorm2.running_mean.detach().numpy(),

"bn2_running_var": self.batchnorm2.running_var.detach().numpy(),

"conv_skip_W": np.moveaxis(

self.conv_skip.weight.detach().numpy(), orig, W_swap

"conv_skip_b": self.conv_skip.bias.detach().numpy().reshape(1, 1, 1, -1),

"bn_skip_intercept": self.batchnorm_skip.bias.detach().numpy(),

"bn_skip_scaler": self.batchnorm_skip.weight.detach().numpy(),

"bn_skip_running_mean": self.batchnorm_skip.running_mean.detach().numpy(),

"bn_skip_running_var": self.batchnorm_skip.running_var.detach().numpy(),

# layer inputs/outputs (forward step)

"X": np.moveaxis(self.X.detach().numpy(), orig, X_swap),

"conv1_out": np.moveaxis(self.conv1_out.detach().numpy(), orig, X_swap),

"conv2_out": np.moveaxis(self.conv2_out.detach().numpy(), orig, X_swap),

"conv_skip_out": np.moveaxis(

self.c_skip_out.detach().numpy(), orig, X_swap

"bn_skip_out": np.moveaxis(self.bn_skip_out.detach().numpy(), orig, X_swap),

"add_out": np.moveaxis(self.layer3_in.detach().numpy(), orig, X_swap),

"Y": np.moveaxis(self.Y.detach().numpy(), orig, X_swap),

# layer gradients (backward step)

"dLdY": np.moveaxis(self.Y.grad.numpy(), orig, X_swap),

"dLdAdd": np.moveaxis(self.layer3_in.grad.numpy(), orig, X_swap),

"dLdBnSkip_out": np.moveaxis(self.bn_skip_out.grad.numpy(), orig, X_swap),

"dLdConvSkip_out": np.moveaxis(self.c_skip_out.grad.numpy(), orig, X_swap),

"dLdConv2_out": np.moveaxis(self.conv2_out.grad.numpy(), orig, X_swap),

"dLdX": np.moveaxis(self.X.grad.numpy(), orig, X_swap),

# layer parameter gradients (backward step)

"dLdBnSkip_intercept": self.batchnorm_skip.bias.grad.numpy(),

"dLdBnSkip_scaler": self.batchnorm_skip.weight.grad.numpy(),

"dLdConvSkip_W": np.moveaxis(

self.conv_skip.weight.grad.numpy(), orig, W_swap

"dLdConvSkip_b": self.conv_skip.bias.grad.numpy().reshape(1, 1, 1, -1),

"dLdBn2_intercept": self.batchnorm2.bias.grad.numpy(),

"dLdBn2_scaler": self.batchnorm2.weight.grad.numpy(),

"dLdConv2_W": np.moveaxis(self.conv2.weight.grad.numpy(), orig, W_swap),

"dLdConv2_b": self.conv2.bias.grad.numpy().reshape(1, 1, 1, -1),

"dLdBn1_intercept": self.batchnorm1.bias.grad.numpy(),

"dLdBn1_scaler": self.batchnorm1.weight.grad.numpy(),

"dLdConv1_W": np.moveaxis(self.conv1.weight.grad.numpy(), orig, W_swap),

"dLdConv1_b": self.conv1.bias.grad.numpy().reshape(1, 1, 1, -1),

return grads

class TorchBidirectionalLSTM(nn.Module):

def __init__(self, n_in, n_out, params, **kwargs):

super(TorchBidirectionalLSTM, self).__init__()

self.layer1 = nn.LSTM(

bidirectional=True,

bias=True,

Wiu = params["components"]["cell_fwd"]["Wu"][n_out:, :].T

Wif = params["components"]["cell_fwd"]["Wf"][n_out:, :].T

Wic = params["components"]["cell_fwd"]["Wc"][n_out:, :].T

Wio = params["components"]["cell_fwd"]["Wo"][n_out:, :].T

W_ih_f = np.vstack([Wiu, Wif, Wic, Wio])

Whu = params["components"]["cell_fwd"]["Wu"][:n_out, :].T

Whf = params["components"]["cell_fwd"]["Wf"][:n_out, :].T

Whc = params["components"]["cell_fwd"]["Wc"][:n_out, :].T

Who = params["components"]["cell_fwd"]["Wo"][:n_out, :].T

W_hh_f = np.vstack([Whu, Whf, Whc, Who])

assert self.layer1.weight_ih_l0.shape == W_ih_f.shape

assert self.layer1.weight_hh_l0.shape == W_hh_f.shape

self.layer1.weight_ih_l0 = nn.Parameter(torch.FloatTensor(W_ih_f))

self.layer1.weight_hh_l0 = nn.Parameter(torch.FloatTensor(W_hh_f))

Wiu = params["components"]["cell_bwd"]["Wu"][n_out:, :].T

Wif = params["components"]["cell_bwd"]["Wf"][n_out:, :].T

Wic = params["components"]["cell_bwd"]["Wc"][n_out:, :].T

Wio = params["components"]["cell_bwd"]["Wo"][n_out:, :].T

W_ih_b = np.vstack([Wiu, Wif, Wic, Wio])

Whu = params["components"]["cell_bwd"]["Wu"][:n_out, :].T

Whf = params["components"]["cell_bwd"]["Wf"][:n_out, :].T

Whc = params["components"]["cell_bwd"]["Wc"][:n_out, :].T

Who = params["components"]["cell_bwd"]["Wo"][:n_out, :].T

W_hh_b = np.vstack([Whu, Whf, Whc, Who])

assert self.layer1.weight_ih_l0_reverse.shape == W_ih_b.shape

assert self.layer1.weight_hh_l0_reverse.shape == W_hh_b.shape

self.layer1.weight_ih_l0_reverse = nn.Parameter(torch.FloatTensor(W_ih_b))

self.layer1.weight_hh_l0_reverse = nn.Parameter(torch.FloatTensor(W_hh_b))

b_f = np.concatenate(

).flatten()

assert self.layer1.bias_ih_l0.shape == b_f.shape

assert self.layer1.bias_hh_l0.shape == b_f.shape

self.layer1.bias_ih_l0 = nn.Parameter(torch.FloatTensor(b_f))

self.layer1.bias_hh_l0 = nn.Parameter(torch.FloatTensor(b_f))

b_b = np.concatenate(

).flatten()

assert self.layer1.bias_ih_l0_reverse.shape == b_b.shape

assert self.layer1.bias_hh_l0_reverse.shape == b_b.shape

self.layer1.bias_ih_l0_reverse = nn.Parameter(torch.FloatTensor(b_b))

self.layer1.bias_hh_l0_reverse = nn.Parameter(torch.FloatTensor(b_b))

def forward(self, X):

# (batch, input_size, seq_len) -> (seq_len, batch, input_size)

self.X = np.moveaxis(X, [0, 1, 2], [-2, -1, -3])

if not isinstance(self.X, torch.Tensor):

self.X = torchify(self.X)

# initial hidden state is 0

n_ex, n_in, n_timesteps = self.X.shape

n_out, n_out = self.layer1.weight_hh_l0.shape

# forward pass

self.A, (At, Ct) = self.layer1(self.X)

return self.A

def extract_grads(self, X):

self.forward(X)

self.loss = self.A.sum()

self.loss.backward()

# forward

w_ii, w_if, w_ic, w_io = self.layer1.weight_ih_l0.chunk(4, 0)

w_hi, w_hf, w_hc, w_ho = self.layer1.weight_hh_l0.chunk(4, 0)

bu_f, bf_f, bc_f, bo_f = self.layer1.bias_ih_l0.chunk(4, 0)

Wu_f = torch.cat([torch.t(w_hi), torch.t(w_ii)], dim=0)

Wf_f = torch.cat([torch.t(w_hf), torch.t(w_if)], dim=0)

Wc_f = torch.cat([torch.t(w_hc), torch.t(w_ic)], dim=0)

Wo_f = torch.cat([torch.t(w_ho), torch.t(w_io)], dim=0)

dw_ii, dw_if, dw_ic, dw_io = self.layer1.weight_ih_l0.grad.chunk(4, 0)

dw_hi, dw_hf, dw_hc, dw_ho = self.layer1.weight_hh_l0.grad.chunk(4, 0)

dbu_f, dbf_f, dbc_f, dbo_f = self.layer1.bias_ih_l0.grad.chunk(4, 0)

dWu_f = torch.cat([torch.t(dw_hi), torch.t(dw_ii)], dim=0)

dWf_f = torch.cat([torch.t(dw_hf), torch.t(dw_if)], dim=0)

dWc_f = torch.cat([torch.t(dw_hc), torch.t(dw_ic)], dim=0)

dWo_f = torch.cat([torch.t(dw_ho), torch.t(dw_io)], dim=0)

# backward

w_ii, w_if, w_ic, w_io = self.layer1.weight_ih_l0_reverse.chunk(4, 0)

w_hi, w_hf, w_hc, w_ho = self.layer1.weight_hh_l0_reverse.chunk(4, 0)

bu_b, bf_b, bc_b, bo_b = self.layer1.bias_ih_l0_reverse.chunk(4, 0)

Wu_b = torch.cat([torch.t(w_hi), torch.t(w_ii)], dim=0)

Wf_b = torch.cat([torch.t(w_hf), torch.t(w_if)], dim=0)

Wc_b = torch.cat([torch.t(w_hc), torch.t(w_ic)], dim=0)

Wo_b = torch.cat([torch.t(w_ho), torch.t(w_io)], dim=0)

dw_ii, dw_if, dw_ic, dw_io = self.layer1.weight_ih_l0_reverse.grad.chunk(4, 0)

dw_hi, dw_hf, dw_hc, dw_ho = self.layer1.weight_hh_l0_reverse.grad.chunk(4, 0)

dbu_b, dbf_b, dbc_b, dbo_b = self.layer1.bias_ih_l0_reverse.grad.chunk(4, 0)

dWu_b = torch.cat([torch.t(dw_hi), torch.t(dw_ii)], dim=0)

dWf_b = torch.cat([torch.t(dw_hf), torch.t(dw_if)], dim=0)

dWc_b = torch.cat([torch.t(dw_hc), torch.t(dw_ic)], dim=0)

dWo_b = torch.cat([torch.t(dw_ho), torch.t(dw_io)], dim=0)

orig, X_swap = [0, 1, 2], [-1, -3, -2]

grads = {

"X": np.moveaxis(self.X.detach().numpy(), orig, X_swap),

"Wu_f": Wu_f.detach().numpy(),

"Wf_f": Wf_f.detach().numpy(),

"Wc_f": Wc_f.detach().numpy(),

"Wo_f": Wo_f.detach().numpy(),

"bu_f": bu_f.detach().numpy().reshape(-1, 1),

"bf_f": bf_f.detach().numpy().reshape(-1, 1),

"bc_f": bc_f.detach().numpy().reshape(-1, 1),

"bo_f": bo_f.detach().numpy().reshape(-1, 1),

"Wu_b": Wu_b.detach().numpy(),

"Wf_b": Wf_b.detach().numpy(),

"Wc_b": Wc_b.detach().numpy(),

"Wo_b": Wo_b.detach().numpy(),

"bu_b": bu_b.detach().numpy().reshape(-1, 1),

"bf_b": bf_b.detach().numpy().reshape(-1, 1),

"bc_b": bc_b.detach().numpy().reshape(-1, 1),

"bo_b": bo_b.detach().numpy().reshape(-1, 1),

"y": np.moveaxis(self.A.detach().numpy(), orig, X_swap),

"dLdA": self.A.grad.numpy(),

"dLdWu_f": dWu_f.numpy(),

"dLdWf_f": dWf_f.numpy(),

"dLdWc_f": dWc_f.numpy(),

"dLdWo_f": dWo_f.numpy(),

"dLdBu_f": dbu_f.numpy().reshape(-1, 1),

"dLdBf_f": dbf_f.numpy().reshape(-1, 1),

"dLdBc_f": dbc_f.numpy().reshape(-1, 1),

"dLdBo_f": dbo_f.numpy().reshape(-1, 1),

"dLdWu_b": dWu_b.numpy(),

"dLdWf_b": dWf_b.numpy(),

"dLdWc_b": dWc_b.numpy(),

"dLdWo_b": dWo_b.numpy(),

"dLdBu_b": dbu_b.numpy().reshape(-1, 1),

"dLdBf_b": dbf_b.numpy().reshape(-1, 1),

"dLdBc_b": dbc_b.numpy().reshape(-1, 1),

"dLdBo_b": dbo_b.numpy().reshape(-1, 1),

"dLdX": np.moveaxis(self.X.grad.numpy(), orig, X_swap),

return grads

class TorchPool2DLayer(nn.Module):

def __init__(self, in_channels, hparams, **kwargs):

super(TorchPool2DLayer, self).__init__()

if hparams["mode"] == "max":

self.layer1 = nn.MaxPool2d(

padding=hparams["pad"],

stride=hparams["stride"],

elif hparams["mode"] == "average":

self.layer1 = nn.AvgPool2d(

padding=hparams["pad"],

stride=hparams["stride"],

def forward(self, X):

# (N, H, W, C) -> (N, C, H, W)

self.X = np.moveaxis(X, [0, 1, 2, 3], [0, -2, -1, -3])

if not isinstance(self.X, torch.Tensor):

self.X = torchify(self.X)

self.Y = self.layer1(self.X)

return self.Y

def extract_grads(self, X):

self.forward(X)

self.loss = self.Y.sum()

self.loss.backward()

# W (theirs): (n_out, n_in, f[0], f[1]) -> W (mine): (f[0], f[1], n_in, n_out)

# X (theirs): (N, C, H, W)              -> X (mine): (N, H, W, C)

# Y (theirs): (N, C, H, W)              -> Y (mine): (N, H, W, C)

orig, X_swap = [0, 1, 2, 3], [0, -1, -3, -2]

grads = {

"X": np.moveaxis(self.X.detach().numpy(), orig, X_swap),

"y": np.moveaxis(self.Y.detach().numpy(), orig, X_swap),

"dLdY": np.moveaxis(self.Y.grad.numpy(), orig, X_swap),

return grads

class TorchConv2DLayer(nn.Module):

def __init__(self, in_channels, out_channels, act_fn, params, hparams, **kwargs):

super(TorchConv2DLayer, self).__init__()

W = params["W"]

b = params["b"]

self.act_fn = act_fn

self.layer1 = nn.Conv2d(

padding=hparams["pad"],

stride=hparams["stride"],

dilation=hparams["dilation"] + 1,

bias=True,

# (f[0], f[1], n_in, n_out) -> (n_out, n_in, f[0], f[1])

W = np.moveaxis(W, [0, 1, 2, 3], [-2, -1, -3, -4])

assert self.layer1.weight.shape == W.shape

assert self.layer1.bias.shape == b.flatten().shape

self.layer1.weight = nn.Parameter(torch.FloatTensor(W))

self.layer1.bias = nn.Parameter(torch.FloatTensor(b.flatten()))

def forward(self, X):

# (N, H, W, C) -> (N, C, H, W)

self.X = np.moveaxis(X, [0, 1, 2, 3], [0, -2, -1, -3])

if not isinstance(self.X, torch.Tensor):

self.X = torchify(self.X)

self.Z = self.layer1(self.X)

self.Y = self.act_fn(self.Z)

return self.Y

def extract_grads(self, X):

self.forward(X)

self.loss = self.Y.sum()

self.loss.backward()

# X (theirs): (N, C, H, W)              -> X (mine): (N, H, W, C)

# Y (theirs): (N, C, H, W)              -> Y (mine): (N, H, W, C)

orig, X_swap, W_swap = [0, 1, 2, 3], [0, -1, -3, -2], [-1, -2, -4, -3]

grads = {

"X": np.moveaxis(self.X.detach().numpy(), orig, X_swap),

"W": np.moveaxis(self.layer1.weight.detach().numpy(), orig, W_swap),

"b": self.layer1.bias.detach().numpy().reshape(1, 1, 1, -1),

"y": np.moveaxis(self.Y.detach().numpy(), orig, X_swap),

"dLdY": np.moveaxis(self.Y.grad.numpy(), orig, X_swap),

"dLdZ": np.moveaxis(self.Z.grad.numpy(), orig, X_swap),

"dLdW": np.moveaxis(self.layer1.weight.grad.numpy(), orig, W_swap),

"dLdB": self.layer1.bias.grad.numpy().reshape(1, 1, 1, -1),

"dLdX": np.moveaxis(self.X.grad.numpy(), orig, X_swap),

return grads

class TorchConv1DLayer(nn.Module):

super(TorchConv1DLayer, self).__init__()

W = params["W"]

b = params["b"]

self.act_fn = act_fn

self.layer1 = nn.Conv1d(

padding=hparams["pad"],

stride=hparams["stride"],

dilation=hparams["dilation"] + 1,

bias=True,

# (f[0], n_in, n_out) -> (n_out, n_in, f[0])

W = np.moveaxis(W, [0, 1, 2], [-1, -2, -3])

assert self.layer1.weight.shape == W.shape

assert self.layer1.bias.shape == b.flatten().shape

self.layer1.weight = nn.Parameter(torch.FloatTensor(W))

self.layer1.bias = nn.Parameter(torch.FloatTensor(b.flatten()))

def forward(self, X):

# (N, W, C) -> (N, C, W)

self.X = np.moveaxis(X, [0, 1, 2], [0, -1, -2])

if not isinstance(self.X, torch.Tensor):

self.X = torchify(self.X)

self.Z = self.layer1(self.X)

self.Y = self.act_fn(self.Z)

return self.Y

def extract_grads(self, X):

self.forward(X)

self.loss = self.Y.sum()

self.loss.backward()

# W (theirs): (n_out, n_in, f[0]) -> W (mine): (f[0], n_in, n_out)

# X (theirs): (N, C, W)              -> X (mine): (N, W, C)

# Y (theirs): (N, C, W)              -> Y (mine): (N, W, C)

orig, X_swap, W_swap = [0, 1, 2], [0, -1, -2], [-1, -2, -3]

grads = {

"X": np.moveaxis(self.X.detach().numpy(), orig, X_swap),

"W": np.moveaxis(self.layer1.weight.detach().numpy(), orig, W_swap),

"b": self.layer1.bias.detach().numpy().reshape(1, 1, -1),

"y": np.moveaxis(self.Y.detach().numpy(), orig, X_swap),

"dLdY": np.moveaxis(self.Y.grad.numpy(), orig, X_swap),

"dLdZ": np.moveaxis(self.Z.grad.numpy(), orig, X_swap),

"dLdW": np.moveaxis(self.layer1.weight.grad.numpy(), orig, W_swap),

"dLdB": self.layer1.bias.grad.numpy().reshape(1, 1, -1),

"dLdX": np.moveaxis(self.X.grad.numpy(), orig, X_swap),

return grads

class TorchDeconv2DLayer(nn.Module):

super(TorchDeconv2DLayer, self).__init__()

W = params["W"]

b = params["b"]

self.act_fn = act_fn

self.layer1 = nn.ConvTranspose2d(

padding=hparams["pad"],

stride=hparams["stride"],

bias=True,

# (f[0], f[1], n_in, n_out) -> (n_in, n_out, f[0], f[1])

W = np.moveaxis(W, [0, 1, 2, 3], [-2, -1, -4, -3])

assert self.layer1.weight.shape == W.shape

assert self.layer1.bias.shape == b.flatten().shape

self.layer1.weight = nn.Parameter(torch.FloatTensor(W))

self.layer1.bias = nn.Parameter(torch.FloatTensor(b.flatten()))

def forward(self, X):

# (N, H, W, C) -> (N, C, H, W)

self.X = np.moveaxis(X, [0, 1, 2, 3], [0, -2, -1, -3])

if not isinstance(self.X, torch.Tensor):

self.X = torchify(self.X)

self.Z = self.layer1(self.X)

self.Y = self.act_fn(self.Z)

return self.Y

def extract_grads(self, X):

self.forward(X)

self.loss = self.Y.sum()

self.loss.backward()

# W (theirs): (n_in, n_out, f[0], f[1]) -> W (mine): (f[0], f[1], n_in, n_out)

# X (theirs): (N, C, H, W)              -> X (mine): (N, H, W, C)

# Y (theirs): (N, C, H, W)              -> Y (mine): (N, H, W, C)

orig, X_swap, W_swap = [0, 1, 2, 3], [0, -1, -3, -2], [-2, -1, -4, -3]

grads = {

"X": np.moveaxis(self.X.detach().numpy(), orig, X_swap),

"W": np.moveaxis(self.layer1.weight.detach().numpy(), orig, W_swap),

"b": self.layer1.bias.detach().numpy().reshape(1, 1, 1, -1),

"y": np.moveaxis(self.Y.detach().numpy(), orig, X_swap),

"dLdY": np.moveaxis(self.Y.grad.numpy(), orig, X_swap),

"dLdZ": np.moveaxis(self.Z.grad.numpy(), orig, X_swap),

"dLdW": np.moveaxis(self.layer1.weight.grad.numpy(), orig, W_swap),

"dLdB": self.layer1.bias.grad.numpy().reshape(1, 1, 1, -1),

"dLdX": np.moveaxis(self.X.grad.numpy(), orig, X_swap),

return grads

class TorchLSTMCell(nn.Module):

def __init__(self, n_in, n_out, params, **kwargs):

super(TorchLSTMCell, self).__init__()

Wiu = params["Wu"][n_out:, :].T

Wif = params["Wf"][n_out:, :].T

Wic = params["Wc"][n_out:, :].T

Wio = params["Wo"][n_out:, :].T

W_ih = np.vstack([Wiu, Wif, Wic, Wio])

Whu = params["Wu"][:n_out, :].T

Whf = params["Wf"][:n_out, :].T

Whc = params["Wc"][:n_out, :].T

Who = params["Wo"][:n_out, :].T

W_hh = np.vstack([Whu, Whf, Whc, Who])

self.layer1 = nn.LSTMCell(input_size=n_in, hidden_size=n_out, bias=True)

assert self.layer1.weight_ih.shape == W_ih.shape

assert self.layer1.weight_hh.shape == W_hh.shape

self.layer1.weight_ih = nn.Parameter(torch.FloatTensor(W_ih))

self.layer1.weight_hh = nn.Parameter(torch.FloatTensor(W_hh))

b = np.concatenate(

[params["bu"], params["bf"], params["bc"], params["bo"]], axis=-1

).flatten()

assert self.layer1.bias_ih.shape == b.shape

assert self.layer1.bias_hh.shape == b.shape

self.layer1.bias_ih = nn.Parameter(torch.FloatTensor(b))

self.layer1.bias_hh = nn.Parameter(torch.FloatTensor(b))

def forward(self, X):

self.X = X

if not isinstance(self.X, torch.Tensor):

self.X = torchify(self.X)

# initial hidden state is 0

n_ex, n_in, n_timesteps = self.X.shape

n_out, n_out = self.layer1.weight_hh.shape

# initialize hidden states

a0 = torchify(np.zeros((n_ex, n_out)))

c0 = torchify(np.zeros((n_ex, n_out)))

# forward pass

A, C = [], []

at = a0

ct = c0

for t in range(n_timesteps):

A.append(at)

C.append(ct)

at1, ct1 = self.layer1(self.X[:, :, t], (at, ct))

at = at1

ct = ct1

A.append(at)

C.append(ct)

# don't inclue a0 in our outputs

self.A = A[1:]

self.C = C[1:]

return self.A, self.C

def extract_grads(self, X):

self.forward(X)

self.loss = torch.stack(self.A).sum()

self.loss.backward()

w_ii, w_if, w_ic, w_io = self.layer1.weight_ih.chunk(4, 0)

w_hi, w_hf, w_hc, w_ho = self.layer1.weight_hh.chunk(4, 0)

bu, bf, bc, bo = self.layer1.bias_ih.chunk(4, 0)

Wu = torch.cat([torch.t(w_hi), torch.t(w_ii)], dim=0)

Wf = torch.cat([torch.t(w_hf), torch.t(w_if)], dim=0)

Wc = torch.cat([torch.t(w_hc), torch.t(w_ic)], dim=0)

Wo = torch.cat([torch.t(w_ho), torch.t(w_io)], dim=0)

dw_ii, dw_if, dw_ic, dw_io = self.layer1.weight_ih.grad.chunk(4, 0)

dw_hi, dw_hf, dw_hc, dw_ho = self.layer1.weight_hh.grad.chunk(4, 0)

dbu, dbf, dbc, dbo = self.layer1.bias_ih.grad.chunk(4, 0)

dWu = torch.cat([torch.t(dw_hi), torch.t(dw_ii)], dim=0)

dWf = torch.cat([torch.t(dw_hf), torch.t(dw_if)], dim=0)

dWc = torch.cat([torch.t(dw_hc), torch.t(dw_ic)], dim=0)

dWo = torch.cat([torch.t(dw_ho), torch.t(dw_io)], dim=0)

grads = {

"X": self.X.detach().numpy(),

"Wu": Wu.detach().numpy(),

"Wf": Wf.detach().numpy(),

"Wc": Wc.detach().numpy(),

"Wo": Wo.detach().numpy(),

"bu": bu.detach().numpy().reshape(-1, 1),

"bf": bf.detach().numpy().reshape(-1, 1),

"bc": bc.detach().numpy().reshape(-1, 1),

"bo": bo.detach().numpy().reshape(-1, 1),

"C": torch.stack(self.C).detach().numpy(),

"y": np.swapaxes(

np.swapaxes(torch.stack(self.A).detach().numpy(), 1, 0), 1, 2

"dLdA": np.array([a.grad.numpy() for a in self.A]),

"dLdWu": dWu.numpy(),

"dLdWf": dWf.numpy(),

"dLdWc": dWc.numpy(),

"dLdWo": dWo.numpy(),

"dLdBu": dbu.numpy().reshape(-1, 1),

"dLdBf": dbf.numpy().reshape(-1, 1),

"dLdBc": dbc.numpy().reshape(-1, 1),

"dLdBo": dbo.numpy().reshape(-1, 1),

"dLdX": self.X.grad.numpy(),

return grads

class TorchRNNCell(nn.Module):

def __init__(self, n_in, n_hid, params, **kwargs):

super(TorchRNNCell, self).__init__()

self.layer1 = nn.RNNCell(n_in, n_hid, bias=True, nonlinearity="tanh")

# set weights and bias to match those of RNNCell

# NB: we pass the *transpose* of the RNNCell weights and biases to

# pytorch, meaning we need to check against the *transpose* of our

# outputs for any function of the weights

self.layer1.weight_ih = nn.Parameter(torch.FloatTensor(params["Wax"].T))

self.layer1.weight_hh = nn.Parameter(torch.FloatTensor(params["Waa"].T))

self.layer1.bias_ih = nn.Parameter(torch.FloatTensor(params["bx"].T))

self.layer1.bias_hh = nn.Parameter(torch.FloatTensor(params["ba"].T))

def forward(self, X):

self.X = X

if not isinstance(self.X, torch.Tensor):

self.X = torchify(self.X)

# initial hidden state is 0

n_ex, n_in, n_timesteps = self.X.shape

n_out, n_out = self.layer1.weight_hh.shape

# initialize hidden states

a0 = torchify(np.zeros((n_ex, n_out)))

# forward pass

A = []

at = a0

for t in range(n_timesteps):

A += [at]

at1 = self.layer1(self.X[:, :, t], at)

at = at1

A += [at]

# don't inclue a0 in our outputs

self.A = A[1:]

return self.A

def extract_grads(self, X):

self.forward(X)

self.loss = torch.stack(self.A).sum()

self.loss.backward()

grads = {

"X": self.X.detach().numpy(),

"ba": self.layer1.bias_hh.detach().numpy(),

"bx": self.layer1.bias_ih.detach().numpy(),

"Wax": self.layer1.weight_ih.detach().numpy(),

"Waa": self.layer1.weight_hh.detach().numpy(),

"y": torch.stack(self.A).detach().numpy(),

"dLdA": np.array([a.grad.numpy() for a in self.A]),

"dLdWaa": self.layer1.weight_hh.grad.numpy(),

"dLdWax": self.layer1.weight_ih.grad.numpy(),

"dLdBa": self.layer1.bias_hh.grad.numpy(),

"dLdBx": self.layer1.bias_ih.grad.numpy(),

"dLdX": self.X.grad.numpy(),

return grads

class TorchFCLayer(nn.Module):

def __init__(self, n_in, n_hid, act_fn, params, **kwargs):

super(TorchFCLayer, self).__init__()

self.layer1 = nn.Linear(n_in, n_hid)

# explicitly set weights and bias

# NB: we pass the *transpose* of the weights to pytorch, meaning

# we'll need to check against the *transpose* of our outputs for

# any function of the weights

self.layer1.weight = nn.Parameter(torch.FloatTensor(params["W"].T))

self.layer1.bias = nn.Parameter(torch.FloatTensor(params["b"]))

self.act_fn = act_fn

self.model = nn.Sequential(self.layer1, self.act_fn)

def forward(self, X):

self.X = X

if not isinstance(X, torch.Tensor):

self.X = torchify(X)

self.z1 = self.layer1(self.X)

self.out1 = self.act_fn(self.z1)

def extract_grads(self, X):

self.forward(X)

self.loss1 = self.out1.sum()

grads = {

"X": self.X.detach().numpy(),

"b": self.layer1.bias.detach().numpy(),

"W": self.layer1.weight.detach().numpy(),

"y": self.out1.detach().numpy(),

"dLdy": self.out1.grad.numpy(),

"dLdZ": self.z1.grad.numpy(),

"dLdB": self.layer1.bias.grad.numpy(),

"dLdW": self.layer1.weight.grad.numpy(),

"dLdX": self.X.grad.numpy(),

return grads

class TorchEmbeddingLayer(nn.Module):

def __init__(self, vocab_size, n_out, params, **kwargs):

super(TorchEmbeddingLayer, self).__init__()

self.layer1 = nn.Embedding(vocab_size, n_out)

# explicitly set embedding weights

self.layer1.weight = nn.Parameter(torch.FloatTensor(params["W"]))

self.model = nn.Sequential(self.layer1)

def forward(self, X):

self.X = X

if not isinstance(X, torch.Tensor):

self.X = torch.from_numpy(X)

self.out1 = self.layer1(self.X)

def extract_grads(self, X):

self.forward(X)

self.loss1 = self.out1.sum()

grads = {

"X": self.X.detach().numpy(),

"y": self.out1.detach().numpy(),

"dLdy": self.out1.grad.numpy(),

return grads

class TorchSDPAttentionLayer(nn.Module):

def __init__(self):

super(TorchSDPAttentionLayer, self).__init__()

def forward(self, Q, K, V, mask=None):

self.Q = Q

self.K = K

self.V = V

if not isinstance(self.Q, torch.Tensor):

self.Q = torchify(self.Q)

if not isinstance(self.K, torch.Tensor):

self.K = torchify(self.K)

if not isinstance(self.V, torch.Tensor):

self.V = torchify(self.V)

self.d_k = self.Q.size(-1)

self.scores = torch.matmul(self.Q, self.K.transpose(-2, -1)) / np.sqrt(self.d_k)

if mask is not None:

self.scores = self.scores.masked_fill(mask == 0, -1e9)

self.weights = F.softmax(self.scores, dim=-1)

self.Y = torch.matmul(self.weights, self.V)

return self.Y, self.weights

def extract_grads(self, Q, K, V, mask=None):

self.forward(Q, K, V, mask=mask)

self.loss1 = self.Y.sum()

grads = {

"Q": self.Q.detach().numpy(),

"K": self.K.detach().numpy(),

"V": self.V.detach().numpy(),

"d_k": self.d_k,

"scores": self.scores.detach().numpy(),

"weights": self.weights.detach().numpy(),

"Y": self.Y.detach().numpy(),

"dLdV": self.V.grad.numpy(),

"dWeights": self.weights.grad.numpy(),

"dScores": self.scores.grad.numpy(),

"dLdQ": self.Q.grad.numpy(),

"dLdK": self.K.grad.numpy(),

return grads

class TorchMultiHeadedAttentionModule(nn.Module):

def __init__(self, params, hparams):

"Take in model size and number of heads."

super(TorchMultiHeadedAttentionModule, self).__init__()

assert hparams["kqv_dim"] % hparams["n_heads"] == 0

self.n_heads = hparams["n_heads"]

self.latent_dim = hparams["kqv_dim"] // hparams["n_heads"]

self.p_dropout = hparams["dropout_p"]

self.projections = {

"Q": nn.Linear(hparams["kqv_dim"], hparams["kqv_dim"]),

"K": nn.Linear(hparams["kqv_dim"], hparams["kqv_dim"]),

"V": nn.Linear(hparams["kqv_dim"], hparams["kqv_dim"]),

"O": nn.Linear(hparams["kqv_dim"], hparams["kqv_dim"]),

self.projections["Q"].weight = nn.Parameter(

torch.FloatTensor(params["components"]["Q"]["W"].T)

self.projections["Q"].bias = nn.Parameter(

torch.FloatTensor(params["components"]["Q"]["b"])

self.projections["K"].weight = nn.Parameter(

torch.FloatTensor(params["components"]["K"]["W"].T)

self.projections["K"].bias = nn.Parameter(

torch.FloatTensor(params["components"]["K"]["b"])

self.projections["V"].weight = nn.Parameter(

torch.FloatTensor(params["components"]["V"]["W"].T)

self.projections["V"].bias = nn.Parameter(

torch.FloatTensor(params["components"]["V"]["b"])

self.projections["O"].weight = nn.Parameter(

torch.FloatTensor(params["components"]["O"]["W"].T)

self.projections["O"].bias = nn.Parameter(

torch.FloatTensor(params["components"]["O"]["b"])

self.attn = None

self.dropout = nn.Dropout(p=hparams["dropout_p"])

def forward(self, Q, K, V, mask=None):

self.Q = Q

self.K = K

self.V = V

if not isinstance(self.Q, torch.Tensor):

self.Q = torchify(self.Q)

if not isinstance(self.K, torch.Tensor):

self.K = torchify(self.K)

if not isinstance(self.V, torch.Tensor):

self.V = torchify(self.V)

if mask is not None:

# Same mask applied to all h heads.

mask = mask.unsqueeze(1)

n_ex = self.Q.size(0)

self.Q_proj = (

self.projections["Q"](self.Q)

.view(n_ex, -1, self.n_heads, self.latent_dim)

.transpose(1, 2)

self.K_proj = (

self.projections["K"](self.K)

.transpose(1, 2)

self.V_proj = (

self.projections["V"](self.V)

.transpose(1, 2)

# 2) Apply attention on all the projected vectors in batch.

self.attn_out, self.attn = TorchSDPAttentionLayer().forward(

self.Q_proj, self.K_proj, self.V_proj, mask=mask

# 3) "Concat" using a view and apply a final linear transformation

self.attn_out_reshaped = (

self.attn_out.transpose(1, 2)

.contiguous()

.view(n_ex, -1, self.n_heads * self.latent_dim)

self.Y = self.projections["O"](self.attn_out_reshaped)

print(self.Y.shape)

def extract_grads(self, Q, K, V, mask=None):

self.forward(Q, K, V, mask=mask)

self.loss1 = self.Y.sum()

grads = {

"Q": self.Q.detach().numpy(),

"K": self.K.detach().numpy(),

"V": self.V.detach().numpy(),

"O_W": self.projections["O"].weight.detach().numpy().T,

"V_W": self.projections["V"].weight.detach().numpy().T,

"K_W": self.projections["K"].weight.detach().numpy().T,

"Q_W": self.projections["Q"].weight.detach().numpy().T,

"O_b": self.projections["O"].bias.detach().numpy(),

"V_b": self.projections["V"].bias.detach().numpy(),

"K_b": self.projections["K"].bias.detach().numpy(),

"Q_b": self.projections["Q"].bias.detach().numpy(),

"latent_dim": self.latent_dim,

"n_heads": self.n_heads,

"Q_proj": self.Q_proj.detach().numpy(),  # .reshape(self.Q_proj.shape[0], -1),

"K_proj": self.K_proj.detach().numpy(),  # .reshape(self.K_proj.shape[0], -1),

"V_proj": self.V_proj.detach().numpy(),  # .reshape(self.V_proj.shape[0], -1),

"weights": self.attn.detach().numpy(),

"attn_out": self.attn_out_reshaped.detach().numpy(),  # .squeeze(),

#  .reshape(self.attn_out_reshaped.shape[0], -1),

"Y": self.Y.detach().numpy(),

"dO_W": self.projections["O"].weight.grad.numpy().T,

"dV_W": self.projections["V"].weight.grad.numpy().T,

"dK_W": self.projections["K"].weight.grad.numpy().T,

"dQ_W": self.projections["Q"].weight.grad.numpy().T,

"dO_b": self.projections["O"].bias.grad.numpy(),

"dV_b": self.projections["V"].bias.grad.numpy(),

"dK_b": self.projections["K"].bias.grad.numpy(),

"dQ_b": self.projections["Q"].bias.grad.numpy(),

"dLdy": self.Y.grad.numpy(),

"dAttn_out": self.attn_out_reshaped.grad.numpy(),

"dWeights": self.attn.grad.numpy(),

"dQ_proj": self.Q_proj.grad.numpy(),

"dK_proj": self.K_proj.grad.numpy(),

"dV_proj": self.V_proj.grad.numpy(),

"dQ": self.Q.grad.numpy(),

"dK": self.K.grad.numpy(),

"dV": self.V.grad.numpy(),

return grads

#              TF WGAN GP Gold Standard Implementation                #

#  adapted from: https://github.com/igul222/improved_wgan_training/   #

_params = {}

_param_aliases = {}

def param(name, *args, **kwargs):

A wrapper for `tf.Variable` which enables parameter sharing in models.

Creates and returns theano shared variables similarly to `tf.Variable`,

except if you try to create a param with the same name as a

previously-created one, `param(...)` will just return the old one instead of

making a new one.

This constructor also adds a `param` attribute to the shared variables it

creates, so that you can easily search a graph for all params.

if name not in _params:

kwargs["name"] = name

param = tf.Variable(*args, **kwargs)

param.param = True

_params[name] = param

result = _params[name]

i = 0

while result in _param_aliases:

i += 1

result = _param_aliases[result]

return result

def params_with_name(name):

return [p for n, p in _params.items() if name in n]

def ReLULayer(name, n_in, n_out, inputs, w_initialization):

if isinstance(w_initialization, np.ndarray):

weight_values = w_initialization.astype("float32")

W = param(name + ".W", weight_values)

result = tf.matmul(inputs, W)

output = tf.nn.bias_add(

result, param(name + ".b", np.zeros((n_out,), dtype="float32"))

output = tf.nn.relu(output)

return output, W

def LinearLayer(name, n_in, n_out, inputs, w_initialization):

W = param(name + ".W", weight_values)

result = tf.matmul(inputs, W)

output = tf.nn.bias_add(

return output, W

def Generator(n_samples, X_real, params=None):

n_feats = 2

W1 = W2 = W3 = W4 = "he"

noise = tf.random.normal([n_samples, 2])

if params is not None:

noise = tf.convert_to_tensor(params["noise"], dtype="float32")

W1 = params["generator"]["FC1"]["W"]

W2 = params["generator"]["FC2"]["W"]

W3 = params["generator"]["FC3"]["W"]

W4 = params["generator"]["FC4"]["W"]

DIM = params["g_hidden"]

n_feats = params["n_in"]

outs = {}

weights = {}

output, W = ReLULayer("Generator.1", n_feats, DIM, noise, w_initialization=W1)

outs["FC1"] = output

weights["FC1"] = W

output, W = ReLULayer("Generator.2", DIM, DIM, output, w_initialization=W2)

outs["FC2"] = output

weights["FC2"] = W

output, W = ReLULayer("Generator.3", DIM, DIM, output, w_initialization=W3)

outs["FC3"] = output

weights["FC3"] = W

output, W = LinearLayer("Generator.4", DIM, n_feats, output, w_initialization=W4)

outs["FC4"] = output

weights["FC4"] = W

return output, outs, weights

def Discriminator(inputs, params=None):

n_feats = 2

W1 = W2 = W3 = W4 = "he"

if params is not None:

W1 = params["critic"]["FC1"]["W"]

W2 = params["critic"]["FC2"]["W"]

W3 = params["critic"]["FC3"]["W"]

W4 = params["critic"]["FC4"]["W"]

DIM = params["g_hidden"]

n_feats = params["n_in"]

outs = {}

weights = {}

output, W = ReLULayer("Discriminator.1", n_feats, DIM, inputs, w_initialization=W1)

outs["FC1"] = output

weights["FC1"] = W

output, W = ReLULayer("Discriminator.2", DIM, DIM, output, w_initialization=W2)

outs["FC2"] = output

weights["FC2"] = W

output, W = ReLULayer("Discriminator.3", DIM, DIM, output, w_initialization=W3)

outs["FC3"] = output

weights["FC3"] = W

output, W = LinearLayer("Discriminator.4", DIM, 1, output, w_initialization=W4)

outs["FC4"] = output

weights["FC4"] = W

# get bias

for var in params_with_name("Discriminator"):

if "1.b:" in var.name:

weights["FC1_b"] = var

elif "2.b:" in var.name:

weights["FC2_b"] = var

elif "3.b:" in var.name:

weights["FC3_b"] = var

elif "4.b:" in var.name:

weights["FC4_b"] = var

return tf.reshape(output, [-1]), outs, weights

def WGAN_GP_tf(X, lambda_, params, batch_size):

batch_size = X.shape[0]

# get alpha value

n_steps = params["n_steps"]

c_updates_per_epoch = params["c_updates_per_epoch"]

alpha = tf.convert_to_tensor(params["alpha"], dtype="float32")

X_real = tf.compat.v1.placeholder(tf.float32, shape=[None, params["n_in"]])

X_fake, G_out_X_fake, G_weights = Generator(batch_size, X_real, params)

Y_real, C_out_Y_real, C_Y_real_weights = Discriminator(X_real, params)

Y_fake, C_out_Y_fake, C_Y_fake_weights = Discriminator(X_fake, params)

# WGAN loss

mean_fake = tf.reduce_mean(Y_fake)

mean_real = tf.reduce_mean(Y_real)

C_loss = tf.reduce_mean(Y_fake) - tf.reduce_mean(Y_real)

G_loss = -tf.reduce_mean(Y_fake)

# WGAN gradient penalty

X_interp = alpha * X_real + ((1 - alpha) * X_fake)

Y_interp, C_out_Y_interp, C_Y_interp_weights = Discriminator(X_interp, params)

gradInterp = tf.gradients(Y_interp, [X_interp])[0]

norm_gradInterp = tf.sqrt(

tf.compat.v1.reduce_sum(tf.square(gradInterp), reduction_indices=[1])

gradient_penalty = tf.reduce_mean((norm_gradInterp - 1) ** 2)

C_loss += lambda_ * gradient_penalty

# extract gradient of Y_interp wrt. each layer output in critic

C_bwd_Y_interp = {}

for k, v in C_out_Y_interp.items():

C_bwd_Y_interp[k] = tf.gradients(Y_interp, [v])[0]

C_bwd_W = {}

for k, v in C_Y_interp_weights.items():

C_bwd_W[k] = tf.gradients(C_loss, [v])[0]

# get gradients

dC_Y_fake = tf.gradients(C_loss, [Y_fake])[0]

dC_Y_real = tf.gradients(C_loss, [Y_real])[0]

dC_gradInterp = tf.gradients(C_loss, [gradInterp])[0]

dG_Y_fake = tf.gradients(G_loss, [Y_fake])[0]

with tf.compat.v1.Session() as session:

session.run(tf.compat.v1.global_variables_initializer())

for iteration in range(n_steps):

# Train critic

for i in range(c_updates_per_epoch):

_data = X

) = session.run(

alpha,

gradInterp,

feed_dict={X_real: _data},

_G_loss = session.run(G_loss, feed_dict={X_real: _data})

grads = {

"X_real": _data,

"X_interp": _X_interp,

"G_weights_FC1": _G_weights_FC1,

"G_weights_FC2": _G_weights_FC2,

"G_weights_FC3": _G_weights_FC3,

"G_weights_FC4": _G_weights_FC4,

"G_fwd_X_fake_FC1": _G_fwd_X_fake_FC1,

"G_fwd_X_fake_FC2": _G_fwd_X_fake_FC2,

"G_fwd_X_fake_FC3": _G_fwd_X_fake_FC3,

"G_fwd_X_fake_FC4": _G_fwd_X_fake_FC4,

"X_fake": _X_fake,

"C_weights_Y_fake_FC1": _C_weights_Y_fake_FC1,

"C_weights_Y_fake_FC2": _C_weights_Y_fake_FC2,

"C_weights_Y_fake_FC3": _C_weights_Y_fake_FC3,

"C_weights_Y_fake_FC4": _C_weights_Y_fake_FC4,

"C_fwd_Y_fake_FC1": _C_fwd_Y_fake_FC1,

"C_fwd_Y_fake_FC2": _C_fwd_Y_fake_FC2,

"C_fwd_Y_fake_FC3": _C_fwd_Y_fake_FC3,

"C_fwd_Y_fake_FC4": _C_fwd_Y_fake_FC4,

"Y_fake": _Y_fake,

"C_weights_Y_real_FC1": _C_weights_Y_real_FC1,

"C_weights_Y_real_FC2": _C_weights_Y_real_FC2,

"C_weights_Y_real_FC3": _C_weights_Y_real_FC3,

"C_weights_Y_real_FC4": _C_weights_Y_real_FC4,

"C_fwd_Y_real_FC1": _C_fwd_Y_real_FC1,

"C_fwd_Y_real_FC2": _C_fwd_Y_real_FC2,

"C_fwd_Y_real_FC3": _C_fwd_Y_real_FC3,

"C_fwd_Y_real_FC4": _C_fwd_Y_real_FC4,

"Y_real": _Y_real,

"C_weights_Y_interp_FC1": _C_weights_Y_interp_FC1,

"C_weights_Y_interp_FC2": _C_weights_Y_interp_FC2,

"C_weights_Y_interp_FC3": _C_weights_Y_interp_FC3,

"C_weights_Y_interp_FC4": _C_weights_Y_interp_FC4,

"C_fwd_Y_interp_FC1": _C_fwd_Y_interp_FC1,

"C_fwd_Y_interp_FC2": _C_fwd_Y_interp_FC2,

"C_fwd_Y_interp_FC3": _C_fwd_Y_interp_FC3,

"C_fwd_Y_interp_FC4": _C_fwd_Y_interp_FC4,

"Y_interp": _Y_interp,

"dY_interp_wrt_FC1": _C_dY_interp_wrt_FC1,

"dY_interp_wrt_FC2": _C_dY_interp_wrt_FC2,

"dY_interp_wrt_FC3": _C_dY_interp_wrt_FC3,

"dY_interp_wrt_FC4": _C_dY_interp_wrt_FC4,

"gradInterp": _gradInterp,

"gradInterp_norm": _norm_gradInterp,

"G_loss": _G_loss,

"C_loss": _C_loss,

"dC_loss_dW_FC1": _C_dW_FC1,

"dC_loss_db_FC1": _C_db_FC1,

"dC_loss_dW_FC2": _C_dW_FC2,

"dC_loss_db_FC2": _C_db_FC2,

"dC_loss_dW_FC3": _C_dW_FC3,

"dC_loss_db_FC3": _C_db_FC3,

"dC_loss_dW_FC4": _C_dW_FC4,

"dC_loss_db_FC4": _C_db_FC4,

"dC_Y_fake": _dC_Y_fake,

"dC_Y_real": _dC_Y_real,

"dC_gradInterp": _dC_gradInterp,

"dG_Y_fake": _dG_Y_fake,

return grads

def TFNCELoss(X, target_word, L):

from tensorflow.python.ops.nn_impl import _compute_sampled_logits

from tensorflow.python.ops.nn_impl import sigmoid_cross_entropy_with_logits

in_embed = tf.compat.v1.placeholder(tf.float32, shape=X.shape)

in_bias = tf.compat.v1.placeholder(

tf.float32, shape=L.parameters["b"].flatten().shape

in_weights = tf.compat.v1.placeholder(tf.float32, shape=L.parameters["W"].shape)

in_target_word = tf.compat.v1.placeholder(tf.int64)

in_neg_samples = tf.compat.v1.placeholder(tf.int32)

in_target_prob = tf.compat.v1.placeholder(tf.float32)

in_neg_samp_prob = tf.compat.v1.placeholder(tf.float32)

#  in_embed = tf.keras.Input(dtype=tf.float32, shape=X.shape)

#  in_bias = tf.keras.Input(dtype=tf.float32, shape=L.parameters["b"].flatten().shape)

#  in_weights = tf.keras.Input(dtype=tf.float32, shape=L.parameters["W"].shape)

#  in_target_word = tf.keras.Input(dtype=tf.int64, shape=())

#  in_neg_samples = tf.keras.Input(dtype=tf.int32, shape=())

#  in_target_prob = tf.keras.Input(dtype=tf.float32, shape=())

#  in_neg_samp_prob = tf.keras.Input(dtype=tf.float32, shape=())

feed = {

in_embed: X,

in_weights: L.parameters["W"],

in_target_word: target_word,

in_bias: L.parameters["b"].flatten(),

in_neg_samples: L.derived_variables["noise_samples"][0],

in_target_prob: L.derived_variables["noise_samples"][1],

in_neg_samp_prob: L.derived_variables["noise_samples"][2],

# Compute the NCE loss, using a sample of the negative labels each time.

nce_unreduced = tf.nn.nce_loss(

sampled_values=(in_neg_samples, in_target_prob, in_neg_samp_prob),

loss = tf.reduce_sum(nce_unreduced)

dLdW = tf.gradients(loss, [in_weights])[0]

dLdb = tf.gradients(loss, [in_bias])[0]

dLdX = tf.gradients(loss, [in_embed])[0]

sampled_logits, sampled_labels = _compute_sampled_logits(

sampled_losses = sigmoid_cross_entropy_with_logits(

labels=sampled_labels, logits=sampled_logits

with tf.compat.v1.Session() as session:

session.run(tf.compat.v1.global_variables_initializer())

) = session.run(

loss,

dLdW,

dLdb,

dLdX,

return {

"final_loss": _final_loss,

"nce_unreduced": _nce_unreduced,

"dLdW": _dLdW,

"dLdb": _dLdb,

"dLdX": _dLdX,

"out_logits": _sampled_logits,

"out_labels": _sampled_labels,

"sampled_loss": _sampled_losses,


# flake8: noqa

import time

import numpy as np

from numpy.testing import assert_almost_equal

from scipy.special import expit

import torch

import torch.nn.functional as F

from numpy_ml.utils.testing import random_stochastic_matrix, random_tensor

def torch_gradient_generator(fn, **kwargs):

def get_grad(z):

z1 = torch.autograd.Variable(torch.from_numpy(z), requires_grad=True)

z2 = fn(z1, **kwargs).sum()

grad = z1.grad.numpy()

return grad

return get_grad

#                           Debug Formatter                           #

def err_fmt(params, golds, ix, warn_str=""):

mine, label = params[ix]

err_msg = "-" * 25 + " DEBUG " + "-" * 25 + "\n"

prev_mine, prev_label = params[max(ix - 1, 0)]

err_msg += "Mine (prev) [{}]:\n{}\n\nTheirs (prev) [{}]:\n{}".format(

prev_label, prev_mine, prev_label, golds[prev_label]

err_msg += "\n\nMine [{}]:\n{}\n\nTheirs [{}]:\n{}".format(

label, mine, label, golds[label]

err_msg += warn_str

err_msg += "\n" + "-" * 23 + " END DEBUG " + "-" * 23

return err_msg

#                            Test Suite                               #

#  def test_activations(N=50):

#      print("Testing Sigmoid activation")

#      time.sleep(1)

#      test_sigmoid_activation(N)

#      test_sigmoid_grad(N)

#      #  print("Testing Softmax activation")

#      #  time.sleep(1)

#      #  test_softmax_activation(N)

#      #  test_softmax_grad(N)

#      print("Testing Tanh activation")

#      time.sleep(1)

#      test_tanh_grad(N)

#      print("Testing ReLU activation")

#      time.sleep(1)

#      test_relu_activation(N)

#      test_relu_grad(N)

#      print("Testing ELU activation")

#      time.sleep(1)

#      test_elu_activation(N)

#      test_elu_grad(N)

#      print("Testing SELU activation")

#      time.sleep(1)

#      test_selu_activation(N)

#      test_selu_grad(N)

#      print("Testing LeakyRelu activation")

#      time.sleep(1)

#      test_leakyrelu_activation(N)

#      test_leakyrelu_grad(N)

#      print("Testing SoftPlus activation")

#      time.sleep(1)

#      test_softplus_activation(N)

#      test_softplus_grad(N)

#                          Activations                                #

def test_sigmoid_activation(N=50):

from numpy_ml.neural_nets.activations import Sigmoid

N = np.inf if N is None else N

mine = Sigmoid()

gold = expit

i = 0

while i < N:

n_dims = np.random.randint(1, 100)

z = random_tensor((1, n_dims))

assert_almost_equal(mine.fn(z), gold(z))

print("PASSED")

i += 1

def test_softplus_activation(N=50):

from numpy_ml.neural_nets.activations import SoftPlus

N = np.inf if N is None else N

mine = SoftPlus()

gold = lambda z: F.softplus(torch.FloatTensor(z)).numpy()

i = 0

while i < N:

n_dims = np.random.randint(1, 100)

z = random_stochastic_matrix(1, n_dims)

print("PASSED")

i += 1

def test_elu_activation(N=50):

from numpy_ml.neural_nets.activations import ELU

N = np.inf if N is None else N

i = 0

while i < N:

n_dims = np.random.randint(1, 10)

z = random_tensor((1, n_dims))

alpha = np.random.uniform(0, 10)

mine = ELU(alpha)

gold = lambda z, a: F.elu(torch.from_numpy(z), alpha).numpy()

assert_almost_equal(mine.fn(z), gold(z, alpha))

print("PASSED")

i += 1

def test_relu_activation(N=50):

from numpy_ml.neural_nets.activations import ReLU

N = np.inf if N is None else N

mine = ReLU()

gold = lambda z: F.relu(torch.FloatTensor(z)).numpy()

i = 0

while i < N:

n_dims = np.random.randint(1, 100)

z = random_stochastic_matrix(1, n_dims)

print("PASSED")

i += 1

def test_selu_activation(N=50):

from numpy_ml.neural_nets.activations import SELU

N = np.inf if N is None else N

mine = SELU()

gold = lambda z: F.selu(torch.FloatTensor(z)).numpy()

i = 0

while i < N:

n_dims = np.random.randint(1, 100)

z = random_stochastic_matrix(1, n_dims)

print("PASSED")

i += 1

def test_leakyrelu_activation(N=50):

from numpy_ml.neural_nets.activations import LeakyReLU

N = np.inf if N is None else N

i = 0

while i < N:

n_dims = np.random.randint(1, 100)

z = random_stochastic_matrix(1, n_dims)

alpha = np.random.uniform(0, 10)

mine = LeakyReLU(alpha=alpha)

gold = lambda z: F.leaky_relu(torch.FloatTensor(z), alpha).numpy()

print("PASSED")

i += 1

def test_gelu_activation(N=50):

from numpy_ml.neural_nets.activations import GELU

N = np.inf if N is None else N

i = 0

while i < N:

n_dims = np.random.randint(1, 100)

z = random_stochastic_matrix(1, n_dims)

approx = np.random.choice([True, False])

mine = GELU(approximate=False)

mine_approx = GELU(approximate=True)

gold = lambda z: F.gelu(torch.FloatTensor(z)).numpy()

np.testing.assert_allclose(mine.fn(z), gold(z), rtol=1e-3)

assert_almost_equal(mine.fn(z), mine_approx.fn(z))

print("PASSED")

i += 1

#                      Activation Gradients                           #

def test_sigmoid_grad(N=50):

N = np.inf if N is None else N

mine = Sigmoid()

gold = torch_gradient_generator(torch.sigmoid)

i = 0

while i < N:

n_ex = np.random.randint(1, 100)

n_dims = np.random.randint(1, 100)

z = random_tensor((n_ex, n_dims))

assert_almost_equal(mine.grad(z), gold(z))

print("PASSED")

i += 1

def test_elu_grad(N=50):

N = np.inf if N is None else N

i = 0

while i < N:

n_ex = np.random.randint(1, 10)

n_dims = np.random.randint(1, 10)

alpha = np.random.uniform(0, 10)

z = random_tensor((n_ex, n_dims))

mine = ELU(alpha)

gold = torch_gradient_generator(F.elu, alpha=alpha)

assert_almost_equal(mine.grad(z), gold(z), decimal=6)

print("PASSED")

i += 1

def test_tanh_grad(N=50):

from numpy_ml.neural_nets.activations import Tanh

N = np.inf if N is None else N

mine = Tanh()

gold = torch_gradient_generator(torch.tanh)

i = 0

while i < N:

n_ex = np.random.randint(1, 100)

n_dims = np.random.randint(1, 100)

z = random_tensor((n_ex, n_dims))

print("PASSED")

i += 1

def test_relu_grad(N=50):

N = np.inf if N is None else N

mine = ReLU()

gold = torch_gradient_generator(F.relu)

i = 0

while i < N:

n_ex = np.random.randint(1, 100)

n_dims = np.random.randint(1, 100)

z = random_tensor((n_ex, n_dims))

print("PASSED")

i += 1

def test_gelu_grad(N=50):

N = np.inf if N is None else N

mine = GELU(approximate=False)

mine_approx = GELU(approximate=True)

gold = torch_gradient_generator(F.gelu)

i = 0

while i < N:

n_ex = np.random.randint(1, 100)

n_dims = np.random.randint(1, 100)

z = random_tensor((n_ex, n_dims))

assert_almost_equal(mine.grad(z), gold(z), decimal=3)

assert_almost_equal(mine.grad(z), mine_approx.grad(z))

print("PASSED")

i += 1

def test_selu_grad(N=50):

N = np.inf if N is None else N

mine = SELU()

gold = torch_gradient_generator(F.selu)

i = 0

while i < N:

n_ex = np.random.randint(1, 100)

n_dims = np.random.randint(1, 100)

z = random_tensor((n_ex, n_dims))

print("PASSED")

i += 1

def test_leakyrelu_grad(N=50):

N = np.inf if N is None else N

i = 0

while i < N:

n_ex = np.random.randint(1, 10)

n_dims = np.random.randint(1, 10)

alpha = np.random.uniform(0, 10)

z = random_tensor((n_ex, n_dims))

mine = LeakyReLU(alpha)

gold = torch_gradient_generator(F.leaky_relu, negative_slope=alpha)

print("PASSED")

i += 1

def test_softplus_grad(N=50):

N = np.inf if N is None else N

mine = SoftPlus()

gold = torch_gradient_generator(F.softplus)

i = 0

while i < N:

n_ex = np.random.randint(1, 100)

n_dims = np.random.randint(1, 100)

z = random_tensor((n_ex, n_dims), standardize=True)

print("PASSED")

i += 1

if __name__ == "__main__":


"""Miscellaneous utility functions"""

import numpy as np

def logsumexp(log_probs, axis=None):

Redefine scipy.special.logsumexp

see: http://bayesjumping.net/log-sum-exp-trick/

_max = np.max(log_probs)

ds = log_probs - _max

exp_sum = np.exp(ds).sum(axis=axis)

return _max + np.log(exp_sum)

def log_gaussian_pdf(x_i, mu, sigma):

"""Compute log N(x_i | mu, sigma)"""

n = len(mu)

a = n * np.log(2 * np.pi)

_, b = np.linalg.slogdet(sigma)

y = np.linalg.solve(sigma, x_i - mu)

c = np.dot(x_i - mu, y)

return -0.5 * (a + b + c)


"""Utilities module"""
from . import testing
from . import data_structures
from . import distance_metrics
from . import kernels
from . import windows
from . import graphs
from . import misc


import numpy as np

def euclidean(x, y):

Compute the Euclidean (`L2`) distance between two real vectors

Notes

The Euclidean distance between two vectors **x** and **y** is

.. math::

d(\mathbf{x}, \mathbf{y}) = \sqrt{ \sum_i (x_i - y_i)^2  }

Parameters

x,y : :py:class:`ndarray <numpy.ndarray>` s of shape `(N,)`

Returns

d : float

The L2 distance between **x** and **y**.

return np.sqrt(np.sum((x - y) ** 2))

def manhattan(x, y):

Compute the Manhattan (`L1`) distance between two real vectors

Notes

The Manhattan distance between two vectors **x** and **y** is

.. math::

d(\mathbf{x}, \mathbf{y}) = \sum_i |x_i - y_i|

Parameters

Returns

d : float

The L1 distance between **x** and **y**.

return np.sum(np.abs(x - y))

def chebyshev(x, y):

Compute the Chebyshev (:math:`L_\infty`) distance between two real vectors

Notes

The Chebyshev distance between two vectors **x** and **y** is

.. math::

d(\mathbf{x}, \mathbf{y}) = \max_i |x_i - y_i|

Parameters

Returns

d : float

The Chebyshev distance between **x** and **y**.

return np.max(np.abs(x - y))

def minkowski(x, y, p):

Compute the Minkowski-`p` distance between two real vectors.

Notes

The Minkowski-`p` distance between two vectors **x** and **y** is

.. math::

d(\mathbf{x}, \mathbf{y}) = \left( \sum_i |x_i - y_i|^p \\right)^{1/p}

Parameters

p : float > 1

The parameter of the distance function. When `p = 1`, this is the `L1`

distance, and when `p=2`, this is the `L2` distance. For `p < 1`,

Minkowski-`p` does not satisfy the triangle inequality and hence is not

a valid distance metric.

Returns

d : float

The Minkowski-`p` distance between **x** and **y**.

return np.sum(np.abs(x - y) ** p) ** (1 / p)

def hamming(x, y):

Compute the Hamming distance between two integer-valued vectors.

Notes

The Hamming distance between two vectors **x** and **y** is

.. math::

d(\mathbf{x}, \mathbf{y}) = \\frac{1}{N} \sum_i \mathbb{1}_{x_i \\neq y_i}

Parameters

The two vectors to compute the distance between. Both vectors should be

integer-valued.

Returns

d : float

The Hamming distance between **x** and **y**.

return np.sum(x != y) / len(x)


import re

from abc import ABC, abstractmethod

import numpy as np

class KernelBase(ABC):

def __init__(self):

self.parameters = {}

self.hyperparameters = {}

@abstractmethod

def _kernel(self, X, Y):

raise NotImplementedError

def __call__(self, X, Y=None):

"""Refer to documentation for the `_kernel` method"""

return self._kernel(X, Y)

def __str__(self):

P, H = self.parameters, self.hyperparameters

p_str = ", ".join(["{}={}".format(k, v) for k, v in P.items()])

return "{}({})".format(H["id"], p_str)

def summary(self):

"""Return the dictionary of model parameters, hyperparameters, and ID"""

return {

"id": self.hyperparameters["id"],

"parameters": self.parameters,

"hyperparameters": self.hyperparameters,

def set_params(self, summary_dict):

Set the model parameters and hyperparameters using the settings in

Parameters

summary_dict : dict

A dictionary with keys 'parameters' and 'hyperparameters',

structured as would be returned by the :meth:`summary` method. If

a particular (hyper)parameter is not included in this dict, the

current value will be used.

Returns

new_kernel : :doc:`Kernel <numpy_ml.utils.kernels>` instance

A kernel with parameters and hyperparameters adjusted to those

specified in `summary_dict`.

kr, sd = self, summary_dict

# collapse `parameters` and `hyperparameters` nested dicts into a single

# merged dictionary

flatten_keys = ["parameters", "hyperparameters"]

for k in flatten_keys:

if k in sd:

entry = sd[k]

sd.update(entry)

del sd[k]

for k, v in sd.items():

if k in self.parameters:

kr.parameters[k] = v

if k in self.hyperparameters:

kr.hyperparameters[k] = v

return kr

class LinearKernel(KernelBase):

def __init__(self, c0=0):

The linear (i.e., dot-product) kernel.

Notes

For input vectors :math:`\mathbf{x}` and :math:`\mathbf{y}`, the linear

kernel is:

.. math::

k(\mathbf{x}, \mathbf{y}) = \mathbf{x}^\\top \mathbf{y} + c_0

Parameters

c0 : float

An "inhomogeneity" parameter. When `c0` = 0, the kernel is said to be

homogenous. Default is 1.

self.hyperparameters = {"id": "LinearKernel"}

self.parameters = {"c0": c0}

def _kernel(self, X, Y=None):

Compute the linear kernel (i.e., dot-product) between all pairs of rows in

`X` and `Y`.

Parameters

X : :py:class:`ndarray <numpy.ndarray>` of shape `(N, C)`

Collection of `N` input vectors

Y : :py:class:`ndarray <numpy.ndarray>` of shape `(M, C)` or None

Default is None.

Returns

out : :py:class:`ndarray <numpy.ndarray>` of shape `(N, M)`

Similarity between `X` and `Y`, where index (`i`, `j`) gives

:math:`k(x_i, y_j)`.

X, Y = kernel_checks(X, Y)

return X @ Y.T + self.parameters["c0"]

class PolynomialKernel(KernelBase):

def __init__(self, d=3, gamma=None, c0=1):

The degree-`d` polynomial kernel.

Notes

For input vectors :math:`\mathbf{x}` and :math:`\mathbf{y}`, the polynomial

kernel is:

.. math::

k(\mathbf{x}, \mathbf{y}) = (\gamma \mathbf{x}^\\top \mathbf{y} + c_0)^d

In contrast to the linear kernel, the polynomial kernel also computes

similarities *across* dimensions of the **x** and **y** vectors,

allowing it to account for interactions between features.  As an

instance of the dot product family of kernels, the polynomial kernel is

invariant to a rotation of the coordinates about the origin, but *not*

to translations.

Parameters

d : int

Degree of the polynomial kernel. Default is 3.

gamma : float or None

A scaling parameter for the dot product between `x` and `y`,

determining the amount of smoothing/resonlution of the kernel.

Larger values result in greater smoothing. If None, defaults to 1 /

`C`.  Sometimes referred to as the kernel bandwidth.  Default is

None.

c0 : float

Parameter trading off the influence of higher-order versus lower-order

terms in the polynomial. If `c0` = 0, the kernel is said to be

homogenous. Default is 1.

self.hyperparameters = {"id": "PolynomialKernel"}

self.parameters = {"d": d, "c0": c0, "gamma": gamma}

def _kernel(self, X, Y=None):

Compute the degree-`d` polynomial kernel between all pairs of rows in `X`

and `Y`.

Parameters

X : :py:class:`ndarray <numpy.ndarray>` of shape `(N, C)`

Collection of `N` input vectors

Collection of `M` input vectors. If None, assume `Y = X`. Default

is None.

Returns

Similarity between `X` and `Y` where index (`i`, `j`) gives

:math:`k(x_i, y_j)` (i.e., the kernel's Gram-matrix).

P = self.parameters

X, Y = kernel_checks(X, Y)

gamma = 1 / X.shape[1] if P["gamma"] is None else P["gamma"]

return (gamma * (X @ Y.T) + P["c0"]) ** P["d"]

class RBFKernel(KernelBase):

def __init__(self, sigma=None):

Radial basis function (RBF) / squared exponential kernel.

Notes

For input vectors :math:`\mathbf{x}` and :math:`\mathbf{y}`, the radial

basis function kernel is:

.. math::

k(\mathbf{x}, \mathbf{y}) = \exp \left\{ -0.5

\left\lVert \\frac{\mathbf{x} -

\mathbf{y}}{\sigma} \\right\\rVert_2^2 \\right\}

The RBF kernel decreases with distance and ranges between zero (in the

limit) to one (when **x** = **y**). Notably, the implied feature space

of the kernel has an infinite number of dimensions.

Parameters

sigma : float or array of shape `(C,)` or None

A scaling parameter for the vectors **x** and **y**, producing an

isotropic kernel if a float, or an anistropic kernel if an array of

length `C`.  Larger values result in higher resolution / greater

smoothing. If None, defaults to :math:`\sqrt(C / 2)`. Sometimes

referred to as the kernel 'bandwidth'. Default is None.

self.hyperparameters = {"id": "RBFKernel"}

self.parameters = {"sigma": sigma}

def _kernel(self, X, Y=None):

Computes the radial basis function (RBF) kernel between all pairs of

rows in `X` and `Y`.

Parameters

X : :py:class:`ndarray <numpy.ndarray>` of shape `(N, C)`

Collection of `N` input vectors, each with dimension `C`.

Y : :py:class:`ndarray <numpy.ndarray>` of shape `(M, C)`

is None.

Returns

Similarity between `X` and `Y` where index (i, j) gives :math:`k(x_i, y_j)`.

P = self.parameters

X, Y = kernel_checks(X, Y)

sigma = np.sqrt(X.shape[1] / 2) if P["sigma"] is None else P["sigma"]

return np.exp(-0.5 * pairwise_l2_distances(X / sigma, Y / sigma) ** 2)

class KernelInitializer(object):

def __init__(self, param=None):

A class for initializing learning rate schedulers. Valid inputs are:

(a) __str__ representations of `KernelBase` instances

(b) `KernelBase` instances

(c) Parameter dicts (e.g., as produced via the :meth:`summary` method in

`KernelBase` instances)

If `param` is None, return `LinearKernel`.

self.param = param

def __call__(self):

param = self.param

if param is None:

kernel = LinearKernel()

elif isinstance(param, KernelBase):

kernel = param

elif isinstance(param, str):

kernel = self.init_from_str()

elif isinstance(param, dict):

kernel = self.init_from_dict()

return kernel

def init_from_str(self):

r = r"([a-zA-Z0-9]*)=([^,)]*)"

kr_str = self.param.lower()

kwargs = dict([(i, eval(j)) for (i, j) in re.findall(r, self.param)])

if "linear" in kr_str:

kernel = LinearKernel(**kwargs)

elif "polynomial" in kr_str:

kernel = PolynomialKernel(**kwargs)

elif "rbf" in kr_str:

kernel = RBFKernel(**kwargs)

else:

raise NotImplementedError("{}".format(kr_str))

return kernel

def init_from_dict(self):

S = self.param

sc = S["hyperparameters"] if "hyperparameters" in S else None

if sc is None:

raise ValueError("Must have `hyperparameters` key: {}".format(S))

if sc and sc["id"] == "LinearKernel":

scheduler = LinearKernel().set_params(S)

elif sc and sc["id"] == "PolynomialKernel":

scheduler = PolynomialKernel().set_params(S)

elif sc and sc["id"] == "RBFKernel":

scheduler = RBFKernel().set_params(S)

elif sc:

raise NotImplementedError("{}".format(sc["id"]))

return scheduler

def kernel_checks(X, Y):

X = X.reshape(-1, 1) if X.ndim == 1 else X

Y = X if Y is None else Y

Y = Y.reshape(-1, 1) if Y.ndim == 1 else Y

assert X.ndim == 2, "X must have 2 dimensions, but got {}".format(X.ndim)

assert Y.ndim == 2, "Y must have 2 dimensions, but got {}".format(Y.ndim)

assert X.shape[1] == Y.shape[1], "X and Y must have the same number of columns"

return X, Y

def pairwise_l2_distances(X, Y):

A fast, vectorized way to compute pairwise l2 distances between rows in `X`

and `Y`.

Notes

An entry of the pairwise Euclidean distance matrix for two vectors is

.. math::

d[i, j]  &=  \sqrt{(x_i - y_i) @ (x_i - y_i)} \\\\

&=  \sqrt{sum (x_i - y_j)^2} \\\\

&=  \sqrt{sum (x_i)^2 - 2 x_i y_j + (y_j)^2}

The code below computes the the third line using numpy broadcasting

fanciness to avoid any for loops.

Parameters

X : :py:class:`ndarray <numpy.ndarray>` of shape `(N, C)`

Collection of `N` input vectors

Y : :py:class:`ndarray <numpy.ndarray>` of shape `(M, C)`

Collection of `M` input vectors. If None, assume `Y` = `X`. Default is

None.

Returns

dists : :py:class:`ndarray <numpy.ndarray>` of shape `(N, M)`

Pairwise distance matrix. Entry (i, j) contains the `L2` distance between

:math:`x_i` and :math:`y_j`.

D = -2 * X @ Y.T + np.sum(Y ** 2, axis=1) + np.sum(X ** 2, axis=1)[:, np.newaxis]

D[D < 0] = 0  # clip any value less than 0 (a result of numerical imprecision)

return np.sqrt(D)


import heapq

from copy import copy

from collections import Hashable

import numpy as np

from .distance_metrics import euclidean

#                           Priority Queue                            #

class PQNode(object):

def __init__(self, key, val, priority, entry_id, **kwargs):

"""A generic node object for holding entries in :class:`PriorityQueue`"""

self.key = key

self.val = val

self.entry_id = entry_id

self.priority = priority

def __repr__(self):

fstr = "PQNode(key={}, val={}, priority={}, entry_id={})"

return fstr.format(self.key, self.val, self.priority, self.entry_id)

def to_dict(self):

"""Return a dictionary representation of the node's contents"""

d = self.__dict__

d["id"] = "PQNode"

return d

def __gt__(self, other):

if not isinstance(other, PQNode):

return -1

if self.priority == other.priority:

return self.entry_id > other.entry_id

return self.priority > other.priority

def __ge__(self, other):

if not isinstance(other, PQNode):

return -1

return self.priority >= other.priority

def __lt__(self, other):

if not isinstance(other, PQNode):

return -1

if self.priority == other.priority:

return self.entry_id < other.entry_id

return self.priority < other.priority

def __le__(self, other):

if not isinstance(other, PQNode):

return -1

return self.priority <= other.priority

class PriorityQueue:

def __init__(self, capacity, heap_order="max"):

A priority queue implementation using a binary heap.

Notes

A priority queue is a data structure useful for storing the top

`capacity` largest or smallest elements in a collection of values. As a

result of using a binary heap, ``PriorityQueue`` offers `O(log N)`

:meth:`push` and :meth:`pop` operations.

Parameters

capacity: int

The maximum number of items that can be held in the queue.

heap_order: {"max", "min"}

Whether the priority queue should retain the items with the

`capacity` smallest (`heap_order` = 'min') or `capacity` largest

(`heap_order` = 'max') priorities.

assert heap_order in ["max", "min"], "heap_order must be either 'max' or 'min'"

self.capacity = capacity

self.heap_order = heap_order

self._pq = []

self._count = 0

self._entry_counter = 0

def __repr__(self):

fstr = "PriorityQueue(capacity={}, heap_order={}) with {} items"

return fstr.format(self.capacity, self.heap_order, self._count)

def __len__(self):

return self._count

def __iter__(self):

return iter(self._pq)

def push(self, key, priority, val=None):

Add a new (key, value) pair with priority `priority` to the queue.

Notes

If the queue is at capacity and `priority` exceeds the priority of the

item with the largest/smallest priority currently in the queue, replace

the current queue item with (`key`, `val`).

Parameters

key : hashable object

The key to insert into the queue.

priority : comparable

The priority for the `key`, `val` pair.

val : object

The value associated with `key`. Default is None.

if self.heap_order == "max":

priority = -1 * priority

item = PQNode(key=key, val=val, priority=priority, entry_id=self._entry_counter)

heapq.heappush(self._pq, item)

self._count += 1

self._entry_counter += 1

while self._count > self.capacity:

self.pop()

def pop(self):

Remove the item with the largest/smallest (depending on

``self.heap_order``) priority from the queue and return it.

Notes

In contrast to :meth:`peek`, this operation is `O(log N)`.

Returns

item : :class:`PQNode` instance or None

Item with the largest/smallest priority, depending on

item = heapq.heappop(self._pq).to_dict()

if self.heap_order == "max":

item["priority"] = -1 * item["priority"]

self._count -= 1

return item

def peek(self):

Return the item with the largest/smallest (depending on

``self.heap_order``) priority *without* removing it from the queue.

Notes

In contrast to :meth:`pop`, this operation is O(1).

Returns

item : :class:`PQNode` instance or None

item = None

if self._count > 0:

item = copy(self._pq[0].to_dict())

if self.heap_order == "max":

return item

#                              Ball Tree                              #

class BallTreeNode:

def __init__(self, centroid=None, X=None, y=None):

self.left = None

self.right = None

self.radius = None

self.is_leaf = False

self.data = X

self.targets = y

self.centroid = centroid

def __repr__(self):

fstr = "BallTreeNode(centroid={}, is_leaf={})"

return fstr.format(self.centroid, self.is_leaf)

def to_dict(self):

d = self.__dict__

d["id"] = "BallTreeNode"

return d

class BallTree:

def __init__(self, leaf_size=40, metric=None):

A ball tree data structure.

Notes

A ball tree is a binary tree in which every node defines a

`D`-dimensional hypersphere ("ball") containing a subset of the points

to be searched. Each internal node of the tree partitions the data

points into two disjoint sets which are associated with different

balls. While the balls themselves may intersect, each point is assigned

to one or the other ball in the partition according to its distance

from the ball's center. Each leaf node in the tree defines a ball and

enumerates all data points inside that ball.

Parameters

leaf_size : int

The maximum number of datapoints at each leaf. Default is 40.

metric : :doc:`Distance metric <numpy_ml.utils.distance_metrics>` or None

The distance metric to use for computing nearest neighbors. If

None, use the :func:`~numpy_ml.utils.distance_metrics.euclidean`

metric. Default is None.

References

.. [1] Omohundro, S. M. (1989). "Five balltree construction algorithms". *ICSI

Technical Report TR-89-063*.

.. [2] Liu, T., Moore, A., & Gray A. (2006). "New algorithms for efficient

high-dimensional nonparametric classification". *J. Mach. Learn. Res.,

7*, 1135-1158.

self.root = None

self.leaf_size = leaf_size

self.metric = metric if metric is not None else euclidean

def fit(self, X, y=None):

Build a ball tree recursively using the O(M log N) `k`-d construction

algorithm.

Notes

Recursively divides data into nodes defined by a centroid `C` and radius

`r` such that each point below the node lies within the hyper-sphere

defined by `C` and `r`.

Parameters

X : :py:class:`ndarray <numpy.ndarray>` of shape `(N, M)`

An array of `N` examples each with `M` features.

y : :py:class:`ndarray <numpy.ndarray>` of shape `(N, \*)` or None

An array of target values / labels associated with the entries in

`X`. Default is None.

centroid, left_X, left_y, right_X, right_y = self._split(X, y)

self.root = BallTreeNode(centroid=centroid)

self.root.radius = np.max([self.metric(centroid, x) for x in X])

self.root.left = self._build_tree(left_X, left_y)

self.root.right = self._build_tree(right_X, right_y)

def _build_tree(self, X, y):

if X.shape[0] <= self.leaf_size:

leaf = BallTreeNode(centroid=centroid, X=X, y=y)

leaf.radius = np.max([self.metric(centroid, x) for x in X])

leaf.is_leaf = True

return leaf

node = BallTreeNode(centroid=centroid)

node.radius = np.max([self.metric(centroid, x) for x in X])

node.left = self._build_tree(left_X, left_y)

node.right = self._build_tree(right_X, right_y)

return node

def _split(self, X, y=None):

# find the dimension with greatest variance

split_dim = np.argmax(np.var(X, axis=0))

# sort X and y along split_dim

sort_ixs = np.argsort(X[:, split_dim])

X, y = X[sort_ixs], y[sort_ixs] if y is not None else None

# divide at median value of split_dim

med_ix = X.shape[0] // 2

centroid = X[med_ix]  # , split_dim

# split data into two halves at the centroid (median always appears on

# the right split)

left_X, left_y = X[:med_ix], y[:med_ix] if y is not None else None

right_X, right_y = X[med_ix:], y[med_ix:] if y is not None else None

return centroid, left_X, left_y, right_X, right_y

def nearest_neighbors(self, k, x):

Find the `k` nearest neighbors in the ball tree to a query vector `x`

using the KNS1 algorithm.

Parameters

k : int

The number of closest points in `X` to return

x : :py:class:`ndarray <numpy.ndarray>` of shape `(1, M)`

The query vector.

Returns

nearest : list of :class:`PQNode` s of length `k`

List of the `k` points in `X` to closest to the query vector. The

``key`` attribute of each :class:`PQNode` contains the point itself, the

``val`` attribute contains its target, and the ``distance``

attribute contains its distance to the query vector.

# maintain a max-first priority queue with priority = distance to x

PQ = PriorityQueue(capacity=k, heap_order="max")

nearest = self._knn(k, x, PQ, self.root)

for n in nearest:

n.distance = self.metric(x, n.key)

return nearest

def _knn(self, k, x, PQ, root):

dist = self.metric

dist_to_ball = dist(x, root.centroid) - root.radius

dist_to_farthest_neighbor = dist(x, PQ.peek()["key"]) if len(PQ) > 0 else np.inf

if dist_to_ball >= dist_to_farthest_neighbor and len(PQ) == k:

return PQ

if root.is_leaf:

targets = [None] * len(root.data) if root.targets is None else root.targets

for point, target in zip(root.data, targets):

dist_to_x = dist(x, point)

if len(PQ) == k and dist_to_x < dist_to_farthest_neighbor:

PQ.push(key=point, val=target, priority=dist_to_x)

else:

else:

l_closest = dist(x, root.left.centroid) < dist(x, root.right.centroid)

PQ = self._knn(k, x, PQ, root.left if l_closest else root.right)

PQ = self._knn(k, x, PQ, root.right if l_closest else root.left)

return PQ

#                         Multinomial Sampler                         #

class DiscreteSampler:

def __init__(self, probs, log=False, with_replacement=True):

Sample from an arbitrary multinomial PMF over the first `N` nonnegative

integers using Vose's algorithm for the alias method.

Notes

Vose's algorithm takes `O(n)` time to initialize, requires `O(n)` memory,

and generates samples in constant time.

References

.. [1] Walker, A. J. (1977) "An efficient method for generating discrete

random variables with general distributions". *ACM Transactions on

Mathematical Software, 3(3)*, 253-256.

.. [2] Vose, M. D. (1991) "A linear algorithm for generating random numbers

with a given distribution". *IEEE Trans. Softw. Eng., 9*, 972-974.

.. [3] Schwarz, K (2011) "Darts, dice, and coins: sampling from a discrete

distribution". http://www.keithschwarz.com/darts-dice-coins/

Parameters

probs : :py:class:`ndarray <numpy.ndarray>` of length `(N,)`

A list of probabilities of the `N` outcomes in the sample space.

`probs[i]` returns the probability of outcome `i`.

log : bool

Whether the probabilities in `probs` are in logspace. Default is

False.

with_replacement : bool

Whether to generate samples with or without replacement. Default is

True.

if not isinstance(probs, np.ndarray):

probs = np.array(probs)

self.log = log

self.N = len(probs)

self.probs = probs

self.with_replacement = with_replacement

alias = np.zeros(self.N)

prob = np.zeros(self.N)

scaled_probs = self.probs + np.log(self.N) if log else self.probs * self.N

selector = scaled_probs < 0 if log else scaled_probs < 1

small, large = np.where(selector)[0].tolist(), np.where(~selector)[0].tolist()

while len(small) and len(large):

l, g = small.pop(), large.pop()

alias[l] = g

prob[l] = scaled_probs[l]

if log:

pg = np.log(np.exp(scaled_probs[g]) + np.exp(scaled_probs[l]) - 1)

else:

pg = scaled_probs[g] + scaled_probs[l] - 1

scaled_probs[g] = pg

to_small = pg < 0 if log else pg < 1

if to_small:

small.append(g)

else:

large.append(g)

while len(large):

prob[large.pop()] = 0 if log else 1

while len(small):

prob[small.pop()] = 0 if log else 1

self.prob_table = prob

self.alias_table = alias

def __call__(self, n_samples=1):

Generate random draws from the `probs` distribution over integers in

[0, N).

Parameters

n_samples: int

The number of samples to generate. Default is 1.

Returns

sample : :py:class:`ndarray <numpy.ndarray>` of shape `(n_samples,)`

A collection of draws from the distribution defined by `probs`.

Each sample is an int in the range `[0, N)`.

return self.sample(n_samples)

def sample(self, n_samples=1):

[0, N).

Parameters

n_samples: int

Returns

ixs = np.random.randint(0, self.N, n_samples)

p = np.exp(self.prob_table[ixs]) if self.log else self.prob_table[ixs]

flips = np.random.binomial(1, p)

samples = [ix if f else self.alias_table[ix] for ix, f in zip(ixs, flips)]

# do recursive rejection sampling to sample without replacement

if not self.with_replacement:

unique = list(set(samples))

while len(samples) != len(unique):

n_new = len(samples) - len(unique)

samples = unique + self.sample(n_new).tolist()

unique = list(set(samples))

return np.array(samples, dtype=int)

#                                Dict                                 #

class Dict(dict):

def __init__(self, encoder=None):

A dictionary subclass which returns the key value if it is not in the

dict.

Parameters

encoder : function or None

A function which is applied to a key before adding / retrieving it

from the dictionary. If None, the function defaults to the

identity. Default is None.

super(Dict, self).__init__()

self._encoder = encoder

self._id_max = 0

def __setitem__(self, key, value):

if self._encoder is not None:

key = self._encoder(key)

elif not isinstance(key, Hashable):

key = tuple(key)

super(Dict, self).__setitem__(key, value)

def _encode_key(self, key):

D = super(Dict, self)

enc_key = self._encoder(key)

if D.__contains__(enc_key):

val = D.__getitem__(enc_key)

else:

val = self._id_max

D.__setitem__(enc_key, val)

self._id_max += 1

return val

def __getitem__(self, key):

self._key = copy.deepcopy(key)

if self._encoder is not None:

return self._encode_key(key)

elif not isinstance(key, Hashable):

key = tuple(key)

return super(Dict, self).__getitem__(key)

def __missing__(self, key):

return self._key


import numpy as np

def blackman_harris(window_len, symmetric=False):

The Blackman-Harris window.

Notes

The Blackman-Harris window is an instance of the more general class of

cosine-sum windows where `K=3`. Additional coefficients extend the Hamming

window to further minimize the magnitude of the nearest side-lobe in the

frequency response.

.. math::

\\text{bh}(n) = a_0 - a_1 \cos\left(\\frac{2 \pi n}{N}\\right) +

a_2 \cos\left(\\frac{4 \pi n }{N}\\right) -

a_3 \cos\left(\\frac{6 \pi n}{N}\\right)

where `N` = `window_len` - 1, :math:`a_0` = 0.35875, :math:`a_1` = 0.48829,

:math:`a_2` = 0.14128, and :math:`a_3` = 0.01168.

Parameters

window_len : int

The length of the window in samples. Should be equal to the

`frame_width` if applying to a windowed signal.

symmetric : bool

If False, create a 'periodic' window that can be used in with an FFT /

in spectral analysis.  If True, generate a symmetric window that can be

used in, e.g., filter design. Default is False.

Returns

window : :py:class:`ndarray <numpy.ndarray>` of shape `(window_len,)`

The window

return generalized_cosine(

window_len, [0.35875, 0.48829, 0.14128, 0.01168], symmetric

def hamming(window_len, symmetric=False):

The Hamming window.

Notes

The Hamming window is an instance of the more general class of cosine-sum

windows where `K=1` and :math:`a_0 = 0.54`. Coefficients selected to

minimize the magnitude of the nearest side-lobe in the frequency response.

.. math::

\\text{hamming}(n) = 0.54 -

0.46 \cos\left(\\frac{2 \pi n}{\\text{window_len} - 1}\\right)

Parameters

window_len : int

symmetric : bool

Returns

The window

return generalized_cosine(window_len, [0.54, 1 - 0.54], symmetric)

def hann(window_len, symmetric=False):

The Hann window.

Notes

The Hann window is an instance of the more general class of cosine-sum

windows where `K=1` and :math:`a_0` = 0.5. Unlike the Hamming window, the

end points of the Hann window touch zero.

.. math::

\\text{hann}(n) = 0.5 - 0.5 \cos\left(\\frac{2 \pi n}{\\text{window_len} - 1}\\right)

Parameters

window_len : int

symmetric : bool

Returns

The window

return generalized_cosine(window_len, [0.5, 0.5], symmetric)

def generalized_cosine(window_len, coefs, symmetric=False):

The generalized cosine family of window functions.

Notes

The generalized cosine window is a simple weighted sum of cosine terms.

For :math:`n \in \{0, \ldots, \\text{window_len} \}`:

.. math::

\\text{GCW}(n) = \sum_{k=0}^K (-1)^k a_k \cos\left(\\frac{2 \pi k n}{\\text{window_len}}\\right)

Parameters

window_len : int

coefs: list of floats

The :math:`a_k` coefficient values

symmetric : bool

Returns

The window

window_len += 1 if not symmetric else 0

entries = np.linspace(-np.pi, np.pi, window_len)  # (-1)^k * 2pi*n / window_len

window = np.sum([ak * np.cos(k * entries) for k, ak in enumerate(coefs)], axis=0)

return window[:-1] if not symmetric else window

class WindowInitializer:

def __call__(self, window):

if window == "hamming":

return hamming

elif window == "blackman_harris":

return blackman_harris

elif window == "hann":

return hann

elif window == "generalized_cosine":

return generalized_cosine

else:

raise NotImplementedError("{}".format(window))