# Fits small MLPs to 8-bit parity (the full 256-row truth table) with
  # full-batch AdamW, written out in numpy (no PyTorch, so it won't match a
  # PyTorch run bit for bit), and reports for each shape how many of 20 seeds
  # reach 100% accuracy and the median accuracy. Needs Python 3.9+ and numpy.
  import itertools

  import numpy as np

  X = np.array(list(itertools.product([0, 1], repeat=8)), float)  # all 256 inputs
  y = X.sum(1) % 2


  def nparams(widths: list[int]) -> int:
      """Parameter count of a dense MLP with these layer widths, biases included."""
      return sum(a * b + b for a, b in zip(widths, widths[1:]))


  def train(
      widths: list[int],
      seed: int,
      steps: int = 4000,
      lr: float = 1e-2,
      wd: float = 1e-4,
  ) -> float:
      """Full-batch AdamW on all 256 inputs; returns the final training accuracy."""
      rng = np.random.default_rng(seed)
      P = []    # parameters: W0, b0, W1, b1, ...
      for a, b in zip(widths, widths[1:]):
          bound = 1 / np.sqrt(a)    # same as PyTorch nn.Linear's default init
          P += [rng.uniform(-bound, bound, (a, b)), rng.uniform(-bound, bound, b)]

      m = [np.zeros_like(p) for p in P]    # Adam first moments
      v = [np.zeros_like(p) for p in P]    # Adam second moments
      nlayers = len(widths) - 1
      for step in range(1, steps + 1):
          hs = [X]    # hs[i] is the input to layer i
          zs = []     # pre-activations
          h = X
          for i in range(nlayers):
              z = h @ P[2 * i] + P[2 * i + 1]
              zs.append(z)
              h = np.maximum(z, 0) if i < nlayers - 1 else z    # ReLU; linear output
              hs.append(h)

          logit = h[:, 0]
          prob = 0.5 * (1 + np.tanh(logit / 2))    # sigmoid without overflow
          g = ((prob - y) / len(y))[:, None]    # d(mean cross-entropy)/d(logit)
          G = [None] * len(P)    # gradients, parallel to P
          for i in reversed(range(nlayers)):
              G[2 * i] = hs[i].T @ g
              G[2 * i + 1] = g.sum(0)
              if i > 0:
                  g = (g @ P[2 * i].T) * (zs[i - 1] > 0)

          # AdamW, using PyTorch's default values for betas (0.9, 0.999) and eps (1e-8)
          for k in range(len(P)):
              P[k] -= lr * wd * P[k]    # decoupled weight decay
              m[k] = 0.9 * m[k] + 0.1 * G[k]
              v[k] = 0.999 * v[k] + 0.001 * G[k] ** 2
              P[k] -= lr * (m[k] / (1 - 0.9 ** step)) / (
                  np.sqrt(v[k] / (1 - 0.999 ** step)) + 1e-8)

      # accuracy of the last forward pass, i.e. before the final update
      return ((logit > 0) == (y > 0.5)).mean()


  for w in [[8, 1, 2, 2, 2, 1], [8, 8, 4, 2, 1], [8, 8, 1], [8, 16, 1], [8, 64, 1]]:
      accs = [train(w, seed) for seed in range(20)]
      print(w, nparams(w), "params; seeds at 100%:", sum(a == 1 for a in accs),
            "/ 20; median acc", np.median(accs))