Skip to content

Export Uzu traces as safetensors#450

Draft
ry2009 wants to merge 1 commit into
mainfrom
ryan/trace-safetensors-export
Draft

Export Uzu traces as safetensors#450
ry2009 wants to merge 1 commit into
mainfrom
ryan/trace-safetensors-export

Conversation

@ry2009
Copy link
Copy Markdown

@ry2009 ry2009 commented May 26, 2026

Summary

Reworks the Uzu tracing integration around activation/tokenization export instead of local comparison:

  • removes the old trace comparator/metrics implementation from the Uzu integration tracer
  • exports Uzu activation traces to traces.safetensors-compatible safetensors output
  • exports Uzu tokenization traces with activation_trace.token_ids / activation_trace.token_positions plus rendered request metadata
  • keeps trace input loading strict for token id/position shape
  • adds a minimal safetensors writer for exported trace tensors and metadata
  • keeps the reusable writer byte/schema oriented instead of tying it to backend arrays

This pairs with the Lalamo comparator PR: trymirai/lalamo#285

Validation

  • cargo test -p backend-uzu test_safetensors_metadata_writer --test integration
  • cargo check -p backend-uzu --tests
  • cargo fmt --check -p backend-uzu
  • git diff --check

cargo fmt --check passes but prints existing stable-rustfmt warnings about nightly-only config keys.

@ry2009 ry2009 force-pushed the ryan/trace-safetensors-export branch from c2ef717 to 296b06c Compare May 26, 2026 23:29
@ry2009 ry2009 force-pushed the ryan/trace-safetensors-export branch from 296b06c to 0aa5d08 Compare May 28, 2026 06:24
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant