text: import numpy as np
  import torch
  from torch import nn
  from core import *
  from collections import namedtuple
  from itertools import count
  
  torch.backends.cudnn.benchmark = True
  device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
  cpu = torch.device("cpu")
  
  
  @cat.register(torch.Tensor)
  def _(*xs):
      return torch.cat(xs)
  
  
  @to_numpy.register(torch.Tensor)
  def _(x):
      return x.detach().cpu().numpy()
  
  
  @pad.register(torch.Tensor)
  def _(x, border):
      return nn.ReflectionPad2d(border)(x)
  
  
  @transpose.register(torch.Tensor)
  def _(x, source, target):
      return x.permute([source.index(d) for d in target])
  
  
  def to(*args, **kwargs):
      return lambda x: x.to(*args, **kwargs)
  
  
  @flip_lr.register(torch.Tensor)
  def _(x):
      return torch.flip(x, [-1])
  
  
  #####################
  ## dataset
  #####################
  from functools import lru_cache as cache
  
  
  @cache(None)
  def cifar10(root='./data'):
      try:
          import torchvision
          download = lambda train: torchvision.datasets.CIFAR10(
              root=root, train=train, download=True)
          return {
              k: {
                  'data': v.data,
                  'targets': v.targets
              }
              for k, v in [('train', download(
                  train=True)), ('valid', download(train=False))]
          }
      except ImportError:
          from tensorflow.keras import datasets
          (train_images,
           train_labels), (valid_images,
                           valid_labels) = datasets.cifar10.load_data()
          return {
              'train': {
                  'data': train_images,
                  'targets': train_labels.squeeze()
              },
              'valid': {
                  'data': valid_images,
                  'targets': valid_labels.squeeze()
              }
          }
  
  
  cifar10_mean, cifar10_std = [
      (125.31, 122.95,
       113.87),  # equals np.mean(cifar10()['train']['data'], axis=(0,1,2)) 
      (62.99, 62.09,
       66.70),  # equals np.std(cifar10()['train']['data'], axis=(0,1,2))
  ]
  cifar10_classes = 'airplane, automobile, bird, cat, deer, dog, frog, horse, ship, truck'.split(
      ', ')
  
  #####################
  ## data loading
  #####################
  
  
  class DataLoader():
      def __init__(self,
                   dataset,
                   batch_size,
                   shuffle,
                   set_random_choices=False,
                   num_workers=0,
                   drop_last=False):
          self.dataset = dataset
          self.batch_size = batch_size
          self.set_random_choices = set_random_choices
          self.dataloader = torch.utils.data.DataLoader(dataset,
                                                        batch_size=batch_size,
                                                        num_workers=num_workers,
                                                        pin_memory=True,
                                                        shuffle=shuffle,
                                                        drop_last=drop_last)
  
      def __iter__(self):
          if self.set_random_choices:
              self.dataset.set_random_choices()
          return ({
              'input': x.to(device).half(),
              'target': y.to(device).long()
          } for (x, y) in self.dataloader)
  
      def __len__(self):
          return len(self.dataloader)
  
  
  #GPU dataloading
  chunks = lambda data, splits: (data[start:end]
                                 for (start, end) in zip(splits, splits[1:]))
  
  even_splits = lambda N, num_chunks: np.cumsum([0] + [(N // num_chunks) + 1] * (
      N % num_chunks) + [N // num_chunks] * (num_chunks - (N % num_chunks)))
  
  
  def shuffled(xs, inplace=False):
      xs = xs if inplace else copy.copy(xs)
      np.random.shuffle(xs)
      return xs
  
  
  def transformed(data, targets, transform, max_options=None, unshuffle=False):
      i = torch.randperm(len(data), device=device)
      data = data[i]
      options = shuffled(transform.options(data.shape),
                         inplace=True)[:max_options]
      data = torch.cat([
          transform(x, **choice) for choice, x in zip(
              options, chunks(data, even_splits(len(data), len(options))))
      ])
      return (data[torch.argsort(i)], targets) if unshuffle else (data,
                                                                  targets[i])
  
  
  class GPUBatches():
      def __init__(self,
                   batch_size,
                   transforms=(),
                   dataset=None,
                   shuffle=True,
                   drop_last=False,
                   max_options=None,
                   mixup_count=1,
                   noise_channels=0,
                   alpha=1.):
          self.dataset = dataset
          self.transforms = transforms
          self.shuffle = shuffle
          self.max_options = max_options
          self.mixup_count = mixup_count
          self.noise_channels = noise_channels
          self.alpha = alpha
  
          N = len(dataset['data']) // mixup_count
  
          self.samples = N * mixup_count
  
          self.dataset['data'] = self.dataset['data'][:self.samples]
  
          noise_size = [*self.dataset['data'].size()]
          noise_size[1] = noise_channels
  
          noise_channels = torch.randn(*noise_size,
                                       dtype=self.dataset['data'].dtype,
                                       device=device)
  
          self.dataset['data'] = torch.cat(
              (self.dataset['data'], noise_channels), dim=1)
  
          self.dataset['targets'] = self.dataset['targets'][:self.samples]
  
          self.splits = list(range(0, N + 1, batch_size))
  
          self.N = N
  
          if not drop_last and self.splits[-1] != N:
              self.splits.append(N)
  
      def __iter__(self):
          data, targets = self.dataset['data'], self.dataset['targets']
          for transform in self.transforms:
              data, targets = transformed(data,
                                          targets,
                                          transform,
                                          max_options=self.max_options,
                                          unshuffle=not self.shuffle)
  
          if self.shuffle:
              i = torch.randperm(len(data), device=device)
              data, targets = data[i], targets[i]
  
          N = self.N
  
          with torch.no_grad():
              mixed_data = data[:N]
              targets_list = [targets[:N]]
              weights_list = []
  
              for i in range(self.mixup_count - 1):
                  t_alpha = torch.tensor(self.alpha,
                                         dtype=mixed_data.dtype,
                                         device=device)
                  weights = torch.distributions.beta.Beta(t_alpha,
                                                          t_alpha).sample(
                                                              (N, 1, 1, 1))
                  start = N * (i + 1)
                  end = N * (i + 2)
                  next_data, next_targets = data[start:end], targets[start:end]
  
                  mixed_data = weights * mixed_data + (1 - weights) * next_data
                  targets_list.append(next_targets)
                  weights_list.append(weights)
  
              data = mixed_data
  
          data_chunks = chunks(data, self.splits)
  
          flip_dim = lambda x: zip(*x)
          l_chunks = lambda xs: zip(*[chunks(x, self.splits) for x in xs])
  
          targets_chunks = l_chunks(targets_list)
  
          if self.mixup_count != 1:
              weights_chunks = l_chunks(weights_list)
          else:
              weights_chunks = [[] for _ in self.splits]
  
          return ({
              'input': x.clone(),
              'targets': y,
              'weights': w,
          } for (x, y, w) in zip(data_chunks, targets_chunks, weights_chunks))
  
      def __len__(self):
          return len(self.splits) - 1
  
  
  #####################
  ## Layers
  #####################
  
  
  #Network
  class Network(nn.Module):
      def __init__(self, net):
          super().__init__()
          self.graph = build_graph(net)
          for path, (val, _) in self.graph.items():
              setattr(self, path.replace('/', '_'), val)
  
      def nodes(self):
          return (node for node, _ in self.graph.values())
  
      def forward(self, inputs):
          outputs = dict(inputs)
          for k, (node, ins) in self.graph.items():
              #only compute nodes that are not supplied as inputs.
              if k not in outputs:
                  outputs[k] = node(*[outputs[x] for x in ins])
          return outputs
  
      def half(self):
          for node in self.nodes():
              if isinstance(node,
                            nn.Module) and not isinstance(node, nn.BatchNorm2d):
                  node.half()
          return self
  
  
  class Identity(namedtuple('Identity', [])):
      def __call__(self, x):
          return x
  
  
  class Add(namedtuple('Add', [])):
      def __call__(self, x, y):
          return x + y
  
  
  class AddWeighted(namedtuple('AddWeighted', ['wx', 'wy'])):
      def __call__(self, x, y):
          return self.wx * x + self.wy * y
  
  
  class Mul(nn.Module):
      def __init__(self, weight):
          super().__init__()
          self.weight = weight
  
      def __call__(self, x):
          return x * self.weight
  
  
  class Flatten(nn.Module):
      def forward(self, x):
          return x.view(x.size(0), x.size(1))
  
  class Concat(nn.Module):
      def forward(self, *xs):
          return torch.cat(xs, 1)
  
  
  class BatchNorm(nn.BatchNorm2d):
      def __init__(self,
                   num_features,
                   eps=1e-05,
                   momentum=0.1,
                   weight_freeze=False,
                   bias_freeze=False,
                   weight_init=1.0,
                   bias_init=0.0):
          super().__init__(num_features, eps=eps, momentum=momentum)
          if weight_init is not None: self.weight.data.fill_(weight_init)
          if bias_init is not None: self.bias.data.fill_(bias_init)
          self.weight.requires_grad = not weight_freeze
          self.bias.requires_grad = not bias_freeze
  
  
  class GhostBatchNorm(BatchNorm):
      def __init__(self, num_features, num_splits, **kw):
          super().__init__(num_features, **kw)
          self.num_splits = num_splits
          self.register_buffer('running_mean',
                               torch.zeros(num_features * self.num_splits))
          self.register_buffer('running_var',
                               torch.ones(num_features * self.num_splits))
  
      def train(self, mode=True):
          if (self.training is True) and (
                  mode is
                  False):  #lazily collate stats when we are going to use them
              self.running_mean = torch.mean(self.running_mean.view(
                  self.num_splits, self.num_features),
                                             dim=0).repeat(self.num_splits)
              self.running_var = torch.mean(self.running_var.view(
                  self.num_splits, self.num_features),
                                            dim=0).repeat(self.num_splits)
          return super().train(mode)
  
      def forward(self, input):
          N, C, H, W = input.shape
          if self.training or not self.track_running_stats:
              return nn.functional.batch_norm(
                  input.view(-1, C * self.num_splits, H,
                             W), self.running_mean, self.running_var,
                  self.weight.repeat(self.num_splits),
                  self.bias.repeat(self.num_splits), True, self.momentum,
                  self.eps).view(N, C, H, W)
          else:
              return nn.functional.batch_norm(
                  input, self.running_mean[:self.num_features],
                  self.running_var[:self.num_features], self.weight, self.bias,
                  False, self.momentum, self.eps)
  
  
  # Losses
  class CrossEntropyLoss(namedtuple('CrossEntropyLoss', [])):
      def __call__(self, log_probs, target):
          return torch.nn.functional.nll_loss(log_probs,
                                              target,
                                              reduction='none')
  
  
  class KLLoss(namedtuple('KLLoss', [])):
      def __call__(self, log_probs):
          return -log_probs.mean(dim=1)
  
  
  class Correct(namedtuple('Correct', [])):
      def __call__(self, classifier, target):
          return classifier.max(dim=1)[1] == target
  
  
  class LogSoftmax(namedtuple('LogSoftmax', ['dim'])):
      def __call__(self, x):
          return torch.nn.functional.log_softmax(x, self.dim, _stacklevel=5)
  
  
  x_ent_loss = Network({
      'loss': (nn.CrossEntropyLoss(reduction='none'), ['logits', 'target']),
      'acc': (Correct(), ['logits', 'target'])
  })
  
  
  class MixupLoss(namedtuple('MixupLoss', [])):
      def __init__(self):
          super().__init__()
  
          self.x_loss = nn.CrossEntropyLoss(reduction='none')
  
      def __call__(self, classifier, targets, weights):
          total = self.x_loss(classifier, targets[0])
  
          for (target, weight) in zip(targets[1:], weights):
              total = weight * total + (1 - weight) * self.x_loss(
                  classifier, target)
  
          return total
  
  
  class MixupCorrect(namedtuple('MixupCorrect', [])):
      def __call__(self, classifier, targets, weights):
          return classifier.max(dim=1)[1] == targets[0]  # TODO: maybe max?
  
  
  mixup_loss = Network({
      'loss': (MixupLoss(), ['logits', 'targets', 'weights']),
      'acc': (MixupCorrect(), ['logits', 'targets', 'weights'])
  })
  
  label_smoothing_loss = lambda alpha: Network({
      'logprobs': (LogSoftmax(dim=1), ['logits']),
      'KL': (KLLoss(), ['logprobs']),
      'xent': (CrossEntropyLoss(), ['logprobs', 'target']),
      'loss': (AddWeighted(wx=1 - alpha, wy=alpha), ['xent', 'KL']),
      'acc': (Correct(), ['logits', 'target']),
  })
  
  trainable_params = lambda model: {
      k: p
      for k, p in model.named_parameters() if p.requires_grad
  }
  
  #####################
  ## Optimisers
  #####################
  
  from functools import partial
  
  
  def nesterov_update(w, dw, v, lr, weight_decay, momentum):
      dw.add_(weight_decay, w).mul_(-lr)
      v.mul_(momentum).add_(dw)
      w.add_(dw.add_(momentum, v))
  
  
  norm = lambda x: torch.norm(x.reshape(x.size(0), -1).float(), dim=1
                              )[:, None, None, None]
  
  
  def LARS_update(w, dw, v, lr, weight_decay, momentum):
      nesterov_update(w, dw, v,
                      lr * (norm(w) / (norm(dw) + 1e-2)).to(w.dtype),
                      weight_decay, momentum)
  
  
  def zeros_like(weights):
      return [torch.zeros_like(w) for w in weights]
  
  
  def optimiser(weights, param_schedule, update, state_init):
      weights = list(weights)
      return {
          'update': update,
          'param_schedule': param_schedule,
          'step_number': 0,
          'weights': weights,
          'opt_state': state_init(weights)
      }
  
  
  def opt_step(update, param_schedule, step_number, weights, opt_state):
      step_number += 1
      param_values = {k: f(step_number) for k, f in param_schedule.items()}
      for w, v in zip(weights, opt_state):
          if w.requires_grad:
              update(w.data, w.grad.data, v, **param_values)
      return {
          'update': update,
          'param_schedule': param_schedule,
          'step_number': step_number,
          'weights': weights,
          'opt_state': opt_state
      }
  
  
  LARS = partial(optimiser, update=LARS_update, state_init=zeros_like)
  SGD = partial(optimiser, update=nesterov_update, state_init=zeros_like)
  
  #####################
  ## training
  #####################
  from itertools import chain
  
  
  def reduce(batches, state, steps):
      #state: is a dictionary
      #steps: are functions that take (batch, state)
      #and return a dictionary of updates to the state (or None)
  
      for batch in chain(batches, [None]):
          #we send an extra batch=None at the end for steps that
          #need to do some tidying-up (e.g. log_activations)
          for step in steps:
              updates = step(batch, state)
              if updates:
                  for k, v in updates.items():
                      state[k] = v
      return state
  
  
  #define keys in the state dict as constants
  MODEL = 'model'
  LOSS = 'loss'
  VALID_MODEL = 'valid_model'
  OUTPUT = 'output'
  OPTS = 'optimisers'
  ACT_LOG = 'activation_log'
  WEIGHT_LOG = 'weight_log'
  
  
  #step definitions
  def forward(training_mode):
      def step(batch, state):
          if not batch: return
          model = state[MODEL] if training_mode or (
              VALID_MODEL not in state) else state[VALID_MODEL]
          if model.training != training_mode:  #without the guard it's slow!
              model.train(training_mode)
          return {OUTPUT: state[LOSS](model(batch))}
  
      return step
  
  
  def forward_tta(tta_transforms):
      def step(batch, state):
          if not batch: return
          model = state[MODEL] if (
              VALID_MODEL not in state) else state[VALID_MODEL]
          if model.training:
              model.train(False)
          logits = torch.mean(torch.stack([
              model({'input': transform(batch['input'].clone())
                     })['logits'].detach() for transform in tta_transforms
          ],
                                          dim=0),
                              dim=0)
          return {OUTPUT: state[LOSS](dict(batch, logits=logits))}
  
      return step
  
  
  def backward(dtype=None):
      def step(batch, state):
          state[MODEL].zero_grad()
          if not batch: return
          loss = state[OUTPUT][LOSS]
          if dtype is not None:
              loss = loss.to(dtype)
          loss.sum().backward()
  
      return step
  
  
  def opt_steps(batch, state):
      if not batch: return
      return {OPTS: [opt_step(**opt) for opt in state[OPTS]]}
  
  
  def log_activations(node_names=('loss', 'acc')):
      def step(batch, state):
          if '_tmp_logs_' not in state:
              state['_tmp_logs_'] = []
          if batch:
              state['_tmp_logs_'].extend(
                  (k, state[OUTPUT][k].detach()) for k in node_names)
          else:
              res = {
                  k: to_numpy(torch.cat(xs)).astype(np.float)
                  for k, xs in group_by_key(state['_tmp_logs_']).items()
              }
              del state['_tmp_logs_']
              return {ACT_LOG: res}
  
      return step
  
  
  epoch_stats = lambda state: {k: np.mean(v) for k, v in state[ACT_LOG].items()}
  
  
  def update_ema(momentum, update_freq=1):
      n = iter(count())
      rho = momentum**update_freq
  
      def step(batch, state):
          if not batch: return
          if (next(n) % update_freq) != 0: return
          for v, ema_v in zip(state[MODEL].state_dict().values(),
                              state[VALID_MODEL].state_dict().values()):
              if not v.dtype.is_floating_point:
                  continue  #skip things like num_batches_tracked.
              ema_v *= rho
              ema_v += (1 - rho) * v
  
      return step
  
  
  default_train_steps = (forward(training_mode=True),
                         log_activations(('loss', 'acc')), backward(), opt_steps)
  default_valid_steps = (forward(training_mode=False),
                         log_activations(('loss', 'acc')))
  
  
  def train_epoch(state,
                  timer,
                  train_batches,
                  valid_batches,
                  train_steps=default_train_steps,
                  valid_steps=default_valid_steps,
                  on_epoch_end=(lambda state: state)):
      train_summary, train_time = epoch_stats(
          on_epoch_end(reduce(train_batches, state, train_steps))), timer()
      valid_summary, valid_time = epoch_stats(
          reduce(valid_batches, state,
                 valid_steps)), timer(include_in_total=False)  #DAWNBench rules
      return {
          'train': union({'time': train_time}, train_summary),
          'valid': union({'time': valid_time}, valid_summary),
          'total time': timer.total_time
      }
  
  
  #on_epoch_end
  def log_weights(state, weights):
      state[WEIGHT_LOG] = state.get(WEIGHT_LOG, [])
      state[WEIGHT_LOG].append({k: to_numpy(v.data) for k, v in weights.items()})
      return state
  
  
  def fine_tune_bn_stats(state, batches, model_key=VALID_MODEL):
      reduce(batches, {MODEL: state[model_key]}, [forward(True)])
      return state
  
  
  #misc
  def warmup_cudnn(model, loss, batch):
      #run forward and backward pass of the model
      #to allow benchmarking of cudnn kernels
      reduce([batch], {MODEL: model, LOSS: loss}, [forward(True), backward()])
      torch.cuda.synchronize()
  
  
  #####################
  ## input whitening
  #####################
  
  
  def cov(X):
      X = X / np.sqrt(X.size(0) - 1)
      return X.t() @ X
  
  
  def patches(data, patch_size=(3, 3), dtype=torch.float32):
      h, w = patch_size
      c = data.size(1)
      return data.unfold(2, h, 1).unfold(3, w,
                                         1).transpose(1, 3).reshape(-1, c, h,
                                                                    w).to(dtype)
  
  
  def eigens(patches):
      n, c, h, w = patches.shape
      Σ = cov(patches.reshape(n, c * h * w))
      Λ, V = torch.symeig(Σ, eigenvectors=True)
      return Λ.flip(0), V.t().reshape(c * h * w, c, h, w).flip(0)
  
  
  def whitening_filter(c_in, Λ, V, eps=1e-2):
      filt = nn.Conv2d(c_in, 27, kernel_size=(3, 3), padding=(1, 1), bias=False)
      filt.weight.data = (V / torch.sqrt(Λ + eps)[:, None, None, None])
      filt.weight.requires_grad = False
      return filt
  