[associative_scan] Lifted arguments #140043

bohnstingl · 2024-11-07T20:21:53Z

This PR implements lifted arguments for associative_scan

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @chauhang @aakhundov @ydwu4

…e xs

…an_73

pytorch-bot · 2024-11-07T20:21:56Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/140043

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

❌ 2 New Failures, 1 Unrelated Failure

As of commit 9cbf5fb with merge base d100e9a ():

NEW FAILURES - The following jobs have failed:

linux-binary-manywheel / manywheel-py3_9-cuda12_4-test / test (gh)
Process completed with exit code 1.
linux-binary-manywheel / manywheel-py3_9-cuda12_6-test / test (gh)
Process completed with exit code 1.

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

linux-binary-manywheel / manywheel-py3_9-cuda11_8-test / test (gh) (trunk failure)
Process completed with exit code 1.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

bohnstingl · 2024-11-07T20:23:59Z

@pytorchbot label "topic: not user facing"

WIP: export tests

ydwu4 · 2025-01-31T21:37:31Z

torch/_dynamo/variables/higher_order_ops.py

        with discard_graph_changes(tx):
+            # See NOTE [unspecialize int carry with unbacked symints]
+            # Note: this must be run under discard graph changes.
+            def create_unbacked_sym_node_var(tx) -> SymNodeVariable:


We don't want to do auto unspecialize for all control flow yet. Additional inputs is read-only so we don't need to create new unbacked symbols.

I reverted the handling back to the original implementation using cloning. In addition, there is now a check that all additional_inputs are TensorVariables.

ydwu4 · 2025-01-31T21:45:13Z

test/export/test_export.py

+            def forward(self, x):
+                return associative_scan(self.combine_fn, x, 1)
+
+        ep = export(Foo(), (xs,), dynamic_shapes={"x": {1: dim0}})


If we make dim 1 dynamic in this case (3, s0, 2), we then associative_scan over dim 1, the subgraph will work on a static shaped input (3, 2). If we want to make the subgraph dynamic, we'll need to scan over 0-th or 2-th dim or we change to first dim being dynamic. and in that case, i'm expecting symbol s0 to be lifted as additional_inputs.

As we discussed offline, during the export tests the symbolic shapes are not lifted and thus we don't see the int in the additional_inputs.

ydwu4 · 2025-02-04T18:53:22Z

test/export/test_export.py

+                return associative_scan(self.combine_fn, x, 1, combine_mode="generic")
+
+        inp = torch.randn(3, 10, 2, device=torch.device("cuda"))
+        ep = torch.export.export(M(), (inp,))


can replace "torch.export.export" with "export", there're some test patching behind the scene if we use export.

The failure is probably caused by the randn call of the buffer (not sure why the randn is called inside vmap rather than outside at module initialization time) used inside vmap. We can put a constant tensor in the buffer.

I changed the torch.randn to torch.ones and the torch.export.export to export. However, these export tests are giving me quite some headache. The logs of the three tests are attached below:

The two last tests fail with a vmap issue. Is this because the associative_scan does not yet have a vmap implementation?

ERROR: test_export_associative_scan_lifted_buffers (__main__.TestExport.test_export_associative_scan_lifted_buffers) ---------------------------------------------------------------------- Traceback (most recent call last): File "/data_malta3_ssd/pytorch/torch/testing/_internal/common_utils.py", line 3120, in wrapper method(*args, **kwargs) File "/data_malta3_ssd/pytorch/test/export/test_export.py", line 6346, in test_export_associative_scan_lifted_buffers ep = export(M(), (inp,)) ^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/export/__init__.py", line 368, in export return _export( ^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/export/_trace.py", line 1044, in wrapper raise e File "/data_malta3_ssd/pytorch/torch/export/_trace.py", line 1017, in wrapper ep = fn(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/export/exported_program.py", line 117, in wrapper return fn(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/export/_trace.py", line 2079, in _export return _export_for_training( ^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/export/_trace.py", line 1044, in wrapper raise e File "/data_malta3_ssd/pytorch/torch/export/_trace.py", line 1017, in wrapper ep = fn(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/export/exported_program.py", line 117, in wrapper return fn(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/export/_trace.py", line 1944, in _export_for_training export_artifact = export_func( # type: ignore[operator] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/export/_trace.py", line 1384, in _strict_export_lower_to_aten_ir aten_export_artifact = lower_to_aten_callback( ^^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/export/_trace.py", line 1665, in _export_to_aten_ir_make_fx gm, graph_signature = transform(_make_fx_helper)( ^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/export/_trace.py", line 1585, in _make_fx_helper gm = make_fx( ^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/fx/experimental/proxy_tensor.py", line 2194, in wrapped return make_fx_tracer.trace(f, *args) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/fx/experimental/proxy_tensor.py", line 2132, in trace return self._trace_inner(f, *args) ^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/fx/experimental/proxy_tensor.py", line 2103, in _trace_inner t = dispatch_trace( ^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/_compile.py", line 51, in inner return disable_fn(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/_dynamo/eval_frame.py", line 749, in _fn return fn(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/fx/experimental/proxy_tensor.py", line 1136, in dispatch_trace graph = tracer.trace(root, concrete_args) # type: ignore[arg-type] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/fx/experimental/proxy_tensor.py", line 1692, in trace res = super().trace(root, concrete_args) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/_dynamo/eval_frame.py", line 749, in _fn return fn(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/fx/_symbolic_trace.py", line 832, in trace (self.create_arg(fn(*args)),), ^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/fx/experimental/proxy_tensor.py", line 1191, in wrapped out = f(*tensors) # type:ignore[call-arg] ^^^^^^^^^^^ File "<string>", line 1, in <lambda> File "/data_malta3_ssd/pytorch/torch/export/_trace.py", line 1488, in wrapped_fn return tuple(flat_fn(*args)) ^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/_functorch/_aot_autograd/utils.py", line 184, in flat_fn tree_out = fn(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 875, in functional_call out = PropagateUnbackedSymInts(mod).run( ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/fx/interpreter.py", line 171, in run self.env[node] = self.run_node(node) ^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/fx/experimental/symbolic_shapes.py", line 6944, in run_node result = super().run_node(n) ^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/fx/interpreter.py", line 234, in run_node return getattr(self, n.op)(n.target, args, kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/fx/interpreter.py", line 314, in call_function return target(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/fx/experimental/proxy_tensor.py", line 1239, in __torch_function__ return func(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/fx/experimental/proxy_tensor.py", line 1286, in __torch_function__ return func(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/_ops.py", line 866, in handler return torch._library.utils.handle_dispatch_mode( ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/_library/utils.py", line 296, in handle_dispatch_mode return curr_mode.__torch_dispatch__(op_overload, overload_types, args, kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/utils/_stats.py", line 27, in wrapper return fn(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/fx/experimental/proxy_tensor.py", line 1341, in __torch_dispatch__ return proxy_call(self, func, self.pre_dispatch, args, kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/fx/experimental/proxy_tensor.py", line 973, in proxy_call track_tensor_tree(out, proxy_out, constant=constant, tracer=tracer) File "/data_malta3_ssd/pytorch/torch/fx/experimental/proxy_tensor.py", line 674, in track_tensor_tree wrap_with_proxy(inner_res, proxy_res, constant) File "/data_malta3_ssd/pytorch/torch/fx/experimental/proxy_tensor.py", line 621, in wrap_with_proxy set_meta(proxy, e) File "/data_malta3_ssd/pytorch/torch/fx/experimental/proxy_tensor.py", line 491, in set_meta proxy.node.meta["tensor_meta"] = _extract_tensor_metadata(val) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/fx/passes/shape_prop.py", line 55, in _extract_tensor_metadata if result.is_contiguous(memory_format=query_format): ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ RuntimeError: NYI: querying is_contiguous inside of vmap for memory_format other than torch.contiguous_format While executing %add_1 : [num_users=1] = call_function[target=operator.add](args = (%_add_batch_dim, %_add_batch_dim_1), kwargs = {}) Original traceback: File "/data_malta3_ssd/pytorch/test/export/test_export.py", line 6343, in forward return associative_scan(self.combine_fn, x, 1, combine_mode="generic") File "/data_malta3_ssd/pytorch/torch/_higher_order_ops/associative_scan.py", line 229, in associative_scan result_flat = generic_associative_scan(combine_fn, leaves, additional_inputs=()) File "/data_malta3_ssd/pytorch/torch/_higher_order_ops/associative_scan.py", line 342, in generic_associative_scan scans = _scan(leaves) File "/data_malta3_ssd/pytorch/torch/_higher_order_ops/associative_scan.py", line 303, in _scan reduced_elems = operator( File "/data_malta3_ssd/pytorch/torch/_higher_order_ops/associative_scan.py", line 36, in wrap_combine_fn_flat combined = combine_fn(lhs, rhs) File "/data_malta3_ssd/pytorch/torch/_functorch/apis.py", line 202, in wrapped return vmap_impl( File "/data_malta3_ssd/pytorch/torch/_functorch/vmap.py", line 331, in vmap_impl return _flat_vmap( File "/data_malta3_ssd/pytorch/torch/_functorch/vmap.py", line 481, in _flat_vmap batched_outputs = func(*batched_inputs, **kwargs) File "/data_malta3_ssd/pytorch/test/export/test_export.py", line 6340, in combine_fn return (x + y) * self.buffer To execute this test, run the following from the base repo dir: python test/export/test_export.py TestExport.test_export_associative_scan_lifted_buffers This message can be suppressed by setting PYTORCH_PRINT_REPRO_ON_FAILURE=0 ====================================================================== ERROR: test_export_associative_scan_symbol_dim (__main__.TestExport.test_export_associative_scan_symbol_dim) ---------------------------------------------------------------------- Traceback (most recent call last): File "/data_malta3_ssd/pytorch/torch/testing/_internal/common_utils.py", line 3120, in wrapper method(*args, **kwargs) File "/data_malta3_ssd/pytorch/test/export/test_export.py", line 6281, in test_export_associative_scan_symbol_dim ep = export(Foo(), (xs,), dynamic_shapes={"x": {1: dim1}}) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/export/__init__.py", line 368, in export return _export( ^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/export/_trace.py", line 1044, in wrapper raise e File "/data_malta3_ssd/pytorch/torch/export/_trace.py", line 1017, in wrapper ep = fn(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/export/exported_program.py", line 117, in wrapper return fn(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/export/_trace.py", line 2079, in _export return _export_for_training( ^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/export/_trace.py", line 1044, in wrapper raise e File "/data_malta3_ssd/pytorch/torch/export/_trace.py", line 1017, in wrapper ep = fn(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/export/exported_program.py", line 117, in wrapper return fn(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/export/_trace.py", line 1944, in _export_for_training export_artifact = export_func( # type: ignore[operator] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/export/_trace.py", line 1296, in _strict_export_lower_to_aten_ir gm_torch_level = _export_to_torch_ir( ^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/export/_trace.py", line 693, in _export_to_torch_ir gm_torch_level, _ = torch._dynamo.export( ^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/_dynamo/eval_frame.py", line 1579, in inner result_traced = opt_f(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/nn/modules/module.py", line 1749, in _wrapped_call_impl return self._call_impl(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/nn/modules/module.py", line 1760, in _call_impl return forward_call(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/_dynamo/eval_frame.py", line 570, in _fn return fn(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/nn/modules/module.py", line 1749, in _wrapped_call_impl return self._call_impl(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/nn/modules/module.py", line 1760, in _call_impl return forward_call(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/_dynamo/convert_frame.py", line 1400, in __call__ return self._torchdynamo_orig_callable( ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/_dynamo/convert_frame.py", line 565, in __call__ return _compile( ^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/_dynamo/convert_frame.py", line 997, in _compile guarded_code = compile_inner(code, one_graph, hooks, transform) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/_utils_internal.py", line 95, in wrapper_function return function(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/_dynamo/convert_frame.py", line 726, in compile_inner return _compile_inner(code, one_graph, hooks, transform) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/_dynamo/convert_frame.py", line 760, in _compile_inner out_code = transform_code_object(code, transform) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/_dynamo/bytecode_transformation.py", line 1404, in transform_code_object transformations(instructions, code_options) File "/data_malta3_ssd/pytorch/torch/_dynamo/convert_frame.py", line 236, in _fn return fn(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/_dynamo/convert_frame.py", line 660, in transform tracer = InstructionTranslator( ^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/_dynamo/symbolic_convert.py", line 2780, in __init__ self._throw_if_in_functorch() File "/data_malta3_ssd/pytorch/torch/_dynamo/symbolic_convert.py", line 2896, in _throw_if_in_functorch unimplemented(msg) File "/data_malta3_ssd/pytorch/torch/_dynamo/exc.py", line 380, in unimplemented raise Unsupported(msg, case_name=case_name) torch._dynamo.exc.Unsupported: If you are reaching here, it means dynamo failed for one of the following reasons: - Calling torch.func.vmap(compiled_fn) function from eager mode is not supported. Ensure that torch.func.vmap is also wrapped within a torch.compile function. For more information, see PyTorch issue #128711. - torch.func.vmap(fn) requires the function to be inlined by dynamo Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information To execute this test, run the following from the base repo dir: python test/export/test_export.py TestExport.test_export_associative_scan_symbol_dim This message can be suppressed by setting PYTORCH_PRINT_REPRO_ON_FAILURE=0 ====================================================================== ERROR: test_export_associative_scan_symbol_scandim (__main__.TestExport.test_export_associative_scan_symbol_scandim) ---------------------------------------------------------------------- Traceback (most recent call last): File "/data_malta3_ssd/pytorch/torch/testing/_internal/common_utils.py", line 3120, in wrapper method(*args, **kwargs) File "/data_malta3_ssd/pytorch/test/export/test_export.py", line 6313, in test_export_associative_scan_symbol_scandim ep = export(Foo(), (xs,), dynamic_shapes={"x": {1: dim1}}) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/export/__init__.py", line 368, in export return _export( ^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/export/_trace.py", line 1044, in wrapper raise e File "/data_malta3_ssd/pytorch/torch/export/_trace.py", line 1017, in wrapper ep = fn(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/export/exported_program.py", line 117, in wrapper return fn(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/export/_trace.py", line 2079, in _export return _export_for_training( ^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/export/_trace.py", line 1044, in wrapper raise e File "/data_malta3_ssd/pytorch/torch/export/_trace.py", line 1017, in wrapper ep = fn(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/export/exported_program.py", line 117, in wrapper return fn(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/export/_trace.py", line 1944, in _export_for_training export_artifact = export_func( # type: ignore[operator] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/export/_trace.py", line 1296, in _strict_export_lower_to_aten_ir gm_torch_level = _export_to_torch_ir( ^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/export/_trace.py", line 693, in _export_to_torch_ir gm_torch_level, _ = torch._dynamo.export( ^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/_dynamo/eval_frame.py", line 1579, in inner result_traced = opt_f(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/nn/modules/module.py", line 1749, in _wrapped_call_impl return self._call_impl(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/nn/modules/module.py", line 1760, in _call_impl return forward_call(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/_dynamo/eval_frame.py", line 570, in _fn return fn(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/nn/modules/module.py", line 1749, in _wrapped_call_impl return self._call_impl(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/nn/modules/module.py", line 1760, in _call_impl return forward_call(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/_dynamo/convert_frame.py", line 1400, in __call__ return self._torchdynamo_orig_callable( ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/_dynamo/convert_frame.py", line 565, in __call__ return _compile( ^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/_dynamo/convert_frame.py", line 997, in _compile guarded_code = compile_inner(code, one_graph, hooks, transform) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/_utils_internal.py", line 95, in wrapper_function return function(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/_dynamo/convert_frame.py", line 726, in compile_inner return _compile_inner(code, one_graph, hooks, transform) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/_dynamo/convert_frame.py", line 760, in _compile_inner out_code = transform_code_object(code, transform) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/_dynamo/bytecode_transformation.py", line 1404, in transform_code_object transformations(instructions, code_options) File "/data_malta3_ssd/pytorch/torch/_dynamo/convert_frame.py", line 236, in _fn return fn(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/_dynamo/convert_frame.py", line 660, in transform tracer = InstructionTranslator( ^^^^^^^^^^^^^^^^^^^^^^ File "/data_malta3_ssd/pytorch/torch/_dynamo/symbolic_convert.py", line 2780, in __init__ self._throw_if_in_functorch() File "/data_malta3_ssd/pytorch/torch/_dynamo/symbolic_convert.py", line 2896, in _throw_if_in_functorch unimplemented(msg) File "/data_malta3_ssd/pytorch/torch/_dynamo/exc.py", line 380, in unimplemented raise Unsupported(msg, case_name=case_name) torch._dynamo.exc.Unsupported: If you are reaching here, it means dynamo failed for one of the following reasons: - Calling torch.func.vmap(compiled_fn) function from eager mode is not supported. Ensure that torch.func.vmap is also wrapped within a torch.compile function. For more information, see PyTorch issue #128711. - torch.func.vmap(fn) requires the function to be inlined by dynamo Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information To execute this test, run the following from the base repo dir: python test/export/test_export.py TestExport.test_export_associative_scan_symbol_scandim This message can be suppressed by setting PYTORCH_PRINT_REPRO_ON_FAILURE=0 ---------------------------------------------------------------------- Ran 3 tests in 2.202s FAILED (errors=3)```

For the first failure, it looks like an export x vmap problem: we probably have to use the point-wise mode and fix export issues. Since export doesn't require inductor compilation, i'm expecting export to succeed.

For the second failure, i'm actually not sure why we'll trigger this error because the exported model is not under vmap. Does run "python test/export/test_export.py TestExport.test_export_associative_scan_symbol_scandim" alone also fail?

Using pointwise for the first testcase doesn't work. I think inductor is involved and in current lowering additional_inputs are not supported in inductor.

LoweringException: RuntimeError: Unable to generate code for associative_scan op, because there are lifted arguments

For the second issue python test/export/test_export.py TestExport.test_export_associative_scan_symbol_scandim does pass. In fact both of the latter tests pass with that. I.e, using:
python test/export/test_export.py TestExport.test_export_associative_scan_symbol_dim and
python test/export/test_export.py TestExport.test_export_associative_scan_symbol_scandim
both pass

We need to update the accociative_scan's backend="eager": https://github.com/pytorch/pytorch/blob/main/torch/_higher_order_ops/associative_scan.py#L137

edit: we can do this change in next PR. Can exp fail this test and remove the exp in next PR.

Right, I've done the changes except the backend. Indeed, the backend does resolve the remaining issues.

However, the problem with switching the backend to "eager" is that the pointwise check needs to be reworked. This is because the pointwise check is invoked only in the tracing function for inductor and by switching the backend, this check is not invoked anymore.
Could you please start the CI tests?

The backend change is pursued in a separate PR here #146973

…m different calls

ydwu4 · 2025-02-10T22:10:48Z

@pytorchbot merge

pytorchmergebot · 2025-02-10T22:12:53Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

pytorchmergebot · 2025-02-11T00:15:36Z

Merge failed

Reason: 2 jobs have failed, first few of them are: linux-binary-manywheel / manywheel-py3_9-cuda12_4-test / test, linux-binary-manywheel / manywheel-py3_9-cuda12_6-test / test

Details for Dev Infra team

Raised by workflow job

ydwu4 · 2025-02-11T18:04:06Z

@pytorchbot merge -i

pytorchmergebot · 2025-02-11T18:05:56Z

Merge started

Your change will be merged while ignoring the following 3 checks: linux-binary-manywheel / manywheel-py3_9-cuda12_4-test / test, linux-binary-manywheel / manywheel-py3_9-cuda12_6-test / test, linux-binary-manywheel / manywheel-py3_9-cuda11_8-test / test

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

This PR fixes some issues with torch export discussed here: #140043 (comment) However, this backend change does still not resolve the failure for specific shapes mentioned here: #137943 (comment) Pull Request resolved: #146973 Approved by: https://github.com/ydwu4

bohnstingl added 18 commits October 26, 2024 22:59

Ensure that the combine_fn is only called with the proper slice of th…

460f753

…e xs

Fixed shape check

1419a79

WIP: nested associative_scan

944649a

Incorporated first review round

f974cf3

Implemented better and more unified testing procedures

ab0e515

Rebase to main

59b164b

Lintrunner cleanup

6dc7811

WIP: new _run_test interface

308e89c

Integrated comments from PR and updated testcases

0a902eb

Integrated nested tuple for the vmap used in generic_associative_scan

022a454

Integrated nit changes

8aeef66

Fixed minor issue with testcase parameters

9e01fff

Rebased to associative_scan_70

90e9ac3

Fixed rebasing issues

ce619ea

Implemented lifted arguments

4d8247c

Merge branch 'main' of github.com:pytorch/pytorch into associative_sc…

8eba2d1

…an_73

Merge branch 'main' of github.com:pytorch/pytorch into associative_sc…

014220e

…an_73

Added additional testcases

55266fb

bohnstingl requested a review from zou3519 as a code owner November 7, 2024 20:21

pytorch-bot bot added module: dynamo module: inductor labels Nov 7, 2024

bohnstingl mentioned this pull request Nov 7, 2024

Improvements for associative_scan - Lifted arguments bohnstingl/pytorch#5

Closed

pytorch-bot bot added the topic: not user facing topic category label Nov 7, 2024

pytorchbot added the open source label Nov 7, 2024

bdhirsh added the triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module label Nov 8, 2024

zou3519 requested review from ydwu4 and removed request for zou3519 November 11, 2024 17:59

bohnstingl changed the title ~~Improvements for associative_scan - Lifted arguments~~ [associative_scan] Lifted arguments Nov 19, 2024

Integrated comments from code review

34c6b56

WIP: export tests

bohnstingl mentioned this pull request Jan 31, 2025

[associative_scan] Support lifted arguments for inductor #146108

Open

ydwu4 reviewed Jan 31, 2025

View reviewed changes

Integrated review comments and minor fixes

b9c5609

ydwu4 reviewed Feb 4, 2025

View reviewed changes

bohnstingl added 5 commits February 4, 2025 23:27

Integrated review comments

d3a44a8

Fixed lintrunner issues

8df0871

CI test issue fix

5eaef86

Removed graph code in export tests, because of conflicting graphs fro…

1d96a79

…m different calls

Fix Serdes issue

9cbf5fb

ydwu4 approved these changes Feb 10, 2025

View reviewed changes

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Feb 10, 2025

pytorchmergebot added the merging label Feb 10, 2025

pytorchmergebot removed the merging label Feb 11, 2025

pytorchmergebot added the merging label Feb 11, 2025

pytorchmergebot added the Merged label Feb 11, 2025

pytorchmergebot closed this in 3a29992 Feb 11, 2025

pytorchmergebot removed the merging label Feb 11, 2025

bohnstingl mentioned this pull request Feb 12, 2025

[associative_scan] compile backend change to "eager" #146973

Closed

[associative_scan] Lifted arguments #140043

[associative_scan] Lifted arguments #140043

Uh oh!

Conversation

bohnstingl commented Nov 7, 2024 • edited by pytorch-bot bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Nov 7, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/140043

❌ 2 New Failures, 1 Unrelated Failure

Uh oh!

bohnstingl commented Nov 7, 2024

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ydwu4 Jan 31, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bohnstingl Feb 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ydwu4 Feb 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bohnstingl Feb 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ydwu4 commented Feb 10, 2025

Uh oh!

pytorchmergebot commented Feb 10, 2025

Merge started

Uh oh!

pytorchmergebot commented Feb 11, 2025

Merge failed

Uh oh!

ydwu4 commented Feb 11, 2025

Uh oh!

pytorchmergebot commented Feb 11, 2025

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

bohnstingl commented Nov 7, 2024 •

edited by pytorch-bot bot

Loading

pytorch-bot bot commented Nov 7, 2024 •

edited

Loading

ydwu4 Jan 31, 2025 •

edited

Loading

bohnstingl Feb 4, 2025 •

edited

Loading

ydwu4 Feb 4, 2025 •

edited

Loading

bohnstingl Feb 4, 2025 •

edited

Loading