Is it clear how to use the fancy new `adapt` kwarg with `grad`, `jvp`, and the like? Or do we need to extend them to cope with it?