mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
abbd71d29d
Reproduce command: ```bash ghstack checkout https://github.com/pytorch/pytorch/pull/138443 git checkout HEAD~1 torch/ lintrunner -a --take "PYFMT" --all-files ``` Pull Request resolved: https://github.com/pytorch/pytorch/pull/138443 Approved by: https://github.com/ezyang
57 lines
1.8 KiB
Python
57 lines
1.8 KiB
Python
import unittest
|
|
|
|
from ..pass_manager import (
|
|
inplace_wrapper,
|
|
PassManager,
|
|
these_before_those_pass_constraint,
|
|
this_before_that_pass_constraint,
|
|
)
|
|
|
|
|
|
class TestPassManager(unittest.TestCase):
|
|
def test_pass_manager_builder(self) -> None:
|
|
passes = [lambda x: 2 * x for _ in range(10)]
|
|
pm = PassManager(passes)
|
|
pm.validate()
|
|
|
|
def test_this_before_that_pass_constraint(self) -> None:
|
|
passes = [lambda x: 2 * x for _ in range(10)]
|
|
pm = PassManager(passes)
|
|
|
|
# add unfulfillable constraint
|
|
pm.add_constraint(this_before_that_pass_constraint(passes[-1], passes[0]))
|
|
|
|
self.assertRaises(RuntimeError, pm.validate)
|
|
|
|
def test_these_before_those_pass_constraint(self) -> None:
|
|
passes = [lambda x: 2 * x for _ in range(10)]
|
|
constraint = these_before_those_pass_constraint(passes[-1], passes[0])
|
|
pm = PassManager([inplace_wrapper(p) for p in passes])
|
|
|
|
# add unfulfillable constraint
|
|
pm.add_constraint(constraint)
|
|
|
|
self.assertRaises(RuntimeError, pm.validate)
|
|
|
|
def test_two_pass_managers(self) -> None:
|
|
"""Make sure we can construct the PassManager twice and not share any
|
|
state between them"""
|
|
|
|
passes = [lambda x: 2 * x for _ in range(3)]
|
|
constraint = these_before_those_pass_constraint(passes[0], passes[1])
|
|
pm1 = PassManager()
|
|
for p in passes:
|
|
pm1.add_pass(p)
|
|
pm1.add_constraint(constraint)
|
|
output1 = pm1(1)
|
|
self.assertEqual(output1, 2**3)
|
|
|
|
passes = [lambda x: 3 * x for _ in range(3)]
|
|
constraint = these_before_those_pass_constraint(passes[0], passes[1])
|
|
pm2 = PassManager()
|
|
for p in passes:
|
|
pm2.add_pass(p)
|
|
pm2.add_constraint(constraint)
|
|
output2 = pm2(1)
|
|
self.assertEqual(output2, 3**3)
|