Differentiation wrt system parameters

Hey,

I can see how brax can be used to differentiate with respect to the system state. I wonder if there is a nice way to also diff wrt, e.g., the mass of a body. Taking the basic tutorial as an example, I would like to have something like `step(sys, ..., ball_mass=10.)` which is a pure jax function.

I have found `brax.physics.bodies.Body`, which could be adapted via a `.replace()` call. However, I don't see how I can a) find it starting out with a `System` instance and b) how I could update that instance with a replaced version without potentially breaking things.