Jax: Optix documentation

[joaogui1]

The optix optimization library has docstrings for most of its classes already, but it is not included in the docs, is that by design?
Edit: @tomhennigan @gnecula

[tomhennigan]

I think optix is in a slightly odd space where it is in jax.experimental to make it easy for folks to try but it is also an orthogonal project to JAX (to me it is clearly a library that sits on top of JAX).

@gnecula @mtthss should we consider putting optix in its own repo? It is a very simple project of course, but it's not clear to me that it should live in JAX core and be documented alongside it.

[gnecula]

I agree that optix being in jax.experimental is not ideal. It is indeed a library on top of JAX. Putting it into its own repo feels costly though (will need its own Travis, Readthedocs, copybara, ... settings).

(How about optimizers.py, should be part of the same discussion?)

@jekbradbury WDYT?

[gnecula]

Meanwhile, I am adding the documentation

[mattjj]

Optix is a great little library! But my weak opinion is that it should be in its own repo, for these weak reasons:

1. no JAX core devs are very familiar with the code and we have no vision for it, and so it's not something we could easily maintain or develop (i.e. we're relying on @mtthss) which makes it different from the other code in the JAX repo;
2. fixes and improvements by @mtthss, @tomhennigan, and others have to be merged into this repo by JAX core devs, which slows progress when we're slow to respond (@mtthss is relying on us);
3. issues like this one incur a small but nonzero time cost for the JAX core team.

Those are only weak reasons because the code is so concise and currently low-velocity. If the code were big, or being developed rapidly, we'd have already felt the pain and decided to move it out (or never merged it at all). But perhaps regardless of size the same principles could be applied.

If it were deemed too small to be in its own repo (though it's not clear to me there is such a minimum size!), it's not clear to me whether it's better to be bundled with JAX or with, say, Haiku.

Overall, though, I want to do whatever is best for @tomhennigan and @mtthss because this isn't costing us much. Even though my weak opinion is that it should be taken out of JAX, I defer to your judgement. (Though if you're ambivalent, we'll take it out!)

(By the way, #1620 is where it got merged, but there doesn't seem to be a relevant discussion in that PR thread.)

[mattjj]

IIUC this specific issue was solved by #2312, so I'm going to close it. I'll start an email thread about taking optix out of jax.