code: def diet_expert(x, hidden_size, params):
      """A two-layer feed-forward network with relu activation on hidden layer.
  
      Uses diet variables.
      Recomputes hidden layer on backprop to save activation memory.
  
      Args:
        x: a Tensor with shape [batch, io_size]
        hidden_size: an integer
        params: a diet variable HParams object.
  
      Returns:
        a Tensor with shape [batch, io_size]
      """
  
      @fn_with_diet_vars(params)
      def diet_expert_internal(x):
          dim = x.get_shape().as_list()[-1]
          h = tf.layers.dense(
              x, hidden_size, activation=tf.nn.relu, use_bias=False
          )
          y = tf.layers.dense(h, dim, use_bias=False)
          y *= tf.rsqrt(tf.to_float(dim * hidden_size))
          return y
  
      return diet_expert_internal(x)
  
imports: from collections import defaultdict
  import copy
  import math
  from tensor2tensor.layers import common_layers
  from tensor2tensor.utils import hparam
  import tensorflow.compat.v1
label_desc: Use image node_id 2 for calling a global function with example usage: diet_expert(x, hidden_size, params) and returns: diet_expert_internal, y