Trace callbacks

RxInfer provides a built-in callback structure called RxInferTraceCallbacks for recording all callback events during the inference procedure. Each event is stored as a TracedEvent containing the event name (as a Symbol) and the event object itself. This is useful for debugging, understanding the inference flow, and inspecting what happens at each stage. For general information about the callbacks system, see Callbacks.

Basic usage

using RxInfer

@model function iid_normal(y)
    μ  ~ Normal(mean = 0.0, variance = 100.0)
    γ  ~ Gamma(shape = 1.0, rate = 1.0)
    y .~ Normal(mean = μ, precision = γ)
end

init = @initialization begin
    q(μ) = vague(NormalMeanVariance)
end

# Create a trace callbacks instance
trace = RxInferTraceCallbacks()

result = infer(
    model = iid_normal(),
    data = (y = randn(10),),
    constraints = MeanField(),
    iterations = 3,
    initialization = init,
    callbacks = trace,
)

events = RxInfer.tracedevents(trace)

println("Recorded $(length(events)) events")
for i in 1:10
    println("  ", events[i])
end
println("...")

Recorded 302 events
  TracedEvent(:before_model_creation)
  TracedEvent(:after_model_creation)
  TracedEvent(:before_inference)
  TracedEvent(:before_iteration)
  TracedEvent(:before_data_update)
  TracedEvent(:before_marginal_computation)
  TracedEvent(:before_product_of_messages)
  TracedEvent(:before_message_rule_call)
  TracedEvent(:after_message_rule_call)
  TracedEvent(:before_message_rule_call)
...

Using trace = true

Instead of creating a RxInferTraceCallbacks instance manually, you can use the trace = true keyword argument in the infer function. This automatically merges a RxInferTraceCallbacks instance with any user-provided callbacks and saves it to the model's metadata:

result = infer(
    model = iid_normal(),
    data = (y = randn(10),),
    constraints = MeanField(),
    iterations = 3,
    initialization = init,
    trace = true,
)


trace = result.model.metadata[:trace]
events = RxInfer.tracedevents(trace)
println("Recorded $(length(events)) events via trace = true")

Recorded 302 events via trace = true

Accessing from model metadata

After model creation, the trace callbacks instance is automatically saved into the model's metadata under the :trace key. This makes it accessible from the inference result without needing to hold onto the callbacks object separately:

result = infer(
    model = iid_normal(),
    data = (y = randn(10),),
    constraints = MeanField(),
    iterations = 3,
    initialization = init,
    callbacks = RxInferTraceCallbacks(),
)


trace = result.model.metadata[:trace]
events = RxInfer.tracedevents(trace)
println("Recorded $(length(events)) events via trace = true")

Recorded 302 events via trace = true

Inspecting traced events

Each TracedEvent has a single field:

• event::ReactiveMP.Event — the original event object that was passed to the callback

You can retrieve the event name via ReactiveMP.event_name(typeof(traced_event.event)) and access event-specific fields directly on traced_event.event.

using RxInfer.ReactiveMP: event_name
events = RxInfer.tracedevents(trace)

# Filter for specific events
iteration_events = filter(e -> event_name(typeof(e.event)) === :before_iteration, events)
println("Number of iterations: ", length(iteration_events))

Number of iterations: 3

Combining with other callbacks

trace = true is compatible with other callbacks, including benchmark = true and custom callbacks:

result = infer(
    model = iid_normal(),
    data = (y = randn(10),),
    constraints = MeanField(),
    iterations = 3,
    initialization = init,
    trace = true,
    benchmark = true,
)


println("Trace included: ", haskey(result.model.metadata, :trace))
println("Benchmark included: ", haskey(result.model.metadata, :benchmark))

Trace included: true
Benchmark included: true

Viewing traces in Perfetto

A recorded trace can be inspected interactively with the Perfetto trace viewer. Use perfetto_view to embed the viewer inside a Pluto, VS Code or Jupyter notebook cell, or perfetto_open to open it in your default browser.

Inside Perfetto, you can navigate (zoom and pan) using the WASD keys. You can select with the mouse, and inspect individual events. Press ? for a quick help menu.

result = infer(model = iid_normal(), data = (y = randn(10),), iterations = 3, trace = true)
traces = RxInfer.tracedevents(result.model.metadata[:trace])

perfetto_view(traces)   # show directly in your IDE (Pluto, VS Code, Jupyter)
perfetto_open(traces)   # open in the browser

In the screenshot above, the first ProductOfMessages event is selected, showing the event details in the bottom panel. Here you see the duration (5ms), and the event arguments, including the result distribution.

If you are interested in debugging the performance of your inference call, take note that runtimes can vary greatly between runs due to Julia features like GC and JIT compilation. Try running your inference multiple times to get a better picture. You can also try to use Julia's built-in profiler.

perfetto_view(traces::Vector{TracedEvent}; name = "$(Time(now())) RxInfer trace")

result = infer(model = my_model(), data = my_data, trace = true)
traces = RxInfer.tracedevents(result.model.metadata[:trace])
perfetto_view(traces)   # display in a notebook cell

perfetto_open(traces::Vector{TracedEvent}; name = "$(Time(now())) RxInfer trace")

result = infer(model = my_model(), data = my_data, trace = true)
traces = RxInfer.tracedevents(result.model.metadata[:trace])
perfetto_open(traces)

PerfettoDisplay

Exporting to TensorBoard

When TensorBoardLogger.jl is loaded, the TensorBoardLoggerExt extension activates and provides RxInfer.convert_to_tensorboard, which converts a recorded trace into TensorFlow event files readable by TensorBoard.

using RxInfer
using TensorBoardLogger  # activates the extension

result = infer(
    model = iid_normal(),
    data = (y = randn(10),),
    constraints = MeanField(),
    iterations = 5,
    initialization = init,
    trace = true,
)

trace = result.model.metadata[:trace]

log_dir = RxInfer.convert_to_tensorboard(trace; log_distributions = true)
# Then run: tensorboard --logdir="<log_dir>"

What gets logged

Posterior scalar tags by family

Each posterior is logged under posteriors/<variable>/<tag> with one step per inference iteration. Specific dispatch is selected by the marginal's distribution type — the most-specific method wins, with the generic mean/var fallback catching anything not listed.

Marginals that are not univariate (e.g. multivariate Normal, matrix-variate posteriors) are silently skipped on the scalar path and produce no posteriors/... scalar tags. The histogram path is also gated on <: UnivariateDistribution, so the same families above are the ones that contribute to the Distributions / Histograms tabs when log_distributions = true.

Run summary

The Summary text tag is a one-shot snapshot of the run, written at step 1 alongside EventCounts. Each row is a key: value line; rows whose underlying measurement is missing (no matching Before*/After* event seen, or no iteration durations recorded) are silently skipped, so partial runs still produce a useful table instead of zero-valued placeholders.

Summary complements rather than replaces existing outputs: per-iteration durations remain in the iteration_time_ms scalar series, and the per-event-type breakdown stays in EventCounts. The Summary tag is always emitted for any trace that contains at least one event — it is not gated by log_text_events, log_posteriors, or log_distributions.

Filtering posteriors

The log_posteriors keyword controls which marginals reach the posteriors/<var>/* tags. It is independent of log_distributions, which controls what is emitted (scalars only vs. scalars + per-iteration histogram). Think of log_posteriors as the row filter and log_distributions as the column filter.

trace = result.model.metadata[:trace]

# Log only μ (the τ posterior is skipped on both scalar and histogram paths).
RxInfer.convert_to_tensorboard(
    trace;
    log_posteriors    = ["μ"],
    log_distributions = true,
)

# Suppress every posterior tag while keeping iteration timing visible.
RxInfer.convert_to_tensorboard(trace; log_posteriors = false)

When log_posteriors is an allow-list, the per-event text breadcrumb under on_marginal_update/<var> (gated by log_text_events) still fires for every variable — so you can keep visibility on which marginals updated without paying the scalar/histogram cost.

convert_to_tensorboard(trace::RxInferTraceCallbacks; output_file::Union{String, Nothing} = nothing,
                       log_distributions::Bool = false, log_text_events::Bool = false,
                       n_samples::Int = 1024, verbose::Bool = true)

results = infer(
    model = my_model(),
    data = my_data,
    trace = true
)

trace = results.model.metadata[:trace]

# Create TensorBoard logs (writes to tensorboard_logs/ in the current directory)
log_dir = convert_to_tensorboard(trace; log_distributions = true)

# Then run: tensorboard --logdir=$log_dir

API Reference

RxInferTraceCallbacks()

# Create a trace callbacks instance that records all events
trace = RxInferTraceCallbacks()

# Or record only specific events
trace = RxInferTraceCallbacks((:before_iteration, :after_iteration))

result = infer(
    model = my_model(),
    data = my_data,
    callbacks = trace,
)

# Access the traced events
events = RxInfer.tracedevents(trace)
for event in events
    println(event_name(event.event))
end

# Or access via model metadata
result.model.metadata[:trace] === trace # true

TracedEvent

tracedevents(callbacks::RxInferTraceCallbacks)

tracedevents(event::Symbol, callbacks::RxInferTraceCallbacks)

is_trace_event_included(callbacks::RxInferTraceCallbacks, event_name::Symbol)

julia> callbacks = RxInfer.RxInferTraceCallbacks((:event1, :event2));

julia> RxInfer.is_trace_event_included(callbacks, :event1)
true

julia> RxInfer.is_trace_event_included(callbacks, :event2)
true

julia> RxInfer.is_trace_event_included(callbacks, :event3)
false