Add usage examples for the documentations

[amir-saadat]

This may be a duplicate. It would be very useful if we add example usage together with the docstrings and/or the documentation.

[fabianp]

hey @amir-saadat , could you elaborate a bit more on where do you think examples should be added or linked?

[amir-saadat]

hey @fabianp thanks for looking into this. a couple of suggestions:

• Within the documentation at https://jaxopt.github.io/, we have examples for a specific category. For example, for stochastic optimization, we give examples for optax, adaptive and manual solvers. As I was searching for a solver usage, say for optax solver, I realized that it would require the user to traverse back to the main category to search for an example. It might not be entirely obvious that the examples can be found in the main category, so I wonder if it would also make sense to either backlink to the examples in the main category, or even provide the examples within the description of the solvers.

• It would be useful IMO to provide some examples in the docstring in addition to relying only on the documentation. Sometimes I find them useful and handy to just hover over the solver and take a quick look at the usage examples. I should start this for nonlinear lsq perhaps ;). For example, for jax.grad at the end there is:

  >>> import jax
  >>>
  >>> grad_tanh = jax.grad(jax.numpy.tanh)
  >>> print(grad_tanh(0.2))
  0.961043

[fabianp]

thanks for the clarification. I agree that both of these points would be nice improvements.

For the first point, for example in the stochastic optimization page , I don't think the examples should be at the top of the page, but rather inside each section, so that you can easily see which solvers use (say) the OptaxSolver or the ArmijoSGD. I think we're saying the same thing, just wanted to state a specific example :-)

[fabianp]

@amir-saadat : would you be willing to submit a pull request implementing (some of) this?