Optimizer(  # 5.11 MB state
  step=0,
  # Transforms:
  ClipByGlobalNorm(max_norm=1.0),
  AdamW(learning_rate=0.0003, b1=0.9, b2=0.999, eps=1e-08, weight_decay=0.1),
  # State:
  mu=MLP(...),  # 669,706 params, 2.55 MB
  nu=MLP(...),  # 669,706 params, 2.55 MB
)