code: def _build(self, inputs):
      
      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
docs: 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.