Assembles the module network and adds it to the graph.

    The internal computation graph is assembled according to the set of
    constraints provided at construction time.

    Args:
      inputs: Tensor containing a batch of transformation parameters.

    Returns:
      A batch of warped grids.

    Raises:
      Error: If the input tensor size is not consistent with the constraints
        passed at construction time.

def _build(self, inputs):
    """Assembles the module network and adds it to the graph.

    The internal computation graph is assembled according to the set of
    constraints provided at construction time.

    Args:
      inputs: Tensor containing a batch of transformation parameters.

    Returns:
      A batch of warped grids.

    Raises:
      Error: If the input tensor size is not consistent with the constraints
        passed at construction time.
    """
    input_shape = tf.shape(inputs)
    input_dtype = inputs.dtype.as_numpy_dtype
    batch_size = tf.expand_dims(input_shape[0], 0)
    number_of_params = inputs.get_shape()[1]
    if number_of_params != self._constraints.num_free_params:
      raise base.Error('Input size is not consistent with constraint '
                       'definition: {} parameters expected, {} provided.'
                       .format(self._constraints.num_free_params,
                               number_of_params))
    num_output_dimensions = len(self._psi) // 3
    def get_input_slice(start, size):
      """Extracts a subset of columns from the input 2D Tensor."""
      return basic.SliceByDim([1], [start], [size])(inputs)

    warped_grid = []
    var_index_offset = 0
    number_of_points = np.prod(self._output_shape)
    for i in xrange(num_output_dimensions):
      if self._psi[i] is not None:
        # The i-th output dimension is not fully specified by the constraints,
        # the graph is setup to perform matrix multiplication in batch mode.
        grid_coord = self._psi[i].astype(input_dtype)

        num_active_vars = self._psi[i].shape[0]
        active_vars = get_input_slice(var_index_offset, num_active_vars)
        warped_coord = tf.matmul(active_vars, grid_coord)
        warped_coord = tf.expand_dims(warped_coord, 1)
        var_index_offset += num_active_vars
        offset = self._psi[num_output_dimensions + i]
        if offset is not None:
          offset = offset.astype(input_dtype)
          # Some entries in the i-th row of the affine matrix were constrained
          # and the corresponding matrix multiplications have been precomputed.
          tiling_params = tf.concat(
              [
                  batch_size, tf.constant(
                      1, shape=(1,)), tf.ones_like(offset.shape)
              ],
              0)
          offset = offset.reshape((1, 1) + offset.shape)
          warped_coord += tf.tile(offset, tiling_params)

      else:
        # The i-th output dimension is fully specified by the constraints, and
        # the corresponding matrix multiplications have been precomputed.
        warped_coord = self._psi[num_output_dimensions + i].astype(input_dtype)
        tiling_params = tf.concat(
            [
                batch_size, tf.constant(
                    1, shape=(1,)), tf.ones_like(warped_coord.shape)
            ],
            0)
        warped_coord = warped_coord.reshape((1, 1) + warped_coord.shape)
        warped_coord = tf.tile(warped_coord, tiling_params)

      warped_coord += self._psi[i + 2 * num_output_dimensions]
      # Need to help TF figuring out shape inference since tiling information
      # is held in Tensors which are not known until run time.
      warped_coord.set_shape([None, 1, number_of_points])
      warped_grid.append(warped_coord)

    # Reshape all the warped coordinates tensors to match the specified output
    # shape and concatenate  into a single matrix.
    grid_shape = self._output_shape + (1,)
    warped_grid = [basic.BatchReshape(grid_shape)(grid) for grid in warped_grid]
    return tf.concat(warped_grid, len(grid_shape))