`__getitem__` method for `Trace`

When using importance resampling and other methods eventually calling  `jax.random.categorical(key, weights)` to get an array of indecies, it would be increadibly helpful to be able to call `trace[indecies]` which produces another trace with the sliced/indexed out traces.