Bug: torch.compile fails to compile torch.func.vmap with reduction functions and raw python numbers

### 🐛 Describe the bug

`torch.compile` fails to compile vmap transformation with reduction functions and native python numbers. This bug was only found when using reduction functions, and there are several workarounds as shown in the following examples:

```python
import torch

torch._dynamo.reset()
torch._dynamo.config.capture_func_transforms=True

def foo(x):
  return torch.vmap(lambda x: torch.sum(x) + 1e-2)(x)  # Error
  # return torch.vmap(lambda x: torch.mean(x) + 1e-2)(x)  # Error
  # return torch.vmap(lambda x: torch.std(x) + 1e-2)(x)  # Error
  # return torch.vmap(lambda x: torch.sum(x) + torch.tensor(1e-2))(x) # OK
  # return torch.vmap(lambda x: torch.sum(x, 0, keepdim=True) + 1e-2)(x) # OK
  # return torch.vmap(lambda x: torch.square(x) + 1e-2)(x) # OK
  # return torch.vmap(lambda x: x + 1e-2)(x) # OK


torch.compile(foo, fullgraph=True)(torch.randn((3, 3), device='cuda:0'))
# foo(torch.randn((3, 3), device='cuda:0'))   # OK
```

Error messages:

```
BackendCompilerFailed: backend='inductor' raised:
AssertionError: While executing %call : [num_users=1] = call_method[target=__call__](args = (%vmap_proxy, %l_x_), kwargs = {})
Original traceback:
  File "/tmp/ipykernel_672249/1649664715.py", line 7, in foo
    return torch.vmap(lambda x: torch.sum(x) + 1e-2)(x)  # Error
  File "/data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_functorch/apis.py", line 188, in wrapped
    return vmap_impl(func, in_dims, out_dims, randomness, chunk_size, *args, **kwargs)



You can suppress this exception and fall back to eager by setting:
    import torch._dynamo
    torch._dynamo.config.suppress_errors = True
```

Traceback:

```
---------------------------------------------------------------------------
BackendCompilerFailed                     Traceback (most recent call last)
Cell In[66], line 16
      7   return torch.vmap(lambda x: torch.sum(x) + 1e-2)(x)  # Error
      8   # return torch.vmap(lambda x: torch.mean(x) + 1e-2)(x)  # Error
      9   # return torch.vmap(lambda x: torch.std(x) + 1e-2)(x)  # Error
     10   # return torch.vmap(lambda x: torch.sum(x) + torch.tensor(1e-2))(x) # OK
     11   # return torch.vmap(lambda x: torch.sum(x, 0, keepdim=True) + 1e-2)(x) # OK
     12   # return torch.vmap(lambda x: torch.square(x) + 1e-2)(x) # OK
     13   # return torch.vmap(lambda x: x + 1e-2)(x) # OK
---> 16 torch.compile(foo, fullgraph=True)(torch.randn((3, 3), device='cuda:0'))
     17 # foo(torch.randn((3, 3), device='cuda:0'))   # OK

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py:328, in _TorchDynamoContext.__call__.<locals>._fn(*args, **kwargs)
    326 dynamic_ctx.__enter__()
    327 try:
--> 328     return fn(*args, **kwargs)
    329 finally:
    330     set_eval_frame(prior)

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py:490, in catch_errors_wrapper.<locals>.catch_errors(frame, cache_entry, frame_state)
    487             return hijacked_callback(frame, cache_entry, hooks, frame_state)
    489 with compile_lock, _disable_current_modes():
--> 490     return callback(frame, cache_entry, hooks, frame_state)

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py:133, in wrap_convert_context.<locals>._fn(*args, **kwargs)
    131 cleanup = setup_compile_debug()
    132 try:
--> 133     return fn(*args, **kwargs)
    134 finally:
    135     cleanup.close()

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py:389, in convert_frame_assert.<locals>._convert_frame_assert(frame, cache_entry, hooks, frame_state)
    376 compile_id = CompileId(frame_id, frame_compile_id)
    378 signpost_event(
    379     "dynamo",
    380     "_convert_frame_assert._compile",
   (...)
    386     },
    387 )
--> 389 return _compile(
    390     frame.f_code,
    391     frame.f_globals,
    392     frame.f_locals,
    393     frame.f_builtins,
    394     compiler_fn,
    395     one_graph,
    396     export,
    397     export_constraints,
    398     hooks,
    399     cache_size,
    400     frame,
    401     frame_state=frame_state,
    402     compile_id=compile_id,
    403 )

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py:569, in _compile(code, globals, locals, builtins, compiler_fn, one_graph, export, export_constraints, hooks, cache_size, frame, frame_state, compile_id)
    567 with compile_context(CompileContext(compile_id)):
    568     try:
--> 569         guarded_code = compile_inner(code, one_graph, hooks, transform)
    570         return guarded_code
    571     except (
    572         Unsupported,
    573         TorchRuntimeError,
   (...)
    578         ValidationException,
    579     ) as e:

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_dynamo/utils.py:189, in dynamo_timed.<locals>.dynamo_timed_inner.<locals>.time_wrapper(*args, **kwargs)
    187 with torch.profiler.record_function(f"{key} (dynamo_timed)"):
    188     t0 = time.time()
--> 189     r = func(*args, **kwargs)
    190     time_spent = time.time() - t0
    191 compilation_time_metrics[key].append(time_spent)

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py:491, in _compile.<locals>.compile_inner(code, one_graph, hooks, transform)
    489 for attempt in itertools.count():
    490     try:
--> 491         out_code = transform_code_object(code, transform)
    492         orig_code_map[out_code] = code
    493         break

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_dynamo/bytecode_transformation.py:1028, in transform_code_object(code, transformations, safe)
   1025 instructions = cleaned_instructions(code, safe)
   1026 propagate_line_nums(instructions)
-> 1028 transformations(instructions, code_options)
   1029 return clean_and_assemble_instructions(instructions, keys, code_options)[1]

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py:458, in _compile.<locals>
.transform(instructions, code_options)
    456 try:
    457     with tracing(tracer.output.tracing_context):
--> 458         tracer.run()
    459 except (exc.RestartAnalysis, exc.SkipFrame):
    460     raise

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py:2074, in InstructionTranslator.run(self)
   2073 def run(self):
-> 2074     super().run()

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py:724, in InstructionTranslatorBase.run(self)
    719 try:
    720     self.output.push_tx(self)
    721     while (
    722         self.instruction_pointer is not None
    723         and not self.output.should_exit
--> 724         and self.step()
    725     ):
    726         pass
    727 except BackendCompilerFailed:

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py:688, in InstructionTranslatorBase.step(self)
    684         unimplemented(f"missing: {inst.opname}")
    685     TracingContext.set_current_loc(
    686         self.f_code.co_filename, self.lineno, self.f_code.co_name
    687     )
--> 688     getattr(self, inst.opname)(inst)
    690     return inst.opname != "RETURN_VALUE"
    691 except Unsupported:

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py:2162, in InstructionTranslator.RETURN_VALUE(self, inst)
   2157 _step_logger()(
   2158     logging.INFO,
   2159     f"torchdynamo done tracing {self.f_code.co_name} (RETURN_VALUE)",
   2160 )
   2161 log.debug("RETURN_VALUE triggered compile")
-> 2162 self.output.compile_subgraph(
   2163     self,
   2164     reason=GraphCompileReason(
   2165         "return_value", [self.frame_summary()], graph_break=False
   2166     ),
   2167 )
   2168 self.output.add_output_instructions([create_instruction("RETURN_VALUE")])

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_dynamo/output_graph.py:833, in OutputGraph.compile_subgraph(self, tx, partial_convert, reason)
    830     append_prefix_insts()
    831     # optimization to generate better code in a common case
    832     self.add_output_instructions(
--> 833         self.compile_and_call_fx_graph(tx, list(reversed(stack_values)), root)
    834         + [create_instruction("UNPACK_SEQUENCE", arg=len(stack_values))]
    835     )
    836 else:
    837     graph_output_var = self.new_var("graph_out")

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/contextlib.py:81, in ContextDecorator.__call__.<locals>.inner(*args, **kwds)
     78 @wraps(func)
     79 def inner(*args, **kwds):
     80     with self._recreate_cm():
---> 81         return func(*args, **kwds)

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_dynamo/output_graph.py:957, in OutputGraph.compile_and_call_fx_graph(self, tx, rv, root)
    952 graph_tabular_log.debug("%s", lazy_format_graph_tabular(name, gm))
953 graph_sizes_log.debug(
    954     "%s", LazyString(lambda: self.get_graph_sizes_log_str(name))
    955 )
--> 957 compiled_fn = self.call_user_compiler(gm)
    958 compiled_fn = disable(compiled_fn)
    960 counters["stats"]["unique_graphs"] += 1

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_dynamo/utils.py:189, in dynamo_timed.<locals>.dynamo_timed_inner.<locals>.time_wrapper(*args, **kwargs)
    187 with torch.profiler.record_function(f"{key} (dynamo_timed)"):
    188     t0 = time.time()
--> 189     r = func(*args, **kwargs)
    190     time_spent = time.time() - t0
    191 compilation_time_metrics[key].append(time_spent)

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_dynamo/output_graph.py:1024, in OutputGraph.call_user_compiler(self, gm)
   1022     unimplemented_with_warning(e, self.root_tx.f_code, msg)
   1023 except Exception as e:
-> 1024     raise BackendCompilerFailed(self.compiler_fn, e).with_traceback(
   1025         e.__traceback__
   1026     ) from None
   1028 signpost_event(
   1029     "dynamo",
   1030     "OutputGraph.call_user_compiler",
   (...)
   1036     },
   1037 )
   1039 return compiled_fn

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_dynamo/output_graph.py:1009, in OutputGraph.call_user_compiler(self, gm)
   1007 if config.verify_correctness:
   1008     compiler_fn = WrapperBackend(compiler_fn)
-> 1009 compiled_fn = compiler_fn(gm, self.example_inputs())
   1010 _step_logger()(logging.INFO, f"done compiler function {name}")
   1011 assert callable(compiled_fn), "compiler_fn did not return callable"

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_dynamo/repro/after_dynamo.py:117, in wrap_backend_debug.<locals>.debug_wrapper(gm, example_inputs, **kwargs)
    115             raise
    116 else:
--> 117     compiled_gm = compiler_fn(gm, example_inputs)
    119 return compiled_gm

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_dynamo/repro/after_dynamo.py:117, in wrap_backend_debug.<locals>.debug_wrapper(gm, example_inputs, **kwargs)
    115             raise
    116 else:
--> 117     compiled_gm = compiler_fn(gm, example_inputs)
    119 return compiled_gm

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/__init__.py:1568, in _TorchCompileInductorWrapper.__call__(self, model_, inputs_)
   1565 def __call__(self, model_, inputs_):
   1566     from torch._inductor.compile_fx import compile_fx
-> 1568     return compile_fx(model_, inputs_, config_patches=self.config)

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_inductor/compile_fx.py:1150, in compile_fx(model_, example_inputs_, inner_compile, config_patches, decompositions)
   1143 tracing_context = (
   1144     torch._guards.TracingContext.get() or torch._guards.TracingContext(fake_mode)
   1145 )
   1147 with V.set_fake_mode(fake_mode), torch._guards.tracing(  # type: ignore[call-arg]
   1148     tracing_context
   1149 ), compiled_autograd.disable():
-> 1150     return aot_autograd(
   1151         fw_compiler=fw_compiler,
   1152         bw_compiler=bw_compiler,
   1153         inference_compiler=inference_compiler,
   1154         decompositions=decompositions,
   1155         partition_fn=partition_fn,
   1156         keep_inference_input_mutations=True,
   1157     )(model_, example_inputs_)

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_dynamo/backends/common.py:55, in aot_autograd
.<locals>.compiler_fn(gm, example_inputs)
     52 try:
     53     # NB: NOT cloned!
     54     with enable_aot_logging(), patch_config:
---> 55         cg = aot_module_simplified(gm, example_inputs, **kwargs)
     56         counters["aot_autograd"]["ok"] += 1
     57         return disable(cg)

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_functorch/aot_autograd.py:3891, in aot_module_simplified(mod, args, fw_compiler, bw_compiler, partition_fn, decompositions, keep_inference_input_mutations, inference_compiler)
   3875 aot_config = AOTConfig(
   3876     fw_compiler=fw_compiler,
   3877     bw_compiler=bw_compiler,
   (...)
   3887     no_tangents=False,
   3888 )
   3890 with compiled_autograd.disable():
-> 3891     compiled_fn = create_aot_dispatcher_function(
   3892         functional_call,
   3893         full_args,
   3894         aot_config,
   3895     )
   3897 # TODO: There is something deeply wrong here; compiled_fn running with
   3898 # the boxed calling convention, but aot_module_simplified somehow
   3899 # historically returned a function that was not the boxed calling
   3900 # convention.  This should get fixed...
   3901 def forward(*runtime_args):

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_dynamo/utils.py:189, in dynamo_timed.<locals>.dynamo_timed_inner.<locals>.time_wrapper(*args, **kwargs)
    187 with torch.profiler.record_function(f"{key} (dynamo_timed)"):
    188     t0 = time.time()
--> 189     r = func(*args, **kwargs)
    190     time_spent = time.time() - t0
    191 compilation_time_metrics[key].append(time_spent)

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_functorch/aot_autograd.py:3429, in create_aot_dispatcher_function(flat_fn, flat_args, aot_config)
   3426 compiler_fn = partial(aot_wrapper_dedupe, compiler_fn=compiler_fn)
   3427 # You can put more passes here
-> 3429 compiled_fn = compiler_fn(flat_fn, fake_flat_args, aot_config, fw_metadata=fw_metadata)
   3430 if aot_config.is_export:
   3432     mutated_user_inp_locs = [
   3433         idx - aot_config.num_params_buffers
   3434         for idx in fw_metadata.mutated_inp_indices
   3435         if idx >= aot_config.num_params_buffers
   3436     ]

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_functorch/aot_autograd.py:2212, in aot_wrapper_dedupe(flat_fn, flat_args, aot_config, compiler_fn, fw_metadata)
   2209         break
   2211 if ok:
-> 2212     return compiler_fn(flat_fn, leaf_flat_args, aot_config, fw_metadata=fw_metadata)
   2214 # export path: ban duplicate inputs for now, add later if requested.
   2215 if aot_config.is_export:

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_functorch/aot_autograd.py:2392, in aot_wrapper_synthetic_base(flat_fn, flat_args, aot_config, fw_metadata, needs_autograd, compiler_fn)
   2390 # Happy path: we don't need synthetic bases
   2391 if synthetic_base_info is None:
-> 2392     return compiler_fn(flat_fn, flat_args, aot_config, fw_metadata=fw_metadata)
   2394 # export path: ban synthetic bases for now, add later if requested.
   2395 if aot_config.is_export:

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_functorch/aot_autograd.py:1558, in aot_dispatch_base(flat_fn, flat_args, aot_config, fw_metadata)
   1557 def aot_dispatch_base(flat_fn, flat_args: List[Tensor], aot_config: AOTConfig, *, fw_metadata: ViewAndMutationMeta):
-> 1558     fw_module = aot_dispatch_base_graph(flat_fn, flat_args, aot_config, fw_metadata=fw_metadata)
   1560     disable_amp = torch._C._is_any_autocast_enabled()
   1561     context = torch._C._DisableAutocast if disable_amp else nullcontext

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_functorch/aot_autograd.py:1533, in aot_dispatch_base_graph(flat_fn, flat_args, aot_config, fw_metadata)
   1526 keep_mutations = aot_config.keep_inference_input_mutations
   1527 fn_to_trace = fn_input_mutations_to_outputs(
   1528     flat_fn,
   1529     fw_metadata,
   1530     keep_data_input_mutations=aot_config.keep_inference_input_mutations,
   1531 )
-> 1533 fw_module = create_functionalized_graph(
   1534     fn_to_trace,
   1535     flat_args,
   1536     meta=fw_metadata,
   1537     aot_config=aot_config,
   1538     trace_joint=False,
   1539 )
   1541 # As long as we opted to remove input mutations, then
   1542 # there should be *NO* mutating ops in the graph at this point.
   1543 copy_count = assert_functional_graph(fw_module.graph, allow_input_mutations=aot_config.keep_inference_input_mutations)

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_functorch/aot_autograd.py:1420, in create_functionalized_graph(fn, args, meta, aot_config, trace_joint)
   1417     helper, args = create_functionalized_rng_ops_wrapper(helper, args, trace_joint)
   1419 with enable_python_dispatcher():
-> 1420     fx_g = make_fx(helper, decomposition_table=aot_config.decompositions)(*args)
   1422 return fx_g

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:809, in make_fx.<locals>.wrapped(*args)
    801 # We disable the autocast cache as the autocast cache causes type conversions on parameters to
    802 # check a cache, which introduces untracked tensors into the graph
    803 #
    804 # We also disable tracing by any other tensor proxy-based tracers except the current. The
    805 # purpose of `make_fx` is to produce graphmodules as a side effect; its internal execution is
    806 # thus irrelevant to any external functional trace.
    807 with decompose(decomposition_table), fake_tensor_mode, python_dispatcher_mode, pre_dispatch_mode, proxy_function_mode, \
    808      sym_mode, proxy_mode, disable_autocast_cache(), disable_proxy_modes_tracing(enable_current=True):
--> 809     t = dispatch_trace(wrap_key(func, args, fx_tracer, pre_dispatch), tracer=fx_tracer, concrete_args=tuple(phs))
    811 # TODO: kind of a bad way to do it, should maybe figure out a better way
    812 if tracing_mode == "symbolic":

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_compile.py:24, in _disable_dynamo.<locals>.inner(*args, **kwargs)
     20 @functools.wraps(fn)
     21 def inner(*args, **kwargs):
     22     import torch._dynamo
---> 24     return torch._dynamo.disable(fn, recursive)(*args, **kwargs)

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py:328, in _TorchDynamoContext.__call__.<locals>._fn(*args, **kwargs)
    326 dynamic_ctx.__enter__()
    327 try:
--> 328     return fn(*args, **kwargs)
    329 finally:
    330     set_eval_frame(prior)

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_dynamo/external_utils.py:17, in wrap_inline.<locals>.inner(*args, **kwargs)
     15 @functools.wraps(fn)
     16 def inner(*args, **kwargs):
---> 17     return fn(*args, **kwargs)

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:468, in dispatch_trace(root, tracer, concrete_args)
    462 @torch._disable_dynamo
    463 def dispatch_trace(
    464         root: Union[torch.nn.Module, Callable],
    465         tracer: Tracer,
    466         concrete_args: Optional[Tuple[Any, ...]] = None,
    467 ) -> GraphModule:
--> 468     graph = tracer.trace(root, concrete_args)
    469     name = root.__class__.__name__ if isinstance(root, torch.nn.Module) else root.__name__
    470     return GraphModule(tracer.root, graph, name)

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py:328, in _TorchDynamoContext.__call__.<locals>._fn(*args, **kwargs)
    326 dynamic_ctx.__enter__()
    327 try:
--> 328     return fn(*args, **kwargs)
    329 finally:
    330     set_eval_frame(prior)

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_dynamo/external_utils.py:17, in wrap_inline.<locals>.inner(*args, **kwargs)
     15 @functools.wraps(fn)
     16 def inner(*args, **kwargs):
---> 17     return fn(*args, **kwargs)

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/fx/_symbolic_trace.py:817, in Tracer.trace(self, root, concrete_args)
    810         for module in self._autowrap_search:
    811             _autowrap_check(
    812                 patcher, module.__dict__, self._autowrap_function_ids
    813             )
    814         self.create_node(
    815             "output",
    816             "output",
--> 817             (self.create_arg(fn(*args)),),
    818             {},
    819             type_expr=fn.__annotations__.get("return", None),
    820         )
    822     self.submodule_paths = None
    823 finally:

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:485, in wrap_key.<locals>.wrapped(*proxies)
    482     assert isinstance(m, ProxyTorchDispatchMode)
    483     track_tensor_tree(flat_tensors, flat_proxies, constant=None, tracer=tracer)
--> 485 out = f(*tensors)
    486 out = pytree.tree_map_only(
    487     torch.Tensor,
    488     lambda t: get_proxy_slot(t, tracer, t, lambda x: x.proxy),
    489     out
    490 )
    491 out = pytree.tree_map_only(
    492     (SymInt, SymFloat, SymBool),
    493     lambda t: get_proxy_slot(t.node, tracer)(),
    494     out
    495 )

File <string>:1, in <lambda>(arg0)

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_functorch/aot_autograd.py:1412, in create_functionalized_graph.<locals>.fwd_helper(*args)
   1411 def fwd_helper(*args):
-> 1412     return functionalized_f_helper(*args)

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_functorch/aot_autograd.py:1363, in create_functionalized_graph.<locals>.functionalized_f_helper(*args)
   1360 torch._enable_functionalization(reapply_views=True)
   1361 try:
   1362     # Run the joint
-> 1363     f_outs = fn(*f_args)
   1364 finally:
   1365     torch._disable_functionalization()

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_functorch/aot_autograd.py:1165, in fn_input_mutations_to_outputs.<locals>.inner_fn(*args)
   1164 def inner_fn(*args):
-> 1165     outs = fn(*args)
   1166     assert len(meta.output_info) == len(outs)
   1167     # The compiled fw will return mutated input tensors, *including* metadata-only mutation.
   1168     # However, if keep_data_input_mutations is set, the compiled fw only needs to return metadata-mutated inputs.
   1169     # (because data-only input mutations are handled directly in the compiled graph)

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_functorch/aot_autograd.py:3496, in create_functional_call.<locals>.functional_call(*args, **kwargs)
   3492         warnings.filterwarnings(
   3493             "ignore", "Anomaly Detection has been enabled."
   3494         )
   3495         with torch.autograd.detect_anomaly(check_nan=False):
-> 3496             out = Interpreter(mod).run(*args[params_len:], **kwargs)
   3497 else:
   3498     out = mod(*args[params_len:], **kwargs)

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/fx/interpreter.py:138, in Interpreter.run(self, initial_env, enable_io_processing, *args)
    135     continue
    137 try:
--> 138     self.env[node] = self.run_node(node)
    139 except Exception as e:
    140     if self.extra_traceback:

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/fx/interpreter.py:195, in Interpreter.run_node(self, n)
    193 assert isinstance(args, tuple)
    194 assert isinstance(kwargs, dict)
--> 195 return getattr(self, n.op)(n.target, args, kwargs)

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/fx/interpreter.py:289, in Interpreter.call_method(self, target, args, kwargs)
    287 # Execute the method and return the result
    288 assert isinstance(target, str)
--> 289 return getattr(self_obj, target)(*args_tail, **kwargs)

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_functorch/apis.py:188, in vmap.<locals>.wrapped(*args, **kwargs)
    187 def wrapped(*args, **kwargs):
--> 188     return vmap_impl(func, in_dims, out_dims, randomness, chunk_size, *args, **kwargs)

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_functorch/vmap.py:266, in vmap_impl(func, in_dims, out_dims, randomness, chunk_size, *args, **kwargs)
    262     return _chunked_vmap(func, flat_in_dims, chunks_flat_args,
    263                          args_spec, out_dims, randomness, **kwargs)
    265 # If chunk_size is not specified.
--> 266 return _flat_vmap(
    267     func, batch_size, flat_in_dims, flat_args, args_spec, out_dims, randomness, **kwargs
    268 )

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_functorch/vmap.py:38, in doesnt_support_saved_tensors_hooks.<locals>.fn(*args, **kwargs)
     35 @functools.wraps(f)
     36 def fn(*args, **kwargs):
     37     with torch.autograd.graph.disable_saved_tensors_hooks(message):
---> 38         return f(*args, **kwargs)

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_functorch/vmap.py:379, in _flat_vmap(func, batch_size, flat_in_dims, flat_args, args_spec, out_dims, randomness, **kwargs)
    377 try:
    378     batched_inputs = _create_batched_inputs(flat_in_dims, flat_args, vmap_level, args_spec)
--> 379     batched_outputs = func(*batched_inputs, **kwargs)
    380     return _unwrap_batched(batched_outputs, out_dims, vmap_level, batch_size, func)
    381 finally:

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/fx/graph_module.py:678, in GraphModule.recompile.<locals>.call_wrapped(self, *args, **kwargs)
    677 def call_wrapped(self, *args, **kwargs):
--> 678     return self._wrapped_call(self, *args, **kwargs)

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/fx/graph_module.py:284, in _WrappedCall.__call__(self, obj, *args, **kwargs)
    282     raise e.with_traceback(None)
    283 else:
--> 284     raise e

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/fx/graph_module.py:274, in _WrappedCall.__call__(self, obj, *args, **kwargs)
    272         return self.cls_call(obj, *args, **kwargs)
    273     else:
--> 274         return super(self.cls, obj).__call__(*args, **kwargs)  # type: ignore[misc]
    275 except Exception as e:
    276     assert e.__traceback__

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/fx/_symbolic_trace.py:795, in Tracer.trace.<locals>.module_call_wrapper(mod, *args, **kwargs)
    788     return _orig_module_call(mod, *args, **kwargs)
    790 _autowrap_check(
    791     patcher,
    792     getattr(getattr(mod, "forward", mod), "__globals__", {}),
    793     self._autowrap_function_ids,
    794 )
--> 795 return self.call_module(mod, forward, args, kwargs)

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:425, in PythonKeyTracer.call_module(self, m, forward, args, kwargs)
    422 def call_module(
    423         self, m: torch.nn.Module, forward: Callable[..., Any], args: Tuple[Any, ...], kwargs: Dict[str, Any]
    424 ) -> Any:
--> 425     return forward(*args, **kwargs)

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/fx/_symbolic_trace.py:788, in Tracer.trace.<locals>.module_call_wrapper.<locals>.forward(*args, **kwargs)
    787 def forward(*args, **kwargs):
--> 788     return _orig_module_call(mod, *args, **kwargs)

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/nn/modules/module.py:1518, in Module._wrapped_call_impl(self, *args, **kwargs)
   1516     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1517 else:
-> 1518     return self._call_impl(*args, **kwargs)

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/nn/modules/module.py:1527, in Module._call_impl(self, *args, **kwargs)
   1522 # If we don't have any hooks, we want to skip the rest of the logic in
   1523 # this function, and just call forward.
   1524 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1525         or _global_backward_pre_hooks or _global_backward_hooks
   1526         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1527     return forward_call(*args, **kwargs)
   1529 try:
   1530     result = None

File <eval_with_key>.429:6, in forward(self, select)
      4 def forward(self, select):
      5     sum_1 = torch.sum(select);  select = None
----> 6     add = sum_1 + 0.01;  sum_1 = None
      7     return add

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/utils/_stats.py:20, in count.<locals>.wrapper(*args, **kwargs)
     18     simple_call_counter[fn.__qualname__] = 0
     19 simple_call_counter[fn.__qualname__] = simple_call_counter[fn.__qualname__] + 1
---> 20 return fn(*args, **kwargs)

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:555, in ProxyTorchDispatchMode.__torch_dispatch__(self, func, types, args, kwargs)
    552 @count
    553 def __torch_dispatch__(self, func, types, args=(), kwargs=None):
    554     with self.sym_mode.enable(False), set_original_aten_op(func):
--> 555         return self.inner_torch_dispatch(func, types, args, kwargs)

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:580, in ProxyTorchDispatchMode.inner_torch_dispatch(self, func, types, args, kwargs)
    577 if func in [prim.device.default]:
    578     return func(*args, **kwargs)
--> 580 return proxy_call(self, func, self.pre_dispatch, args, kwargs)

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:262, in proxy_call(proxy_mode, func, pre_dispatch, args, kwargs)
    260 if func in CURRENT_DECOMPOSITION_TABLE:
    261     with proxy_mode:
--> 262         r = CURRENT_DECOMPOSITION_TABLE[func](*args, **kwargs)
    263         if r is not NotImplemented:
    264             return r

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_decomp/decompositions.py:1737, in _to_copy(x, dtype, layout, device, pin_memory, non_blocking, memory_format)
   1735     x = torch._prims.device_put(x, device)
   1736 if dtype is not None and not dtype_converted:
-> 1737     x = torch._prims.convert_element_type(x, dtype)
   1738     dtype_converted = True
   1739 # In case of dtype promotion, faketensor converted into tensor.
   1740 # Need to convert into faketensor if input was a faketensor.

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_ops.py:448, in OpOverload.__call__(self, *args, **kwargs)
    447 def __call__(self, *args, **kwargs):
--> 448     return self._op(*args, **kwargs or {})

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/utils/_stats.py:20, in count.<locals>.wrapper(*args, **kwargs)
     18     simple_call_counter[fn.__qualname__] = 0
     19 simple_call_counter[fn.__qualname__] = simple_call_counter[fn.__qualname__] + 1
---> 20 return fn(*args, **kwargs)

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:555, in ProxyTorchDispatchMode.__torch_dispatch__(self, func, types, args, kwargs)
    552 @count
    553 def __torch_dispatch__(self, func, types, args=(), kwargs=None):
    554     with self.sym_mode.enable(False), set_original_aten_op(func):
--> 555         return self.inner_torch_dispatch(func, types, args, kwargs)

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:580, in ProxyTorchDispatchMode.inner_torch_dispatch(self, func, types, args, kwargs)
    577 if func in [prim.device.default]:
    578     return func(*args, **kwargs)
--> 580 return proxy_call(self, func, self.pre_dispatch, args, kwargs)

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:361, in proxy_call(proxy_mode, func, pre_dispatch, args, kwargs)
    358     else:
    359         args[0].proxy = proxy_out
--> 361 out = func(*args, **kwargs)
    363 # In some circumstances, we will be tracing in a situation where a tensor
    364 # is *statically* known to be a constant (currently, this only happens if
    365 # you run torch.tensor; deterministic factory functions like torch.arange
   (...)
    382 # propagating const-ness.  Similarly, we don't require the constant to
    383 # live on CPU, but we could.
    384 any_constant = pytree.tree_any_only(_ProxyTensor, lambda t: t.constant is not None, (f_args, f_kwargs))

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_ops.py:448, in OpOverload.__call__(self, *args, **kwargs)
    447 def __call__(self, *args, **kwargs):
--> 448     return self._op(*args, **kwargs or {})

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/utils/_stats.py:20, in count.<locals>.wrapper(*args, **kwargs)
     18     simple_call_counter[fn.__qualname__] = 0
     19 simple_call_counter[fn.__qualname__] = simple_call_counter[fn.__qualname__] + 1
---> 20 return fn(*args, **kwargs)

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_subclasses/fake_tensor.py:1250, in FakeTensorMode.__torch_dispatch__(self, func, types, args, kwargs)
   1248 assert self not in _get_current_dispatch_mode_stack(), func
   1249 try:
-> 1250     return self.dispatch(func, types, args, kwargs)
   1251 except TypeError:
   1252     log.exception("fake tensor raised TypeError")

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_subclasses/fake_tensor.py:1470, in FakeTensorMode.dispatch(self, func, types, args, kwargs)
   1464 if (
   1465     "prims::" in func._schema.name
   1466     and hasattr(func, "prim_meta_impl")
   1467     and not stride_incorrect_op(func)
   1468 ):
   1469     with self:
-> 1470         return func.prim_meta_impl(*args, **kwargs)
   1472 # Users can register FakeTensor rules for custom operators
   1473 # Call them if they exist.
   1474 if func.name() in torch._custom_op.impl.global_registry:

File /data/shurui.gui/mambaforge/envs/CSR/lib/python3.11/site-packages/torch/_prims/__init__.py:1993, in _convert_element_type_meta(a, dtype)
   1991 def _convert_element_type_meta(a: TensorLikeType, dtype: torch.dtype) -> TensorLikeType:
   1992     # Type checks
-> 1993     assert isinstance(a, TensorLike)
   1994     assert isinstance(dtype, torch.dtype)
   1996     # dtype conversion preserves dense strides
```

### Versions

PyTorch version: 2.1.0
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.4 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.2) 9.4.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.31

Python version: 3.11.6 | packaged by conda-forge | (main, Oct  3 2023, 10:40:35) [GCC 12.3.0] (64-bit runtime)
Python platform: Linux-5.15.0-69-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 12.3.52
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration:
GPU 0: NVIDIA A100 80GB PCIe
GPU 1: NVIDIA A100 80GB PCIe
GPU 2: NVIDIA A100 80GB PCIe
GPU 3: NVIDIA A100 80GB PCIe
GPU 4: NVIDIA A100 80GB PCIe
GPU 5: NVIDIA A100 80GB PCIe
GPU 6: NVIDIA A100 80GB PCIe
GPU 7: NVIDIA A100 80GB PCIe

Nvidia driver version: 510.60.02
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                    x86_64
CPU op-mode(s):                  32-bit, 64-bit
Byte Order:                      Little Endian
Address sizes:                   46 bits physical, 48 bits virtual
CPU(s):                          112
On-line CPU(s) list:             0-111
Thread(s) per core:              2
Core(s) per socket:              28
Socket(s):                       2
NUMA node(s):                    2
Vendor ID:                       GenuineIntel
CPU family:                      6
Model:                           85
Model name:                      Intel(R) Xeon(R) Gold 6258R CPU @ 2.70GHz
Stepping:                        7
CPU MHz:                         3399.999
CPU max MHz:                     4000.0000
CPU min MHz:                     1000.0000
BogoMIPS:                        5400.00
Virtualization:                  VT-x
L1d cache:                       1.8 MiB
L1i cache:                       1.8 MiB
L2 cache:                        56 MiB
L3 cache:                        77 MiB
NUMA node0 CPU(s):               0,2,4,6,8,10,12,14,16,18,20,22,24,26,28,30,32,34,36,38,40,42,44,46,48,50,52,54,56,58,60,62,64,66,68,70,72,74,76,78,80,82,84,86,88,90,92,94,96,98,100,102,104,106,108,110
NUMA node1 CPU(s):               1,3,5,7,9,11,13,15,17,19,21,23,25,27,29,31,33,35,37,39,41,43,45,47,49,51,53,55,57,59,61,63,65,67,69,71,73,75,77,79,81,83,85,87,89,91,93,95,97,99,101,103,105,107,109,111
Vulnerability Itlb multihit:     KVM: Mitigation: VMX disabled
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Mmio stale data:   Mitigation; Clear CPU buffers; SMT vulnerable
Vulnerability Retbleed:          Mitigation; Enhanced IBRS
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; Enhanced IBRS, IBPB conditional, RSB filling, PBRSB-eIBRS SW sequence
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Mitigation; TSX disabled
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf pni pclmulqdq dtes64 monitor ds_cpl vmx smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid dca sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault epb cat_l3 cdp_l3 invpcid_single intel_ppin ssbd mba ibrs ibpb stibp ibrs_enhanced tpr_shadow vnmi flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid cqm mpx rdt_a avx512f avx512dq rdseed adx smap clflushopt clwb intel_pt avx512cd avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local dtherm ida arat pln pts pku ospke avx512_vnni md_clear flush_l1d arch_capabilities

Versions of relevant libraries:
[pip3] numpy==1.26.0
[pip3] torch==2.1.0
[pip3] torch_geometric==2.4.0
[pip3] torchaudio==2.1.0
[pip3] torchvision==0.16.0
[pip3] triton==2.1.0
[conda] blas                      2.116                       mkl    conda-forge
[conda] blas-devel                3.9.0            16_linux64_mkl    conda-forge
[conda] cudatoolkit               11.8.0              h4ba93d1_12    conda-forge
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] libblas                   3.9.0            16_linux64_mkl    conda-forge
[conda] libcblas                  3.9.0            16_linux64_mkl    conda-forge
[conda] libjpeg-turbo             2.0.0                h9bf148f_0    pytorch
[conda] liblapack                 3.9.0            16_linux64_mkl    conda-forge
[conda] liblapacke                3.9.0            16_linux64_mkl    conda-forge
[conda] mkl                       2022.1.0           h84fe81f_915    conda-forge
[conda] mkl-devel                 2022.1.0           ha770c72_916    conda-forge
[conda] mkl-include               2022.1.0           h84fe81f_915    conda-forge
[conda] numpy                     1.26.0          py311h64a7726_0    conda-forge
[conda] pyg                       2.4.0           py311_torch_2.1.0_cu118    pyg
[conda] pytorch                   2.1.0           py3.11_cuda11.8_cudnn8.7.0_0    pytorch
[conda] pytorch-cuda              11.8                 h7e8668a_5    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torchaudio                2.1.0               py311_cu118    pytorch
[conda] torchtriton               2.1.0                     py311    pytorch
[conda] torchvision               0.16.0              py311_cu118    pytorch