parameters: (num_blocks, n_in, n_out, mask_type=MASK_EXCLUSIVE)
return_statement: return slices
docstring: Generate the slices for building an autoregressive mask.
docstring_summary: Generate the slices for building an autoregressive mask.
docstring_tokens: Generate the slices for building an autoregressive mask .
function: def _gen_slices(num_blocks, n_in, n_out, mask_type=MASK_EXCLUSIVE):
    """Generate the slices for building an autoregressive mask."""
    # TODO(b/67594795): Better support of dynamic shape.
    slices = []
    col = 0
    d_in = n_in // num_blocks
    d_out = n_out // num_blocks
    row = d_out if mask_type == MASK_EXCLUSIVE else 0
    for _ in range(num_blocks):
      row_slice = slice(row, None)
      col_slice = slice(col, col + d_in)
      slices.append([row_slice, col_slice])
      col += d_in
      row += d_out
    return slices
function_tokens: def _gen_slices ( num_blocks , n_in , n_out , mask_type = MASK_EXCLUSIVE ) : # TODO(b/67594795): Better support of dynamic shape. slices = [ ] col = 0 d_in = n_in // num_blocks d_out = n_out // num_blocks row = d_out if mask_type == MASK_EXCLUSIVE else 0 for _ in range ( num_blocks ) : row_slice = slice ( row , None ) col_slice = slice ( col , col + d_in ) slices . append ( [ row_slice , col_slice ] ) col += d_in row += d_out return slices