docs: def model_inference_fn(features, training, params):
      """Builds just the inference part of the model graph.
  
      Args:
          features: input features tensor.
          training: True if the model is training.
          params: A dictionary
  
      Returns:
          (policy_output, value_output, logits) tuple of tensors.
      """
  
      mg_batchn = functools.partial(
          tf.layers.batch_normalization,
          axis=-1,
          momentum=.95,
          epsilon=1e-5,
          center=True,
          scale=True,
          fused=True,
          training=training)
  
      mg_conv2d = functools.partial(
          tf.layers.conv2d,
          filters=params['conv_width'],
          kernel_size=3,
          padding="same",
          data_format="channels_last",
          use_bias=False)
  
      mg_global_avgpool2d = functools.partial(
          tf.layers.average_pooling2d,
          pool_size=go.N,
          strides=1,
          padding="valid",
          data_format="channels_last")
  
      def mg_activation(inputs):
          if FLAGS.use_swish:
              return tf.nn.swish(inputs)
  
          return tf.nn.relu(inputs)
  
  
      def residual_inner(inputs):
          conv_layer1 = mg_batchn(mg_conv2d(inputs))
          initial_output = mg_activation(conv_layer1)
          conv_layer2 = mg_batchn(mg_conv2d(initial_output))
          return conv_layer2
  
      def mg_res_layer(inputs):
          residual = residual_inner(inputs)
          output = mg_activation(inputs + residual)
          return output
  
      def mg_squeeze_excitation_layer(inputs):
          # Hu, J., Shen, L., & Sun, G. (2018). Squeeze-and-Excitation Networks.
          # 2018 IEEE/CVF Conference on Computer Vision, 7132-7141.
          # arXiv:1709.01507 [cs.CV]
  
          channels = params['conv_width']
          ratio = FLAGS.SE_ratio
          assert channels % ratio == 0
  
          residual = residual_inner(inputs)
          pool = mg_global_avgpool2d(residual)
          fc1 = tf.layers.dense(pool, units=channels // ratio)
          squeeze = mg_activation(fc1)
  
          if FLAGS.use_SE_bias:
              fc2 = tf.layers.dense(squeeze, units=2*channels)
              # Channels_last so axis = 3 = -1
              gamma, bias = tf.split(fc2, 2, axis=3)
          else:
              gamma = tf.layers.dense(squeeze, units=channels)
              bias = 0
  
          sig = tf.nn.sigmoid(gamma)
          # Explicitly signal the broadcast.
          scale = tf.reshape(sig, [-1, 1, 1, channels])
  
          excitation = tf.multiply(scale, residual) + bias
          return mg_activation(inputs + excitation)
  
      initial_block = mg_activation(mg_batchn(mg_conv2d(features)))
  
      # the shared stack
      shared_output = initial_block
      for _ in range(params['trunk_layers']):
          if FLAGS.use_SE or FLAGS.use_SE_bias:
              shared_output = mg_squeeze_excitation_layer(shared_output)
          else:
              shared_output = mg_res_layer(shared_output)
  
      # Policy head
      policy_conv = mg_conv2d(
          shared_output, filters=params['policy_conv_width'], kernel_size=1)
      policy_conv = mg_activation(mg_batchn(policy_conv, center=False, scale=False))
      logits = tf.layers.dense(
          tf.reshape(
              policy_conv, [-1, params['policy_conv_width'] * go.N * go.N]),
          go.N * go.N + 1)
  
      policy_output = tf.nn.softmax(logits, name='policy_output')
  
      # Value head
      value_conv = mg_conv2d(
          shared_output, filters=params['value_conv_width'], kernel_size=1)
      value_conv = mg_activation(mg_batchn(value_conv, center=False, scale=False))
  
      value_fc_hidden = mg_activation(tf.layers.dense(
          tf.reshape(value_conv, [-1, params['value_conv_width'] * go.N * go.N]),
          params['fc_width']))
      value_output = tf.nn.tanh(
          tf.reshape(tf.layers.dense(value_fc_hidden, 1), [-1]),
          name='value_output')
  
      return policy_output, value_output, logits