Repository: pytorch
Issue: [TorchScript] Support ScriptFunction arguments in torch.jit.script calls.
Labels: oncall: jit

### 🚀 The feature, motivation and pitch

I'm trying to use TorchScript with code that uses an ordinary differential equation solver. This code basically has a loop that iteratively evaluates an ODE step function, which itself takes another function as part of its inputs and evaluates it multiple times at different points. The particular function to use in the ODE steps is decided at runtime from a set based on a condition.

Putting things together, it looks roughly like this:

```python
def solver(..., step_f1: callable, step_f2: callable):
    if x:
        for i in range(...):
            y = step_f1(y, ...)
        return y
    else:
        z = ...
        for i in range(...):
            y = step_f2(y, z, ...)
        return y

def f1(...):
    ...

def f2(...):
    ...

def ode_step(func, ...):
  # Invokes func as part of its logic.
  ...
```

This solver function would be normally used with some `func_f1`, `func_f2` like this:
```python
step_f1 = lambda *args: ode_step(f1, *args)
step_f2 = lambda *args: ode_step(f2, *args)
```

When exporting the model, because of the if condition and the for loops, the solver function itself cannot be traced and requires using `torch.jit.script`. However, since `f1` and `f2` would actually be expensive model evaluations I trace these:
```python
step_f1 = torch.jit.trace(lambda *args: ode_step(f1, *args))
step_f2 = torch.jit.trace(lambda *args: ode_step(f2, *args))
```

Ideally I would like to be able to use the same solver function for both use cases: one without `torch.jit.script` and with no tracing, and another when exporting the model where I manually call `torch.jit.script` and use traced functions. However, if I try to use `torch.jit.script` I get the runtime error "function cannot be used as a value" despite the fact that these are not really python functions but `torch.jit.ScriptFunction` from tracing.

I also found that this seems to be functionally possible, since I found a way to get `torch.jit.script` to use my traced functions, but it comes with nasty consequences. See the alternatives section below for more details.

Since this does affect script/trace composability and apparently it is already functionally possible, albeit somewhat broken, would it make sense to allow using `torch.jit.ScriptFunction` objects as valid argument types for torch.jit.script?

### Alternatives

Since `torch.jit.script` complains if I pass any `torch.jit.ScriptFunction` argument directly or in any nested function it calls, I tried something like this:

```python
def trace_functions(...):
    step_f1 = torch.jit.trace(lambda *args: ode_step(f1, *args))
    step_f2 = torch.jit.trace(lambda *args: ode_step(f2, *args))

    def solver_local(...):
        if x:
            for i in range(...):
                y = func_f1(y, ...)
            return y
        else:
            z = ...
            for i in range(...):
                y = func_f2(y, z, ...)
            return y

    solver_script = torch.jit.script(solver_local)
```

In this case, the `torch.jit.ScriptFunction` objects `func_f1` and `func_f2` are captured in `solver_local` instead of passed as arguments. This version actually doesn't raise any exceptions, but it forces me to create an inner copy of `solver` that I cannot reuse either in other parts of the code, effectively duplicating the code maintenance of this function and making it error-prone.

Trying to use `torch.jit.script` with some other function, lambda, or result from `functools.partial` that calls `solver` and sets `func_f1` and `func_f2` does not work.

### Additional context

_No response_

cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel