ProbProg: Docs

[sbrantq]

Assumes #2831

[sbrantq]

https://enzymead.github.io/Reactant.jl/previews/PR2834/api/probprog

Thanks changes are in preview

[penelopeysm]

Hi @sbrantq many congratulations on your work here! By way of introduction, I've worked on Turing.jl for the past couple of years-ish. I hope it's not too rude of me to comment here, @wsmoses asked me to take a look.

I think at this early stage your docs are looking really thorough in terms of coverage. Docs are definitely something that can & will be iterated on over time, but I think it's most important to get a really good first impression, which is going to come via this index page! On the other pages you'll have a lot more leeway to not be so engaging.

For the examples on this page, I wonder if it would be possible to break it up into separate code blocks plus add some explanatory text to show at each stage what these things do? For example, that would let you move the mention of generate() closer to the point where it's actually called. Another benefit of splitting them up is that it lets a reader see what objects these functions return. For example, in this bit,

trace_tensor, _ = compiled_fn(
    rng, linreg, xs, obs_tensor, constrained_addresses,
    selection, step_size, inverse_mass_matrix,
)

selected_entries = ProbProg.filter_entries_by_selection(tt.entries, selection)
trace            = ProbProg.unflatten_trace(trace_tensor, 0.0, selected_entries, nothing)

I'd really like to see what things like selection, trace_tensor, and trace really are (and I don't want to have to run it myself to find out 😄)

If you'd like to see an example of what I mean, here are the docs for a package I wrote: https://pysm.dev/FlexiChains.jl/stable/turing/. I think that this is in spirit similar to the idea of literate programming, or what we'd nowadays call a 'notebook'.

Alternatively, or in addition to that, I wonder to what extent the examples can be simplified? For example, the linreg example in Turing is much shorter:

using Turing
@model function linreg(x, y)
    slope ~ Normal(0, 2)
    intercept ~ Normal(0, 10)
    for i in eachindex(y)
        y[i] ~ Normal(slope * x[i] + intercept)
    end
end
x = ...
y = ...
sample(linreg(x, y), NUTS(), 500)

Now I can tell that there is a lot of irreducible complexity in the present example and it's unlikely that you'll be able to get it down to 10 lines, but I would ask whether there is a need to specify things like the initial mass matrix and step size, or whether the library could just choose or find reasonable defaults for the user. If this isn't possible yet then that's definitely something worth exploring as a future direction 🙂.

More broadly, and not for this PR, I'd suggest thinking about who the likely audience for the docs is, which will help shape the general approach to writing docs. Feel free to come up with your own ideas, but my suspicion is that you're initially going to get a bunch of people who know some other PPL (indeed most likely Turing, given that we're in Julia land) but have no clue how to use Reactant. I think you'd have to bring those people on board and hand hold them through some of the Reactant-specific stuff, for example explain what @compile does (in 2026 our attention spans are very short 😅 so I wouldn't assume that they're going to click to the rest of the Reactant docs to read it).

[sbrantq]

Hi Penelope, nice to meet you and really appreciate your valuable insights! Completely agree with your points here. I will do a few more iterations to restructure the index into smaller cells with more explanations in between (the FlexiChains page is a great reference, thanks for the pointer!), trim down the reducible complexity, and add Reactant-land explanation to better target our audience!