frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

Update to TorchFix 0.4.0 (#119424)

Browse Source 
`torch.library.Library` updated to `torch.library._scoped_library` in files with many tests where it seems obvious to do, otherwise `noqa: TOR901` added - see https://github.com/pytorch/pytorch/pull/118318 for more context.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/119424
Approved by: https://github.com/zou3519
This commit is contained in:
Sergii Dymchenko
2024-02-12 23:30:08 +00:00
committed by PyTorch MergeBot
parent 5acd1f0f7d
commit bd9db6a9c7
22 changed files with 348 additions and 355 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
.flake8
View File

@ -2,7 +2,7 @@
# NOTE: **Mirror any changes** to this file the [tool.ruff] config in pyproject.toml
# before we can fully move to use ruff
enable-extensions = G
select = B,C,E,F,G,P,SIM1,T4,W,B9,TOR0,TOR1,TOR2
select = B,C,E,F,G,P,SIM1,T4,W,B9,TOR0,TOR1,TOR2,TOR9
max-line-length = 120
# C408 ignored because we like the dict keyword argument syntax
# E501 is not flexible enough, we're using B950 instead
@ -27,6 +27,9 @@ ignore =
# TODO(kit1980): fix all TOR102 issues
# `torch.load` without `weights_only` parameter is unsafe
TOR102,
# TODO(kit1980): resolve all TOR003 issues
# pass `use_reentrant` explicitly to `checkpoint`.
TOR003
per-file-ignores =
__init__.py: F401
test/**: F821
@ -36,6 +39,21 @@ per-file-ignores =
torchgen/executorch/api/types/__init__.py: F401,F403
test/dynamo/test_higher_order_ops.py: B950
torch/testing/_internal/dynamo_test_failures.py: B950
# TOR901 is only for test, we want to ignore it for everything else.
# It's not easy to configure this without affecting other per-file-ignores,
# so we explicitly list every file where it's violated outside of test.
torch/__init__.py: F401,TOR901
torch/_custom_op/impl.py: TOR901
torch/_export/serde/upgrade.py: TOR901
torch/_functorch/vmap.py: TOR901
torch/_inductor/test_operators.py: TOR901
torch/_library/abstract_impl.py: TOR901
torch/_meta_registrations.py: TOR901
torch/_prims/__init__.py: F401,TOR901
torch/_prims/rng_prims.py: TOR901
torch/ao/quantization/fx/_decomposed.py: TOR901
torch/distributed/_functional_collectives.py: TOR901
torch/distributed/_spmd/data_parallel.py: TOR901
optional-ascii-coding = True
exclude =
./.git,

2
.lintrunner.toml
View File

@ -46,7 +46,7 @@ init_command = [
'mccabe==0.7.0',
'pycodestyle==2.11.1',
'pyflakes==3.1.0',
'torchfix==0.2.0',
'torchfix==0.4.0',
]

2
requirements-flake8.txt
View File

@ -8,4 +8,4 @@ flake8-pyi==20.5.0
mccabe==0.6.1
pycodestyle==2.6.0
pyflakes==2.2.0
torchfix==0.2.0
torchfix==0.4.0

2
test/distributed/_spmd/test_tracing.py
View File

@ -256,7 +256,7 @@ def ddm_backward(grad: torch.Tensor) -> torch.Tensor:
return grad
dummy_lib = torch.library.Library("dummy", "DEF")
dummy_lib = torch.library.Library("dummy", "DEF") # noqa: TOR901
dummy_lib.define("ddm(Tensor x) -> Tensor")
dummy_lib.impl("ddm", ddm, "CompositeExplicitAutograd")
dummy_lib.define("ddm_backward(Tensor x) -> Tensor")

2
test/dynamo/test_aot_autograd.py
View File

@ -25,7 +25,7 @@ def maybe_dupe_op(x):
aten = torch.ops.aten
lib = torch.library.Library("custom", "DEF")
lib = torch.library.Library("custom", "DEF") # noqa: TOR901
lib.define("maybe_dupe_op(Tensor a) -> (Tensor, Tensor)")
lib.impl("maybe_dupe_op", maybe_dupe_op, "CPU")
lib.impl("maybe_dupe_op", maybe_dupe_op, "Meta")

2
test/dynamo/test_decorators.py
View File

@ -39,7 +39,7 @@ class DecoratorTests(torch._dynamo.test_case.TestCase):
import torch.library
from torch.library import Library
foo = Library("foo", "DEF")
foo = Library("foo", "DEF") # noqa: TOR901
foo.define("custom(Tensor self) -> Tensor")
# Dynamic shape data dependent operator. For static shape compilation, Dynamo

2
test/dynamo/test_repros.py
View File

@ -43,7 +43,7 @@ from torch.testing._internal.common_utils import (
_orig_module_call = torch.nn.Module.__call__
# Custom operator that only supports CPU and Meta
lib = torch.library.Library("test_sample", "DEF")
lib = torch.library.Library("test_sample", "DEF") # noqa: TOR901
lib.define("foo(Tensor self) -> Tensor")
lib.impl("foo", torch.sin, "CPU")

4
test/export/test_serialize.py
View File

@ -326,7 +326,7 @@ class TestDeserialize(TestCase):
def test_auto_functionalize(self):
try:
lib = torch.library.Library("mylib", "FRAGMENT")
lib = torch.library.Library("mylib", "FRAGMENT") # noqa: TOR901
torch.library.define(
"mylib::foo1",
"(Tensor(a!) x, Tensor[] y, Tensor(b!) z, SymInt w, Tensor n) -> Tensor",
@ -522,7 +522,7 @@ class TestDeserialize(TestCase):
def test_tensor_tensor_list(self):
try:
from torch.library import Library
lib = Library("_export", "FRAGMENT")
lib = Library("_export", "FRAGMENT") # noqa: TOR901
lib.define(
"_test_tensor_tensor_list_output(Tensor x, Tensor y) -> (Tensor, Tensor[])",
tags=torch.Tag.pt2_compliant_tag)

12
test/inductor/test_custom_lowering.py
View File

@ -17,9 +17,15 @@ class TestCustomLowering(TorchTestCase):
@classmethod
def setUpClass(cls):
super().setUpClass()
cls.test_inductor_ops = torch.library.Library("test_inductor_ops", "DEF")
cls.impl_cuda = torch.library.Library("test_inductor_ops", "IMPL", "CUDA")
cls.impl_meta = torch.library.Library("test_inductor_ops", "IMPL", "Meta")
cls.test_inductor_ops = torch.library.Library( # noqa: TOR901
"test_inductor_ops", "DEF"
)
cls.impl_cuda = torch.library.Library( # noqa: TOR901
"test_inductor_ops", "IMPL", "CUDA"
)
cls.impl_meta = torch.library.Library( # noqa: TOR901
"test_inductor_ops", "IMPL", "Meta"
)
cls._register_jagged_to_padded_dense()
@classmethod

2
test/inductor/test_torchinductor.py
View File

@ -123,7 +123,7 @@ skip_if_x86_mac = functools.partial(
)
vec_dtypes = [torch.float, torch.bfloat16, torch.float16]
libtest = torch.library.Library("test", "FRAGMENT")
libtest = torch.library.Library("test", "FRAGMENT") # noqa: TOR901
ids = set()
f32 = torch.float32

6
test/inductor/test_torchinductor_dynamic_shapes.py
View File

@ -275,8 +275,7 @@ class TestInductorDynamic(TestCase):
@torch._dynamo.config.patch(capture_scalar_outputs=True)
@torch._inductor.config.patch(implicit_fallbacks=True)
def test_item_to_inputs_kernel_nobreak(self, device):
lib = torch.library.Library("test", "DEF")
with torch.library._scoped_library("test", "DEF") as lib:
try:
@custom_ops.custom_op("test::foo")
@ -359,8 +358,7 @@ class TestInductorDynamic(TestCase):
)
@torch._inductor.config.patch(implicit_fallbacks=True)
def test_dynamic_stride_nobreak(self, device):
lib = torch.library.Library("test", "DEF")
with torch.library._scoped_library("test", "DEF") as lib:
try:
@custom_ops.custom_op("test::foo")

2
test/quantization/pt2e/test_quantize_pt2e.py
View File

@ -1744,7 +1744,7 @@ class TestQuantizePT2E(PT2EQuantizationTestCase):
def test_observer_callback(self):
from torch.library import Library, impl
test_lib = Library("test_int4", "DEF")
test_lib = Library("test_int4", "DEF") # noqa: TOR901
test_lib.define("quantize_per_tensor_int4(Tensor input, float scale, int zero_point) -> Tensor")
@impl(test_lib, "quantize_per_tensor_int4", "CompositeExplicitAutograd")

19
test/test_autograd_fallback.py
View File

@ -5,7 +5,7 @@ import warnings
import numpy as np
import torch
from torch.library import Library
from torch.library import _scoped_library, Library
from torch.testing._internal.common_utils import (
instantiate_parametrized_tests,
parametrize,
@ -38,7 +38,7 @@ class TestAutogradFallback(TestCase):
return getattr(getattr(torch.ops, self.test_ns), name).default
def get_lib(self):
lib = Library(self.test_ns, "FRAGMENT")
lib = Library(self.test_ns, "FRAGMENT") # noqa: TOR901
self.lib = lib
return lib
@ -146,7 +146,7 @@ class TestAutogradFallback(TestCase):
# To be clear, none of these situations are OK and will lead
# to other problems down the line. We're testing them because
# it is fairly common to actually do these things.
lib = Library(self.test_ns, "FRAGMENT")
with _scoped_library(self.test_ns, "FRAGMENT") as lib:
lib.define("foo(Tensor self) -> Tensor")
lib.impl("foo", lambda x: x, "CPU")
op = self.get_op("foo")
@ -170,7 +170,7 @@ class TestAutogradFallback(TestCase):
@parametrize("mode", ("nothing", "warn"))
def test_composite_registered_to_cpu(self, mode):
with autograd_fallback_mode(mode):
lib = Library(self.test_ns, "FRAGMENT")
with _scoped_library(self.test_ns, "FRAGMENT") as lib:
lib.define("foo(Tensor self) -> Tensor")
lib.impl("foo", lambda x: x.sin().sum(), "CPU")
op = self.get_op("foo")
@ -184,7 +184,7 @@ class TestAutogradFallback(TestCase):
@parametrize("mode", ("nothing", "warn"))
def test_autograd_function_registered_to_cpu(self, mode):
with autograd_fallback_mode(mode):
lib = Library(self.test_ns, "FRAGMENT")
with _scoped_library(self.test_ns, "FRAGMENT") as lib:
lib.define("foo(Tensor self) -> Tensor")
class NumpySin(torch.autograd.Function):
@ -210,7 +210,7 @@ class TestAutogradFallback(TestCase):
@parametrize("mode", ("nothing", "warn"))
def test_inplace_autograd_function_registered_to_cpu(self, mode):
with autograd_fallback_mode(mode):
lib = Library(self.test_ns, "FRAGMENT")
with _scoped_library(self.test_ns, "FRAGMENT") as lib:
lib.define("foo(Tensor(a!) self) -> Tensor(a!)")
class NumpySin_(torch.autograd.Function):
@ -238,7 +238,9 @@ class TestAutogradFallback(TestCase):
expected = torch.zeros_like(x)
expected[0] = x[0].cos()
with self._check_ctx(mode):
(gx,) = torch.autograd.grad(y, x, torch.ones_like(y), retain_graph=True)
(gx,) = torch.autograd.grad(
y, x, torch.ones_like(y), retain_graph=True
)
self.assertEqual(gx, expected)
expected = torch.ones_like(x)
@ -252,8 +254,7 @@ class TestAutogradFallback(TestCase):
# We don't do anything special (that is, we don't rebase history).
# See NOTE [autograd fallback and in-place operations] for why
with autograd_fallback_mode(mode):
lib = Library(self.test_ns, "FRAGMENT")
with _scoped_library(self.test_ns, "FRAGMENT") as lib:
# Correct usage of (a!)
lib.define("foo(Tensor(a!) self, Tensor other) -> Tensor(a!)")

2
test/test_custom_ops.py
View File

@ -46,7 +46,7 @@ class CustomOpTestCaseBase(TestCase):
return getattr(torch.ops, self.test_ns)
def lib(self):
result = torch.library.Library(self.test_ns, "FRAGMENT")
result = torch.library.Library(self.test_ns, "FRAGMENT") # noqa: TOR901
self.libraries.append(result)
return result

2
test/test_fake_tensor.py
View File

@ -87,7 +87,7 @@ class FakeTensorTest(TestCase):
def test_custom_op_fallback(self):
from torch.library import Library, impl
test_lib = Library("my_test_op", "DEF")
test_lib = Library("my_test_op", "DEF") # noqa: TOR901
test_lib.define('foo(Tensor self) -> Tensor')
@impl(test_lib, 'foo', 'CPU')

2
test/test_fx_passes.py
View File

@ -743,7 +743,7 @@ class MultiOutputWithWithInvalidMatches:
class QuantizationFp8Pattern:
@classmethod
def setup(cls):
cls.quantization = torch.library.Library("fp8_quantization", "DEF")
cls.quantization = torch.library.Library("fp8_quantization", "DEF") # noqa: TOR901
cls.quantization.define("quantize_per_tensor_affine_fp8(Tensor self, int dtype, float scale) -> Tensor")
cls.quantization.define("dequantize_per_tensor_affine_fp8(Tensor self, int dtype, float scale) -> Tensor")

10
test/test_meta.py
View File

@ -1325,10 +1325,9 @@ class TestMeta(TestCase):
@onlyCPU
def test_meta_autograd_no_error(self):
lib = torch.library.Library("meta_test", "DEF")
impl_cpu = torch.library.Library("meta_test", "IMPL", "CPU")
impl_meta = torch.library.Library("meta_test", "IMPL", "Meta")
with torch.library._scoped_library("meta_test", "DEF") as lib:
with torch.library._scoped_library("meta_test", "IMPL", "CPU") as impl_cpu:
with torch.library._scoped_library("meta_test", "IMPL", "Meta") as impl_meta:
def foo_impl(x):
return x + 1
@ -1342,9 +1341,6 @@ class TestMeta(TestCase):
# key for custom ops, so it's fine that `foo()` doesn't have
# an autograd kernel.
b = torch.ops.meta_test.foo.default(a)
del impl_meta
del impl_cpu
del lib
def test_huber_loss_backward(self):
inps = [torch.rand(2**52, device='meta') for _ in range(3)]

2
test/test_proxy_tensor.py
View File

@ -955,7 +955,7 @@ class TestSymbolicTracing(TestCase):
import torch.library
from torch.library import Library
foo = Library("foo", "DEF")
foo = Library("foo", "DEF") # noqa: TOR901
foo.define("foo(Tensor self) -> Tensor")
# Operator where meta and cpu disagree on strides

75
test/test_python_dispatch.py
View File

@ -63,10 +63,9 @@ class TestPythonRegistration(TestCase):
# RuntimeError: impl("aten::neg", ...):
# Explicitly provided namespace (aten) in operator name does not match ...
with self.assertRaisesRegex(RuntimeError, "operator name does not match namespace"):
my_lib3 = Library("foo", "DEF")
with _scoped_library("foo", "DEF") as my_lib3:
my_lib3.define("neg(Tensor self) -> Tensor")
my_lib3.impl(torch.ops.aten.neg.default, my_neg, "AutogradCPU")
del my_lib3
# Example 2
def my_mul(*args, **kwargs):
@ -92,12 +91,12 @@ class TestPythonRegistration(TestCase):
def test_error_if_fn_not_callable(self):
with self.assertRaisesRegex(TypeError, "Input function is required to be a callable"):
my_lib = Library("aten", "IMPL")
with _scoped_library("aten", "IMPL") as my_lib:
my_lib.impl(torch.ops.aten.neg.default, [], "AutogradCPU")
def test_finalizer(self):
impls_refcnt = sys.getrefcount(torch.library._impls)
lib = Library(self.test_ns, "FRAGMENT")
lib = Library(self.test_ns, "FRAGMENT") # noqa: TOR901
lib.define("foo123(Tensor x) -> Tensor")
# 1 for `lib`, 1 for sys.getrefcount
@ -142,12 +141,11 @@ class TestPythonRegistration(TestCase):
run[0] = True
return args[0].clone()
my_lib1 = Library("aten", "IMPL")
with _scoped_library("aten", "IMPL") as my_lib1:
my_lib1.impl('aten::sum', my_sum, "CPU")
x = torch.tensor([1, 2])
self.assertEqual(torch.sum(x), x)
self.assertTrue(run[0])
del my_lib1
# Validate that the old behavior is restored for sum
self.assertEqual(torch.sum(x), torch.tensor(3))
@ -168,7 +166,7 @@ class TestPythonRegistration(TestCase):
return jitted_where(*args, **kwargs)
# overriding where's cuda kernel with Jiterator generated kernel
my_lib = Library("aten", "IMPL")
with _scoped_library("aten", "IMPL") as my_lib:
my_lib.impl('aten::where.self', inverted_where, "CUDA")
device = 'cuda'
@ -178,7 +176,6 @@ class TestPythonRegistration(TestCase):
self.assertEqual(torch.where(cond, x, y), torch.tensor([-1, -2, 3]))
self.assertTrue(CALLED[0])
del my_lib
# behavior restored after deregistration
self.assertEqual(torch.where(cond, x, y), torch.tensor([1, 2, -3]))
@ -199,13 +196,12 @@ class TestPythonRegistration(TestCase):
return jitted_gelu(*args, **kwargs)
# overriding gelu's cuda kernel with Jiterator generated relu kernel
my_lib = Library("aten", "IMPL")
with _scoped_library("aten", "IMPL") as my_lib:
my_lib.impl('aten::gelu', fast_gelu, "CUDA")
x = torch.rand([3, 3], device='cuda', dtype=torch.float)
self.assertEqual(torch.nn.functional.gelu(x), torch.nn.functional.relu(x))
self.assertTrue(CALLED[0])
del my_lib
# behavior restored after deregistration
self.assertNotEqual(torch.nn.functional.gelu(x), torch.nn.functional.relu(x))
@ -226,13 +222,12 @@ class TestPythonRegistration(TestCase):
return jitted_exp(*args, **kwargs)
# overriding exp's cuda kernel with clipped_exp kernel
my_lib = Library("aten", "IMPL")
with _scoped_library("aten", "IMPL") as my_lib:
my_lib.impl('aten::exp', clipped_exp, "CUDA")
x = torch.tensor([0.0, 100.0], device='cuda', dtype=torch.float16)
self.assertEqual(torch.exp(x), torch.tensor([1.0, 22026.4657948], dtype=torch.float16))
self.assertTrue(CALLED[0])
del my_lib
# behavior restored after deregistration
self.assertEqual(torch.exp(x), torch.tensor([1.0, torch.inf], dtype=torch.float16))
@ -252,7 +247,7 @@ class TestPythonRegistration(TestCase):
CALLED[0] = True
return jitted_add(*args, **kwargs)
my_lib = Library("aten", "IMPL")
with _scoped_library("aten", "IMPL") as my_lib:
my_lib.impl('aten::add.Tensor', buggy_add, "CUDA")
x_cpu = torch.rand([3, 3], device='cpu')
@ -263,7 +258,6 @@ class TestPythonRegistration(TestCase):
self.assertEqual(x_cuda + y_cuda, x_cpu + y_cpu + 1)
self.assertTrue(CALLED[0])
del my_lib
# behavior restored after deregistration
self.assertEqual(x_cuda + y_cuda, x_cpu + y_cpu)
@ -277,8 +271,7 @@ class TestPythonRegistration(TestCase):
def test_extend_library_with_dispatch_key_arg(self):
def my_sum(*args, **kwargs):
return args[0].clone()
my_lib1 = Library("aten", "IMPL", dispatch_key="CPU")
with _scoped_library("aten", "IMPL", dispatch_key="CPU") as my_lib1:
# RuntimeError: Explicitly provided dispatch key (Conjugate) is
# inconsistent with the dispatch key of the enclosing TORCH_LIBRARY_IMPL block
with self.assertRaisesRegex(RuntimeError, "inconsistent with the dispatch key"):
@ -286,11 +279,9 @@ class TestPythonRegistration(TestCase):
my_lib1.impl('aten::sum', my_sum)
x = torch.tensor([1, 2])
self.assertEqual(torch.sum(x), x)
del my_lib1
def test_create_new_library(self) -> None:
my_lib1 = Library(self.test_ns, "DEF")
with _scoped_library(self.test_ns, "DEF") as my_lib1:
my_lib1.define("sum(Tensor self) -> Tensor")
# Example 1
@ -302,8 +293,7 @@ class TestPythonRegistration(TestCase):
op = getattr(torch.ops, self.test_ns).sum
self.assertEqual(op(x), x)
my_lib2 = Library(self.test_ns, "IMPL")
with _scoped_library(self.test_ns, "IMPL") as my_lib2:
# Example 2
@torch.library.impl(my_lib2, op.default, "ZeroTensor")
def my_sum_zt(*args, **kwargs):
@ -316,12 +306,8 @@ class TestPythonRegistration(TestCase):
self.assertTrue(op(y)._is_zerotensor())
self.assertEqual(op(x), x)
del my_lib2
del my_lib1
def test_create_new_library_fragment_no_existing(self):
my_lib = Library(self.test_ns, "FRAGMENT")
with _scoped_library(self.test_ns, "FRAGMENT") as my_lib:
my_lib.define("sum2(Tensor self) -> Tensor")
@torch.library.impl(my_lib, "sum2", "CPU")
@ -331,14 +317,10 @@ class TestPythonRegistration(TestCase):
x = torch.tensor([1, 2])
self.assertEqual(getattr(torch.ops, self.test_ns).sum2(x), x)
del my_lib
def test_create_new_library_fragment_with_existing(self):
my_lib1 = Library(self.test_ns, "DEF")
with _scoped_library(self.test_ns, "DEF") as my_lib1:
# Create a fragment
my_lib2 = Library(self.test_ns, "FRAGMENT")
with _scoped_library(self.test_ns, "FRAGMENT") as my_lib2:
my_lib2.define("sum4(Tensor self) -> Tensor")
@torch.library.impl(my_lib2, "sum4", "CPU")
@ -349,8 +331,7 @@ class TestPythonRegistration(TestCase):
self.assertEqual(getattr(torch.ops, self.test_ns).sum4(x), x)
# Create another fragment
my_lib3 = Library(self.test_ns, "FRAGMENT")
with _scoped_library(self.test_ns, "FRAGMENT") as my_lib3:
my_lib3.define("sum3(Tensor self) -> Tensor")
@torch.library.impl(my_lib3, "sum3", "CPU")
@ -360,14 +341,10 @@ class TestPythonRegistration(TestCase):
x = torch.tensor([1, 2])
self.assertEqual(getattr(torch.ops, self.test_ns).sum3(x), x)
del my_lib1
del my_lib2
del my_lib3
@unittest.skipIf(IS_WINDOWS, "Skipped under Windows")
def test_alias_analysis(self):
def test_helper(alias_analysis=""):
my_lib1 = Library(self.test_ns, "DEF")
my_lib1 = Library(self.test_ns, "DEF") # noqa: TOR901
called = [0]
@ -388,11 +365,11 @@ class TestPythonRegistration(TestCase):
def test_error_for_unsupported_ns_or_kind(self) -> None:
with self.assertRaisesRegex(ValueError, "Unsupported kind"):
my_lib1 = Library("myns", "BLA")
my_lib1 = Library("myns", "BLA") # noqa: TOR901
for kind in ('DEF', 'FRAGMENT'):
with self.assertRaisesRegex(ValueError, "reserved namespace"):
my_lib1 = Library("prim", kind)
my_lib1 = Library("prim", kind) # noqa: TOR901
def test_returning_symint(self) -> None:
shape_env = ShapeEnv()
@ -402,7 +379,7 @@ class TestPythonRegistration(TestCase):
s0, s1 = ft.shape
tlib = Library(self.test_ns, "DEF")
with _scoped_library(self.test_ns, "DEF") as tlib:
tlib.define("sqsum(SymInt a, SymInt b) -> SymInt")
@impl(tlib, "sqsum", "CompositeExplicitAutograd")
@ -566,8 +543,7 @@ class TestPythonRegistration(TestCase):
getattr(torch.ops, self.test_ns).foo_functional.default, (x, y, z, w))
def test_register_fallthrough(self):
try:
my_lib = Library('aten', 'IMPL')
with _scoped_library('aten', 'IMPL') as my_lib:
my_lib.impl("mm", fallthrough_kernel, "AutocastCPU")
a = torch.randn(2, 3, device='cpu', dtype=torch.float32)
@ -577,8 +553,6 @@ class TestPythonRegistration(TestCase):
self.assertEqual(torch.mm(a, b).dtype, torch.float32)
# ops that don't have a fallthrough registered should not be affected
self.assertEqual(torch.matmul(a, b).dtype, torch.bfloat16)
finally:
del my_lib
with torch.autocast(device_type="cpu", dtype=torch.bfloat16):
# default behavior should have been restored
@ -694,7 +668,7 @@ $5: f32[2] = torch._ops.aten.clone.default($4, memory_format=torch.contiguous_fo
print("woof")
return torch.empty(())
my_lib = Library("my_lib", "DEF")
with _scoped_library("my_lib", "DEF") as my_lib:
my_lib.define("weird(Tensor?[] self) -> Tensor")
my_lib.impl("weird", weird, "CPU")
with capture_logs() as logs:
@ -1485,10 +1459,11 @@ $0: f32[] = torch._ops.aten.empty.memory_format([], device=device(type='cpu'), p
t.record_stream(s)
def test_return_stream(self) -> None:
l_def = torch.library.Library("test_return_stream", "DEF")
with _scoped_library("test_return_stream", "DEF") as l_def:
l_def.define("return_stream(Tensor self) -> Stream")
l_impl = torch.library.Library("test_return_stream", "IMPL", "CPU")
l_impl.impl("return_stream", lambda _: torch.Stream(stream_id=0, device_index=1, device_type=2))
with _scoped_library("test_return_stream", "IMPL", "CPU") as l_impl:
l_impl.impl("return_stream",
lambda _: torch.Stream(stream_id=0, device_index=1, device_type=2))
class TestMode(TorchDispatchMode):
def __torch_dispatch__(self, func, types, args=(), kwargs=None):

6
test/test_schema_check.py
View File

@ -26,17 +26,17 @@ def secretly_mutating(x):
def output_is_input(x):
return x
custom_lib = torch.library.Library("bad_schemas", "DEF")
custom_lib = torch.library.Library("bad_schemas", "DEF") # noqa: TOR901
custom_lib.define("secretly_aliasing(Tensor x) -> Tensor")
custom_lib.define("secretly_mutating(Tensor x) -> Tensor")
custom_lib.define("output_is_input(Tensor(a) x) -> Tensor(a)")
custom_lib_cpu = torch.library.Library("bad_schemas", "IMPL", "CPU")
custom_lib_cpu = torch.library.Library("bad_schemas", "IMPL", "CPU") # noqa: TOR901
custom_lib_cpu.impl("secretly_aliasing", secretly_aliasing)
custom_lib_cpu.impl("secretly_mutating", secretly_mutating)
custom_lib_cpu.impl("output_is_input", output_is_input)
custom_lib_meta = torch.library.Library("bad_schemas", "IMPL", "Meta")
custom_lib_meta = torch.library.Library("bad_schemas", "IMPL", "Meta") # noqa: TOR901
custom_lib_meta.impl("secretly_aliasing", secretly_aliasing)
custom_lib_meta.impl("secretly_mutating", secretly_mutating)
custom_lib_meta.impl("output_is_input", output_is_input)

2
test/test_torch.py
View File

@ -10205,7 +10205,7 @@ tensor([[[1.+1.j, 1.+1.j, 1.+1.j, ..., 1.+1.j, 1.+1.j, 1.+1.j],
from torch.library import Library, impl
global _my_storage
my_lib = Library("my_lib", "DEF")
my_lib = Library("my_lib", "DEF") # noqa: TOR901
my_lib.define('my_func() -> None')
a = torch.tensor([1.])

1
torch/testing/_internal/dynamo_test_failures.py
View File

@ -2089,7 +2089,6 @@ dynamo_expected_failures = {
"TestPythonRegistration.test_alias_analysis", # test_python_dispatch
"TestWrapperSubclassAliasingCPU.test_wrapper_subclass_aliasing_conv2d_cpu", # test_python_dispatch
"TestPythonRegistration.test_finalizer", # test_python_dispatch
"TestPythonRegistration.test_override_cpu_sum", # test_python_dispatch
"TestPythonDispatch.test_subclass_autograd_device_check", # test_python_dispatch
"TestPythonDispatch.test_make_subclass_with_modes", # test_python_dispatch
"LoggingTests.test_trace_source_nested", # dynamo/test_logging