code: def test_mpi_adam():
      """
      tests the MpiAdam object's functionality
      """
      np.random.seed(0)
      tf.set_random_seed(0)
  
      a_var = tf.Variable(np.random.randn(3).astype("float32"))
      b_var = tf.Variable(np.random.randn(2, 5).astype("float32"))
      loss = tf.reduce_sum(tf.square(a_var)) + tf.reduce_sum(
          tf.sin(b_var)
      )
  
      learning_rate = 1e-2
      update_op = tf.train.AdamOptimizer(learning_rate).minimize(loss)
      do_update = tf_utils.function([], loss, updates=[update_op])
  
      tf.get_default_session().run(tf.global_variables_initializer())
      for step in range(10):
          print(step, do_update())
  
      tf.set_random_seed(0)
      tf.get_default_session().run(tf.global_variables_initializer())
  
      var_list = [a_var, b_var]
      lossandgrad = tf_utils.function(
          [],
          [loss, tf_utils.flatgrad(loss, var_list)],
          updates=[update_op],
      )
      adam = MpiAdam(var_list)
  
      for step in range(10):
          loss, grad = lossandgrad()
          adam.update(grad, learning_rate)
          print(step, loss)
  
file_path: stable-baselines/stable_baselines/common/mpi_adam.py
imports: import tensorflow
  import numpy
  import mpi4py
  import stable_baselines.common.tf_util
label_desc: Use image node_id 1 for calling a global function with example usage: test_mpi_adam() without return types