Dynamic unroll across input objects.

    Args:
      inputs: tensor (batch x num_objects x feature). Objects to sort.

    Returns:
      Tensor (batch x num_objects); logits indicating the reference objects.

def _build(self, inputs):
    """Dynamic unroll across input objects.

    Args:
      inputs: tensor (batch x num_objects x feature). Objects to sort.

    Returns:
      Tensor (batch x num_objects); logits indicating the reference objects.
    """
    batch_size = inputs.get_shape()[0]
    output_sequence, _ = tf.nn.dynamic_rnn(
        cell=self._core,
        inputs=inputs,
        time_major=False,
        initial_state=self._core.initial_state(
            batch_size, trainable=False)
    )
    outputs = snt.BatchFlatten()(output_sequence[:, -1, :])
    outputs = self._final_mlp(outputs)
    logits = snt.Linear(self._target_size)(outputs)
    return logits