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)
  init_fn = tensorflow_computation.tf_computation(
      lambda: optimizer.initialize(spec)
  )
  next_fn = tensorflow_computation.tf_computation(optimizer.next)
  state = init_fn()
  state_history = [state]
  weights_history = [weights]
  for _ in range(3):
    state, weights = next_fn(state, weights, gradients)
    state_history.append(state)
    weights_history.append(weights)
  return state_history, weights_history