parameters: (self, update_row_factors,
                              sp_input=None, transpose_input=False)
return_statement: return (new_left_values,
              self.scatter_update(left,
                                  update_indices,
                                  new_left_values,
                                  sharding_func))
docstring: Creates the graph for processing a sparse slice of input.
  
      Args:
        update_row_factors: if True, update the row_factors, else update the
          column factors.
        sp_input: Please refer to comments for update_row_factors and
          update_col_factors.
        transpose_input: If true, logically transpose the input.
  
      Returns:
        A tuple consisting of the following two elements:
        new_values: New values for the row/column factors.
        update_op: An op that assigns the newly computed values to the row/column
          factors.
docstring_summary: Creates the graph for processing a sparse slice of input.
docstring_tokens: Creates the graph for processing a sparse slice of input .
function: def _process_input_helper(self, update_row_factors,
                              sp_input=None, transpose_input=False):
      """Creates the graph for processing a sparse slice of input.
  
      Args:
        update_row_factors: if True, update the row_factors, else update the
          column factors.
        sp_input: Please refer to comments for update_row_factors and
          update_col_factors.
        transpose_input: If true, logically transpose the input.
  
      Returns:
        A tuple consisting of the following two elements:
        new_values: New values for the row/column factors.
        update_op: An op that assigns the newly computed values to the row/column
          factors.
      """
      assert isinstance(sp_input, ops.SparseTensor)
  
      if update_row_factors:
        left = self._row_factors
        right = self._col_factors_cache
        row_weights = self._row_wt_cache
        col_weights = self._col_wt_cache
        sharding_func = WALSModel._get_sharding_func(self._input_rows,
                                                     self._num_row_shards)
        right_length = self._input_cols
      else:
        left = self._col_factors
        right = self._row_factors_cache
        row_weights = self._col_wt_cache
        col_weights = self._row_wt_cache
        sharding_func = WALSModel._get_sharding_func(self._input_cols,
                                                     self._num_col_shards)
        right_length = self._input_rows
        transpose_input = not transpose_input
  
      # Note that the row indices of sp_input are based on the original full input
      # Here we reindex the rows and give them contiguous ids starting at 0.
      # We use tf.unique to achieve this reindexing. Note that this is done so
      # that the downstream kernel can assume that the input is "dense" along the
      # row dimension.
      row_ids, col_ids = tf.split(1, 2, sp_input.indices)
  
      if transpose_input:
        update_indices, all_ids = tf.unique(col_ids[:, 0])
        col_ids = tf.expand_dims(tf.cast(all_ids, tf.int64), 1)
      else:
        update_indices, all_ids = tf.unique(row_ids[:, 0])
        row_ids = tf.expand_dims(tf.cast(all_ids, tf.int64), 1)
  
      num_rows = tf.cast(tf.shape(update_indices)[0], tf.int64)
      row_shape = tf.constant([right_length], tf.int64)
      col_shape = [num_rows]
  
      new_sp_indices = tf.concat(1, [row_ids, col_ids])
      new_sp_shape = (tf.concat(0, [row_shape, col_shape]) if transpose_input
                      else tf.concat(0, [col_shape, row_shape]))
      new_sp_input = tf.SparseTensor(indices=new_sp_indices,
                                     values=sp_input.values, shape=new_sp_shape)
  
      # Compute lhs and rhs of the normal equations
      total_lhs = (self._unobserved_weight *
                   tf.matmul(right, right, transpose_a=True))
      if self._regularization is not None:
        total_lhs += self._regularization
      if self._row_weights is None:
        # Special case of ALS. Use a much simpler update rule.
        total_rhs = (self._unobserved_weight *
                     tf.sparse_tensor_dense_matmul(new_sp_input, right,
                                                   adjoint_a=transpose_input))
        # TODO(rmlarsen): handle transposing in tf.matrix_solve instead of
        # transposing explicitly.
        # TODO(rmlarsen): multi-thread tf.matrix_solve.
        new_left_values = tf.transpose(tf.matrix_solve(total_lhs,
                                                       tf.transpose(total_rhs)))
      else:
        row_weights_slice = tf.gather(row_weights, update_indices)
        partial_lhs, total_rhs = wals_compute_partial_lhs_and_rhs(
            right,
            col_weights,
            self._unobserved_weight,
            row_weights_slice,
            new_sp_input.indices,
            new_sp_input.values,
            num_rows,
            transpose_input,
            name="wals_compute_partial_lhs_rhs")
        total_lhs = tf.expand_dims(total_lhs, 0) + partial_lhs
        total_rhs = tf.expand_dims(total_rhs, -1)
        new_left_values = tf.squeeze(tf.batch_matrix_solve(total_lhs, total_rhs),
                                     [2])
  
      return (new_left_values,
              self.scatter_update(left,
                                  update_indices,
                                  new_left_values,
                                  sharding_func))
function_tokens: def _process_input_helper ( self , update_row_factors , sp_input = None , transpose_input = False ) : assert isinstance ( sp_input , ops . SparseTensor ) if update_row_factors : left = self . _row_factors right = self . _col_factors_cache row_weights = self . _row_wt_cache col_weights = self . _col_wt_cache sharding_func = WALSModel . _get_sharding_func ( self . _input_rows , self . _num_row_shards ) right_length = self . _input_cols else : left = self . _col_factors right = self . _row_factors_cache row_weights = self . _col_wt_cache col_weights = self . _row_wt_cache sharding_func = WALSModel . _get_sharding_func ( self . _input_cols , self . _num_col_shards ) right_length = self . _input_rows transpose_input = not transpose_input # Note that the row indices of sp_input are based on the original full input # Here we reindex the rows and give them contiguous ids starting at 0. # We use tf.unique to achieve this reindexing. Note that this is done so # that the downstream kernel can assume that the input is "dense" along the # row dimension. row_ids , col_ids = tf . split ( 1 , 2 , sp_input . indices ) if transpose_input : update_indices , all_ids = tf . unique ( col_ids [ : , 0 ] ) col_ids = tf . expand_dims ( tf . cast ( all_ids , tf . int64 ) , 1 ) else : update_indices , all_ids = tf . unique ( row_ids [ : , 0 ] ) row_ids = tf . expand_dims ( tf . cast ( all_ids , tf . int64 ) , 1 ) num_rows = tf . cast ( tf . shape ( update_indices ) [ 0 ] , tf . int64 ) row_shape = tf . constant ( [ right_length ] , tf . int64 ) col_shape = [ num_rows ] new_sp_indices = tf . concat ( 1 , [ row_ids , col_ids ] ) new_sp_shape = ( tf . concat ( 0 , [ row_shape , col_shape ] ) if transpose_input else tf . concat ( 0 , [ col_shape , row_shape ] ) ) new_sp_input = tf . SparseTensor ( indices = new_sp_indices , values = sp_input . values , shape = new_sp_shape ) # Compute lhs and rhs of the normal equations total_lhs = ( self . _unobserved_weight * tf . matmul ( right , right , transpose_a = True ) ) if self . _regularization is not None : total_lhs += self . _regularization if self . _row_weights is None : # Special case of ALS. Use a much simpler update rule. total_rhs = ( self . _unobserved_weight * tf . sparse_tensor_dense_matmul ( new_sp_input , right , adjoint_a = transpose_input ) ) # TODO(rmlarsen): handle transposing in tf.matrix_solve instead of # transposing explicitly. # TODO(rmlarsen): multi-thread tf.matrix_solve. new_left_values = tf . transpose ( tf . matrix_solve ( total_lhs , tf . transpose ( total_rhs ) ) ) else : row_weights_slice = tf . gather ( row_weights , update_indices ) partial_lhs , total_rhs = wals_compute_partial_lhs_and_rhs ( right , col_weights , self . _unobserved_weight , row_weights_slice , new_sp_input . indices , new_sp_input . values , num_rows , transpose_input , name = "wals_compute_partial_lhs_rhs" ) total_lhs = tf . expand_dims ( total_lhs , 0 ) + partial_lhs total_rhs = tf . expand_dims ( total_rhs , - 1 ) new_left_values = tf . squeeze ( tf . batch_matrix_solve ( total_lhs , total_rhs ) , [ 2 ] ) return ( new_left_values , self . scatter_update ( left , update_indices , new_left_values , sharding_func ) )