Outsource tensor functionality into mixins - #27
Merged
Conversation
added 10 commits
September 18, 2025 09:31
Move tensor shape, math, type, and device operations from `TensorContainer` into dedicated mixin classes. This refactoring reduces code duplication across `TensorContainer`, `TensorDataClass`, `TensorDict`, and `TensorDistribution` by centralizing common functionalities. It improves modularity and maintainability, making the codebase easier to extend and reason about. `TensorContainer` now explicitly implements `TensorContainerProtocol` and its subclasses inherit these operations via the new mixins.
Replaced implicit pytree operations triggered by .copy() with an explicit pytree.tree_flatten and pytree.tree_unflatten roundtrip.
Refactor and expand the test suite for `TensorAnnotated` subclassing. This significantly improves coverage and robustness by: - Introducing dedicated test fixtures and helper assertion functions. - Restructuring positive and negative test cases for better clarity. - Adding new tests for cross-device handling, method inheritance, and shape consistency. - Covering complex edge cases like deep inheritance chains, diamond patterns, and empty base classes. - Ensuring proper preservation of tensor identity and properties through PyTree operations. - Validating correct exclusion of non-annotated attributes in mixed inheritance.
Removed outdated and redundant sections from the TensorContainer class docstring, such as "Usage Patterns" and "Limitations". Consolidated and rephrased core concepts like Shape Management, Device Management, PyTree Integration, and Torch Function Override for clarity. Added a new section explaining the concept of "Metadata" within the container. Significantly improved the docstring for `unsafe_construction` to provide a detailed explanation of its purpose and usage, including a concrete example. Simplified the "Subclassing Guide" for better readability.
Adds detailed Google-style docstrings to the internal PyTree integration methods (`_pytree_flatten`, `_pytree_flatten_with_keys_fn`, `_pytree_unflatten`, `_tree_map`, `tree_map_with_path`) and validation methods (`_is_shape_compatible`, `_is_device_compatible`, `_validate_shape`, `_validate_device`, `_validate`).
Add comprehensive Google-style docstrings to the module, classes, and methods within `TensorAnnotated` to enhance readability and maintainability.
…efault slots=True The extensive docstrings for `TensorDataClass` and its methods have been streamlined for improved readability and conciseness.
…prove readability
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Refactor: Extract Common Tensor Operations to Mixins
Summary
• Extract tensor operations into modular mixins - Decomposed TensorContainer's monolithic operation set into focused, reusable mixins for shape, math, type, and device
operations
• Establish protocol-based mixin constraints - Introduced TensorContainerProtocol to ensure type safety and proper interface contracts for mixin usage
• Improve code organization and maintainability - Operations are now grouped by functionality, making the codebase easier to understand and extend
Changes
Core Architecture Refactoring
New mixin system (
src/tensorcontainer/mixins/) with four focused operation mixins:TensorShapeOperationsMixin- Shape transformations (view, reshape, permute, etc.)TensorMathOperationsMixin- Mathematical operations (add, sub, mul, sqrt, etc.)TensorTypeOperationsMixin- Type conversions (float, int, double, etc.)TensorDeviceOperationsMixin- Device/memory operations (to, cpu, cuda, clone, etc.)Protocol-based type safety (
src/tensorcontainer/protocols.py):TensorContainerProtocoldefines required interface for mixin compatibilityUpdated Container Implementations
API Compatibility