def METHOD_NAME(optimizer, spec):
  weights = tf.nest.map_structure(lambda s: tf.ones(s.shape, s.dtype), spec)
  gradients = tf.nest.map_structure(lambda s: tf.ones(s.shape, s.dtype), spec)
  state = optimizer.initialize(spec)
  state_history = [state]
  weights_history = [weights]
  for _ in range(3):
    state, weights = optimizer.next(state, weights, gradients)
    state_history.append(state)
    weights_history.append(weights)
  return state_history, weights_history