original_docstring: Return a 1D integer tensor of counts such that
  
      - len(counts) == n_rows
      - 0 <= counts.min()
      - counts.max() <= n_cols
      - counts.sum() == nnz
      - counts[-1] - counts[0] is as large as possible value
      - the distribution of different count values is as uniform as
        possible
  
      Application
      -----------
   
      cumsum(cat([0], distribute_column_indices(n_rows, n_cols, nnz))
      gives a valid crow indices of the CSR tensor format.
  
      Description of the algoritm
      ---------------------------
  
      The counts can be represented graphically as the following
      histogram:
  
        ^ number of columns per row
        |
        |
  
          *   *   *   *   ###
         **  **  **  **+ o###
        *** *** *** ***+oo###
        @@@@@@@@@@@@@@@@@@###
        @@@@@@@@@@@@@@@@@@###      --> row indices
  
      where regions are named and marked as follows
  
      - no counts                    : space
      - final correction             : +
      - an incomplete sawtooth       : o
      - a sequence of full sawteeth  : *
      - lower rectangle              : @
      - right rectangle              : #
  
      The sizes of these regions depend on the input parameter and are
      found such that the following invariants hold:
  
      - the number of all non-space symbols is equal to nnz
      - the width of the histogram is equal to n_rows
      - the height of the histogram is equal to n_cols
  
      For very small nnz values, only the incomplete sawtooth is
      generated. For increasing nnz and when n_rows < n_cols, the lower
      rectangle is introduced. Further increasing of nnz leads to
      generation of full sawteeth followed by the incomplete
      sawtooth. For larger values of nnz, the right rectangle will
      emerge which will decrease the width of the lower rectangle as
      well as the sequence of sawteeth. Finally, for very large nnz
      values, the right rectangle will squash the incomplete sawtooth to
      zero-width so that at nnz == n_rows * n_cols the histogram
      contains only the right rectangle. During the whole process, the
      first position of the incomplete rectangle, denoted as the final
      correction, is most volatile to increments of nnz. Notice that the
      lower rectangle and the sequence of full sawteeth never appear in
      the same histogram.
  
      Performance notes
      -----------------
  
      Computing
        counts = distribute_column_indices(...)
      takes about the same time as computing the tuple
        (counts.sum(), counts.max(), counts.min())
  
      Using numpy backend is about twice faster than torch backend.
  
      Using python backend is about 1.5 times faster than numpy backend
      for small nnz sizes but 4 or more times slower for large nnz
      because of explicit python for loop.
  
      For maximum performance, distribute_column_indices should be
      implemented in C/C++.
  
      These performance results are obtained by calling
      distribute_column_indices for all allowed nnz values. For specific
      nnz values, the performance results may vary creatly.
      
docstring: 
  
  Application
  
  cumsum(cat([0], distribute_column_indices(n_rows, n_cols, nnz))
  gives a valid crow indices of the CSR tensor format.
  
  Description of the algoritm
  
  The counts can be represented graphically as the following
  histogram.
  
  ^ number of columns per row
  |
  |
  
  + o
  +oo
  
  > row indices
  
  
  
  
  
  The sizes of these regions depend on the input parameter and are
  found such that the following invariants hold.
  
  the number of all non-space symbols is equal to nnz
  the width of the histogram is equal to n_rows
  the height of the histogram is equal to n_cols
  
  For very small nnz values, only the incomplete sawtooth is
  generated. For increasing nnz and when n_rows < n_cols, the lower
  rectangle is introduced. Further increasing of nnz leads to
  generation of full sawteeth followed by the incomplete
  sawtooth. For larger values of nnz, the right rectangle will
  emerge which will decrease the width of the lower rectangle as
  well as the sequence of sawteeth. Finally, for very large nnz
  values, the right rectangle will squash the incomplete sawtooth to
  zero-width so that at nnz == n_rows * n_cols the histogram
  contains only the right rectangle. During the whole process, the
  first position of the incomplete rectangle, denoted as the final
  correction, is most volatile to increments of nnz. Notice that the
  lower rectangle and the sequence of full sawteeth never appear in
  the same histogram.
  
  Performance notes
  
  Computing
  counts = distribute_column_indices(...)
  takes about the same time as computing the tuple
  (counts.sum(), counts.max(), counts.min())
  
  Using numpy backend is about twice faster than torch backend.
  
  Using python backend is about 1.5 times faster than numpy backend
  for small nnz sizes but 4 or more times slower for large nnz
  because of explicit python for loop.
  
  For maximum performance, distribute_column_indices should be
  implemented in C/C++.
  
  These performance results are obtained by calling
  distribute_column_indices for all allowed nnz values. For specific
  nnz values, the performance results may vary creatly.
docstring_tokens: Application cumsum ( cat ( [ 0 ] distribute_column_indices ( n_rows n_cols nnz )) gives a valid crow indices of the CSR tensor format . Description of the algoritm The counts can be represented graphically as the following histogram . ^ number of columns per row | | + o + oo > row indices The sizes of these regions depend on the input parameter and are found such that the following invariants hold . the number of all non - space symbols is equal to nnz the width of the histogram is equal to n_rows the height of the histogram is equal to n_cols For very small nnz values only the incomplete sawtooth is generated . For increasing nnz and when n_rows < n_cols the lower rectangle is introduced . Further increasing of nnz leads to generation of full sawteeth followed by the incomplete sawtooth . For larger values of nnz the right rectangle will emerge which will decrease the width of the lower rectangle as well as the sequence of sawteeth . Finally for very large nnz values the right rectangle will squash the incomplete sawtooth to zero - width so that at nnz == n_rows * n_cols the histogram contains only the right rectangle . During the whole process the first position of the incomplete rectangle denoted as the final correction is most volatile to increments of nnz . Notice that the lower rectangle and the sequence of full sawteeth never appear in the same histogram . Performance notes Computing counts = distribute_column_indices ( ... ) takes about the same time as computing the tuple ( counts . sum () counts . max () counts . min () ) Using numpy backend is about twice faster than torch backend . Using python backend is about 1 . 5 times faster than numpy backend for small nnz sizes but 4 or more times slower for large nnz because of explicit python for loop . For maximum performance distribute_column_indices should be implemented in C / C ++ . These performance results are obtained by calling distribute_column_indices for all allowed nnz values . For specific nnz values the performance results may vary creatly .
code: def distribute_column_indices(n_rows, n_cols, nnz, backend='numpy'):
      """Return a 1D integer tensor of counts such that
  
      - len(counts) == n_rows
      - 0 <= counts.min()
      - counts.max() <= n_cols
      - counts.sum() == nnz
      - counts[-1] - counts[0] is as large as possible value
      - the distribution of different count values is as uniform as
        possible
  
      Application
      -----------
   
      cumsum(cat([0], distribute_column_indices(n_rows, n_cols, nnz))
      gives a valid crow indices of the CSR tensor format.
  
      Description of the algoritm
      ---------------------------
  
      The counts can be represented graphically as the following
      histogram:
  
        ^ number of columns per row
        |
        |
  
          *   *   *   *   ###
         **  **  **  **+ o###
        *** *** *** ***+oo###
        @@@@@@@@@@@@@@@@@@###
        @@@@@@@@@@@@@@@@@@###      --> row indices
  
      where regions are named and marked as follows
  
      - no counts                    : space
      - final correction             : +
      - an incomplete sawtooth       : o
      - a sequence of full sawteeth  : *
      - lower rectangle              : @
      - right rectangle              : #
  
      The sizes of these regions depend on the input parameter and are
      found such that the following invariants hold:
  
      - the number of all non-space symbols is equal to nnz
      - the width of the histogram is equal to n_rows
      - the height of the histogram is equal to n_cols
  
      For very small nnz values, only the incomplete sawtooth is
      generated. For increasing nnz and when n_rows < n_cols, the lower
      rectangle is introduced. Further increasing of nnz leads to
      generation of full sawteeth followed by the incomplete
      sawtooth. For larger values of nnz, the right rectangle will
      emerge which will decrease the width of the lower rectangle as
      well as the sequence of sawteeth. Finally, for very large nnz
      values, the right rectangle will squash the incomplete sawtooth to
      zero-width so that at nnz == n_rows * n_cols the histogram
      contains only the right rectangle. During the whole process, the
      first position of the incomplete rectangle, denoted as the final
      correction, is most volatile to increments of nnz. Notice that the
      lower rectangle and the sequence of full sawteeth never appear in
      the same histogram.
  
      Performance notes
      -----------------
  
      Computing
        counts = distribute_column_indices(...)
      takes about the same time as computing the tuple
        (counts.sum(), counts.max(), counts.min())
  
      Using numpy backend is about twice faster than torch backend.
  
      Using python backend is about 1.5 times faster than numpy backend
      for small nnz sizes but 4 or more times slower for large nnz
      because of explicit python for loop.
  
      For maximum performance, distribute_column_indices should be
      implemented in C/C++.
  
      These performance results are obtained by calling
      distribute_column_indices for all allowed nnz values. For specific
      nnz values, the performance results may vary creatly.
      """
      assert n_rows > 0 and n_cols > 0 and nnz >= 0 and nnz <= n_rows * n_cols
  
      if backend == 'torch':
          index_dtype = torch.int32
          counts = torch.zeros(n_rows, dtype=index_dtype)
      elif backend == 'numpy':
          index_dtype = numpy.int32
          counts = numpy.zeros(n_rows, dtype=index_dtype)
      elif backend == 'python':
          counts = [0] * n_rows
      else:
          raise NotImplementedError(backend)
          
      # the number of counts per single full sawtooth:
      M = n_cols * (n_cols + 1) // 2
      # the width of an incomplete sawtooth:
      K = n_rows % (n_cols + 1)
      # the total number of counts in the sequence of sawteeth
      N = M * (n_rows // (n_cols + 1)) + K * (K - 1) // 2
  
      # Right rectangle:
      if N == 0 or nnz < max(N, n_cols):
          n = 0
      else:
          # Find the width of the right rectangle
          left, n, N = 0, n_rows-1, 0
          while n - left > 1:
              middle = (n + left) // 2
              K = (n_rows - middle) % (n_cols + 1)
              N2 = M * ((n_rows - middle) // (n_cols + 1)) + K * (K - 1) // 2
              if N2 == 0 or middle * n_cols > nnz - max(N2, n_cols):
                  n, N = middle, N2
              else:
                  left = middle
          # Fill the right rectangle:
          if backend == 'torch':
              counts[-n:].fill_(n_cols)
          elif backend == 'python':
              counts[-n:] = [n_cols] * n
          else:
              counts[-n:] = n_cols
  
      nnz2 = nnz - n * n_cols
      n_rows2 = n_rows - n
  
      # Lower rectangle:
      if N == 0 or nnz2 < max(N, n_rows2):
          m = 0
      else:
          # Find the height of lower rectangle
          left, m, M, N = 0, n_cols-1, 0, 0
          while m - left > 1:
              middle = (m + left) // 2
              K = n_rows2 % (n_cols - middle + 1)
              M2 = (n_cols - middle) * (n_cols - middle + 1) // 2
              N2 = M2 * (n_rows2 // (n_cols - middle + 1)) + K * (K - 1) // 2
              if N2 == 0 or middle * n_rows2 > nnz2 - max(N2, n_rows2):
                  m, M, N = middle, M2, N2
              else:
                  left = middle
          # Fill the lower rectangle:
          if backend == 'torch':
              counts[:n_rows2].fill_(m)
          elif backend == 'python':
              counts[:n_rows2] = [m] * n_rows2
          else:
              counts[:n_rows2] = m
  
      n_cols2 = n_cols - m
      nnz2 = nnz2 - m * n_rows2
  
      if N == 0:
          # no sawteeth
          counts[0] = nnz2
      else:
          offset = (nnz2 // M) * (n_cols2 + 1)
          k = math.isqrt(2 * (nnz2 % M))
          if k * (k + 1) > 2 * (nnz2 % M):
              k -= 1
          #correction = nnz2 - k * (k + 1) // 2 - (nnz2 // M) * M
          correction = nnz2 % M - k * (k + 1) // 2
          if backend == 'torch':
              tmp = torch.arange(max(offset, k + 1), dtype=counts.dtype)
          elif backend == 'numpy':
              tmp = numpy.arange(max(offset, k + 1), dtype=counts.dtype)
          elif backend == 'python':
              pass
          else:
              raise NotImplementedError(backend)
          if offset:
              # Fill the sequence of full sawteeth:
              if backend == 'python':
                  for i in range(offset):
                      counts[i] = i % (n_cols2 + 1)
              else:
                  counts[:offset] = tmp[:offset] % (n_cols2 + 1)
              # Since the height of the full sawteeth is n_cols, the
              # lower rectangle must have zero-height.
          if correction:
              # Final correction to the incomplete sawtooth
              if backend != 'python':
                  tmp[0] += correction
          if k:
              # The incomplete sawtooth added on top of the bottom
              # rectangle:
              if m:
                  if backend == 'python':
                      for i in range(k + 1):
                          counts[offset + i] += i
                      counts[offset] += correction
                  else:
                      counts[offset:offset+k+1] += tmp[:k + 1]
              else:
                  if backend == 'python':
                      for i in range(k + 1):
                          counts[offset + i] = i
                      counts[offset] += correction
                  else:
                      counts[offset:offset+k+1] = tmp[:k + 1]
  
      if 1:
          # Check the invariants (it will increase the computation time by 2x or more):
          assert sum(counts) == nnz, (counts, nnz)
          assert min(counts) >= 0
          assert max(counts) <= n_cols
  
      return counts
code_tokens: def distribute_column_indices ( n_rows , n_cols , nnz , backend = 'numpy' ) : assert n_rows > 0 and n_cols > 0 and nnz >= 0 and nnz <= n_rows * n_cols if backend == 'torch' : index_dtype = torch . int32 counts = torch . zeros ( n_rows , dtype = index_dtype ) elif backend == 'numpy' : index_dtype = numpy . int32 counts = numpy . zeros ( n_rows , dtype = index_dtype ) elif backend == 'python' : counts = [ 0 ] * n_rows else : raise NotImplementedError ( backend ) M = n_cols * ( n_cols + 1 ) // 2 K = n_rows % ( n_cols + 1 ) N = M * ( n_rows // ( n_cols + 1 ) ) + K * ( K - 1 ) // 2 if N == 0 or nnz < max ( N , n_cols ) : n = 0 else : left , n , N = 0 , n_rows - 1 , 0 while n - left > 1 : middle = ( n + left ) // 2 K = ( n_rows - middle ) % ( n_cols + 1 ) N2 = M * ( ( n_rows - middle ) // ( n_cols + 1 ) ) + K * ( K - 1 ) // 2 if N2 == 0 or middle * n_cols > nnz - max ( N2 , n_cols ) : n , N = middle , N2 else : left = middle if backend == 'torch' : counts [ - n : ] . fill_ ( n_cols ) elif backend == 'python' : counts [ - n : ] = [ n_cols ] * n else : counts [ - n : ] = n_cols nnz2 = nnz - n * n_cols n_rows2 = n_rows - n if N == 0 or nnz2 < max ( N , n_rows2 ) : m = 0 else : left , m , M , N = 0 , n_cols - 1 , 0 , 0 while m - left > 1 : middle = ( m + left ) // 2 K = n_rows2 % ( n_cols - middle + 1 ) M2 = ( n_cols - middle ) * ( n_cols - middle + 1 ) // 2 N2 = M2 * ( n_rows2 // ( n_cols - middle + 1 ) ) + K * ( K - 1 ) // 2 if N2 == 0 or middle * n_rows2 > nnz2 - max ( N2 , n_rows2 ) : m , M , N = middle , M2 , N2 else : left = middle if backend == 'torch' : counts [ : n_rows2 ] . fill_ ( m ) elif backend == 'python' : counts [ : n_rows2 ] = [ m ] * n_rows2 else : counts [ : n_rows2 ] = m n_cols2 = n_cols - m nnz2 = nnz2 - m * n_rows2 if N == 0 : counts [ 0 ] = nnz2 else : offset = ( nnz2 // M ) * ( n_cols2 + 1 ) k = math . isqrt ( 2 * ( nnz2 % M ) ) if k * ( k + 1 ) > 2 * ( nnz2 % M ) : k -= 1 correction = nnz2 % M - k * ( k + 1 ) // 2 if backend == 'torch' : tmp = torch . arange ( max ( offset , k + 1 ) , dtype = counts . dtype ) elif backend == 'numpy' : tmp = numpy . arange ( max ( offset , k + 1 ) , dtype = counts . dtype ) elif backend == 'python' : pass else : raise NotImplementedError ( backend ) if offset : if backend == 'python' : for i in range ( offset ) : counts [ i ] = i % ( n_cols2 + 1 ) else : counts [ : offset ] = tmp [ : offset ] % ( n_cols2 + 1 ) if correction : if backend != 'python' : tmp [ 0 ] += correction if k : if m : if backend == 'python' : for i in range ( k + 1 ) : counts [ offset + i ] += i counts [ offset ] += correction else : counts [ offset : offset + k + 1 ] += tmp [ : k + 1 ] else : if backend == 'python' : for i in range ( k + 1 ) : counts [ offset + i ] = i counts [ offset ] += correction else : counts [ offset : offset + k + 1 ] = tmp [ : k + 1 ] if 1 : assert sum ( counts ) == nnz , ( counts , nnz ) assert min ( counts ) >= 0 assert max ( counts ) <= n_cols return counts
short_docstring: Return a 1D integer tensor of counts such that
   len(counts) == n_rows
   0 <= counts.min()
   counts.max() <= n_cols
   counts.sum() == nnz
   counts[-1] - counts[0] is as large as possible value
   the distribution of different count values is as uniform as
  possible
short_docstring_tokens: Return a 1D integer tensor of counts such that len ( counts ) == n_rows 0 < = counts . min () counts . max () < = n_cols counts . sum () == nnz counts [ - 1 ] - counts [ 0 ] is as large as possible value the distribution of different count values is as uniform as possible
comment: """Return a 1D integer tensor of counts such that
  
      - len(counts) == n_rows
      - 0 <= counts.min()
      - counts.max() <= n_cols
      - counts.sum() == nnz
      - counts[-1] - counts[0] is as large as possible value
      - the distribution of different count values is as uniform as
        possible
  
      Application
      -----------
   
      cumsum(cat([0], distribute_column_indices(n_rows, n_cols, nnz))
      gives a valid crow indices of the CSR tensor format.
  
      Description of the algoritm
      ---------------------------
  
      The counts can be represented graphically as the following
      histogram:
  
        ^ number of columns per row
        |
        |
  
          *   *   *   *   ###
         **  **  **  **+ o###
        *** *** *** ***+oo###
        @@@@@@@@@@@@@@@@@@###
        @@@@@@@@@@@@@@@@@@###      --> row indices
  
      where regions are named and marked as follows
  
      - no counts                    : space
      - final correction             : +
      - an incomplete sawtooth       : o
      - a sequence of full sawteeth  : *
      - lower rectangle              : @
      - right rectangle              : #
  
      The sizes of these regions depend on the input parameter and are
      found such that the following invariants hold:
  
      - the number of all non-space symbols is equal to nnz
      - the width of the histogram is equal to n_rows
      - the height of the histogram is equal to n_cols
  
      For very small nnz values, only the incomplete sawtooth is
      generated. For increasing nnz and when n_rows < n_cols, the lower
      rectangle is introduced. Further increasing of nnz leads to
      generation of full sawteeth followed by the incomplete
      sawtooth. For larger values of nnz, the right rectangle will
      emerge which will decrease the width of the lower rectangle as
      well as the sequence of sawteeth. Finally, for very large nnz
      values, the right rectangle will squash the incomplete sawtooth to
      zero-width so that at nnz == n_rows * n_cols the histogram
      contains only the right rectangle. During the whole process, the
      first position of the incomplete rectangle, denoted as the final
      correction, is most volatile to increments of nnz. Notice that the
      lower rectangle and the sequence of full sawteeth never appear in
      the same histogram.
  
      Performance notes
      -----------------
  
      Computing
        counts = distribute_column_indices(...)
      takes about the same time as computing the tuple
        (counts.sum(), counts.max(), counts.min())
  
      Using numpy backend is about twice faster than torch backend.
  
      Using python backend is about 1.5 times faster than numpy backend
      for small nnz sizes but 4 or more times slower for large nnz
      because of explicit python for loop.
  
      For maximum performance, distribute_column_indices should be
      implemented in C/C++.
  
      These performance results are obtained by calling
      distribute_column_indices for all allowed nnz values. For specific
      nnz values, the performance results may vary creatly.
      """$SEP$# the number of counts per single full sawtooth:$SEP$# the width of an incomplete sawtooth:$SEP$# the total number of counts in the sequence of sawteeth$SEP$# Right rectangle:$SEP$# Find the width of the right rectangle$SEP$# Fill the right rectangle:$SEP$# Lower rectangle:$SEP$# Find the height of lower rectangle$SEP$# Fill the lower rectangle:$SEP$# no sawteeth$SEP$#correction = nnz2 - k * (k + 1) // 2 - (nnz2 // M) * M$SEP$# Fill the sequence of full sawteeth:$SEP$# Since the height of the full sawteeth is n_cols, the$SEP$# lower rectangle must have zero-height.$SEP$# Final correction to the incomplete sawtooth$SEP$# The incomplete sawtooth added on top of the bottom$SEP$# rectangle:$SEP$# Check the invariants (it will increase the computation time by 2x or more):