Add aimx trace distribution with tensor (histogram) output modes#15
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for Aim distribution sequences, allowing users to fetch and render histogram weights via the aimx trace distribution command. Key additions include new data models for distributions, collection logic from the Aim repository, and rendering capabilities for terminal tables, CSV, and JSON outputs. The review feedback highlights opportunities to optimize step range filtering into a single pass and suggests refactoring the command execution logic to reduce duplication between metric and distribution processing paths.
| if start is not None: | ||
| points = [point for point in points if point.step >= start] | ||
| if end is not None: | ||
| points = [point for point in points if point.step <= end] |
There was a problem hiding this comment.
The step range filtering can be optimized to a single pass. This avoids creating an intermediate list when both start and end are provided, which improves performance for series with many points.
if start is not None or end is not None:
points = [
point for point in points
if (start is None or point.step >= start) and (end is None or point.step <= end)
]| if invocation.target == "distribution": | ||
| from aimx.aim_bridge.metric_stats import ( | ||
| collect_distribution_series, | ||
| filter_distribution_by_step_range, | ||
| subsample_distribution, | ||
| ) | ||
|
|
||
| if not series_list: | ||
| return QueryCommandResult(exit_status=0, output="No matching metrics found.") | ||
| series_list = collect_distribution_series(invocation.expression, normalized_repo_path) | ||
| if not series_list: | ||
| return QueryCommandResult(exit_status=0, output="No matching distributions found.") | ||
|
|
||
| # Step range filter is a hard constraint applied before density subsampling | ||
| if invocation.step_slice is not None: | ||
| step_start, step_end = parse_step_slice(invocation.step_slice) | ||
| series_list = [filter_by_step_range(s, step_start, step_end) for s in series_list] | ||
| # Drop empty series so they don't clutter plots | ||
| series_list = [s for s in series_list if s.count > 0] | ||
| if invocation.step_slice is not None: | ||
| step_start, step_end = parse_step_slice(invocation.step_slice) | ||
| series_list = [ | ||
| filter_distribution_by_step_range(s, step_start, step_end) for s in series_list | ||
| ] | ||
| series_list = [s for s in series_list if s.count > 0] | ||
|
|
||
| if not series_list: | ||
| return QueryCommandResult(exit_status=0, output="No data in the requested step range.") | ||
| if not series_list: | ||
| return QueryCommandResult(exit_status=0, output="No data in the requested step range.") | ||
|
|
||
| # Density subsampling for visualisation | ||
| needs_sample = any( | ||
| x is not None for x in (invocation.head, invocation.tail, invocation.every) | ||
| ) | ||
| if needs_sample: | ||
| series_list = [ | ||
| subsample(s, head=invocation.head, tail=invocation.tail, every=invocation.every) | ||
| for s in series_list | ||
| ] | ||
|
|
||
| if invocation.mode == "json": | ||
| output = render_trace_json(series_list) | ||
| elif invocation.mode == "csv": | ||
| output = render_csv(series_list) | ||
| elif invocation.mode == "table": | ||
| output = render_trace_table(series_list, no_color=effective_no_color) | ||
| needs_sample = any( | ||
| x is not None for x in (invocation.head, invocation.tail, invocation.every) | ||
| ) | ||
| if needs_sample: | ||
| series_list = [ | ||
| subsample_distribution( | ||
| s, head=invocation.head, tail=invocation.tail, every=invocation.every | ||
| ) | ||
| for s in series_list | ||
| ] | ||
|
|
||
| if invocation.mode == "json": | ||
| output = render_distribution_json(series_list) | ||
| elif invocation.mode == "csv": | ||
| output = render_distribution_csv(series_list) | ||
| else: | ||
| output = render_distribution_table(series_list, no_color=effective_no_color) | ||
| else: | ||
| output = render_plot( | ||
| series_list, | ||
| width=invocation.width, | ||
| height=invocation.height, | ||
| from aimx.aim_bridge.metric_stats import ( | ||
| collect_metric_series, | ||
| filter_by_step_range, | ||
| subsample, | ||
| ) | ||
|
|
||
| series_list = collect_metric_series(invocation.expression, normalized_repo_path) | ||
|
|
||
| if not series_list: | ||
| return QueryCommandResult(exit_status=0, output="No matching metrics found.") | ||
|
|
||
| # Step range filter is a hard constraint applied before density subsampling | ||
| if invocation.step_slice is not None: | ||
| step_start, step_end = parse_step_slice(invocation.step_slice) | ||
| series_list = [filter_by_step_range(s, step_start, step_end) for s in series_list] | ||
| # Drop empty series so they don't clutter plots | ||
| series_list = [s for s in series_list if s.count > 0] | ||
|
|
||
| if not series_list: | ||
| return QueryCommandResult(exit_status=0, output="No data in the requested step range.") | ||
|
|
||
| # Density subsampling for visualisation | ||
| needs_sample = any( | ||
| x is not None for x in (invocation.head, invocation.tail, invocation.every) | ||
| ) | ||
| if needs_sample: | ||
| series_list = [ | ||
| subsample(s, head=invocation.head, tail=invocation.tail, every=invocation.every) | ||
| for s in series_list | ||
| ] | ||
|
|
||
| if invocation.mode == "json": | ||
| output = render_trace_json(series_list) | ||
| elif invocation.mode == "csv": | ||
| output = render_csv(series_list) | ||
| elif invocation.mode == "table": | ||
| output = render_trace_table(series_list, no_color=effective_no_color) | ||
| else: | ||
| output = render_plot( | ||
| series_list, | ||
| width=invocation.width, | ||
| height=invocation.height, | ||
| ) |
There was a problem hiding this comment.
There is significant code duplication between the distribution and metrics execution paths in run_trace_command. Both paths follow the same sequence: collect series, filter by step range, subsample, and render based on the mode.
Consider refactoring this into a unified pipeline. You could define a set of functions (collect, filter, subsample, render) based on the invocation.target and then execute the pipeline once. This would improve maintainability and ensure that improvements to the tracing logic (like new filters or sampling methods) are automatically applied to both metrics and distributions.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 589f405094
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| " Repo defaults to the current directory.", | ||
| " Short run hashes in the expression are transparently expanded.", | ||
| " Example: aimx trace \"metric.name=='loss'\" --repo data --steps 100:500", | ||
| " Example: aimx trace distribution \"distribution.name=='weights'\" --repo data --json", |
There was a problem hiding this comment.
distributions in trace-distribution expressions
This example advertises distribution.name, but Aim’s Repo.query_distributions(...) binds the sequence variable as distributions (plural). When users copy this command, the query variable is undefined and the new trace path reports no matches, making the feature look broken even when data exists. Update the expression to distributions.name == 'weights' (and align related docs/tests) or add alias rewriting before executing the query.
Useful? React with 👍 / 👎.
|
@codex 审核代码 |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 589f405094
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| query_result = repo.query_distributions( | ||
| expression, report_mode=QueryReportMode.DISABLED |
There was a problem hiding this comment.
collect_distribution_series forwards the user expression directly to Repo.query_distributions, but Aim’s distribution sequence variable is distributions (see aim.sdk.sequences.distribution_sequence.Distributions.sequence_name()), while this change’s own examples/tests use distribution.name == .... In that common case, Aim treats distribution as undefined and the query silently yields no matches, so aimx trace distribution reports “No matching distributions found” even when data exists. Please rewrite/alias the singular form before calling query_distributions (or reject it with a clear error) so documented expressions actually return results.
Useful? React with 👍 / 👎.
1 participant
Motivation
trace distributionpath so users can inspect Aim-tracked distribution sequences (histogram/tensor payloads) from the CLI.Description
DistributionPointandDistributionSeriestypes and implementcollect_distribution_series(...)inaimx.aim_bridge.metric_statsthat usesRepo.query_distributions(...)to collect per-stepstep,epoch,weights, andbin_edgesdata.filter_distribution_by_step_rangeandsubsample_distribution, and keep metric helpers unchanged.distributiontarget and route to a distribution-specific execution path inrun_trace_command, preserving existing metric trace behavior.src/aimx/rendering/trace_views.py:render_distribution_table,render_distribution_csv, andrender_distribution_json, plus a small_format_tensorhelper and labeling helper; update metric render helpers to coexist.README.mdwithtrace distributionusage examples and add unit tests for parsing and renderers (tests/unit/test_trace_distribution_views.py) and adjust existing trace parsing tests.Testing
uv run pytest tests/unit/test_trace_helpers.py tests/unit/test_trace_distribution_views.py tests/unit/test_router.py tests/unit/test_owned_commands.py, which completed successfully with32 passed.uv run pytest tests/integration/test_trace_command.py tests/contract/test_trace_contract.py, which completed with1 passed, 12 skippeddue to the sample-repo fixture behavior.Codex Task