Merge branch 'torch_dot' of github.com:HangenYuu/pytensor into torch_dot

HangenYuu · HangenYuu · commit 398383a31585 · 2024-07-10T16:25:19.000+08:00
diff --git a/pyproject.toml b/pyproject.toml
@@ -125,8 +125,8 @@ line-length = 88
 exclude = ["doc/", "pytensor/_version.py"]
 
 [tool.ruff.lint]
-select = ["C", "E", "F", "I", "UP", "W", "RUF"]
-ignore = ["C408", "C901", "E501", "E741", "RUF012"]
+select = ["C", "E", "F", "I", "UP", "W", "RUF", "PERF", "PTH"]
+ignore = ["C408", "C901", "E501", "E741", "RUF012", "PERF203"]
 
 
 [tool.ruff.lint.isort]
@@ -136,6 +136,7 @@ lines-after-imports = 2
 # TODO: Get rid of these:
 "**/__init__.py" = ["F401", "E402", "F403"]
 "pytensor/tensor/linalg.py" = ["F403"]
+"pytensor/link/c/cmodule.py" = ["PTH"]
 # For the tests we skip because `pytest.importorskip` is used:
 "tests/link/jax/test_scalar.py" = ["E402"]
 "tests/link/jax/test_tensor_basic.py" = ["E402"]
diff --git a/pytensor/compile/debugmode.py b/pytensor/compile/debugmode.py
@@ -478,7 +478,7 @@ def _check_inputs(
     """
     destroyed_idx_list = []
     destroy_map = node.op.destroy_map
-    for o_pos, i_pos_list in destroy_map.items():
+    for i_pos_list in destroy_map.values():
         destroyed_idx_list.extend(i_pos_list)
     destroyed_res_list = [node.inputs[i] for i in destroyed_idx_list]
 
@@ -598,7 +598,7 @@ def _check_viewmap(fgraph, node, storage_map):
         # TODO: make sure this is correct
         # According to OB, duplicate inputs are rejected on build graph time
         # if they cause problems. So if they are here it should be ok.
-        for key, val in good_alias.items():
+        for key in good_alias:
             bad_alias.pop(key, None)
         if bad_alias:
             raise BadViewMap(node, oi, outstorage, list(bad_alias.values()))
@@ -756,10 +756,7 @@ def _get_preallocated_maps(
     # TODO: Sparse? Scalar does not really make sense.
 
     # Do not preallocate memory for outputs that actually work inplace
-    considered_outputs = []
-    for r in node.outputs:
-        if r not in inplace_outs:
-            considered_outputs.append(r)
+    considered_outputs = [r for r in node.outputs if r not in inplace_outs]
 
     # Output storage that was initially present in the storage_map
     if "initial" in prealloc_modes or "ALL" in prealloc_modes:
diff --git a/pytensor/compile/profiling.py b/pytensor/compile/profiling.py
@@ -1380,9 +1380,9 @@ def print_stats(stats1, stats2):
         items.sort(key=lambda a: a[1], reverse=True)
         for idx, ((fgraph, node), node_outputs_size) in enumerate(items[:N]):
             code = ["c"] * len(node.outputs)
-            for out, inp in node.op.destroy_map.items():
+            for out in node.op.destroy_map:
                 code[out] = "i"
-            for out, inp in node.op.view_map.items():
+            for out in node.op.view_map:
                 code[out] = "v"
             shapes = str(fct_shapes[fgraph][node])
 
diff --git a/pytensor/configdefaults.py b/pytensor/configdefaults.py
@@ -348,7 +348,7 @@ def add_compile_configvars():
         if sys.platform == "win32":
             mingw_w64_gcc = Path(sys.executable).parent / "Library/mingw-w64/bin/g++"
             try:
-                rc = call_subprocess_Popen([mingw_w64_gcc, "-v"])
+                rc = call_subprocess_Popen([str(mingw_w64_gcc), "-v"])
                 if rc == 0:
                     maybe_add_to_os_environ_pathlist("PATH", mingw_w64_gcc.parent)
             except OSError:
diff --git a/pytensor/graph/destroyhandler.py b/pytensor/graph/destroyhandler.py
@@ -182,7 +182,7 @@ def _build_droot_impact(destroy_handler):
     root_destroyer = {}  # root -> destroyer apply
 
     for app in destroy_handler.destroyers:
-        for output_idx, input_idx_list in app.op.destroy_map.items():
+        for input_idx_list in app.op.destroy_map.values():
             if len(input_idx_list) != 1:
                 raise NotImplementedError()
             input_idx = input_idx_list[0]
@@ -698,7 +698,7 @@ def orderings(self, fgraph, ordered=True):
                 # keep track of clients that should run before the current Apply
                 root_clients = set_type()
                 # for each destroyed input...
-                for output_idx, input_idx_list in app.op.destroy_map.items():
+                for input_idx_list in app.op.destroy_map.values():
                     destroyed_idx = input_idx_list[0]
                     destroyed_variable = app.inputs[destroyed_idx]
                     root = droot[destroyed_variable]
diff --git a/pytensor/link/c/basic.py b/pytensor/link/c/basic.py
@@ -1450,12 +1450,16 @@ def in_sig(i, topological_pos, i_idx):
             if props:
                 version.append(props)
 
-            for i in node.inputs:
-                if isinstance(i.type, CLinkerObject):
-                    version.append(i.type.c_code_cache_version())
-            for o in node.outputs:
-                if isinstance(o.type, CLinkerObject):
-                    version.append(o.type.c_code_cache_version())
+            version.extend(
+                i.type.c_code_cache_version()
+                for i in node.inputs
+                if isinstance(i.type, CLinkerObject)
+            )
+            version.extend(
+                o.type.c_code_cache_version()
+                for o in node.outputs
+                if isinstance(o.type, CLinkerObject)
+            )
 
             # add the signature for this node
             sig.append(
diff --git a/pytensor/link/c/cmodule.py b/pytensor/link/c/cmodule.py
@@ -2131,9 +2131,11 @@ def get_lines(cmd, parse=True):
                             or "-march=native" in line
                         ):
                             continue
-                        for reg in ("-march=", "-mtune=", "-target-cpu", "-mabi="):
-                            if reg in line:
-                                selected_lines.append(line.strip())
+                        selected_lines.extend(
+                            line.strip()
+                            for reg in ("-march=", "-mtune=", "-target-cpu", "-mabi=")
+                            if reg in line
+                        )
                     lines = list(set(selected_lines))  # to remove duplicate
 
                 return lines
diff --git a/pytensor/link/c/lazylinker_c.py b/pytensor/link/c/lazylinker_c.py
@@ -1,6 +1,4 @@
-import errno
 import logging
-import os
 import sys
 import warnings
 from importlib import reload
@@ -43,23 +41,12 @@ def try_reload():
     # compile_str()) but if another lazylinker_ext does exist then it will be
     # imported and compile_str won't get called at all.
     location = config.compiledir / "lazylinker_ext"
-    if not location.exists():
-        try:
-            # Try to make the location
-            os.mkdir(location)
-        except OSError as e:
-            # If we get an error, verify that the error was # 17, the
-            # path already exists, and that it is a directory Note: we
-            # can't check if it exists before making it, because we
-            # are not holding the lock right now, so we could race
-            # another process and get error 17 if we lose the race
-            assert e.errno == errno.EEXIST
-            assert location.is_dir()
+    location.mkdir(exist_ok=True)
 
     init_file = location / "__init__.py"
     if not init_file.exists():
         try:
-            with open(init_file, "w"):
+            with init_file.open("w"):
                 pass
         except OSError as e:
             if init_file.exists():
@@ -129,12 +116,7 @@ def try_reload():
             code = cfile.read_text("utf-8")
 
             loc = config.compiledir / dirname
-            if not loc.exists():
-                try:
-                    os.mkdir(loc)
-                except OSError as e:
-                    assert e.errno == errno.EEXIST
-                    assert loc.exists()
+            loc.mkdir(exist_ok=True)
 
             args = GCC_compiler.compile_args()
             GCC_compiler.compile_str(dirname, code, location=loc, preargs=args)
@@ -147,8 +129,7 @@ def try_reload():
             # imported at the same time: we need to make sure we do not
             # reload the now outdated __init__.pyc below.
             init_pyc = loc / "__init__.pyc"
-            if init_pyc.is_file():
-                os.remove(init_pyc)
+            init_pyc.unlink(missing_ok=True)
 
             try_import()
             try_reload()
diff --git a/pytensor/link/pytorch/dispatch/__init__.py b/pytensor/link/pytorch/dispatch/__init__.py
@@ -7,5 +7,5 @@
 import pytensor.link.pytorch.dispatch.elemwise
 import pytensor.link.pytorch.dispatch.extra_ops
 import pytensor.link.pytorch.dispatch.math
-
+import pytensor.link.pytorch.dispatch.sort
 # isort: on
diff --git a/pytensor/link/pytorch/dispatch/sort.py b/pytensor/link/pytorch/dispatch/sort.py
@@ -0,0 +1,25 @@
+import torch
+
+from pytensor.link.pytorch.dispatch.basic import pytorch_funcify
+from pytensor.tensor.sort import ArgSortOp, SortOp
+
+
+@pytorch_funcify.register(SortOp)
+def pytorch_funcify_Sort(op, **kwargs):
+    stable = op.kind == "stable"
+
+    def sort(arr, axis):
+        sorted, _ = torch.sort(arr, dim=axis, stable=stable)
+        return sorted
+
+    return sort
+
+
+@pytorch_funcify.register(ArgSortOp)
+def pytorch_funcify_ArgSort(op, **kwargs):
+    stable = op.kind == "stable"
+
+    def argsort(arr, axis):
+        return torch.argsort(arr, dim=axis, stable=stable)
+
+    return argsort
diff --git a/pytensor/link/vm.py b/pytensor/link/vm.py
@@ -545,7 +545,7 @@ def __call__(self, output_subset=None):
             # Add the outputs that are needed for the in-place updates of the
             # inputs in `self.update_vars`
             output_subset = list(output_subset)
-            for inp, out in self.update_vars.items():
+            for out in self.update_vars.values():
                 out_idx = self.fgraph.outputs.index(out)
                 if out_idx not in output_subset:
                     output_subset.append(out_idx)
@@ -1055,12 +1055,7 @@ def make_vm(
             for v in self.fgraph.inputs + self.fgraph.outputs:
                 vars_idx.setdefault(v, len(vars_idx))
 
-            nodes_idx_inv = {}
-            vars_idx_inv = {}
-            for node, i in nodes_idx.items():
-                nodes_idx_inv[i] = node
-            for var, i in vars_idx.items():
-                vars_idx_inv[i] = var
+            vars_idx_inv = {i: var for var, i in vars_idx.items()}
 
             # put storage_map and compute_map into a int-based scheme
             storage_map_list = [
@@ -1270,15 +1265,16 @@ def make_all(
         if self.allow_gc:
             post_thunk_clear = []
             for node in order:
-                clear_after_this_thunk = []
-                for input in node.inputs:
+                clear_after_this_thunk = [
+                    storage_map[input]
+                    for input in node.inputs
                     if (
                         input in computed
                         and input not in fgraph.outputs
                         and node == last_user[input]
                         and input not in reallocated_vars
-                    ):
-                        clear_after_this_thunk.append(storage_map[input])
+                    )
+                ]
                 post_thunk_clear.append(clear_after_this_thunk)
         else:
             post_thunk_clear = None
diff --git a/pytensor/scalar/basic.py b/pytensor/scalar/basic.py
@@ -4434,8 +4434,7 @@ def apply(self, fgraph):
             )
             # make sure we don't produce any float16.
             assert not any(o.dtype == "float16" for o in new_node.outputs)
-            for o, no in zip(node.outputs, new_node.outputs):
-                mapping[o] = no
+            mapping.update(zip(node.outputs, new_node.outputs))
 
         new_ins = [mapping[inp] for inp in fgraph.inputs]
         new_outs = [mapping[out] for out in fgraph.outputs]
diff --git a/pytensor/scan/op.py b/pytensor/scan/op.py
@@ -2240,8 +2240,7 @@ def infer_shape(self, fgraph, node, input_shapes):
         # Non-sequences have a direct equivalent from self.inner_inputs in
         # node.inputs
         inner_non_sequences = self.inner_inputs[len(seqs_shape) + len(outs_shape) :]
-        for in_ns, out_ns in zip(inner_non_sequences, node.inputs[offset:]):
-            out_equivalent[in_ns] = out_ns
+        out_equivalent.update(zip(inner_non_sequences, node.inputs[offset:]))
 
         if info.as_while:
             self_outs = self.inner_outputs[:-1]
diff --git a/pytensor/scan/rewriting.py b/pytensor/scan/rewriting.py
@@ -1623,9 +1623,7 @@ def scan_save_mem(fgraph, node):
         (inps, outs, info, node_ins, compress_map) = compress_outs(
             op, not_required, nw_inputs
         )
-        inv_compress_map = {}
-        for k, v in compress_map.items():
-            inv_compress_map[v] = k
+        inv_compress_map = {v: k for k, v in compress_map.items()}
 
         # 3.6 Compose the new scan
         # TODO: currently we don't support scan with 0 step. So
diff --git a/pytensor/tensor/conv/abstract_conv.py b/pytensor/tensor/conv/abstract_conv.py
@@ -585,10 +585,11 @@ def assert_shape(x, expected_shape, msg="Unexpected shape."):
     if expected_shape is None or not config.conv__assert_shape:
         return x
     shape = x.shape
-    tests = []
-    for i in range(x.ndim):
-        if expected_shape[i] is not None:
-            tests.append(pt.eq(shape[i], expected_shape[i]))
+    tests = [
+        pt.eq(shape[i], expected_shape[i])
+        for i in range(x.ndim)
+        if expected_shape[i] is not None
+    ]
     if tests:
         return Assert(msg)(x, *tests)
     else:
diff --git a/pytensor/tensor/subtensor.py b/pytensor/tensor/subtensor.py
@@ -2107,10 +2107,7 @@ def perform(self, node, inp, out_):
         out[0] = x.take(i, axis=0, out=o)
 
     def connection_pattern(self, node):
-        rval = [[True]]
-
-        for ipt in node.inputs[1:]:
-            rval.append([False])
+        rval = [[True], *([False] for _ in node.inputs[1:])]
 
         return rval
 
diff --git a/pytensor/utils.py b/pytensor/utils.py
@@ -123,19 +123,19 @@ def maybe_add_to_os_environ_pathlist(var: str, newpath: Path | str) -> None:
         pass
 
 
-def subprocess_Popen(command, **params):
+def subprocess_Popen(command: str | list[str], **params):
     """
     Utility function to work around windows behavior that open windows.
 
     :see: call_subprocess_Popen and output_subprocess_Popen
     """
     startupinfo = None
     if os.name == "nt":
-        startupinfo = subprocess.STARTUPINFO()
+        startupinfo = subprocess.STARTUPINFO()  # type: ignore[attr-defined]
         try:
-            startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW
+            startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW  # type: ignore[attr-defined]
         except AttributeError:
-            startupinfo.dwFlags |= subprocess._subprocess.STARTF_USESHOWWINDOW
+            startupinfo.dwFlags |= subprocess._subprocess.STARTF_USESHOWWINDOW  # type: ignore[attr-defined]
 
         # Anaconda for Windows does not always provide .exe files
         # in the PATH, they also have .bat files that call the corresponding
@@ -156,7 +156,7 @@ def subprocess_Popen(command, **params):
     # with the default None values.
     stdin = None
     if "stdin" not in params:
-        stdin = open(os.devnull)
+        stdin = Path(os.devnull).open()
         params["stdin"] = stdin.fileno()
 
     try:
diff --git a/tests/link/pytorch/test_sort.py b/tests/link/pytorch/test_sort.py
@@ -0,0 +1,26 @@
+import numpy as np
+import pytest
+
+from pytensor.graph import FunctionGraph
+from pytensor.tensor import matrix
+from pytensor.tensor.sort import argsort, sort
+from tests.link.pytorch.test_basic import compare_pytorch_and_py
+
+
+@pytest.mark.parametrize("func", (sort, argsort))
+@pytest.mark.parametrize(
+    "axis",
+    [
+        pytest.param(0),
+        pytest.param(1),
+        pytest.param(
+            None, marks=pytest.mark.xfail(reason="Reshape Op not implemented")
+        ),
+    ],
+)
+def test_sort(func, axis):
+    x = matrix("x", shape=(2, 2), dtype="float64")
+    out = func(x, axis=axis)
+    fgraph = FunctionGraph([x], [out])
+    arr = np.array([[1.0, 4.0], [5.0, 2.0]])
+    compare_pytorch_and_py(fgraph, [arr])
diff --git a/tests/misc/test_pkl_utils.py b/tests/misc/test_pkl_utils.py
diff --git a/tests/scan/test_basic.py b/tests/scan/test_basic.py
diff --git a/tests/sparse/test_basic.py b/tests/sparse/test_basic.py
diff --git a/tests/tensor/test_math.py b/tests/tensor/test_math.py
diff --git a/tests/tensor/test_subtensor.py b/tests/tensor/test_subtensor.py
diff --git a/tests/tensor/utils.py b/tests/tensor/utils.py
diff --git a/tests/typed_list/test_type.py b/tests/typed_list/test_type.py
