Support for stateful traces

As I understand it, implementing function transformations final-style should be preferred over initial-style whenever possible because the latter (while being more general) usually results in more complex code. I found that the scope of what's possible with final-style transformations can be expanded by attaching mutable state to the master trace:

```python
with new_master(InitTrace) as master:
master.state = InitTraceState(rng)
```

which is then accessible from within the trace class:

```python
rng = self.master.state.next_rng()
```
This approach allowed to massively simplify JAXnet's implementation by
- converting one transformation ([`_init_transform`](https://github.com/JuliusKunze/jaxnet/blob/master/jaxnet/core.py#L413)) from initial- to final-style,
- making another one simpler ([`_apply_transform`](https://github.com/JuliusKunze/jaxnet/blob/master/jaxnet/core.py#L501)) by freeing it from "squeezing" state along data flow trajectories.

I got rid of any custom stateful stacks / context managers, meaning that JAXnet handles state solely through (this extended version of) JAX's tracing mechanisms. Another potential use-case would be fastar, where both [`firstpass`](https://github.com/j-towns/fastar/blob/master/fastar/interpreter.py#L48) and [`fastpass`](https://github.com/j-towns/fastar/blob/master/fastar/interpreter.py#L121) could be converted from initial- to final-style, again simplifying code by getting rid of two near-copies of [a jaxpr interpreter](https://github.com/google/jax/blob/master/jax/core.py#L185).

Before I try to rewrite fastar in this way: Is this unintended for some reason? If not, shouldn't JAX provide a streamlined / less hacky way to have a stateful `MasterTrace`? I. e.

```python
with new_master(InitTrace, state=InitTraceState(rng)) as master:
```

and adding

```python
@property
def state(self):
return self.master.state
```

into the `Trace` base class. To me, it looks like this would encourage writing simpler function transformation code in many cases.

Start by reading the proposed state flow through new_master, MasterTrace, and the Trace base class, then compare it with the existing tracing mechanisms. Review the linked jax/core.py interpreter reference and the cited JAXnet transformations; a complete outcome would need a decided, supported API for stateful master traces and agreement on its intended use.

Written by the indexing model from the issue text.