Handle jitted functions while jitting

Maybe first step:
```
def step(m, inp):
return m(inp)

jm = jit(model)
jfn = jit(fn)

jfn(jm, inp)
```

It would be cool if this worked:
```
def step(m, opt, inps):
logits = m(inps)
opt.zero_grad()
m.backward()
opt.step()

jm = jit(model)
jstep = jit(step)

jstep(jm, opt, inps)
```

This is also a crucial part #1134

Start with the nested-jit examples in the issue and read the related issue #1134 for context. No implementation files or tests are named, so first locate the jit entry point and its existing tests. Done should include jitted functions calling jitted functions, including the optimizer-training example shown here.

Written by the indexing model from the issue text.