Merge pull request #16 from TuringLang/py/hover

Add hoverable model source code

Author: penelopeysm

ad.py

@@ -93,20 +93,40 @@ def run_ad(args):
     append_to_github_output("results", results)
 
 
+def get_model_definition(model_key):
+    """Get the model definition from the Julia script."""
+    lines = []
+    submodels = []
+    record = False
+    with open("models.jl", "r") as file:
+        for line in file:
+            line = line.rstrip()
+            if line.startswith(f"@model function {model_key}"):
+                record = True
+            if record:
+                lines.append(line)
+
+                if "to_submodel" in line:
+                    submodel_name = line.split("to_submodel(")[1].split("(")[0]
+                    submodels.append(submodel_name)
+                if line == "end":
+                    break
+    for submodel in submodels:
+        lines = [get_model_definition(submodel), *lines]
+    return "<br>".join(lines)
+
+
 def html(_args):
     ## Here you can register known errors that have been reported on GitHub /
     ## have otherwise been documented. They will be turned into links in the table.
 
     ENZYME_RVS_ONE_PARAM = "https://github.com/EnzymeAD/Enzyme.jl/issues/2337"
     ENZYME_FWD_BLAS = "https://github.com/EnzymeAD/Enzyme.jl/issues/1995"
+    MOONCAKE_THREADED = "https://github.com/chalk-lab/Mooncake.jl/issues/570"
     KNOWN_ERRORS = {
-        ("assume_beta", "EnzymeReverse"): ENZYME_RVS_ONE_PARAM,
-        ("assume_dirichlet", "EnzymeReverse"): ENZYME_RVS_ONE_PARAM,
-        ("assume_lkjcholu", "EnzymeReverse"): ENZYME_RVS_ONE_PARAM,
-        ("assume_normal", "EnzymeReverse"): ENZYME_RVS_ONE_PARAM,
-        ("assume_wishart", "EnzymeReverse"): ENZYME_RVS_ONE_PARAM,
         ("assume_mvnormal", "EnzymeForward"): ENZYME_FWD_BLAS,
         ("assume_wishart", "EnzymeForward"): ENZYME_FWD_BLAS,
+        ("multithreaded", "Mooncake"): MOONCAKE_THREADED,
     }
 
 
@@ -198,6 +218,8 @@ def html(_args):
 </ul>
 
 <h2>Results</h2>
+
+<p>(New: You can also hover over the model names to see their definitions.)</p>
 """)
 
         # Table header
@@ -211,7 +233,7 @@ def html(_args):
         for model_name in models:
             ad_results = results[model_name]
             f.write("\n<tr>")
-            f.write(f"<td>{model_name}</td>")
+            f.write(f'<td>{model_name}<div class="model-definition"><pre>{get_model_definition(model_name)}</pre></div></td>')
             for adtype in adtypes:
                 ad_result = ad_results[adtype]
                 try:
@@ -272,6 +294,7 @@ def html(_args):
 td, th {
     border: 1px solid black;
     padding: 0px 10px;
+    white-space: nowrap;
 }
 
 th {
@@ -291,6 +314,7 @@ def html(_args):
     font-family: "Fira Sans", sans-serif;
     font-weight: 700;
     background-color: #ececec;
+    position: relative;
 }
 
 tr > th:first-child {
@@ -321,6 +345,29 @@ def html(_args):
 a.issue:visited {
     color: #880000;
 }
+
+div.model-definition {
+    background-color: #f6f6f6;
+    border: 1px solid black;
+    border-radius: 5px;
+    padding: 0 10px;
+    z-index: 5;
+    font-size: 0.9em;
+    text-align: left;
+    font-weight: normal;
+    position: absolute;
+    left: 100%;
+    top: 0;
+    display: none;
+}
+
+td:hover {
+    background-color: #f6f6f6;
+}
+
+td:hover > div.model-definition {
+    display: block;
+}
 """)
 
 

main.jl

@@ -47,6 +47,7 @@ elseif length(ARGS) == 3 && ARGS[1] == "--run"
         # If reached here - nothing went wrong
         @printf("%.3f", result.time_vs_primal)
     catch e
+        @show e
         if e isa ADIncorrectException
             # First check for completely incorrect ones
             for (a, b) in zip(e.grad_expected, e.grad_actual)

models.jl

@@ -2,7 +2,7 @@ module Models
 
 using DynamicPPL
 using Distributions
-using LinearAlgebra: I
+using LinearAlgebra:I
 
 export MODELS
 
@@ -94,8 +94,6 @@ end
 end
 add_model!(MODELS, observe_submodel())
 
-# This one fails with Enzyme ...
-
 @model function dot_assume_observe_index(x=[1.5, 2.0, 2.5], ::Type{TV}=Vector{Float64}) where {TV}
     a = TV(undef, length(x))
     a .~ Normal()
@@ -105,46 +103,67 @@ add_model!(MODELS, observe_submodel())
 end
 add_model!(MODELS, dot_assume_observe_index())
 
-# Add models with different distributions
-
-DISTRIBUTIONS = Dict(
-    # Univariate
-    :assume_normal => Normal(),
-    :assume_beta => Beta(2, 2),
-    # Multivariate
-    :assume_mvnormal => MvNormal([0.0, 0.0], [1.0 0.5; 0.5 1.0]),
-    :assume_dirichlet => Dirichlet([1.0, 5.0]),
-    # Matrixvariate
-    :assume_wishart => Wishart(7, [1.0 0.5; 0.5 1.0]),
-    :assume_lkjcholu => LKJCholesky(5, 1.0, 'U'),
-)
-
-for (name, distribution) in DISTRIBUTIONS
-    @eval begin
-        @model function $name()
-            a ~ $distribution
-        end
-        add_model!(MODELS, $name())
+@model function assume_normal()
+    a ~ Normal()
+end
+add_model!(MODELS, assume_normal())
+
+@model function assume_beta()
+    a ~ Beta(2, 2)
+end
+add_model!(MODELS, assume_beta())
+
+@model function assume_mvnormal()
+    a ~ MvNormal([0.0, 0.0], [1.0 0.5; 0.5 1.0])
+end
+add_model!(MODELS, assume_mvnormal())
+
+@model function assume_dirichlet()
+    a ~ Dirichlet([1.0, 5.0])
+end
+add_model!(MODELS, assume_dirichlet())
+
+@model function assume_wishart()
+    a ~ Wishart(7, [1.0 0.5; 0.5 1.0])
+end
+add_model!(MODELS, assume_wishart())
+
+@model function assume_lkjcholu()
+    a ~ LKJCholesky(5, 1.0, 'U')
+end
+add_model!(MODELS, assume_lkjcholu())
+
+@model function n010(::Type{TV}=Vector{Float64}) where {TV}
+    a = TV(undef, 10)
+    for i in eachindex(a)
+        a[i] ~ Normal()
     end
 end
+add_model!(MODELS, n010())
 
-# Add models with different sizes
+@model function n050(::Type{TV}=Vector{Float64}) where {TV}
+    a = TV(undef, 50)
+    for i in eachindex(a)
+        a[i] ~ Normal()
+    end
+end
+add_model!(MODELS, n050())
 
-NS = [10, 50, 100, 500]
+@model function n100(::Type{TV}=Vector{Float64}) where {TV}
+    a = TV(undef, 100)
+    for i in eachindex(a)
+        a[i] ~ Normal()
+    end
+end
+add_model!(MODELS, n100())
 
-for n in NS
-    # pad to make sure they sort correctly alphabetically
-    name = Symbol("n$(lpad(n, 3, "0"))")
-    @eval begin
-        @model function $name(::Type{TV}=Vector{Float64}) where {TV}
-            a = TV(undef, $n)
-            for i in eachindex(a)
-                a[i] ~ Normal()
-            end
-        end
-        add_model!(MODELS, $name())
+@model function n500(::Type{TV}=Vector{Float64}) where {TV}
+    a = TV(undef, 500)
+    for i in eachindex(a)
+        a[i] ~ Normal()
     end
 end
+add_model!(MODELS, n500())
 
 @model function multithreaded(x)
     a ~ Normal()