ProbProg: Sample & Generate op #1236

    
           create_operation( 
        
               "enzyme.init", location; 
        
               operands, owned_regions, successors, attributes,

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Line 169 in b3ba477

result_inference=false

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Line 173 in b3ba477

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Line 175 in b3ba477

op_ty_results = IR.Type[output, ]

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Lines 180 to 183 in b3ba477

    
           create_operation( 
        
               "enzyme.placeholder", location; 
        
               operands, owned_regions, successors, attributes,

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Line 185 in b3ba477

result_inference=false

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Line 189 in b3ba477

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Lines 191 to 192 in b3ba477

    
           op_ty_results = IR.Type[output, ] 
        
           operands = Value[cache, ]

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Lines 196 to 199 in b3ba477

    
           create_operation( 
        
               "enzyme.pop", location; 
        
               operands, owned_regions, successors, attributes,

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Line 201 in b3ba477

result_inference=false

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Line 205 in b3ba477

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Line 208 in b3ba477

operands = Value[cache, value, ]

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Lines 212 to 215 in b3ba477

    
           create_operation( 
        
               "enzyme.push", location; 
        
               operands, owned_regions, successors, attributes,

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Line 217 in b3ba477

result_inference=false

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Lines 221 to 224 in b3ba477

    
           function sample(inputs::Vector{Value}; outputs::Vector{IR.Type}, fn, name=nothing, location=Location()) 
        
               op_ty_results = IR.Type[outputs..., ] 
        
               operands = Value[inputs..., ]

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Line 227 in b3ba477

attributes = NamedAttribute[namedattribute("fn", fn), ]

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Lines 229 to 232 in b3ba477

    
           create_operation( 
        
               "enzyme.sample", location; 
        
               operands, owned_regions, successors, attributes,

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Line 234 in b3ba477

result_inference=false

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Line 238 in b3ba477

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Line 241 in b3ba477

operands = Value[gradient, value, ]

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Lines 245 to 248 in b3ba477

    
           create_operation( 
        
               "enzyme.set", location; 
        
               operands, owned_regions, successors, attributes,

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Line 250 in b3ba477

result_inference=false

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Lines 261 to 263 in b3ba477

    
           function simulate(inputs::Vector{Value}; trace::IR.Type, fn, name=nothing, location=Location()) 
        
               op_ty_results = IR.Type[trace, ] 
        
               operands = Value[inputs..., ]

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Line 266 in b3ba477

attributes = NamedAttribute[namedattribute("fn", fn), ]

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Lines 268 to 271 in b3ba477

    
           create_operation( 
        
               "enzyme.simulate", location; 
        
               operands, owned_regions, successors, attributes,

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Line 273 in b3ba477

result_inference=false

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Lines 284 to 286 in b3ba477

    
           function trace(inputs::Vector{Value}, oldTrace=nothing::Union{Nothing, Value}; constraints=nothing::Union{Nothing, Value}, newTrace::IR.Type, weights::Vector{IR.Type}, fn, name=nothing, location=Location()) 
        
               op_ty_results = IR.Type[newTrace, weights..., ] 
        
               operands = Value[inputs..., ]

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Line 289 in b3ba477

attributes = NamedAttribute[namedattribute("fn", fn), ]

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Line 292 in b3ba477

    
           push!(attributes, operandsegmentsizes([length(inputs), (oldTrace==nothing) ? 0 : 1(constraints==nothing) ? 0 : 1]))

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Lines 294 to 297 in b3ba477

    
           create_operation( 
        
               "enzyme.trace", location; 
        
               operands, owned_regions, successors, attributes,

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Line 299 in b3ba477

result_inference=false

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Line 310 in b3ba477

operands = Value[trace, sample, ]

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Lines 315 to 318 in b3ba477

    
           create_operation( 
        
               "enzyme.addSampleToTrace", location; 
        
               operands, owned_regions, successors, attributes,

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Line 320 in b3ba477

result_inference=false

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Lines 329 to 331 in b3ba477

    
           function insertChoiceToMap(choiceMap::Value, choice::Value; newChoiceMap::IR.Type, name=nothing, location=Location()) 
        
               op_ty_results = IR.Type[newChoiceMap, ] 
        
               operands = Value[choiceMap, choice, ]

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Lines 336 to 339 in b3ba477

    
           create_operation( 
        
               "enzyme.insertChoiceToMap", location; 
        
               operands, owned_regions, successors, attributes,

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Line 341 in b3ba477

result_inference=false

github-actions

Remaining comments which cannot be posted as a review comment to avoid GitHub Rate Limit

JuliaFormatter

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Lines 284 to 286 in 47e9fe3

    
           function trace(inputs::Vector{Value}, oldTrace=nothing::Union{Nothing, Value}; constraints=nothing::Union{Nothing, Value}, newTrace::IR.Type, weights::Vector{IR.Type}, fn, name=nothing, location=Location()) 
        
               op_ty_results = IR.Type[newTrace, weights..., ] 
        
               operands = Value[inputs..., ]

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Line 289 in 47e9fe3

attributes = NamedAttribute[namedattribute("fn", fn), ]

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Line 292 in 47e9fe3

    
           push!(attributes, operandsegmentsizes([length(inputs), (oldTrace==nothing) ? 0 : 1(constraints==nothing) ? 0 : 1]))

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Lines 294 to 297 in 47e9fe3

    
           create_operation( 
        
               "enzyme.trace", location; 
        
               operands, owned_regions, successors, attributes,

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Line 299 in 47e9fe3

result_inference=false

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Line 310 in 47e9fe3

operands = Value[trace, sample, ]

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Lines 315 to 318 in 47e9fe3

    
           create_operation( 
        
               "enzyme.addSampleToTrace", location; 
        
               operands, owned_regions, successors, attributes,

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Line 320 in 47e9fe3

result_inference=false

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Lines 329 to 331 in 47e9fe3

    
           function insertChoiceToMap(choiceMap::Value, choice::Value; outputs::IR.Type, name=nothing, location=Location()) 
        
               op_ty_results = IR.Type[outputs, ] 
        
               operands = Value[choiceMap, choice, ]

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Lines 336 to 339 in 47e9fe3

    
           create_operation( 
        
               "enzyme.insertChoiceToMap", location; 
        
               operands, owned_regions, successors, attributes,

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/mlir/Dialects/Enzyme.jl

Line 341 in 47e9fe3

result_inference=false

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/test/probprog/generate.jl

Lines 30 to 32 in 47e9fe3

    
           @test Array(model_compiled(seed1, μ1, σ1, shape)) ≈ Array(model_compiled(seed1, μ1, σ1, shape)) 
        
           @test mean(Array(model_compiled(seed1, μ1, σ1, shape))) ≈ 0.0 atol = 0.05 rtol = 0.05 
        
           @test mean(Array(model_compiled(seed2, μ2, σ2, shape))) ≈ 1000.0 atol = 0.05 rtol = 0.05

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/test/probprog/generate.jl

Line 34 in 47e9fe3

    
           Array(model_compiled(seed1, μ1, σ1, shape)) .≈ Array(model_compiled(seed2, μ2, σ2, shape))

github-actions · 2025-05-15T23:11:11Z

src/ProbProg.jl

+    return result
+end
+
+end


[JuliaFormatter] _{reported by reviewdog 🐶}

Suggested change

end

end

…robprog

…ed yet

…robprog

generate

902ced9

sbrantq requested review from wsmoses and abajpai15 May 2, 2025 03:33

sbrantq changed the title ~~Front end for EnzymeMLIR ProbProg pass~~ WIP: Front end for EnzymeMLIR ProbProg pass May 2, 2025

sbrantq added 10 commits May 2, 2025 15:40

refactor

e2c77e4

Merge branch 'main' of https://github.com/EnzymeAD/Reactant.jl into p…

e204d13

…robprog

Merge branch 'main' of https://github.com/EnzymeAD/Reactant.jl into p…

327b10a

…robprog

add probprog pass to :all

d611ae4

improve test

3672d83

only probprog opt mode

b70843e

fix up test

597fa89

move

e6c2c0a

simplify

9b9395e

fix up

b3ba477

github-actions bot reviewed May 14, 2025

View reviewed changes

saving changes

47e9fe3

github-actions bot reviewed May 15, 2025

View reviewed changes

sbrantq added 10 commits May 16, 2025 17:35

Merge branch 'main' of https://github.com/EnzymeAD/Reactant.jl into p…

982b2bf

…robprog

Merge branch 'main' of https://github.com/EnzymeAD/Reactant.jl into p…

06b7464

…robprog

Merge branch 'main' of https://github.com/EnzymeAD/Reactant.jl into p…

bd73c62

…robprog

fix sample op

a6fcca3

save tests

e51e04b

temporarily removing probprog pass from :all as MLIR pass is not merg…

ce68f6a

…ed yet

Merge branch 'main' of https://github.com/EnzymeAD/Reactant.jl into p…

94bbe62

…robprog

undo enzyme binding change

d31bba6

format

573fa02

format

0264a3d

sbrantq changed the title ~~WIP: Front end for EnzymeMLIR ProbProg pass~~ ProbProg: Sample & Generate op May 20, 2025

sbrantq added 2 commits May 20, 2025 18:50

improve

2e18bdf

improve

1f19979

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

ProbProg: Sample & Generate op #1236

ProbProg: Sample & Generate op #1236

sbrantq commented May 2, 2025 •

edited

Loading

github-actions bot left a comment

github-actions bot left a comment

github-actions bot May 15, 2025


	create_operation(
	"enzyme.init", location;
	operands, owned_regions, successors, attributes,


	create_operation(
	"enzyme.placeholder", location;
	operands, owned_regions, successors, attributes,


	create_operation(
	"enzyme.pop", location;
	operands, owned_regions, successors, attributes,


	create_operation(
	"enzyme.push", location;
	operands, owned_regions, successors, attributes,


	function sample(inputs::Vector{Value}; outputs::Vector{IR.Type}, fn, name=nothing, location=Location())
	op_ty_results = IR.Type[outputs..., ]
	operands = Value[inputs..., ]


	create_operation(
	"enzyme.sample", location;
	operands, owned_regions, successors, attributes,


	create_operation(
	"enzyme.set", location;
	operands, owned_regions, successors, attributes,

	function simulate(inputs::Vector{Value}; trace::IR.Type, fn, name=nothing, location=Location())
	op_ty_results = IR.Type[trace, ]
	operands = Value[inputs..., ]


	create_operation(
	"enzyme.simulate", location;
	operands, owned_regions, successors, attributes,

	function trace(inputs::Vector{Value}, oldTrace=nothing::Union{Nothing, Value}; constraints=nothing::Union{Nothing, Value}, newTrace::IR.Type, weights::Vector{IR.Type}, fn, name=nothing, location=Location())
	op_ty_results = IR.Type[newTrace, weights..., ]
	operands = Value[inputs..., ]


	create_operation(
	"enzyme.trace", location;
	operands, owned_regions, successors, attributes,


	create_operation(
	"enzyme.addSampleToTrace", location;
	operands, owned_regions, successors, attributes,

	function insertChoiceToMap(choiceMap::Value, choice::Value; newChoiceMap::IR.Type, name=nothing, location=Location())
	op_ty_results = IR.Type[newChoiceMap, ]
	operands = Value[choiceMap, choice, ]


	create_operation(
	"enzyme.insertChoiceToMap", location;
	operands, owned_regions, successors, attributes,

	function insertChoiceToMap(choiceMap::Value, choice::Value; outputs::IR.Type, name=nothing, location=Location())
	op_ty_results = IR.Type[outputs, ]
	operands = Value[choiceMap, choice, ]

	@test Array(model_compiled(seed1, μ1, σ1, shape)) ≈ Array(model_compiled(seed1, μ1, σ1, shape))
	@test mean(Array(model_compiled(seed1, μ1, σ1, shape))) ≈ 0.0 atol = 0.05 rtol = 0.05
	@test mean(Array(model_compiled(seed2, μ2, σ2, shape))) ≈ 1000.0 atol = 0.05 rtol = 0.05

ProbProg: Sample & Generate op #1236

Are you sure you want to change the base?

ProbProg: Sample & Generate op #1236

Conversation

sbrantq commented May 2, 2025 • edited Loading

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot May 15, 2025

Choose a reason for hiding this comment

sbrantq commented May 2, 2025 •

edited

Loading