code: def ffn_expert_fn(
      input_size,
      hidden_sizes,
      output_size,
      hidden_activation=tf.nn.relu,
  ):
      """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
      """
  
      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
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  import math
  import six
  from six.moves import range
  from six.moves import zip
  from tensor2tensor.layers import common_layers
  from tensor2tensor.layers.vq_discrete import DiscreteBottleneck
  import tensorflow.compat.v1
label_desc: Use image node_id 16 for calling a global function with example usage: ffn_expert_fn(input_size, hidden_sizes, output_size, hidden_activation) and returns: my_fn, x