To use the setup as it is now, one would need to define their neural network (or whatever parametric observable you like that yields histograms) using only jax, which allows the tracing of gradients. This is generally seen as a limitation for two reasons:
- New users will not want to learn yet another way to do neural networks
- Model implementations that are non-trivially portable between frameworks have no way to link with the functionality of the library
This issue outlines the kind of syntax and operations needed to enable this, but jax is still in it's infancy, so there is no good user-facing API for this yet. It seems like it could be possible! (:
To use the setup as it is now, one would need to define their neural network (or whatever parametric observable you like that yields histograms) using only jax, which allows the tracing of gradients. This is generally seen as a limitation for two reasons:
This issue outlines the kind of syntax and operations needed to enable this, but jax is still in it's infancy, so there is no good user-facing API for this yet. It seems like it could be possible! (: