docs: 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))