Repository: pytorch
Issue: [dynamo] Support using UserDefinedFunction as argument (as_proxy).
Labels: feature,triaged,oncall: pt2,module: dynamo

### 🚀 The feature, motivation and pitch

At the moment using UserDefainedFunction as_proxy() fails as unimplemented.

The original reason why it is needed - we want to dynamo trace through Subclasses constructor that receives python Callable as an argument (and it is not used inside the constructor).

Dynamo tries to inline it and fails on UserDefinedFunction. as_proxy()


The straightforward approach will be:
1. trace through UserDefinedFn with speculate_subgraph and register subgraph as attribute of output_graph
At the moment of using UserDefinedFn, arguments for this function are not specified.

One of the ideas - could we put some "stub", empty subgraph at the moment of creating an argument and replace it with real subgraph only when this UserDefinedFunciton is called?


The original testcase:
```
    def test_unwrap_subclass_parameters_with_unused_callable_arg_in_ctor(self):
        def fn(x):
            return x

        _test_fn = fn

        class SC(WrapperSubclass):
            @staticmethod
            def __new__(cls, a, fn, outer_size=None, outer_stride=None):
                return WrapperSubclass.__new__(cls, a, outer_size, outer_stride)

            def __init__(self, a, fn, outer_size=None, outer_stride=None):
                self.a = a
                self.fn = fn

            def __tensor_flatten__(self):
                return ["a"], [self.fn]

            @staticmethod
            def __tensor_unflatten__(inner_tensors, meta, outer_size, outer_stride):
                a = inner_tensors["a"]
                fn = meta[0]
                return SC(a, fn, outer_size, outer_stride)

            @classmethod
            def __torch_dispatch__(cls, func, types, args, kwargs):
                if kwargs is None:
                    kwargs = {}
                args_a = pytree.tree_map_only(cls, lambda x: x.a, args)
                kwargs_a = pytree.tree_map_only(cls, lambda x: x.a, kwargs)

                out_a = func(*args_a, **kwargs_a)
                out_a_flat, spec = pytree.tree_flatten(out_a)
                out_flat = [
                    cls(o_a, _test_fn) if isinstance(o_a, torch.Tensor) else o_a
                    for o_a in out_a_flat
                ]
                out = pytree.tree_unflatten(out_flat, spec)
                from torch._higher_order_ops.cond import cond_op

                if func is cond_op:
                    return out
                else:
                    return return_and_correct_aliasing(func, args, kwargs, out)

        class M(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.p1 = torch.nn.Parameter(torch.ones(3, 4))
                self.p2 = torch.nn.Parameter(SC(torch.ones(3, 4), _test_fn))

            def forward(self, x):
                return x + 2 * self.p1 + self.p2

        m = M()
        from torch._functorch._aot_autograd.subclass_parametrization import (
            unwrap_tensor_subclass_parameters,
        )

        unwrap_tensor_subclass_parameters(m)

        x = torch.randn(3, 4)
        comp_fn = torch.compile(m, backend="aot_eager", fullgraph=True)
        out = comp_fn(x)
```
Error:
```
  File "/data/users/ivankobzarev/a/pytorch/torch/_dynamo/symbolic_convert.py", line 1685, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "/data/users/ivankobzarev/a/pytorch/torch/_dynamo/symbolic_convert.py", line 921, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "/data/users/ivankobzarev/a/pytorch/torch/_dynamo/variables/user_defined.py", line 600, in call_function
    *proxy_args_kwargs(args, kwargs, tx=tx),
  File "/data/users/ivankobzarev/a/pytorch/torch/_dynamo/utils.py", line 976, in proxy_args_kwargs
    unimplemented(
  File "/data/users/ivankobzarev/a/pytorch/torch/_dynamo/exc.py", line 355, in unimplemented
    raise Unsupported(msg, case_name=case_name) from from_exc
torch._dynamo.exc.Unsupported: call_function args: TensorVariable() UserFunctionVariable() ConstantVariable(NoneType: None) ConstantVariable(NoneType: None) 

from user code:
   File "/data/users/ivankobzarev/a/pytorch/test/functorch/test_aotdispatch.py", line 6421, in forward
    return x + 2 * self.p1 + self.p2
  File "/data/users/ivankobzarev/a/pytorch/torch/nn/utils/parametrize.py", line 407, in get_parametrized
    return parametrization()
  File "/data/users/ivankobzarev/a/pytorch/torch/nn/utils/parametrize.py", line 303, in forward
    x = self[0](*originals)
  File "/data/users/ivankobzarev/a/pytorch/torch/_functorch/_aot_autograd/subclass_parametrization.py", line 16, in forward
    rebuilt = tp.__tensor_unflatten__(d, meta, None, None)  # type: ignore[attr-defined]
  File "/data/users/ivankobzarev/a/pytorch/test/functorch/test_aotdispatch.py", line 6391, in __tensor_unflatten__
    return SC(a, fn, outer_size, outer_stride)
```

### Alternatives

_No response_

### Additional context

_No response_

cc @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @amjames