mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
3fe437b24b
Updates flake8 to v6.1.0 and fixes a few lints using sed and some ruff tooling.
- Replace `assert(0)` with `raise AssertionError()`
- Remove extraneous parenthesis i.e.
- `assert(a == b)` -> `assert a == b`
- `if(x > y or y < z):`->`if x > y or y < z:`
- And `return('...')` -> `return '...'`
Co-authored-by: Nikita Shulga <2453524+malfet@users.noreply.github.com>
Pull Request resolved: https://github.com/pytorch/pytorch/pull/116591
Approved by: https://github.com/albanD, https://github.com/malfet
43 lines
1.7 KiB
Python
43 lines
1.7 KiB
Python
import threading
|
|
|
|
__all__ = [
|
|
"LinearBlockSparsePattern"
|
|
]
|
|
|
|
def _is_valid_linear_block_sparse_pattern(row_block_size, col_block_size):
|
|
return (row_block_size == 1 and col_block_size == 4) or \
|
|
(row_block_size == 8 and col_block_size == 1)
|
|
|
|
# This is a stop-gap measure as current flow does not allow module
|
|
# specific block sparse pattern.
|
|
# Infact there is no way to convey sparse pattern via module config
|
|
# of quantization flow. Thus using the global context to convey
|
|
# sparsity pattern.
|
|
# Once the flow supports it, this should be removed.
|
|
class LinearBlockSparsePattern:
|
|
rlock = threading.RLock()
|
|
row_block_size = 1
|
|
col_block_size = 4
|
|
prev_row_block_size = 1
|
|
prev_col_block_size = 4
|
|
|
|
def __init__(self, row_block_size=1, col_block_size=4):
|
|
assert _is_valid_linear_block_sparse_pattern(row_block_size, col_block_size)
|
|
LinearBlockSparsePattern.rlock.acquire()
|
|
LinearBlockSparsePattern.prev_row_block_size = LinearBlockSparsePattern.row_block_size
|
|
LinearBlockSparsePattern.prev_col_block_size = LinearBlockSparsePattern.col_block_size
|
|
LinearBlockSparsePattern.row_block_size = row_block_size
|
|
LinearBlockSparsePattern.col_block_size = col_block_size
|
|
|
|
def __enter__(self):
|
|
pass
|
|
|
|
def __exit__(self, exc_type, exc_value, backtrace):
|
|
LinearBlockSparsePattern.row_block_size = LinearBlockSparsePattern.prev_row_block_size
|
|
LinearBlockSparsePattern.col_block_size = LinearBlockSparsePattern.prev_col_block_size
|
|
LinearBlockSparsePattern.rlock.release()
|
|
|
|
@staticmethod
|
|
def block_size():
|
|
return LinearBlockSparsePattern.row_block_size, LinearBlockSparsePattern.col_block_size
|