code: def lstm_model(nlstm=128, layer_norm=False):
      """
      Builds LSTM (Long-Short Term Memory) network to be used in a policy.
      Note that the resulting function returns not only the output of the LSTM
      (i.e. hidden state of lstm for each step in the sequence), but also a dictionary
      with auxiliary tensors to be set as policy attributes.
  
      Specifically,
          S is a placeholder to feed current state (LSTM state has to be managed outside policy)
          M is a placeholder for the mask (used to mask out observations after the end of the episode, but can be used for other purposes too)
          initial_state is a numpy array containing initial lstm state (usually zeros)
          state is the output LSTM state (to be fed into S at the next call)
  
  
      An example of usage of lstm-based policy can be found here: common/tests/test_doc_examples.py/test_lstm_example
  
      Parameters
      ----------
      nlstm : int
          LSTM hidden state size
      layer_norm : bool
          if True, layer-normalized version of LSTM is used
  
      Returns
      -------
      function that builds LSTM with a given input tensor / placeholder
      """
  
      def network_fn(X, nenv=1, obs_size=-1):
          with tf.variable_scope("emb", reuse=tf.AUTO_REUSE):
              w_emb = tf.get_variable("w_emb", [obs_size + 1, 32])
              X = tf.nn.embedding_lookup(w_emb, X)
  
          nbatch = X.shape[0]
          nsteps = nbatch // nenv
  
          h = tf.layers.flatten(X)
  
          M = tf.placeholder(tf.float32, [nbatch])  # mask (done t-1)
          S = tf.placeholder(tf.float32, [nenv, 2 * nlstm])  # states
  
          xs = batch_to_seq(h, nenv, nsteps)
          ms = batch_to_seq(M, nenv, nsteps)
  
          assert not layer_norm
          h5, snew = lstm(xs, ms, S, scope="lstm", nh=nlstm)
  
          h = seq_to_batch(h5)
          initial_state = np.zeros(S.shape.as_list(), dtype=float)
  
          return h, {
              "S": S,
              "M": M,
              "state": snew,
              "initial_state": initial_state,
          }
  
      return network_fn
  
file_path: nni/nni/algorithms/hpo/ppo_tuner/util.py
imports: import os
  import random
  import multiprocessing
  import numpy
  import tensorflow.compat.v1
  from gym.spaces import Discrete, Box, MultiDiscrete
label_desc: Use image node_id 5 for calling a global function with example usage: lstm_model(nlstm, layer_norm) and returns: network_fn, h,  dict