code: def ffn_expert_fn(input_size,
                    hidden_sizes,
                    output_size,
                    hidden_activation=tf.nn.relu):
    
    def my_fn(x):
      layer_sizes = [input_size] + hidden_sizes + [output_size]
      for i in range(1 + len(hidden_sizes)):
        w = tf.get_variable("w_%d" % i, layer_sizes[i:i+2], tf.float32)
        x = tf.matmul(x, w)
        if i < len(hidden_sizes):
          x = hidden_activation(x)
        if layer_sizes[i] != input_size:
          x *= (layer_sizes[i] / float(input_size))**-0.5
      return x
    return my_fn
docs: Returns a function that creates a feed-forward network.
  
    Use this function to create the expert_fn argument to distributed_moe.
  
    Args:
      input_size: an integer
      hidden_sizes: a list of integers
      output_size: an integer
      hidden_activation: a unary function.
  
    Returns:
      a unary function