Skip to content

feat(nx): Implement lazy views and symbolic shapes #23

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 2 commits into from
Closed

Conversation

tmattio
Copy link
Contributor

@tmattio tmattio commented Jun 2, 2025

This commit introduces two major architectural improvements to nx:

  1. Lazy View Operations: View operations (reshape, permute, expand, etc.) are now lazy - they update metadata instead of copying data. Views are only materialized when operations need to access the underlying data.

  2. Symbolic Shapes: Replace concrete int arrays with symbolic dimensions that can be bound at runtime. This enables shape-polymorphic kernels and dynamic batching.

Key changes:

  • Add Symbolic_shape module supporting Static/Dynamic dimensions
  • Add Lazy_view module tracking sequences of view transformations
  • Update tensor type to use Lazy_view.t instead of View.t
  • Refactor backend interface to accept symbolic shapes
  • Update Native and Metal backends to handle lazy views
  • Add ensure_materializable pattern for on-demand view materialization

Benefits:

  • Memory efficiency: Avoid unnecessary data copies for view operations
  • Performance: Enable view fusion and better memory access patterns
  • Flexibility: Support dynamic shapes for batching and compilation
  • Foundation: Prepare for advanced JIT optimizations in Rune

The changes maintain backward compatibility and NumPy semantics while bringing nx closer to modern tensor frameworks.

We once again draw significant inspiration from Tinygrad.

tmattio added 2 commits June 3, 2025 14:57
This commit introduces two major architectural improvements to nx:

1. **Lazy View Operations**: View operations (reshape, permute, expand, etc.)
   are now lazy - they update metadata instead of copying data. Views are only
   materialized when operations need to access the underlying data.

2. **Symbolic Shapes**: Replace concrete int arrays with symbolic dimensions
   that can be bound at runtime. This enables shape-polymorphic kernels and
   dynamic batching.

Key changes:
- Add `Symbolic_shape` module supporting Static/Dynamic dimensions
- Add `Lazy_view` module tracking sequences of view transformations
- Update tensor type to use `Lazy_view.t` instead of `View.t`
- Refactor backend interface to accept symbolic shapes
- Update Native and Metal backends to handle lazy views
- Add `ensure_materializable` pattern for on-demand view materialization

Benefits:
- Memory efficiency: Avoid unnecessary data copies for view operations
- Performance: Enable view fusion and better memory access patterns
- Flexibility: Support dynamic shapes for batching and compilation
- Foundation: Prepare for advanced JIT optimizations in Rune

The changes maintain backward compatibility and NumPy semantics while
bringing nx closer to modern tensor frameworks.

We once again draw significant inspiration from Tinygrad.
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