text: """
  This module contains machine learning functionality. It is work in
  progress, so you must expect things to change. The only tested
  functionality for training is using consective layers.
  This includes logistic regression. It can be run as
  follows::
  
      sgd = ml.SGD([ml.Dense(n_examples, n_features, 1),
                    ml.Output(n_examples, approx=True)], n_epochs,
                   report_loss=True)
      sgd.layers[0].X.input_from(0)
      sgd.layers[1].Y.input_from(1)
      sgd.reset()
      sgd.run()
  
  This loads measurements from party 0 and labels (0/1) from party
  1. After running, the model is stored in :py:obj:`sgd.layers[0].W` and
  :py:obj:`sgd.layers[1].b`. The :py:obj:`approx` parameter determines
  whether to use an approximate sigmoid function. Setting it to 5 uses
  a five-piece approximation instead of a three-piece one.
  
  A simple network for MNIST using two dense layers can be trained as
  follows::
  
      sgd = ml.SGD([ml.Dense(60000, 784, 128, activation='relu'),
                    ml.Dense(60000, 128, 10),
                    ml.MultiOutput(60000, 10)], n_epochs,
                   report_loss=True)
      sgd.layers[0].X.input_from(0)
      sgd.layers[1].Y.input_from(1)
      sgd.reset()
      sgd.run()
  
  See `this repository <https://github.com/csiro-mlai/mnist-mpc>`_
  for scripts importing MNIST training data and further examples.
  
  Inference can be run as follows::
  
      data = sfix.Matrix(n_test, n_features)
      data.input_from(0)
      res = sgd.eval(data)
      print_ln('Results: %s', [x.reveal() for x in res])
  
  For inference/classification, this module offers the layers necessary
  for neural networks such as DenseNet, ResNet, and SqueezeNet. A
  minimal example using input from player 0 and model from player 1
  looks as follows::
  
      graph = Optimizer()
      graph.layers = layers
      layers[0].X.input_from(0)
      for layer in layers:
          layer.input_from(1)
      graph.forward(1)
      res = layers[-1].Y
  
  See the `readme <https://github.com/data61/MP-SPDZ/#tensorflow-inference>`_ for
  an example of how to run MP-SPDZ on TensorFlow graphs.
  """
  
  import math
  import re
  
  from Compiler import mpc_math, util
  from Compiler.types import *
  from Compiler.types import _unreduced_squant
  from Compiler.library import *
  from Compiler.util import is_zero, tree_reduce
  from Compiler.comparison import CarryOutRawLE
  from Compiler.GC.types import sbitint
  from functools import reduce
  
  def log_e(x):
      return mpc_math.log_fx(x, math.e)
  
  def exp(x):
      return mpc_math.pow_fx(math.e, x)
  
  def get_limit(x):
      exp_limit = 2 ** (x.k - x.f - 1)
      return math.log(exp_limit)
  
  def sanitize(x, raw, lower, upper):
      limit = get_limit(x)
      res = (x > limit).if_else(upper, raw)
      return (x < -limit).if_else(lower, res)
  
  def sigmoid(x):
      """ Sigmoid function.
  
      :param x: sfix """
      return sigmoid_from_e_x(x, exp(-x))
  
  def sigmoid_from_e_x(x, e_x):
      return sanitize(x, 1 / (1 + e_x), 0, 1)
  
  def sigmoid_prime(x):
      """ Sigmoid derivative.
  
      :param x: sfix """
      sx = sigmoid(x)
      return sx * (1 - sx)
  
  @vectorize
  def approx_sigmoid(x, n=3):
      """ Piece-wise approximate sigmoid as in
      `Dahl et al. <https://arxiv.org/abs/1810.08130>`_
  
      :param x: input
      :param n: number of pieces, 3 (default) or 5
      """
      if n == 5:
          cuts = [-5, -2.5, 2.5, 5]
          le = [0] + [x <= cut for cut in cuts] + [1]
          select = [le[i + 1] - le[i] for i in range(5)]
          outputs = [cfix(10 ** -4),
                     0.02776 * x + 0.145,
                     0.17 * x + 0.5,
                     0.02776 * x + 0.85498,
                     cfix(1 - 10 ** -4)]
          return sum(a * b for a, b in zip(select, outputs))
      else:
          a = x < -0.5
          b = x > 0.5
          return a.if_else(0, b.if_else(1, 0.5 + x))
  
  def lse_0_from_e_x(x, e_x):
      return sanitize(-x, log_e(1 + e_x), x + 2 ** -x.f, 0)
  
  def lse_0(x):
      return lse_0_from_e_x(x, exp(x))
  
  def approx_lse_0(x, n=3):
      assert n != 5
      a = x < -0.5
      b = x > 0.5
      return a.if_else(0, b.if_else(x, 0.5 * (x + 0.5) ** 2)) - x
  
  def relu_prime(x):
      """ ReLU derivative. """
      return (0 <= x)
  
  def relu(x):
      """ ReLU function (maximum of input and zero). """
      return (0 < x).if_else(x, 0)
  
  def argmax(x):
      """ Compute index of maximum element.
  
      :param x: iterable
      :returns: sint
      """
      def op(a, b):
          comp = (a[1] > b[1])
          return comp.if_else(a[0], b[0]), comp.if_else(a[1], b[1])
      return tree_reduce(op, enumerate(x))[0]
  
  report_progress = False
  
  def progress(x):
      if report_progress:
          print_ln(x)
          time()
  
  def set_n_threads(n_threads):
      Layer.n_threads = n_threads
      Optimizer.n_threads = n_threads
  
  def _no_mem_warnings(function):
      def wrapper(*args, **kwargs):
          get_program().warn_about_mem.append(False)
          res = function(*args, **kwargs)
          get_program().warn_about_mem.pop()
          return res
      return wrapper
  
  class Tensor(MultiArray):
      def __init__(self, *args, **kwargs):
          kwargs['alloc'] = False
          super(Tensor, self).__init__(*args, **kwargs)
  
      def input_from(self, *args, **kwargs):
          self.alloc()
          super(Tensor, self).input_from(*args, **kwargs)
  
      def __getitem__(self, *args):
          self.alloc()
          return super(Tensor, self).__getitem__(*args)
  
      def assign_vector(self, *args):
          self.alloc()
          return super(Tensor, self).assign_vector(*args)
  
      def assign_vector_by_indices(self, *args):
          self.alloc()
          return super(Tensor, self).assign_vector_by_indices(*args)
  
  class Layer:
      n_threads = 1
      inputs = []
      input_bias = True
      thetas = lambda self: ()
      debug_output = False
      back_batch_size = 128
  
      @property
      def shape(self):
          return list(self._Y.sizes)
  
      @property
      def X(self):
          self._X.alloc()
          return self._X
  
      @X.setter
      def X(self, value):
          self._X = value
  
      @property
      def Y(self):
          self._Y.alloc()
          return self._Y
  
      @Y.setter
      def Y(self, value):
          self._Y = value
  
      def forward(self, batch=None, training=None):
          if batch is None:
              batch = Array.create_from(regint(0))
          self._forward(batch)
  
      def __str__(self):
          return type(self).__name__ + str(self._Y.sizes)
  
  class NoVariableLayer(Layer):
      input_from = lambda *args, **kwargs: None
      output_weights = lambda *args: None
  
      nablas = lambda self: ()
      reset = lambda self: None
  
  class Output(NoVariableLayer):
      """ Fixed-point logistic regression output layer.
  
      :param N: number of examples
      :param approx: :py:obj:`False` (default) or parameter for :py:obj:`approx_sigmoid`
      """
      n_outputs = 2
  
      @classmethod
      def from_args(cls, N, program):
          res = cls(N, approx='approx' in program.args)
          res.compute_loss = not 'no_loss' in program.args
          return res
  
      def __init__(self, N, debug=False, approx=False):
          self.N = N
          self.X = sfix.Array(N)
          self.Y = sfix.Array(N)
          self.nabla_X = sfix.Array(N)
          self.l = MemValue(sfix(-1))
          self.e_x = sfix.Array(N)
          self.debug = debug
          self.weights = None
          self.approx = approx
          self.compute_loss = True
  
      def divisor(self, divisor, size):
          return cfix(1.0 / divisor, size=size)
  
      def _forward(self, batch):
          if self.approx == 5:
              self.l.write(999)
              return
          N = len(batch)
          lse = sfix.Array(N)
          @multithread(self.n_threads, N)
          def _(base, size):
              x = self.X.get_vector(base, size)
              y = self.Y.get(batch.get_vector(base, size))
              if self.approx:
                  if self.compute_loss:
                      lse.assign(approx_lse_0(x, self.approx) + x * (1 - y), base)
                  return
              e_x = exp(-x)
              self.e_x.assign(e_x, base)
              if self.compute_loss:
                  lse.assign(lse_0_from_e_x(-x, e_x) + x * (1 - y), base)
          self.l.write(sum(lse) * \
                       self.divisor(N, 1))
  
      def eval(self, size, base=0):
          if self.approx:
              return approx_sigmoid(self.X.get_vector(base, size), self.approx)
          else:
              return sigmoid_from_e_x(self.X.get_vector(base, size),
                                      self.e_x.get_vector(base, size))
  
      def backward(self, batch):
          N = len(batch)
          @multithread(self.n_threads, N)
          def _(base, size):
              diff = self.eval(size, base) - \
                     self.Y.get(batch.get_vector(base, size))
              assert sfix.f == cfix.f
              if self.weights is not None:
                  assert N == len(self.weights)
                  diff *= self.weights.get_vector(base, size)
                  assert self.weight_total == N
              self.nabla_X.assign(diff, base)
          # @for_range_opt(len(diff))
          # def _(i):
          #     self.nabla_X[i] = self.nabla_X[i] * self.weights[i]
          if self.debug_output:
              print_ln('sigmoid X %s', self.X.reveal_nested())
              print_ln('sigmoid nabla %s', self.nabla_X.reveal_nested())
              print_ln('batch %s', batch.reveal_nested())
  
      def set_weights(self, weights):
          self.weights = cfix.Array(len(weights))
          self.weights.assign(weights)
          self.weight_total = sum(weights)
  
      def average_loss(self, N):
          return self.l.reveal()
  
      def reveal_correctness(self, n=None, Y=None, debug=False):
          if n is None:
              n = self.X.sizes[0]
          if Y is None:
              Y = self.Y
          n_correct = MemValue(0)
          n_printed = MemValue(0)
          @for_range_opt(n)
          def _(i):
              truth = Y[i].reveal()
              b = self.X[i].reveal()
              if debug:
                  nabla = self.nabla_X[i].reveal()
              guess = b > 0
              correct = truth == guess
              n_correct.iadd(correct)
              if debug:
                  to_print = (1 - correct) * (n_printed < 10)
                  n_printed.iadd(to_print)
                  print_ln_if(to_print, '%s: %s %s %s %s',
  	                    i, truth, guess, b, nabla)
          return n_correct
  
  class MultiOutputBase(NoVariableLayer):
      def __init__(self, N, d_out, approx=False, debug=False):
          self.X = sfix.Matrix(N, d_out)
          self.Y = sint.Matrix(N, d_out)
          self.nabla_X = sfix.Matrix(N, d_out)
          self.l = MemValue(sfix(-1))
          self.losses = sfix.Array(N)
          self.approx = None
          self.N = N
          self.d_out = d_out
          self.compute_loss = True
  
      def eval(self, N):
          d_out = self.X.sizes[1]
          res = sfix.Matrix(N, d_out)
          res.assign_vector(self.X.get_part_vector(0, N))
          return res
  
      def average_loss(self, N):
          return sum(self.losses.get_vector(0, N)).reveal() / N
  
      def reveal_correctness(self, n=None, Y=None, debug=False):
          if n is None:
              n = self.X.sizes[0]
          if Y is None:
              Y = self.Y
          n_printed = MemValue(0)
          assert n <= len(self.X)
          assert n <= len(Y)
          Y.address = MemValue.if_necessary(Y.address)
          @map_sum(None if debug else self.n_threads, None, n, 1, regint)
          def _(i):
              a = Y[i].reveal_list()
              b = self.X[i].reveal_list()
              if debug:
                  loss = self.losses[i].reveal()
                  exp = self.get_extra_debugging(i)
                  nabla = self.nabla_X[i].reveal_list()
              truth = argmax(a)
              guess = argmax(b)
              correct = truth == guess
              if debug:
                  to_print = (1 - correct) * (n_printed < 10)
                  n_printed.iadd(to_print)
                  print_ln_if(to_print, '%s: %s %s %s %s %s %s',
  	                    i, truth, guess, loss, b, exp, nabla)
              return correct
          return _()
  
      @property
      def n_outputs(self):
          return self.d_out
  
      def get_extra_debugging(self, i):
          return ''
  
      @staticmethod
      def from_args(program, N, n_output):
          if 'relu_out' in program.args:
              res = ReluMultiOutput(N, n_output)
          else:
              res = MultiOutput(N, n_output, approx='approx' in program.args)
              res.cheaper_loss = 'mse' in program.args
          res.compute_loss = not 'no_loss' in program.args
          for arg in program.args:
              m = re.match('approx=(.*)', arg)
              if m:
                  res.approx = float(m.group(1))
          return res
  
  class MultiOutput(MultiOutputBase):
      """
      Output layer for multi-class classification with softmax and cross entropy.
  
      :param N: number of examples
      :param d_out: number of classes
      :param approx: use ReLU division instead of softmax for the loss
      """
      def __init__(self, N, d_out, approx=False, debug=False):
          MultiOutputBase.__init__(self, N, d_out)
          self.exp = sfix.Matrix(N, d_out)
          self.approx = approx
          self.positives = sint.Matrix(N, d_out)
          self.relus = sfix.Matrix(N, d_out)
          self.cheaper_loss = False
          self.debug = debug
          self.true_X = sfix.Array(N)
  
      def _forward(self, batch):
          N = len(batch)
          d_out = self.X.sizes[1]
          tmp = self.losses
          @for_range_opt_multithread(self.n_threads, N)
          def _(i):
              if self.approx:
                  if self.cheaper_loss or isinstance(self.approx, float):
                      limit = 0
                  else:
                      limit = 0.1
                  positives = self.X[i].get_vector() > limit
                  relus = positives.if_else(self.X[i].get_vector(), 0)
                  self.positives[i].assign_vector(positives)
                  self.relus[i].assign_vector(relus)
                  if self.compute_loss:
                      if self.cheaper_loss:
                          s = sum(relus)
                          tmp[i] = sum((self.Y[batch[i]][j] * s - relus[j]) ** 2
                                       for j in range(d_out)) / s ** 2 * 0.5
                      else:
                          div = relus / sum(relus).expand_to_vector(d_out)
                          self.losses[i] = -sfix.dot_product(
                              self.Y[batch[i]].get_vector(), log_e(div))
              else:
                  m = util.max(self.X[i])
                  mv = m.expand_to_vector(d_out)
                  x = self.X[i].get_vector()
                  e = (x - mv > -get_limit(x)).if_else(exp(x - mv), 0)
                  self.exp[i].assign_vector(e)
                  if self.compute_loss:
                      true_X = sfix.dot_product(self.Y[batch[i]], self.X[i])
                      tmp[i] = m + log_e(sum(e)) - true_X
                      self.true_X[i] = true_X
          self.l.write(sum(tmp.get_vector(0, N)) / N)
  
      def eval(self, N):
          d_out = self.X.sizes[1]
          res = sfix.Matrix(N, d_out)
          if self.approx:
              @for_range_opt_multithread(self.n_threads, N)
              def _(i):
                  relus = (self.X[i].get_vector() > 0).if_else(
                      self.X[i].get_vector(), 0)
                  res[i].assign_vector(relus / sum(relus).expand_to_vector(d_out))
              return res
          @for_range_opt_multithread(self.n_threads, N)
          def _(i):
              x = self.X[i].get_vector() - \
                  util.max(self.X[i].get_vector()).expand_to_vector(d_out)
              e = exp(x)
              res[i].assign_vector(e / sum(e).expand_to_vector(d_out))
          return res
  
      def backward(self, batch):
          d_out = self.X.sizes[1]
          if self.approx:
              @for_range_opt_multithread(self.n_threads, len(batch))
              def _(i):
                  if self.cheaper_loss:
                      s = sum(self.relus[i])
                      ss = s * s * s
                      inv = 1 / ss
                      @for_range_opt(d_out)
                      def _(j):
                          res = 0
                          for k in range(d_out):
                              relu = self.relus[i][k]
                              summand = relu - self.Y[batch[i]][k] * s
                              summand *= (sfix.from_sint(j == k) - relu)
                              res += summand
                          fallback = -self.Y[batch[i]][j]
                          res *= inv
                          self.nabla_X[i][j] = self.positives[i][j].if_else(res, fallback)
                      return
                  relus = self.relus[i].get_vector()
                  if isinstance(self.approx, float):
                      relus += self.approx
                  positives = self.positives[i].get_vector()
                  inv = (1 / sum(relus)).expand_to_vector(d_out)
                  truths = self.Y[batch[i]].get_vector()
                  raw = truths / relus - inv
                  self.nabla_X[i] = -positives.if_else(raw, truths)
              self.maybe_debug_backward(batch)
              return
          @for_range_opt_multithread(self.n_threads, len(batch))
          def _(i):
              for j in range(d_out):
                  dividend = self.exp[i][j]
                  divisor = sum(self.exp[i])
                  div = (divisor > 0.1).if_else(dividend / divisor, 0)
                  self.nabla_X[i][j] = (-self.Y[batch[i]][j] + div)
          self.maybe_debug_backward(batch)
  
      def maybe_debug_backward(self, batch):
          if self.debug:
              @for_range(len(batch))
              def _(i):
                  check = 0
                  for j in range(self.X.sizes[1]):
                      to_check = self.nabla_X[i][j].reveal()
                      check += (to_check > len(batch)) + (to_check < -len(batch))
                  print_ln_if(check, 'X %s', self.X[i].reveal_nested())
                  print_ln_if(check, 'exp %s', self.exp[i].reveal_nested())
                  print_ln_if(check, 'nabla X %s',
                              self.nabla_X[i].reveal_nested())
  
      def get_extra_debugging(self, i):
          if self.approx:
              return self.relus[i].reveal_list()
          else:
              return self.exp[i].reveal_list()
  
  class ReluMultiOutput(MultiOutputBase):
      """
      Output layer for multi-class classification with back-propagation
      based on ReLU division.
  
      :param N: number of examples
      :param d_out: number of classes
      """
      def forward(self, batch, training=None):
          self.l.write(999)
  
      def backward(self, batch):
          N = len(batch)
          d_out = self.X.sizes[1]
          relus = sfix.Matrix(N, d_out)
          @for_range_opt_multithread(self.n_threads, len(batch))
          def _(i):
              positives = self.X[i].get_vector() > 0
              relus = positives.if_else(self.X[i].get_vector(), 0)
              s = sum(relus)
              inv = 1 / s
              prod = relus * inv
              res = prod - self.Y[batch[i]].get_vector()
              self.nabla_X[i].assign_vector(res)
  
  class DenseBase(Layer):
      thetas = lambda self: (self.W, self.b)
      nablas = lambda self: (self.nabla_W, self.nabla_b)
  
      def output_weights(self):
          print_ln('%s', self.W.reveal_nested())
          print_ln('%s', self.b.reveal_nested())
  
      def backward_params(self, f_schur_Y, batch):
          N = len(batch)
          tmp = Matrix(self.d_in, self.d_out, unreduced_sfix)
  
          @multithread(self.n_threads, self.d_in)
          def _(base, size):
              A = sfix.Matrix(self.N, self.d_out, address=f_schur_Y.address)
              B = sfix.Matrix(self.N, self.d_in, address=self.X.address)
              mp = B.direct_trans_mul(A, reduce=False,
                                      indices=(regint.inc(size, base),
                                               batch.get_vector(),
                                               regint.inc(N),
                                               regint.inc(self.d_out)))
              tmp.assign_part_vector(mp, base)
  
          progress('nabla W (matmul)')
  
          if self.d_in * self.d_out < 200000:
              print('reduce at once')
              @multithread(self.n_threads, self.d_in * self.d_out)
              def _(base, size):
                  self.nabla_W.assign_vector(
                      tmp.get_vector(base, size).reduce_after_mul(), base=base)
          else:
              @for_range_opt(self.d_in)
              def _(i):
                  self.nabla_W[i] = tmp[i].get_vector().reduce_after_mul()
  
          progress('nabla W')
  
          self.nabla_b.assign_vector(sum(sum(f_schur_Y[k][j].get_vector()
                                             for k in range(N))
                                         for j in range(self.d)))
  
          progress('nabla b')
  
          if self.debug_output:
              print_ln('dense nabla Y %s', self.nabla_Y.reveal_nested())
              print_ln('dense W %s', self.W.reveal_nested())
              print_ln('dense nabla X %s', self.nabla_X.reveal_nested())
          if self.debug:
              limit = N * self.debug
              @for_range_opt(self.d_in)
              def _(i):
                  @for_range_opt(self.d_out)
                  def _(j):
                      to_check = self.nabla_W[i][j].reveal()
                      check = sum(to_check > limit) + sum(to_check < -limit)
                      @if_(check)
                      def _():
                          print_ln('nabla W %s %s %s: %s', i, j, self.W.sizes, to_check)
                          print_ln('Y %s', [f_schur_Y[k][0][j].reveal()
                                            for k in range(N)])
                          print_ln('X %s', [self.X[k][0][i].reveal()
                                            for k in range(N)])
              @for_range_opt(self.d_out)
              def _(j):
                  to_check = self.nabla_b[j].reveal()
                  check = sum(to_check > limit) + sum(to_check < -limit)
                  @if_(check)
                  def _():
                      print_ln('nabla b %s %s: %s', j, len(self.b), to_check)
                      print_ln('Y %s', [f_schur_Y[k][0][j].reveal()
                                        for k in range(N)])
              @for_range_opt(len(batch))
              def _(i):
                  to_check = self.nabla_X[i].get_vector().reveal()
                  check = sum(to_check > limit) + sum(to_check < -limit)
                  @if_(check)
                  def _():
                      print_ln('X %s %s', i, self.X[i].reveal_nested())
                      print_ln('Y %s %s', i, f_schur_Y[i].reveal_nested())
  
  class Dense(DenseBase):
      """ Fixed-point dense (matrix multiplication) layer.
  
      :param N: number of examples
      :param d_in: input dimension
      :param d_out: output dimension
      """
      def __init__(self, N, d_in, d_out, d=1, activation='id', debug=False):
          if activation == 'id':
              self.activation_layer = None
          elif activation == 'relu':
              self.activation_layer = Relu([N, d, d_out])
          elif activation == 'square':
              self.activation_layer = Square([N, d, d_out])
          else:
              raise CompilerError('activation not supported: %s', activation)
  
          self.N = N
          self.d_in = d_in
          self.d_out = d_out
          self.d = d
  
          self.X = MultiArray([N, d, d_in], sfix)
          self.Y = MultiArray([N, d, d_out], sfix)
          self.W = Tensor([d_in, d_out], sfix)
          self.b = sfix.Array(d_out)
  
          back_N = min(N, self.back_batch_size)
          self.nabla_Y = MultiArray([back_N, d, d_out], sfix)
          self.nabla_X = MultiArray([back_N, d, d_in], sfix)
          self.nabla_W = sfix.Matrix(d_in, d_out)
          self.nabla_b = sfix.Array(d_out)
  
          self.debug = debug
  
          l = self.activation_layer
          if l:
              self.f_input = l.X
              l.Y = self.Y
              l.nabla_Y = self.nabla_Y
          else:
              self.f_input = self.Y
  
      def reset(self):
          d_in = self.d_in
          d_out = self.d_out
          r = math.sqrt(6.0 / (d_in + d_out))
          print('Initializing dense weights in [%f,%f]' % (-r, r))
          self.W.assign_vector(sfix.get_random(-r, r, size=self.W.total_size()))
          self.b.assign_all(0)
  
      def input_from(self, player, raw=False):
          self.W.input_from(player, raw=raw)
          if self.input_bias:
              self.b.input_from(player, raw=raw)
  
      def compute_f_input(self, batch):
          N = len(batch)
          assert self.d == 1
          if self.input_bias:
              prod = MultiArray([N, self.d, self.d_out], sfix)
          else:
              prod = self.f_input
          max_size = program.Program.prog.budget // self.d_out
          @multithread(self.n_threads, N, max_size)
          def _(base, size):
              X_sub = sfix.Matrix(self.N, self.d_in, address=self.X.address)
              prod.assign_part_vector(
                  X_sub.direct_mul(self.W, indices=(
                      batch.get_vector(base, size), regint.inc(self.d_in),
                      regint.inc(self.d_in), regint.inc(self.d_out))), base)
  
          if self.input_bias:
              if self.d_out == 1:
                  @multithread(self.n_threads, N)
                  def _(base, size):
                      v = prod.get_vector(base, size) + self.b.expand_to_vector(0, size)
                      self.f_input.assign_vector(v, base)
              else:
                  @for_range_multithread(self.n_threads, 100, N)
                  def _(i):
                      v = prod[i].get_vector() + self.b.get_vector()
                      self.f_input[i].assign_vector(v)
          progress('f input')
  
      def _forward(self, batch=None):
          if batch is None:
              batch = regint.Array(self.N)
              batch.assign(regint.inc(self.N))
          self.compute_f_input(batch=batch)
          if self.activation_layer:
              self.activation_layer.forward(batch)
          if self.debug_output:
              print_ln('dense X %s', self.X.reveal_nested())
              print_ln('dense W %s', self.W.reveal_nested())
              print_ln('dense b %s', self.b.reveal_nested())
              print_ln('dense Y %s', self.Y.reveal_nested())
          if self.debug:
              limit = self.debug
              @for_range_opt(len(batch))
              def _(i):
                  @for_range_opt(self.d_out)
                  def _(j):
                      to_check = self.Y[i][0][j].reveal()
                      check = to_check > limit
                      @if_(check)
                      def _():
                          print_ln('dense Y %s %s %s %s', i, j, self.W.sizes, to_check)
                          print_ln('X %s', self.X[i].reveal_nested())
                          print_ln('W %s',
                                   [self.W[k][j].reveal() for k in range(self.d_in)])
  
      def backward(self, compute_nabla_X=True, batch=None):
          N = len(batch)
          d = self.d
          d_out = self.d_out
          X = self.X
          Y = self.Y
          W = self.W
          b = self.b
          nabla_X = self.nabla_X
          nabla_Y = self.nabla_Y
          nabla_W = self.nabla_W
          nabla_b = self.nabla_b
  
          if self.activation_layer:
              self.activation_layer.backward(batch)
              f_schur_Y = self.activation_layer.nabla_X
          else:
              f_schur_Y = nabla_Y
  
          if compute_nabla_X:
              @multithread(self.n_threads, N)
              def _(base, size):
                  B = sfix.Matrix(N, d_out, address=f_schur_Y.address)
                  nabla_X.assign_part_vector(
                      B.direct_mul_trans(W, indices=(regint.inc(size, base),
                                                     regint.inc(self.d_out),
                                                     regint.inc(self.d_out),
                                                     regint.inc(self.d_in))),
                      base)
  
              progress('nabla X')
  
          self.backward_params(f_schur_Y, batch=batch)
  
  class QuantizedDense(DenseBase):
      def __init__(self, N, d_in, d_out):
          self.N = N
          self.d_in = d_in
          self.d_out = d_out
          self.d = 1
          self.H = math.sqrt(1.5 / (d_in + d_out))
  
          self.W = sfix.Matrix(d_in, d_out)
          self.nabla_W = self.W.same_shape()
          self.T = sint.Matrix(d_in, d_out)
          self.b = sfix.Array(d_out)
          self.nabla_b = self.b.same_shape()
  
          self.X = MultiArray([N, 1, d_in], sfix)
          self.Y = MultiArray([N, 1, d_out], sfix)
          self.nabla_Y = self.Y.same_shape()
  
      def reset(self):
          @for_range(self.d_in)
          def _(i):
              @for_range(self.d_out)
              def _(j):
                  self.W[i][j] = sfix.get_random(-1, 1)
          self.b.assign_all(0)
  
      def _forward(self):
          @for_range_opt(self.d_in)
          def _(i):
              @for_range_opt(self.d_out)
              def _(j):
                  over = self.W[i][j] > 0.5
                  under = self.W[i][j] < -0.5
                  self.T[i][j] = over.if_else(1, under.if_else(-1, 0))
                  over = self.W[i][j] > 1
                  under = self.W[i][j] < -1
                  self.W[i][j] = over.if_else(1, under.if_else(-1, self.W[i][j]))
          @for_range_opt(self.N)
          def _(i):
              assert self.d_out == 1
              self.Y[i][0][0] = self.b[0] + self.H * sfix._new(
                  sint.dot_product([self.T[j][0] for j in range(self.d_in)],
                                   [self.X[i][0][j].v for j in range(self.d_in)]))
  
      def backward(self, compute_nabla_X=False):
          assert not compute_nabla_X
          self.backward_params(self.nabla_Y)
  
  class Dropout(NoVariableLayer):
      """ Dropout layer.
  
      :param N: number of examples
      :param d1: total dimension
      :param alpha: probability (power of two)
      """
      def __init__(self, N, d1, d2=1, alpha=0.5):
          self.N = N
          self.d1 = d1
          self.d2 = d2
          self.X = MultiArray([N, d1, d2], sfix)
          self.Y = MultiArray([N, d1, d2], sfix)
          self.nabla_Y = MultiArray([N, d1, d2], sfix)
          self.nabla_X = MultiArray([N, d1, d2], sfix)
          self.alpha = alpha
          self.B = MultiArray([N, d1, d2], sint)
  
      def forward(self, batch, training=False):
          if training:
              n_bits = -math.log(self.alpha, 2)
              assert n_bits == int(n_bits)
              n_bits = int(n_bits)
              @for_range_opt_multithread(self.n_threads, len(batch))
              def _(i):
                  size = self.d1 * self.d2
                  self.B[i].assign_vector(util.tree_reduce(
                      util.or_op, (sint.get_random_bit(size=size)
                                   for i in range(n_bits))))
              @for_range_opt_multithread(self.n_threads, len(batch))
              def _(i):
                  self.Y[i].assign_vector(1 / (1 - self.alpha) *
                      self.X[batch[i]].get_vector() * self.B[i].get_vector())
          else:
              @for_range(len(batch))
              def _(i):
                  self.Y[i] = self.X[batch[i]]
          if self.debug_output:
              print_ln('dropout X %s', self.X.reveal_nested())
              print_ln('dropout Y %s', self.Y.reveal_nested())
  
      def backward(self, compute_nabla_X=True, batch=None):
          if compute_nabla_X:
              @for_range_opt_multithread(self.n_threads, len(batch))
              def _(i):
                  self.nabla_X[batch[i]].assign_vector(
                      self.nabla_Y[i].get_vector() * self.B[i].get_vector())
          if self.debug_output:
              print_ln('dropout nabla_Y %s', self.nabla_Y.reveal_nested())
              print_ln('dropout nabla_X %s', self.nabla_X.reveal_nested())
  
  class ElementWiseLayer(NoVariableLayer):
      def __init__(self, shape, inputs=None):
          self.X = Tensor(shape, sfix)
          self.Y = Tensor(shape, sfix)
          backward_shape = list(shape)
          backward_shape[0] = min(shape[0], self.back_batch_size)
          self.nabla_X = Tensor(backward_shape, sfix)
          self.nabla_Y = Tensor(backward_shape, sfix)
          self.inputs = inputs
  
      def f_part(self, base, size):
          return self.f(self.X.get_part_vector(base, size))
  
      def f_prime_part(self, base, size):
          return self.f_prime(self.Y.get_part_vector(base, size))
  
      def _forward(self, batch=[0]):
          n_per_item = reduce(operator.mul, self.X.sizes[1:])
          @multithread(self.n_threads, len(batch), max(1, 1000 // n_per_item))
          def _(base, size):
              self.Y.assign_part_vector(self.f_part(base, size), base)
  
          if self.debug_output:
              name = self
              @for_range(len(batch))
              def _(i):
                  print_ln('%s X %s %s', name, i, self.X[i].reveal_nested())
                  print_ln('%s Y %s %s', name, i, self.Y[i].reveal_nested())
  
      def backward(self, batch):
          f_prime_bit = MultiArray(self.X.sizes, self.prime_type)
          n_elements = len(batch) * reduce(operator.mul, f_prime_bit.sizes[1:])
  
          @multithread(self.n_threads, n_elements)
          def _(base, size):
              f_prime_bit.assign_vector(self.f_prime_part(base, size), base)
  
          progress('f prime')
  
          @multithread(self.n_threads, n_elements)
          def _(base, size):
              self.nabla_X.assign_vector(self.nabla_Y.get_vector(base, size) *
                                         f_prime_bit.get_vector(base, size),
                                         base)
  
          progress('f prime schur Y')
  
          if self.debug_output:
              name = self
              @for_range(len(batch))
              def _(i):
                  print_ln('%s X %s %s', name, i, self.X[i].reveal_nested())
                  print_ln('%s f_prime %s %s', name, i, f_prime_bit[i].reveal_nested())
                  print_ln('%s nabla Y %s %s', name, i, self.nabla_Y[i].reveal_nested())
                  print_ln('%s nabla X %s %s', name, i, self.nabla_X[i].reveal_nested())
  
  class Relu(ElementWiseLayer):
      """ Fixed-point ReLU layer.
  
      :param shape: input/output shape (tuple/list of int)
      """
      f = staticmethod(relu)
      f_prime = staticmethod(relu_prime)
      prime_type = sint
      comparisons = None
  
      def __init__(self, shape, inputs=None):
          super(Relu, self).__init__(shape)
          self.comparisons = MultiArray(shape, sint)
  
      def f_part(self, base, size):
          x = self.X.get_part_vector(base, size)
          c = x > 0
          self.comparisons.assign_part_vector(c, base)
          return c.if_else(x, 0)
  
      def f_prime_part(self, base, size):
          return self.comparisons.get_vector(base, size)
  
  class Square(ElementWiseLayer):
      """ Fixed-point square layer.
  
      :param shape: input/output shape (tuple/list of int)
      """
      f = staticmethod(lambda x: x ** 2)
      f_prime = staticmethod(lambda x: cfix(2, size=x.size) * x)
      prime_type = sfix
  
  class MaxPool(NoVariableLayer):
      """ Fixed-point MaxPool layer.
  
      :param shape: input shape (tuple/list of four int)
      :param strides: strides (tuple/list of four int, first and last must be 1)
      :param ksize: kernel size (tuple/list of four int, first and last must be 1)
      :param padding: :py:obj:`'VALID'` (default) or :py:obj:`'SAME'`
      """
      def __init__(self, shape, strides=(1, 2, 2, 1), ksize=(1, 2, 2, 1),
                   padding='VALID'):
          assert len(shape) == 4
          for x in strides, ksize:
              for i in 0, 3:
                  assert x[i] == 1
          self.X = Tensor(shape, sfix)
          if padding == 'SAME':
              output_shape = [int(math.ceil(shape[i] / strides[i])) for i in range(4)]
          else:
              output_shape = [(shape[i] - ksize[i]) // strides[i] + 1 for i in range(4)]
          self.Y = Tensor(output_shape, sfix)
          self.strides = strides
          self.ksize = ksize
          self.nabla_X = Tensor(shape, sfix)
          self.nabla_Y = Tensor(output_shape, sfix)
          self.N = shape[0]
          self.comparisons = MultiArray([self.N, self.X.sizes[3],
                                         ksize[1] * ksize[2]], sint)
  
      def _forward(self, batch):
          def process(pool, bi, k, i, j):
              def m(a, b):
                  c = a[0] > b[0]
                  l = [c * x for x in a[1]]
                  l += [(1 - c) * x for x in b[1]]
                  return c.if_else(a[0], b[0]), l
              red = util.tree_reduce(m, [(x[0], [1]) for x in pool])
              self.Y[bi][i][j][k] = red[0]
              for i, x in enumerate(red[1]):
                  self.comparisons[bi][k][i] = x
          self.traverse(batch, process)
  
      def backward(self, compute_nabla_X=True, batch=None):
          if compute_nabla_X:
              self.nabla_X.alloc()
              def process(pool, bi, k, i, j):
                  for (x, h_in, w_in, h, w), c in zip(pool,
                                                      self.comparisons[bi][k]):
                      hh = h * h_in
                      ww = w * w_in
                      self.nabla_X[bi][hh][ww][k] = \
                          util.if_else(h_in * w_in, c * self.nabla_Y[bi][i][j][k],
                                       self.nabla_X[bi][hh][ww][k])
          self.traverse(batch, process)
  
      def traverse(self, batch, process):
          need_padding = [self.strides[i] * (self.Y.sizes[i] - 1) + self.ksize[i] >
                          self.X.sizes[i] for i in range(4)]
          @for_range_opt_multithread(self.n_threads,
                                     [len(batch), self.X.sizes[3]])
          def _(l, k):
              bi = batch[l]
              @for_range_opt(self.Y.sizes[1])
              def _(i):
                  h_base = self.strides[1] * i
                  @for_range_opt(self.Y.sizes[2])
                  def _(j):
                      w_base = self.strides[2] * j
                      pool = []
                      for ii in range(self.ksize[1]):
                          h = h_base + ii
                          if need_padding[1]:
                              h_in = h < self.X.sizes[1]
                          else:
                              h_in = True
                          for jj in range(self.ksize[2]):
                              w = w_base + jj
                              if need_padding[2]:
                                  w_in = w < self.X.sizes[2]
                              else:
                                  w_in = True
                              if not is_zero(h_in * w_in):
                                  pool.append([h_in * w_in * self.X[bi][h_in * h]
                                               [w_in * w][k], h_in, w_in, h, w])
                      process(pool, bi, k, i, j)
  
  
  class Argmax(NoVariableLayer):
      """ Fixed-point Argmax layer.
  
      :param shape: input shape (tuple/list of two int)
      """
      def __init__(self, shape):
          assert len(shape) == 2
          self.X = MultiArray(shape, sfix)
          self.Y = Array(shape[0], sint)
  
      def _forward(self, batch=[0]):
          assert len(batch) == 1
          self.Y[batch[0]] = argmax(self.X[batch[0]])
  
  class Concat(NoVariableLayer):
      """ Fixed-point concatentation layer.
  
      :param inputs: two input layers (tuple/list)
      :param dimension: dimension for concatenation (must be 3)
      """
      def __init__(self, inputs, dimension):
          self.inputs = inputs
          self.dimension = dimension
          shapes = [inp.shape for inp in inputs]
          assert dimension == 3
          assert len(shapes) == 2
          assert len(shapes[0]) == len(shapes[1])
          shape = []
          for i in range(len(shapes[0])):
              if i == dimension:
                  shape.append(shapes[0][i] + shapes[1][i])
              else:
                  assert shapes[0][i] == shapes[1][i]
                  shape.append(shapes[0][i])
          self.Y = Tensor(shape, sfix)
  
      def _forward(self, batch=[0]):
          assert len(batch) == 1
          @for_range_multithread(self.n_threads, 1, self.Y.sizes[1:3])
          def _(i, j):
              X = [x.Y[batch[0]] for x in self.inputs]
              self.Y[batch[0]][i][j].assign_vector(X[0][i][j].get_vector())
              self.Y[batch[0]][i][j].assign_part_vector(
                  X[1][i][j].get_vector(),
                  len(X[0][i][j]))
  
  class Add(NoVariableLayer):
      """ Fixed-point addition layer.
  
      :param inputs: two input layers with same shape (tuple/list)
      """
      def __init__(self, inputs):
          assert len(inputs) > 1
          shape = inputs[0].shape
          for inp in inputs:
              assert inp.shape == shape
          self.Y = Tensor(shape, sfix)
          self.inputs = inputs
  
      def _forward(self, batch=[0]):
          assert len(batch) == 1
          @multithread(self.n_threads, self.Y[0].total_size())
          def _(base, size):
              tmp = sum(inp.Y[batch[0]].get_vector(base, size)
                        for inp in self.inputs)
              self.Y[batch[0]].assign_vector(tmp, base)
  
  class FusedBatchNorm(Layer):
      """ Fixed-point fused batch normalization layer.
  
      :param shape: input/output shape (tuple/list of four int)
      """
      def __init__(self, shape, inputs=None):
          assert len(shape) == 4
          self.X = Tensor(shape, sfix)
          self.Y = Tensor(shape, sfix)
          self.weights = sfix.Array(shape[3])
          self.bias = sfix.Array(shape[3])
          self.inputs = inputs
  
      def input_from(self, player, raw=False):
          self.weights.input_from(player, raw=raw)
          self.bias.input_from(player, raw=raw)
          tmp = sfix.Array(len(self.bias))
          tmp.input_from(player, raw=raw)
          tmp.input_from(player, raw=raw)
  
      def _forward(self, batch=[0]):
          assert len(batch) == 1
          @for_range_opt_multithread(self.n_threads, self.X.sizes[1:3])
          def _(i, j):
              self.Y[batch[0]][i][j].assign_vector(
                  self.X[batch[0]][i][j].get_vector() * self.weights.get_vector()
                  + self.bias.get_vector())
  
  class QuantBase(object):
      bias_before_reduction = True
  
      @staticmethod
      def new_squant():
          class _(squant):
              @classmethod
              def get_params_from(cls, player):
                  cls.set_params(sfloat.get_input_from(player),
                                 sint.get_input_from(player))
              @classmethod
              def get_input_from(cls, player, size=None):
                  return cls._new(sint.get_input_from(player, size=size))
          return _
  
      def const_div(self, acc, n):
          logn = int(math.log(n, 2))
          acc = (acc + n // 2)
          if 2 ** logn == n:
              acc = acc.round(self.output_squant.params.k + logn, logn, nearest=True)
          else:
              acc = acc.int_div(sint(n), self.output_squant.params.k + logn)
          return acc
  
  class FixBase:
      bias_before_reduction = False
  
      @staticmethod
      def new_squant():
          class _(sfix):
              params = None
          return _
  
      def input_params_from(self, player):
          pass
  
      def const_div(self, acc, n):
          return (sfix._new(acc) * self.output_squant(1 / n)).v
  
  class BaseLayer(Layer):
      def __init__(self, input_shape, output_shape, inputs=None):
          self.input_shape = input_shape
          self.output_shape = output_shape
  
          self.input_squant = self.new_squant()
          self.output_squant = self.new_squant()
  
          self.X = Tensor(input_shape, self.input_squant)
          self.Y = Tensor(output_shape, self.output_squant)
  
          back_shapes = list(input_shape), list(output_shape)
          for x in back_shapes:
              x[0] = min(x[0], self.back_batch_size)
  
          self.nabla_X = MultiArray(back_shapes[0], self.input_squant)
          self.nabla_Y = MultiArray(back_shapes[1], self.output_squant)
          self.inputs = inputs
  
      def temp_shape(self):
          return [0]
  
      @property
      def N(self):
          return self.input_shape[0]
  
  class ConvBase(BaseLayer):
      fewer_rounds = True
      use_conv2ds = True
      temp_weights = None
      temp_inputs = None
      thetas = lambda self: (self.weights, self.bias)
      nablas = lambda self: (self.nabla_weights, self.nabla_bias)
  
      @classmethod
      def init_temp(cls, layers):
          size = 0
          for layer in layers:
              size = max(size, reduce(operator.mul, layer.temp_shape()))
          cls.temp_weights = sfix.Array(size)
          cls.temp_inputs = sfix.Array(size)
  
      def __init__(self, input_shape, weight_shape, bias_shape, output_shape, stride,
                   padding='SAME', tf_weight_format=False, inputs=None):
          super(ConvBase, self).__init__(input_shape, output_shape, inputs=inputs)
  
          self.weight_shape = weight_shape
          self.bias_shape = bias_shape
          self.stride = stride
          self.tf_weight_format = tf_weight_format
          if padding == 'SAME':
              # https://web.archive.org/web/20171223022012/https://www.tensorflow.org/api_guides/python/nn
              self.padding = []
              for i in 1, 2:
                  s = stride[i - 1]
                  assert output_shape[i] >= input_shape[i] // s
                  if tf_weight_format:
                      w = weight_shape[i - 1]
                  else:
                      w = weight_shape[i]
                  if (input_shape[i] % stride[1] == 0):
                      pad_total = max(w - s, 0)
                  else:
                      pad_total = max(w - (input_shape[i] % s), 0)
                  self.padding.append(pad_total // 2)
          elif padding == 'VALID':
              self.padding = [0, 0]
          else:
              self.padding = padding
  
          self.weight_squant = self.new_squant()
          self.bias_squant = self.new_squant()
  
          self.weights = Tensor(weight_shape, self.weight_squant)
          self.bias = Array(output_shape[-1], self.bias_squant)
  
          self.nabla_weights = Tensor(weight_shape, self.weight_squant)
          self.nabla_bias = Array(output_shape[-1], self.bias_squant)
  
          self.unreduced = Tensor(self.output_shape, sint, address=self.Y.address)
  
          if tf_weight_format:
              weight_in = weight_shape[2]
          else:
              weight_in = weight_shape[3]
          assert(weight_in == input_shape[-1])
          assert(bias_shape[0] == output_shape[-1])
          assert(len(bias_shape) == 1)
          assert(len(input_shape) == 4)
          assert(len(output_shape) == 4)
          assert(len(weight_shape) == 4)
  
      def input_from(self, player, raw=False):
          self.input_params_from(player)
          self.weights.input_from(player, budget=100000, raw=raw)
          if self.input_bias:
              self.bias.input_from(player, raw=raw)
  
      def output_weights(self):
          print_ln('%s', self.weights.reveal_nested())
          print_ln('%s', self.bias.reveal_nested())
  
      def dot_product(self, iv, wv, out_y, out_x, out_c):
          bias = self.bias[out_c]
          acc = self.output_squant.unreduced_dot_product(iv, wv)
          acc.v += bias.v
          acc.res_params = self.output_squant.params
          #self.Y[0][out_y][out_x][out_c] = acc.reduce_after_mul()
          self.unreduced[0][out_y][out_x][out_c] = acc.v
  
      def reduction(self, batch_length=1):
          unreduced = self.unreduced
          n_summands = self.n_summands()
          #start_timer(2)
          n_outputs = batch_length * reduce(operator.mul, self.output_shape[1:])
          @multithread(self.n_threads, n_outputs,
                       1000 if sfix.round_nearest else 10 ** 6)
          def _(base, n_per_thread):
              res = self.input_squant().unreduced(
                  sint.load_mem(unreduced.address + base,
                                size=n_per_thread),
                  self.weight_squant(),
                  self.output_squant.params,
                  n_summands).reduce_after_mul()
              res.store_in_mem(self.Y.address + base)
          #stop_timer(2)
  
      def temp_shape(self):
          return list(self.output_shape[1:]) + [self.n_summands()]
  
      def prepare_temp(self):
          shape = self.temp_shape()
          inputs = MultiArray(shape, self.input_squant,
                              address=self.temp_inputs)
          weights = MultiArray(shape, self.weight_squant,
                               address=self.temp_weights)
          return inputs, weights
  
  class Conv2d(ConvBase):
      def n_summands(self):
          _, weights_h, weights_w, _ = self.weight_shape
          _, inputs_h, inputs_w, n_channels_in = self.input_shape
          return weights_h * weights_w * n_channels_in
  
      def _forward(self, batch):
          if self.tf_weight_format:
              assert(self.weight_shape[3] == self.output_shape[-1])
              weights_h, weights_w, _, _ = self.weight_shape
          else:
              assert(self.weight_shape[0] == self.output_shape[-1])
              _, weights_h, weights_w, _ = self.weight_shape
          _, inputs_h, inputs_w, n_channels_in = self.input_shape
          _, output_h, output_w, n_channels_out = self.output_shape
  
          stride_h, stride_w = self.stride
          padding_h, padding_w = self.padding
  
          if self.use_conv2ds:
              n_parts = max(1, round(self.n_threads / n_channels_out))
              while len(batch) % n_parts != 0:
                  n_parts -= 1
              print('Convolution in %d parts' % n_parts)
              part_size = len(batch) // n_parts
              @for_range_multithread(self.n_threads, 1, [n_parts, n_channels_out])
              def _(i, j):
                  inputs = self.X.get_slice_vector(
                      batch.get_part(i * part_size, part_size))
                  if self.tf_weight_format:
                      weights = self.weights.get_vector_by_indices(None, None, None, j)
                  else:
                      weights = self.weights.get_part_vector(j)
                  inputs = inputs.pre_mul()
                  weights = weights.pre_mul()
                  res = sint(size = output_h * output_w * part_size)
                  conv2ds(res, inputs, weights, output_h, output_w,
                          inputs_h, inputs_w, weights_h, weights_w,
                          stride_h, stride_w, n_channels_in, padding_h, padding_w,
                          part_size)
                  if self.bias_before_reduction:
                      res += self.bias.expand_to_vector(j, res.size).v
                  else:
                      res += self.bias.expand_to_vector(j, res.size).v << \
                          self.input_squant.f
                  addresses = regint.inc(res.size,
                                         self.unreduced[i * part_size].address + j,
                                         n_channels_out)
                  res.store_in_mem(addresses)
              self.reduction(len(batch))
              if self.debug_output:
                  print_ln('%s weights %s', self, self.weights.reveal_nested())
                  print_ln('%s bias %s', self, self.bias.reveal_nested())
                  @for_range(len(batch))
                  def _(i):
                      print_ln('%s X %s %s', self, i, self.X[batch[i]].reveal_nested())
                      print_ln('%s Y %s %s', self, i, self.Y[i].reveal_nested())
              return
          else:
              assert len(batch) == 1
              if self.fewer_rounds:
                  inputs, weights = self.prepare_temp()
  
          @for_range_opt_multithread(self.n_threads,
                                     [output_h, output_w, n_channels_out])
          def _(out_y, out_x, out_c):
                      in_x_origin = (out_x * stride_w) - padding_w
                      in_y_origin = (out_y * stride_h) - padding_h
                      iv = []
                      wv = []
                      for filter_y in range(weights_h):
                          in_y = in_y_origin + filter_y
                          inside_y = (0 <= in_y) * (in_y < inputs_h)
                          for filter_x in range(weights_w):
                              in_x = in_x_origin + filter_x
                              inside_x = (0 <= in_x) * (in_x < inputs_w)
                              inside = inside_y * inside_x
                              if is_zero(inside):
                                  continue
                              for in_c in range(n_channels_in):
                                  iv += [self.X[0][in_y * inside_y]
                                         [in_x * inside_x][in_c]]
                                  wv += [self.weights[out_c][filter_y][filter_x][in_c]]
                                  wv[-1] *= inside
                      if self.fewer_rounds:
                          inputs[out_y][out_x][out_c].assign(iv)
                          weights[out_y][out_x][out_c].assign(wv)
                      else:
                          self.dot_product(iv, wv, out_y, out_x, out_c)
  
          if self.fewer_rounds:
              @for_range_opt_multithread(self.n_threads,
                                         list(self.output_shape[1:]))
              def _(out_y, out_x, out_c):
                  self.dot_product(inputs[out_y][out_x][out_c],
                                   weights[out_y][out_x][out_c],
                                   out_y, out_x, out_c)
  
          self.reduction()
  
  class QuantConvBase(QuantBase):
      def input_params_from(self, player):
          for s in self.input_squant, self.weight_squant, self.bias_squant, self.output_squant:
              s.get_params_from(player)
          print('WARNING: assuming that bias quantization parameters are correct')
          self.output_squant.params.precompute(self.input_squant.params, self.weight_squant.params)
  
  class QuantConv2d(QuantConvBase, Conv2d):
      pass
  
  class FixConv2d(Conv2d, FixBase):
      """ Fixed-point 2D convolution layer.
  
      :param input_shape: input shape (tuple/list of four int)
      :param weight_shape: weight shape (tuple/list of four int)
      :param bias_shape: bias shape (tuple/list of one int)
      :param output_shape: output shape (tuple/list of four int)
      :param stride: stride (tuple/list of two int)
      :param padding: :py:obj:`'SAME'` (default), :py:obj:`'VALID'`, or tuple/list of two int
      :param tf_weight_format: weight shape format is (height, width, input channels, output channels) instead of the default (output channels, height, width, input channels)
      """
  
      def reset(self):
          assert not self.tf_weight_format
          kernel_size = self.weight_shape[1] * self.weight_shape[2]
          r = math.sqrt(6.0 / (kernel_size * sum(self.weight_shape[::3])))
          print('Initializing convolution weights in [%f,%f]' % (-r, r))
          self.weights.assign_vector(
              sfix.get_random(-r, r, size=self.weights.total_size()))
          self.bias.assign_all(0)
  
      def backward(self, compute_nabla_X=True, batch=None):
          assert self.use_conv2ds
  
          assert not self.tf_weight_format
          _, weights_h, weights_w, _ = self.weight_shape
          _, inputs_h, inputs_w, n_channels_in = self.input_shape
          _, output_h, output_w, n_channels_out = self.output_shape
  
          stride_h, stride_w = self.stride
          padding_h, padding_w = self.padding
  
          N = len(batch)
  
          self.nabla_bias.assign_all(0)
  
          @for_range(N)
          def _(i):
              self.nabla_bias.assign_vector(
                  self.nabla_bias.get_vector() + sum(sum(
                      self.nabla_Y[i][j][k].get_vector() for k in range(output_w))
                                                     for j in range(output_h)))
  
          input_size = inputs_h * inputs_w * N
          batch_repeat = regint.Matrix(N, inputs_h * inputs_w)
          batch_repeat.assign_vector(batch.get(
              regint.inc(input_size, 0, 1, 1, N)) *
                                     reduce(operator.mul, self.input_shape[1:]))
  
          @for_range_opt_multithread(self.n_threads, [n_channels_in, n_channels_out])
          def _(i, j):
              a = regint.inc(input_size, self.X.address + i, n_channels_in, N,
                             inputs_h * inputs_w)
              inputs = sfix.load_mem(batch_repeat.get_vector() + a).pre_mul()
              b = regint.inc(N * output_w * output_h, self.nabla_Y.address + j, n_channels_out, N)
              rep_out = regint.inc(output_h * output_w * N, 0, 1, 1, N) * \
                  reduce(operator.mul, self.output_shape[1:])
              nabla_outputs = sfix.load_mem(rep_out + b).pre_mul()
              res = sint(size = weights_h * weights_w)
              conv2ds(res, inputs, nabla_outputs, weights_h, weights_w, inputs_h,
                      inputs_w, output_h, output_w, -stride_h, -stride_w, N,
                      padding_h, padding_w, 1)
              reduced = unreduced_sfix._new(res).reduce_after_mul()
              self.nabla_weights.assign_vector_by_indices(reduced, j, None, None, i)
  
          if compute_nabla_X:
              assert tuple(self.padding) == (0, 0)
              assert tuple(self.stride) == (1, 1)
              reverse_weights = MultiArray(
                  [n_channels_in, weights_h, weights_w, n_channels_out], sfix)
              @for_range(n_channels_out)
              def _(i):
                  @for_range(weights_h)
                  def _(j):
                      @for_range(weights_w)
                      def _(k):
                          @for_range(n_channels_in)
                          def _(l):
                              reverse_weights[l][weights_h-j-1][k][i] = \
                                  self.weights[i][j][weights_w-k-1][l]
              padded_w = inputs_w + 2 * padding_w
              padded_h = inputs_h + 2 * padding_h
              if padding_h or padding_w:
                  output = MultiArray(
                      [N, padded_h, padded_w, n_channels_in], sfix)
              else:
                  output = self.nabla_X
              @for_range_opt_multithread(self.n_threads,
                                         [N, n_channels_in])
              def _(i, j):
                  res = sint(size = (padded_w * padded_h))
                  conv2ds(res, self.nabla_Y[i].get_vector().pre_mul(),
                          reverse_weights[j].get_vector().pre_mul(),
                          padded_h, padded_w, output_h, output_w,
                          weights_h, weights_w, 1, 1, n_channels_out,
                          weights_h - 1, weights_w - 1, 1)
                  output.assign_vector_by_indices(
                      unreduced_sfix._new(res).reduce_after_mul(),
                      i, None, None, j)
              if padding_h or padding_w:
                  @for_range(N)
                  def _(i):
                      @for_range(inputs_h)
                      def _(j):
                          @for_range(inputs_w)
                          def _(k):
                              self.nabla_X[i][j][k].assign_vector(
                                  output[i][j][k].get_vector())
  
          if self.debug_output:
              @for_range(len(batch))
              def _(i):
                  print_ln('%s X %s %s', self, i, list(self.X[i].reveal_nested()))
                  print_ln('%s nabla Y %s %s', self, i, list(self.nabla_Y[i].reveal_nested()))
                  if compute_nabla_X:
                      print_ln('%s nabla X %s %s', self, i, self.nabla_X[batch[i]].reveal_nested())
              print_ln('%s nabla weights %s', self,
                       (self.nabla_weights.reveal_nested()))
              print_ln('%s weights %s', self, (self.weights.reveal_nested()))
              print_ln('%s nabla b %s', self, (self.nabla_bias.reveal_nested()))
              print_ln('%s bias %s', self, (self.bias.reveal_nested()))
  
  class QuantDepthwiseConv2d(QuantConvBase, Conv2d):
      def n_summands(self):
          _, weights_h, weights_w, _ = self.weight_shape
          return weights_h * weights_w
  
      def _forward(self, batch):
          assert len(batch) == 1
          assert(self.weight_shape[-1] == self.output_shape[-1])
          assert(self.input_shape[-1] == self.output_shape[-1])
  
          _, weights_h, weights_w, _ = self.weight_shape
          _, inputs_h, inputs_w, n_channels_in = self.input_shape
          _, output_h, output_w, n_channels_out = self.output_shape
  
          stride_h, stride_w = self.stride
          padding_h, padding_w = self.padding
  
          depth_multiplier = 1
  
          if self.use_conv2ds:
              assert depth_multiplier == 1
              assert self.weight_shape[0] == 1
              @for_range_opt_multithread(self.n_threads, n_channels_in)
              def _(j):
                  inputs = self.X.get_vector_by_indices(0, None, None, j)
                  assert not self.tf_weight_format
                  weights = self.weights.get_vector_by_indices(0, None, None,
                                                               j)
                  inputs = inputs.pre_mul()
                  weights = weights.pre_mul()
                  res = sint(size = output_h * output_w)
                  conv2ds(res, inputs, weights, output_h, output_w,
                          inputs_h, inputs_w, weights_h, weights_w,
                          stride_h, stride_w, 1, padding_h, padding_w, 1)
                  res += self.bias.expand_to_vector(j, res.size).v
                  self.unreduced.assign_vector_by_indices(res, 0, None, None, j)
              self.reduction()
              return
          else:
              if self.fewer_rounds:
                  inputs, weights = self.prepare_temp()
  
          @for_range_opt_multithread(self.n_threads,
                                     [output_h, output_w, n_channels_in])
          def _(out_y, out_x, in_c):
                      for m in range(depth_multiplier):
                          oc = m + in_c * depth_multiplier
                          in_x_origin = (out_x * stride_w) - padding_w
                          in_y_origin = (out_y * stride_h) - padding_h
                          iv = []
                          wv = []
                          for filter_y in range(weights_h):
                              for filter_x in range(weights_w):
                                  in_x = in_x_origin + filter_x
                                  in_y = in_y_origin + filter_y
                                  inside = (0 <= in_x) * (in_x < inputs_w) * \
                                           (0 <= in_y) * (in_y < inputs_h)
                                  if is_zero(inside):
                                      continue
                                  iv += [self.X[0][in_y][in_x][in_c]]
                                  wv += [self.weights[0][filter_y][filter_x][oc]]
                                  wv[-1] *= inside
                          if self.fewer_rounds:
                              inputs[out_y][out_x][oc].assign(iv)
                              weights[out_y][out_x][oc].assign(wv)
                          else:
                              self.dot_product(iv, wv, out_y, out_x, oc)
  
          if self.fewer_rounds:
              @for_range_opt_multithread(self.n_threads,
                                         list(self.output_shape[1:]))
              def _(out_y, out_x, out_c):
                  self.dot_product(inputs[out_y][out_x][out_c],
                                   weights[out_y][out_x][out_c],
                                   out_y, out_x, out_c)
  
          self.reduction()
  
  class AveragePool2d(BaseLayer):
      def __init__(self, input_shape, output_shape, filter_size, strides=(1, 1)):
          super(AveragePool2d, self).__init__(input_shape, output_shape)
          self.filter_size = filter_size
          self.strides = strides
          for i in (0, 1):
              if strides[i] == 1:
                  assert output_shape[1+i] == 1
                  assert filter_size[i] == input_shape[1+i]
              else:
                  assert strides[i] == filter_size[i]
                  assert output_shape[1+i] * strides[i] == input_shape[1+i]
  
      def input_from(self, player, raw=False):
          self.input_params_from(player)
  
      def _forward(self, batch=[0]):
          assert len(batch) == 1
  
          _, input_h, input_w, n_channels_in = self.input_shape
          _, output_h, output_w, n_channels_out = self.output_shape
  
          assert n_channels_in == n_channels_out
  
          padding_h, padding_w = (0, 0)
          stride_h, stride_w = self.strides
          filter_h, filter_w = self.filter_size
          n = filter_h * filter_w
          print('divisor: ', n)
  
          @for_range_opt_multithread(self.n_threads,
                                     [output_h, output_w, n_channels_in])
          def _(out_y, out_x, c):
              in_x_origin = (out_x * stride_w) - padding_w
              in_y_origin = (out_y * stride_h) - padding_h
              fxs = util.max(-in_x_origin, 0)
              #fxe = min(filter_w, input_w - in_x_origin)
              fys = util.max(-in_y_origin, 0)
              #fye = min(filter_h, input_h - in_y_origin)
              acc = 0
              #fc = 0
              for i in range(filter_h):
                  filter_y = fys + i
                  for j in range(filter_w):
                      filter_x = fxs + j
                      in_x = in_x_origin + filter_x
                      in_y = in_y_origin + filter_y
                      acc += self.X[0][in_y][in_x][c].v
                      #fc += 1
              acc = self.const_div(acc, n)
              self.Y[0][out_y][out_x][c] = self.output_squant._new(acc)
  
  class QuantAveragePool2d(QuantBase, AveragePool2d):
      def input_params_from(self, player):
          print('WARNING: assuming that input and output quantization parameters are the same')
          for s in self.input_squant, self.output_squant:
              s.get_params_from(player)
  
  class FixAveragePool2d(FixBase, AveragePool2d):
      """ Fixed-point 2D AvgPool layer.
  
      :param input_shape: input shape (tuple/list of four int)
      :param output_shape: output shape (tuple/list of four int)
      :param filter_size: filter size (tuple/list of two int)
      :param strides: strides (tuple/list of two int)
      """
  
  class QuantReshape(QuantBase, BaseLayer):
      def __init__(self, input_shape, _, output_shape):
          super(QuantReshape, self).__init__(input_shape, output_shape)
  
      def input_from(self, player):
          print('WARNING: assuming that input and output quantization parameters are the same')
          _ = self.new_squant()
          for s in self.input_squant, _, self.output_squant:
              s.set_params(sfloat.get_input_from(player), sint.get_input_from(player))
          for i in range(2):
              sint.get_input_from(player)
  
      def _forward(self, batch):
          assert len(batch) == 1
          # reshaping is implicit
          self.Y.assign(self.X)
  
  class QuantSoftmax(QuantBase, BaseLayer):
      def input_from(self, player):
          print('WARNING: assuming that input and output quantization parameters are the same')
          for s in self.input_squant, self.output_squant:
              s.set_params(sfloat.get_input_from(player), sint.get_input_from(player))
  
      def _forward(self, batch):
          assert len(batch) == 1
          assert(len(self.input_shape) == 2)
  
          # just print the best
          def comp(left, right):
              c = left[1].v.greater_than(right[1].v, self.input_squant.params.k)
              #print_ln('comp %s %s %s', c.reveal(), left[1].v.reveal(), right[1].v.reveal())
              return [c.if_else(x, y) for x, y in zip(left, right)]
          print_ln('guess: %s', util.tree_reduce(comp, list(enumerate(self.X[0])))[0].reveal())
  
  class Optimizer:
      """ Base class for graphs of layers. """
      n_threads = Layer.n_threads
      always_shuffle = True
      time_layers = False
      revealing_correctness = False
  
      @staticmethod
      def from_args(program, layers):
          if 'adam' in program.args or 'adamapprox' in program.args:
              return Adam(layers, 1, approx='adamapprox' in program.args)
          elif 'amsgrad' in program.args:
              return Adam(layers, approx=True, amsgrad=True)
          elif 'quotient' in program.args:
              return Adam(layers, approx=True, amsgrad=True, normalize=True)
          else:
              return SGD(layers, 1)
  
      def __init__(self, report_loss=None):
          self.tol = 0.000
          self.report_loss = report_loss
          self.X_by_label = None
          self.print_update_average = False
          self.print_losses = False
          self.print_loss_reduction = False
          self.i_epoch = MemValue(0)
          self.stopped_on_loss = MemValue(0)
  
      @property
      def layers(self):
          """ Get all layers. """
          return self._layers
  
      @layers.setter
      def layers(self, layers):
          """ Construct linear graph from list of layers. """
          self._layers = layers
          prev = None
          for layer in layers:
              if not layer.inputs and prev is not None:
                  layer.inputs = [prev]
              prev = layer
  
      def set_layers_with_inputs(self, layers):
          """ Construct graph from :py:obj:`inputs` members of list of layers. """
          self._layers = layers
          used = set([None])
          for layer in reversed(layers):
              layer.last_used = list(filter(lambda x: x not in used, layer.inputs))
              used.update(layer.inputs)
  
      def reset(self):
          """ Initialize weights. """
          for layer in self.layers:
              layer.reset()
          self.i_epoch.write(0)
          self.stopped_on_loss.write(0)
  
      def batch_for(self, layer, batch):
          if layer in (self.layers[0], self.layers[-1]):
              return batch
          else:
              batch = regint.Array(len(batch))
              batch.assign(regint.inc(len(batch)))
              return batch
  
      @_no_mem_warnings
      def forward(self, N=None, batch=None, keep_intermediate=True,
                  model_from=None, training=False):
          """ Compute graph.
  
          :param N: batch size (used if batch not given)
          :param batch: indices for computation (:py:class:`~Compiler.types.Array` or list)
          :param keep_intermediate: do not free memory of intermediate results after use
          """
          if batch is None:
              batch = regint.Array(N)
              batch.assign(regint.inc(N))
          for i, layer in enumerate(self.layers):
              if layer.inputs and len(layer.inputs) == 1 and layer.inputs[0] is not None:
                  layer._X.address = layer.inputs[0].Y.address
              layer.Y.alloc()
              if model_from is not None:
                  layer.input_from(model_from)
              break_point()
              if self.time_layers:
                  start_timer(100 + i)
              layer.forward(batch=self.batch_for(layer, batch), training=training)
              if self.time_layers:
                  stop_timer(100 + i)
              break_point()
              if not keep_intermediate:
                  for l in layer.last_used:
                      l.Y.delete()
                  for theta in layer.thetas():
                      theta.delete()
  
      @_no_mem_warnings
      def eval(self, data):
          """ Compute evaluation after training. """
          N = len(data)
          self.layers[0].X.assign(data)
          self.forward(N)
          return self.layers[-1].eval(N)
  
      @_no_mem_warnings
      def backward(self, batch):
          """ Compute backward propagation. """
          for i, layer in reversed(list(enumerate(self.layers))):
              assert len(batch) <= layer.back_batch_size
              if self.time_layers:
                  start_timer(200 + i)
              if not layer.inputs:
                  layer.backward(compute_nabla_X=False,
                                 batch=self.batch_for(layer, batch))
              else:
                  layer.backward(batch=self.batch_for(layer, batch))
                  if len(layer.inputs) == 1:
                      layer.inputs[0].nabla_Y.address = \
                          layer.nabla_X.address
              if self.time_layers:
                  stop_timer(200 + i)
  
      @_no_mem_warnings
      def run(self, batch_size=None, stop_on_loss=0):
          """ Run training.
  
          :param batch_size: batch size (defaults to example size of first layer)
          """
          if self.n_epochs == 0:
              return
          if batch_size is not None:
              N = batch_size
          else:
              N = self.layers[0].N
          i = self.i_epoch
          n_iterations = MemValue(0)
          self.n_correct = MemValue(0)
          @for_range(self.n_epochs)
          def _(_):
              if self.X_by_label is None:
                  self.X_by_label = [[None] * self.layers[0].N]
              assert len(self.X_by_label) in (1, 2)
              assert N % len(self.X_by_label) == 0
              n = N // len(self.X_by_label)
              n_per_epoch = int(math.ceil(1. * max(len(X) for X in
                                                   self.X_by_label) / n))
              print('%d runs per epoch' % n_per_epoch)
              indices_by_label = []
              for label, X in enumerate(self.X_by_label):
                  indices = regint.Array(n * n_per_epoch)
                  indices_by_label.append(indices)
                  indices.assign(regint.inc(len(indices), 0, 1, 1, len(X)))
                  if self.always_shuffle or n_per_epoch > 1:
                      indices.shuffle()
              loss_sum = MemValue(sfix(0))
              self.n_correct.write(0)
              @for_range(n_per_epoch)
              def _(j):
                  n_iterations.iadd(1)
                  batch = regint.Array(N)
                  for label, X in enumerate(self.X_by_label):
                      indices = indices_by_label[label]
                      batch.assign(indices.get_vector(j * n, n) +
                                   regint(label * len(self.X_by_label[0]), size=n),
                                   label * n)
                  self.forward(batch=batch, training=True)
                  self.backward(batch=batch)
                  self.update(i, batch=batch)
                  loss_sum.iadd(self.layers[-1].l)
                  if self.print_loss_reduction:
                      before = self.layers[-1].average_loss(N)
                      self.forward(batch=batch)
                      after = self.layers[-1].average_loss(N)
                      print_ln('loss reduction in batch %s: %s (%s - %s)', j,
                               before - after, before, after)
                  elif self.print_losses:
                      print_str('\rloss in batch %s: %s/%s', j,
                               self.layers[-1].average_loss(N),
                               loss_sum.reveal() / (j + 1))
                  if self.revealing_correctness:
                      part_truth = self.layers[-1].Y.same_shape()
                      part_truth.assign_vector(
                          self.layers[-1].Y.get_slice_vector(batch))
                      self.n_correct.iadd(
                          self.layers[-1].reveal_correctness(batch_size, part_truth))
                  if stop_on_loss:
                      loss = self.layers[-1].average_loss(N)
                      res = (loss < stop_on_loss) * (loss >= -1)
                      self.stopped_on_loss.write(1 - res)
                      return res
              if self.print_losses:
                  print_ln()
              if self.report_loss and self.layers[-1].compute_loss and self.layers[-1].approx != 5:
                  print_ln('loss in epoch %s: %s', i,
                           (loss_sum.reveal() * cfix(1 / n_per_epoch)))
              else:
                  print_ln('done with epoch %s', i)
              time()
              i.iadd(1)
              res = True
              if self.tol > 0:
                  res *= (1 - (loss >= 0) * (loss < self.tol)).reveal()
              return res
  
      def reveal_correctness(self, data, truth, batch_size):
          training_data = self.layers[0].X.address
          training_truth = self.layers[-1].Y.address
          self.layers[0].X.address = data.address
          self.layers[-1].Y.address = truth.address
          N = data.sizes[0]
          batch = regint.Array(batch_size)
          n_correct = MemValue(0)
          loss = MemValue(sfix(0))
          def f(start, batch_size):
              batch.assign_vector(regint.inc(batch_size, start))
              self.forward(batch=batch)
              part_truth = truth.get_part(start, batch_size)
              n_correct.iadd(
                  self.layers[-1].reveal_correctness(batch_size, part_truth))
              loss.iadd(self.layers[-1].l * batch_size)
          @for_range(N // batch_size)
          def _(i):
              start = i * batch_size
              f(start, batch_size)
          batch_size = N % batch_size
          if batch_size:
              start = N - batch_size
              f(start, batch_size)
          self.layers[0].X.address = training_data
          self.layers[-1].Y.address = training_truth
          loss = loss.reveal()
          if cfix.f < 31:
              loss = cfix._new(loss.v << (31 - cfix.f), k=63, f=31)
          return n_correct, loss / N
  
      @_no_mem_warnings
      def run_by_args(self, program, n_runs, batch_size, test_X, test_Y,
                      acc_batch_size=None):
          if acc_batch_size is None:
              acc_batch_size = batch_size
          depreciation = None
          for arg in program.args:
              m = re.match('rate(.*)', arg)
              if m:
                  self.gamma = MemValue(cfix(float(m.group(1))))
              m = re.match('dep(.*)', arg)
              if m:
                  depreciation = float(m.group(1))
          if 'nomom' in program.args:
              self.momentum = 0
          self.print_losses = 'print_losses' in program.args
          self.time_layers = 'time_layers' in program.args
          self.revealing_correctness = not 'no_acc' in program.args
          self.layers[-1].compute_loss = not 'no_loss' in program.args
          model_input = 'model_input' in program.args
          acc_first = model_input and not 'train_first' in program.args
          if model_input:
              for layer in self.layers:
                  layer.input_from(0)
          else:
              self.reset()
          if 'one_iter' in program.args:
              self.output_weights()
              print_ln('loss')
              print_ln('%s', self.eval(
                  self.layers[0].X.get_part(0, batch_size)).reveal_nested())
              for layer in self.layers:
                  print_ln('%s', layer.X.get_part(0, batch_size).reveal_nested())
              print_ln('%s', self.layers[-1].Y.get_part(0, batch_size).reveal_nested())
              batch = Array.create_from(regint.inc(batch_size))
              self.forward(batch=batch, training=True)
              self.backward(batch=batch)
              self.update(0, batch=batch)
              print_ln('loss %s', self.layers[-1].l.reveal())
              self.output_weights()
              return
          @for_range(n_runs)
          def _(i):
              if not acc_first:
                  start_timer(1)
                  self.run(batch_size,
                           stop_on_loss=0 if 'no_loss' in program.args else 100)
                  stop_timer(1)
              if 'no_acc' in program.args:
                  return
              N = self.layers[0].X.sizes[0]
              n_trained = (N + batch_size - 1) // batch_size * batch_size
              print_ln('train_acc: %s (%s/%s)',
                       cfix(self.n_correct, k=63, f=31) / n_trained,
                       self.n_correct, n_trained)
              n_test = len(test_Y)
              n_correct, loss = self.reveal_correctness(test_X, test_Y, acc_batch_size)
              print_ln('test loss: %s', loss)
              print_ln('acc: %s (%s/%s)', cfix(n_correct, k=63, f=31) / n_test,
                       n_correct, n_test)
              if acc_first:
                  start_timer(1)
                  self.run(batch_size)
                  stop_timer(1)
              else:
                  @if_(util.or_op(self.stopped_on_loss, n_correct <
                                  int(n_test // self.layers[-1].n_outputs * 1.2)))
                  def _():
                      self.gamma.imul(.5)
                      self.reset()
                      print_ln('reset after reducing learning rate to %s',
                               self.gamma)
              if depreciation:
                  self.gamma.imul(depreciation)
                  print_ln('reducing learning rate to %s', self.gamma)
          if 'model_output' in program.args:
              self.output_weights()
  
      def output_weights(self):
          print_float_precision(max(6, sfix.f // 3))
          for layer in self.layers:
              layer.output_weights()
  
  class Adam(Optimizer):
      """ Adam/AMSgrad optimizer.
  
      :param layers: layers of linear graph
      :param approx: use approximation for inverse square root (bool)
      :param amsgrad: use AMSgrad (bool)
      """
      def __init__(self, layers, n_epochs=1, approx=False, amsgrad=False,
                   normalize=False):
          self.gamma = MemValue(cfix(.001))
          self.beta1 = 0.9
          self.beta2 = 0.999
          self.beta1_power = MemValue(cfix(1))
          self.beta2_power = MemValue(cfix(1))
          self.epsilon = max(2 ** -((sfix.k - sfix.f - 8) / (1 + approx)), 10 ** -8)
          self.n_epochs = n_epochs
          self.approx = approx
          self.amsgrad = amsgrad
          self.normalize = normalize
          if amsgrad:
              print_str('Using AMSgrad ')
          else:
              print_str('Using Adam ')
          if approx:
              print_ln('with inverse square root approximation')
          else:
              print_ln('with more precise inverse square root')
          if normalize:
              print_ln('Normalize gradient')
  
          self.layers = layers
          self.ms = []
          self.vs = []
          self.gs = []
          self.thetas = []
          self.vhats = []
          for layer in layers:
              for nabla in layer.nablas():
                  self.gs.append(nabla)
                  for x in self.ms, self.vs:
                      x.append(nabla.same_shape())
                  if amsgrad:
                      self.vhats.append(nabla.same_shape())
              for theta in layer.thetas():
                  self.thetas.append(theta)
  
          super(Adam, self).__init__()
  
      def update(self, i_epoch, batch):
          self.beta1_power *= self.beta1
          self.beta2_power *= self.beta2
          m_factor = MemValue(1 / (1 - self.beta1_power))
          v_factor = MemValue(1 / (1 - self.beta2_power))
          for i_layer, (m, v, g, theta) in enumerate(zip(self.ms, self.vs,
                                                         self.gs, self.thetas)):
              if self.normalize:
                  abs_g = g.same_shape()
                  @multithread(self.n_threads, g.total_size())
                  def _(base, size):
                      abs_g.assign_vector(abs(g.get_vector(base, size)), base)
                  max_g = tree_reduce_multithread(self.n_threads,
                                                  util.max, abs_g.get_vector())
                  scale = MemValue(sfix._new(library.AppRcr(
                      max_g.v, max_g.k, max_g.f, simplex_flag=True)))
              @multithread(self.n_threads, m.total_size())
              def _(base, size):
                  m_part = m.get_vector(base, size)
                  v_part = v.get_vector(base, size)
                  g_part = g.get_vector(base, size)
                  if self.normalize:
                      g_part *= scale.expand_to_vector(size)
                  m_part = self.beta1 * m_part + (1 - self.beta1) * g_part
                  v_part = self.beta2 * v_part + (1 - self.beta2) * g_part ** 2
                  m.assign_vector(m_part, base)
                  v.assign_vector(v_part, base)
                  if self.amsgrad:
                      vhat = self.vhats [i_layer].get_vector(base, size)
                      vhat = util.max(vhat, v_part)
                      self.vhats[i_layer].assign_vector(vhat, base)
                      diff = self.gamma.expand_to_vector(size) * m_part
                  else:
                      mhat = m_part * m_factor.expand_to_vector(size)
                      vhat = v_part * v_factor.expand_to_vector(size)
                      diff = self.gamma.expand_to_vector(size) * mhat
                  if self.approx:
                      diff *= mpc_math.InvertSqrt(vhat + self.epsilon ** 2)
                  else:
                      diff /= mpc_math.sqrt(vhat) + self.epsilon
                  theta.assign_vector(theta.get_vector(base, size) - diff, base)
  
  class SGD(Optimizer):
      """ Stochastic gradient descent.
  
      :param layers: layers of linear graph
      :param n_epochs: number of epochs for training
      :param report_loss: disclose and print loss
      """
      def __init__(self, layers, n_epochs, debug=False, report_loss=None):
          self.momentum = 0.9
          self.layers = layers
          self.n_epochs = n_epochs
          self.thetas = []
          self.nablas = []
          self.delta_thetas = []
          for layer in layers:
              self.nablas.extend(layer.nablas())
              self.thetas.extend(layer.thetas())
              for theta in layer.thetas():
                  self.delta_thetas.append(theta.same_shape())
          self.gamma = MemValue(cfix(0.01))
          self.debug = debug
          super(SGD, self).__init__(report_loss)
  
      @_no_mem_warnings
      def reset(self, X_by_label=None):
          """ Reset layer parameters.
  
          :param X_by_label: if given, set training data by public labels for balancing
          """
          self.X_by_label = X_by_label
          if X_by_label is not None:
              for label, X in enumerate(X_by_label):
                  @for_range_multithread(self.n_threads, 1, len(X))
                  def _(i):
                      j = i + label * len(X_by_label[0])
                      self.layers[0].X[j] = X[i]
                      self.layers[-1].Y[j] = label
          for y in self.delta_thetas:
              y.assign_all(0)
          super(SGD, self).reset()
  
      def update(self, i_epoch, batch):
          for nabla, theta, delta_theta in zip(self.nablas, self.thetas,
                                               self.delta_thetas):
              @multithread(self.n_threads, nabla.total_size())
              def _(base, size):
                  old = delta_theta.get_vector(base, size)
                  red_old = self.momentum * old
                  rate = self.gamma.expand_to_vector(size)
                  nabla_vector = nabla.get_vector(base, size)
                  log_batch_size = math.log(len(batch), 2)
                  # divide by len(batch) by truncation
                  # increased rate if len(batch) is not a power of two
                  pre_trunc = nabla_vector.v * rate.v
                  k = nabla_vector.k + rate.k
                  m = rate.f + int(log_batch_size)
                  v = pre_trunc.round(k, m, signed=True,
                                      nearest=sfix.round_nearest)
                  new = nabla_vector._new(v)
                  diff = red_old - new
                  delta_theta.assign_vector(diff, base)
                  theta.assign_vector(theta.get_vector(base, size) +
                                      delta_theta.get_vector(base, size), base)
              if self.print_update_average:
                  vec = abs(delta_theta.get_vector().reveal())
                  print_ln('update average: %s (%s)',
                           sum(vec) * cfix(1 / len(vec)), len(vec))
              if self.debug:
                  limit = int(self.debug)
                  d = delta_theta.get_vector().reveal()
                  aa = [cfix.Array(len(d.v)) for i in range(3)]
                  a = aa[0]
                  a.assign(d)
                  @for_range(len(a))
                  def _(i):
                      x = a[i]
                      print_ln_if((x > limit) + (x < -limit),
                                  'update epoch=%s %s index=%s %s',
                                  i_epoch.read(), str(delta_theta), i, x)
                  a = aa[1]
                  a.assign(nabla.get_vector().reveal())
                  @for_range(len(a))
                  def _(i):
                      x = a[i]
                      print_ln_if((x > len(batch) * limit) + (x < -len(batch) * limit),
                                  'nabla epoch=%s %s index=%s %s',
                                  i_epoch.read(), str(nabla), i, x)
                  a = aa[2]
                  a.assign(theta.get_vector().reveal())
                  @for_range(len(a))
                  def _(i):
                      x = a[i]
                      print_ln_if((x > limit) + (x < -limit),
                                  'theta epoch=%s %s index=%s %s',
                                  i_epoch.read(), str(theta), i, x)
                  index = regint.get_random(64) % len(a)
                  print_ln('%s at %s: nabla=%s update=%s theta=%s', str(theta), index,
                           aa[1][index], aa[0][index], aa[2][index])
          self.gamma.imul(1 - 10 ** - 6)
  