def ff_layer(x, hidden_size, output_size, nonlinearity=tf.nn.relu):
  """Single hidden-layer feed-forward network.

  Args:
    x: <float>[batch_size, length, input_size]
    hidden_size: Integer number of dimensions for hidden layer.
    output_size: Integer number of dimensions for output.
    nonlinearity: Function to use for non-linearity.

  Returns:
    <float>[batch_size, length, output_size]
  """
  x = linear_transform(x, hidden_size, "ffn_1", bias=True)
  x = nonlinearity(x)
  x = linear_transform(x, output_size, "ffn_2", bias=True)
  return x
Single hidden-layer feed-forward network.

  Args:
    x: <float>[batch_size, length, input_size]
    hidden_size: Integer number of dimensions for hidden layer.
    output_size: Integer number of dimensions for output.
    nonlinearity: Function to use for non-linearity.

  Returns:
    <float>[batch_size, length, output_size]
  
language/xsp/model/common_layers.py