[TorchScript, PT2] Add torch._check compatibility support (#159988)

Summary:
Add support for torch._check() in TorchScript jit.script frontend.

* It will be special cased to behave like torch._assert, turned into an if + raise exception.

Test Plan:
Unit tests

Rollback Plan:

Differential Revision: D79744604

Pull Request resolved: https://github.com/pytorch/pytorch/pull/159988
Approved by: https://github.com/davidberard98
This commit is contained in:
Yanan Cao (PyTorch)
2025-08-08 23:14:13 +00:00
committed by PyTorch MergeBot
parent 566c6d52ef
commit 731ee31f7b
5 changed files with 260 additions and 7 deletions

View File

@ -1222,6 +1222,8 @@ std::shared_ptr<SugaredValue> toSugaredValue(
} else if (
obj.ptr() == py::module::import("torch.jit").attr("isinstance").ptr()) {
return SpecialFormValue::create(prim::isinstance);
} else if (obj.ptr() == py::module::import("torch").attr("_check").ptr()) {
return std::make_shared<TorchCheckValue>();
#ifdef USE_RPC
// RPC module is only available when build flag "USE_DISTRIBUTED" is on.
} else if (