[BE]: Remove useless lambdas (pytorch#113602)

Applies PLW0108 which removes useless lambda calls in Python, the rule is in preview so it is not ready to be enabled by default just yet. These are the autofixes from the rule.

Pull Request resolved: pytorch#113602
Approved by: https://github.com/albanD

Author: Skylion007

benchmarks/cpp/tensorexpr/bench_ops.py

@@ -51,7 +51,7 @@ def hardswish(x):
 
 for op in unary_ops:
     x = torch.rand((1024, 1024))
-    traced = torch.jit.trace(lambda x: op(x), (x))
+    traced = torch.jit.trace(op, (x))
 
     # Warmup.
     warmup_iters = 8
@@ -88,7 +88,7 @@ def test_batch_norm():
         x = torch.rand((n, c, h, w))
         y = torch.rand(c)
         z = torch.rand(c)
-        traced = torch.jit.trace(lambda x, y, z: op(x, y, z), (x, y, z))
+        traced = torch.jit.trace(op, (x, y, z))
 
         # Warmup.
         warmup_iters = 8

benchmarks/tensorexpr/broadcast.py

@@ -268,18 +268,18 @@ def register_broadcast_ops():
         ["div", lambda a, b: a / (b + 1e-4)],
         [
             "pow",
-            lambda a, b: torch.pow(a, b),
-            lambda a, b: np.power(a, b),
+            torch.pow,
+            np.power,
         ],  # no fuson triggered
-        ["max", lambda a, b: torch.max(a, b), lambda a, b: np.maximum(a, b)],
-        ["min", lambda a, b: torch.min(a, b), lambda a, b: np.minimum(a, b)],
+        ["max", torch.max, np.maximum],
+        ["min", torch.min, np.minimum],
     ]
 
     unary_op_list = [
-        ["erf", lambda x: torch.erf(x), lambda x: np.erf(x)],
-        ["exp", lambda x: torch.exp(x), lambda x: np.exp(x)],
-        ["sin", lambda x: torch.sin(x), lambda x: np.sin(x)],
-        ["cos", lambda x: torch.cos(x), lambda x: np.cos(x)],
+        ["erf", torch.erf, np.erf],
+        ["exp", torch.exp, np.exp],
+        ["sin", torch.sin, np.sin],
+        ["cos", torch.cos, np.cos],
     ]
 
     for split_input, binary_op in itertools.product([True, False], binary_op_list):

benchmarks/tensorexpr/elementwise.py

@@ -122,19 +122,19 @@ def register_element_ops():
         ["div", lambda a, b: a / (b + 1e-4)],
         [
             "pow",
-            lambda a, b: torch.pow(a, b),
-            lambda a, b: np.power(a, b),
+            torch.pow,
+            np.power,
         ],  # no fuson triggered
-        ["max", lambda a, b: torch.max(a, b), lambda a, b: np.maximum(a, b)],
-        ["min", lambda a, b: torch.min(a, b), lambda a, b: np.minimum(a, b)],
+        ["max", torch.max, np.maximum],
+        ["min", torch.min, np.minimum],
     ]
 
     unary_op_list = [
-        ["erf", lambda x: torch.erf(x), lambda x: scipy.special.erf(x)],
-        ["exp", lambda x: torch.exp(x), lambda x: np.exp(x)],
-        ["sin", lambda x: torch.sin(x), lambda x: np.sin(x)],
-        ["cos", lambda x: torch.cos(x), lambda x: np.cos(x)],
-        ["rand_like", lambda x: torch.rand_like(x), lambda x: np.random.rand(*x.shape)],
+        ["erf", torch.erf, scipy.special.erf],
+        ["exp", torch.exp, np.exp],
+        ["sin", torch.sin, np.sin],
+        ["cos", torch.cos, np.cos],
+        ["rand_like", torch.rand_like, lambda x: np.random.rand(*x.shape)],
     ]
 
     for split_input, binary_op in itertools.product([True, False], binary_op_list):

caffe2/python/memonger.py

@@ -648,7 +648,7 @@ def _find_best(ranges, init_assignment, prev_best_assignment, counter):
         # Try to put 'find_range' in a new assignment
         best_candidates.append(prev_best_assignment + [[find_range]])
 
-        ret = min(best_candidates, key=lambda x: get_memory_usage(x))
+        ret = min(best_candidates, key=get_memory_usage)
         return ret
 
     if not counter:

caffe2/python/net_drawer.py

@@ -374,7 +374,7 @@ def main():
         content = fid.read()
         graphs = utils.GetContentFromProtoString(
             content, {
-                caffe2_pb2.PlanDef: lambda x: GetOperatorMapForPlan(x),
+                caffe2_pb2.PlanDef: GetOperatorMapForPlan,
                 caffe2_pb2.NetDef: lambda x: {x.name: x.op},
             }
         )

caffe2/python/parallel_workers_test.py

@@ -60,7 +60,7 @@ def testParallelWorkers(self):
         workspace.ResetWorkspace()
 
         queue = create_queue()
-        dummy_worker = create_worker(queue, lambda worker_id: str(worker_id))
+        dummy_worker = create_worker(queue, str)
         worker_coordinator = parallel_workers.init_workers(dummy_worker)
         worker_coordinator.start()
 
@@ -102,7 +102,7 @@ def testParallelWorkersShutdownFun(self):
         workspace.ResetWorkspace()
 
         queue = create_queue()
-        dummy_worker = create_worker(queue, lambda worker_id: str(worker_id))
+        dummy_worker = create_worker(queue, str)
         workspace.FeedBlob('data', 'not shutdown')
 
         def shutdown_fun():

tools/autograd/gen_python_functions.py

@@ -405,7 +405,7 @@ def create_python_bindings(
 
     grouped = group_filter_overloads(pairs, pred)
 
-    for name in sorted(grouped.keys(), key=lambda x: str(x)):
+    for name in sorted(grouped.keys(), key=str):
         overloads = grouped[name]
         py_methods.append(
             method_impl(name, module, overloads, method=method, symint=symint)
@@ -443,7 +443,7 @@ def create_python_return_type_bindings(
 
     grouped = group_filter_overloads(pairs, pred)
 
-    for name in sorted(grouped.keys(), key=lambda x: str(x)):
+    for name in sorted(grouped.keys(), key=str):
         overloads = grouped[name]
         definitions, registrations = generate_return_type_definition_and_registrations(
             overloads
@@ -481,7 +481,7 @@ def create_python_return_type_bindings_header(
 
     grouped = group_filter_overloads(pairs, pred)
 
-    for name in sorted(grouped.keys(), key=lambda x: str(x)):
+    for name in sorted(grouped.keys(), key=str):
         overloads = grouped[name]
         declarations = generate_return_type_declarations(overloads)
         py_return_types_declarations.append(

torch/_custom_op/autograd.py

@@ -108,7 +108,7 @@ def forward(ctx, *flat_args):
 
             # We use the info about args to give better error messages in backward
             args_info = namedtuple_args(
-                schema, pytree.tree_map(lambda arg: type(arg), args))
+                schema, pytree.tree_map(type, args))
 
             save_for_backward_fn_inputs = namedtuple_args(schema, args)
             to_save = save_for_backward_fn(save_for_backward_fn_inputs, output)

torch/_dynamo/guards.py

@@ -114,7 +114,7 @@ def uninteresting_files():
     "___tuple_iterator_getitem": tuple_iterator_getitem,
     "__math_isnan": math.isnan,
     "inf": float("inf"),
-    "__load_module": lambda name: importlib.import_module(name),
+    "__load_module": importlib.import_module,
     "utils_device": torch.utils._device,
     "device": torch.device,
     "___from_numpy":

torch/_dynamo/utils.py

@@ -1325,7 +1325,7 @@ class CompileProfiler:
     def __init__(self):
         self.frame_count = 0
         self.op_count = 0
-        self.backend_ctx_ctor = lambda: disable_cache_limit()
+        self.backend_ctx_ctor = disable_cache_limit
 
     def __call__(self, gm: torch.fx.GraphModule, example_inputs):
         self.frame_count += 1

torch/_export/db/case.py

@@ -79,7 +79,7 @@ class ExportCase:
     name: str
     extra_inputs: Optional[InputsType] = None  # For testing graph generalization.
     # Tags associated with the use case. (e.g dynamic-shape, escape-hatch)
-    tags: Set[str] = field(default_factory=lambda: set())
+    tags: Set[str] = field(default_factory=set)
     support_level: SupportLevel = SupportLevel.SUPPORTED
     dynamic_shapes: Optional[Dict[str, Any]] = None
 

torch/_export/unflatten.py

@@ -464,7 +464,7 @@ def finalize_outputs(self):
 
     def copy_node(self, node):
         self.print("copying", node.format_node())
-        self.node_map[node] = self.graph.node_copy(node, lambda n: self.remap_input(n))
+        self.node_map[node] = self.graph.node_copy(node, self.remap_input)
         self.seen_nodes[node.name] = node
 
     def run_outer(self):

torch/_functorch/partitioners.py

@@ -886,7 +886,7 @@ def get_node_weight(node) -> int:
     node_idx = {node: idx for idx, node in enumerate(joint_module.graph.nodes)}
     saved_values = sorted((name_to_node[node] for node in cut_nodes), key=lambda x: node_idx[x])
     # save_for_backward on tensors and stashes symints in autograd .ctx
-    saved_sym_nodes = list(filter(lambda n: is_sym_node(n), saved_values))
+    saved_sym_nodes = list(filter(is_sym_node, saved_values))
     saved_values = list(filter(lambda n: not is_sym_node(n), saved_values))
     # NB: saved_sym_nodes will be mutated to reflect the actual saved symbols
     fw_module, bw_module = _extract_fwd_bwd_modules(

torch/_inductor/autotune_process.py

@@ -243,7 +243,7 @@ def initialize(self) -> None:
             EXIT_HANDLER_REGISTERED = True
             import atexit
 
-            atexit.register(lambda: self.terminate())
+            atexit.register(self.terminate)
 
     def get_device_list(self) -> List[Optional[int]]:
         """

torch/_inductor/codegen/cpp.py

@@ -2720,9 +2720,7 @@ def select_tiling_indices():
                 and contig_vars_sorted[-1] == len(self.itervars) - 1
             ):
                 return contig_vars_sorted
-            return sorted(contig_vars_sorted, key=lambda i: contig_vars_list.count(i))[
-                -1:
-            ]
+            return sorted(contig_vars_sorted, key=contig_vars_list.count)[-1:]
 
         def select_tiling(dtype: torch.dtype = torch.float):
             # TODO(jgong5): support alternative tiling factors and data types

torch/_inductor/graph.py

@@ -715,7 +715,7 @@ def set_current_node(self, node: torch.fx.Node):
 
     def run_node(self, n: torch.fx.Node):
         def debug(msg):
-            log.debug("lowering %s %s", LazyString(lambda: n.format_node()), msg)
+            log.debug("lowering %s %s", LazyString(n.format_node), msg)
 
         origins = {n}
         if n.op == "call_function":

torch/_inductor/inductor_prims.py

@@ -71,7 +71,7 @@ def eager_force_stride(input_tensor: Tensor, stride) -> Tensor:
 )
 force_stride_order = make_prim(
     "inductor_force_stride_order(Tensor input, SymInt[] stride) -> Tensor",
-    lambda input_tensor, stride: eager_force_stride(input_tensor, stride),
+    eager_force_stride,
     doc="Force the stride order for input tensor. No-op if the input tensor already has the stride. Do a copy otherwise",
 )
 masked_scatter_with_index = make_prim(

torch/_inductor/kernel/mm.py

@@ -300,6 +300,6 @@ def tuned_fused_int_mm_mul(mat1, mat2, mat3, out_dtype, *, layout=None):
             layout=layout,
             **dict(mm_options(config, k, layout), **{"ACC_TYPE": "tl.int32"}),
             suffix_args=1,
-            epilogue_fn=lambda acc, mat3: V.ops.mul(acc, mat3),
+            epilogue_fn=V.ops.mul,
         )
     return autotune_select_algorithm("int_mm", choices, [mat1, mat2, mat3], layout)

torch/_inductor/lowering.py

@@ -1968,21 +1968,21 @@ def inner_fn(index):
 
 def require_dense(_, *args, **kwargs):
     args, kwargs = pytree.tree_map_only(
-        ir.IRNode, lambda t: ir.ExternKernel.require_stride1(t), (args, kwargs)
+        ir.IRNode, ir.ExternKernel.require_stride1, (args, kwargs)
     )
     return args, kwargs
 
 
 def require_contiguous(_, *args, **kwargs):
     args, kwargs = pytree.tree_map_only(
-        ir.IRNode, lambda t: ir.ExternKernel.require_contiguous(t), (args, kwargs)
+        ir.IRNode, ir.ExternKernel.require_contiguous, (args, kwargs)
     )
     return args, kwargs
 
 
 def require_channels_last(_, *args, **kwargs):
     args, kwargs = pytree.tree_map_only(
-        ir.IRNode, lambda t: ir.ExternKernel.require_channels_last(t), (args, kwargs)
+        ir.IRNode, ir.ExternKernel.require_channels_last, (args, kwargs)
     )
     return args, kwargs
 

torch/_inductor/triton_heuristics.py

@@ -608,7 +608,7 @@ def __init__(self, *args, regex_filter="", **kwargs):
 
     def run(self, *args, grid, stream):
         possible_names = _find_names(self)
-        kernel_name = f"{max(possible_names, key=lambda x: len(x))}"
+        kernel_name = f"{max(possible_names, key=len)}"
         if not re.match(self.regex_filter, kernel_name):
             return
         super().run(*args, grid=grid, stream=stream)

torch/_prims/__init__.py

@@ -1231,7 +1231,7 @@ def _greater_than_reduce(acc, x):
 
         return x
 
-    reduce(lambda acc, x: _greater_than_reduce(acc, x), broadcast_dimensions, -1)
+    reduce(_greater_than_reduce, broadcast_dimensions, -1)
 
     # shape must be broadcastable to
     for idx, new_idx in enumerate(broadcast_dimensions):

torch/distributed/_spmd/data_parallel.py

@@ -120,7 +120,7 @@ def gradients_tagging(params: Dict[str, torch.Tensor]):
     tagging_hooks = []
     try:
         for p in params.values():
-            h = p.register_hook(lambda grad: torch.ops._spmd.tag_grad(grad))
+            h = p.register_hook(torch.ops._spmd.tag_grad)
             tagging_hooks.append(h)
         yield
     finally:

torch/distributed/checkpoint/_fsspec_filesystem.py

@@ -268,12 +268,12 @@ def _write_files_from_queue(
 
             if torch.cuda.is_available() and inflight_threshhold > 0:
                 loader = _OverlappingCpuLoader(
-                    lambda x: planner.resolve_data(x),
+                    planner.resolve_data,
                     inflight_threshhold=inflight_threshhold,
                 )
             else:
                 loader = _SerialCpuLoader(
-                    lambda x: planner.resolve_data(x),
+                    planner.resolve_data,
                 )
 
             tensor_w = [

torch/distributed/checkpoint/filesystem.py

@@ -265,12 +265,12 @@ def _write_files_from_queue(
 
             if torch.cuda.is_available() and inflight_threshhold > 0:
                 loader = _OverlappingCpuLoader(
-                    lambda x: planner.resolve_data(x),
+                    planner.resolve_data,
                     inflight_threshhold=inflight_threshhold,
                 )
             else:
                 loader = _SerialCpuLoader(
-                    lambda x: planner.resolve_data(x),
+                    planner.resolve_data,
                 )
 
             tensor_w = [

torch/distributed/fsdp/_init_utils.py

@@ -345,18 +345,14 @@ def _check_ignored_states(
         all_modules = all(isinstance(state, nn.Module) for state in ignored_states)
         if not all_params and not all_modules:
             # Sort for consistent ordering for unit test regex matching
-            sorted_types = sorted(
-                {type(state) for state in ignored_states}, key=lambda x: repr(x)
-            )
+            sorted_types = sorted({type(state) for state in ignored_states}, key=repr)
             raise ValueError(
                 "ignored_states expects all nn.Parameter or all nn.Module list "
                 f"elements but got types {sorted_types}"
             )
     else:
         if not all(isinstance(state, nn.Module) for state in ignored_states):
-            sorted_types = sorted(
-                {type(state) for state in ignored_states}, key=lambda x: repr(x)
-            )
+            sorted_types = sorted({type(state) for state in ignored_states}, key=repr)
             raise ValueError(
                 "ignored_modules expects nn.Module list elements but got "
                 f"types {sorted_types}"

torch/distributed/fsdp/fully_sharded_data_parallel.py

@@ -1185,7 +1185,7 @@ def clip_grad_norm_(
             )  # warn since this is generally unexpected
             return total_norm
         total_norm_dtype = functools.reduce(
-            lambda dtype1, dtype2: torch.promote_types(dtype1, dtype2),
+            torch.promote_types,
             [grad.dtype for grad in grads],
         )
         return total_norm.to(total_norm_dtype)

torch/distributed/rpc/rref_proxy.py

@@ -59,7 +59,7 @@ def _complete_op(fut):
             except BaseException as ex:
                 result.set_exception(ex)
 
-        rref_fut.then(lambda fut: _wrap_rref_type_cont(fut))
+        rref_fut.then(_wrap_rref_type_cont)
         return result
 
 # This class manages proxied RPC API calls for RRefs. It is entirely used from

torch/fx/experimental/accelerator_partitioner.py

@@ -583,8 +583,8 @@ def dump_dag(self, module_with_submodules: GraphModule) -> DAG:
             if node.target == operator.__getitem__:
                 continue
             input_nodes: Dict[Node, None] = {}
-            map_arg(node.args, lambda n: input_nodes.setdefault(n))
-            map_arg(node.kwargs, lambda n: input_nodes.setdefault(n))
+            map_arg(node.args, input_nodes.setdefault)
+            map_arg(node.kwargs, input_nodes.setdefault)
             # When a node has two or more output nodes,
             # it outputs its result to 'getitem' nodes.
             # Those 'getitem' nodes are the output node for this node.