Merge pull request #1 from mctigger/CI-init

Create python-package.yml, adds Python 3.9 compatibility and no-cuda pytest compatibility

Author: mctigger

.github/workflows/python-package.yml

@@ -0,0 +1,36 @@
+name: Python package
+
+on:
+  push:
+    branches: [ "main" ]
+  pull_request:
+    branches: [ "main" ]
+
+jobs:
+  build:
+
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.9", "3.10", "3.11"]
+
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v5
+      with:
+        python-version: ${{ matrix.python-version }}
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        python -m pip install .[dev]
+    - name: Lint and Format with Ruff
+      run: |
+        # Run ruff to check for linting errors
+        ruff check .
+        # Run ruff to check for formatting issues (like black --check)
+        ruff format --check .
+    - name: Test with pytest
+      run: |
+        pytest

README.md

@@ -1,14 +1,16 @@
 # Tensor Container
 
-*Modern tensor containers for PyTorch with PyTree compatibility and torch.compile optimization*
+*Tensor containers for PyTorch with PyTree compatibility and torch.compile optimization*
 
-[![Python 3.8+](https://img.shields.io/badge/python-3.8+-blue.svg)](https://www.python.org/downloads/)
+[![Python 3.9+](https://img.shields.io/badge/python-3.9+-blue.svg)](https://www.python.org/downloads/)
 [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
 [![PyTorch](https://img.shields.io/badge/PyTorch-2.0+-red.svg)](https://pytorch.org/)
 
 > **⚠️ Academic Research Project**: This project exists solely for academic purposes to explore and learn PyTorch internals. For production use, please use the official, well-maintained [**torch/tensordict**](https://github.com/pytorch/tensordict) library.
 
-Tensor Container provides efficient, type-safe tensor container implementations designed for modern PyTorch workflows. Built from the ground up with PyTree integration and torch.compile optimization, it enables seamless batched tensor operations with minimal overhead and maximum performance.
+Tensor Container provides efficient, type-safe tensor container implementations for PyTorch workflows. It includes PyTree integration and torch.compile optimization for batched tensor operations.
+
+The library includes tensor containers, probabilistic distributions, and batch/event dimension semantics for machine learning workflows.
 
 ## Table of Contents
 
@@ -18,6 +20,7 @@ Tensor Container provides efficient, type-safe tensor container implementations
 - [API Overview](#api-overview)
 - [torch.compile Compatibility](#torchcompile-compatibility)
 - [Contributing](#contributing)
+- [Documentation](#documentation)
 - [License](#license)
 - [Authors](#authors)
 - [Contact and Support](#contact-and-support)
@@ -40,7 +43,7 @@ pip install -e .[dev]
 
 ### Requirements
 
-- Python 3.8+
+- Python 3.9+
 - PyTorch 2.0+
 
 ## Quick Start
@@ -91,6 +94,39 @@ obs = data.observations
 data.actions = torch.randn(32, 8)  # Type-checked assignment
 ```
 
+### TensorDistribution: Probabilistic Containers
+
+```python
+import torch
+from tensorcontainer import TensorDistribution
+
+# Built-in distribution types
+from tensorcontainer.tensor_distribution import (
+    TensorNormal, TensorBernoulli, TensorCategorical,
+    TensorTruncatedNormal, TensorTanhNormal
+)
+
+# Create probabilistic tensor containers
+normal_dist = TensorNormal(
+    loc=torch.zeros(32, 4),
+    scale=torch.ones(32, 4),
+    shape=(32,),
+    device='cpu'
+)
+
+# Sample and compute probabilities
+samples = normal_dist.sample()  # Shape: (32, 4)
+log_probs = normal_dist.log_prob(samples)
+entropy = normal_dist.entropy()
+
+# Categorical distributions for discrete actions
+categorical = TensorCategorical(
+    logits=torch.randn(32, 6),  # 6 possible actions
+    shape=(32,),
+    device='cpu'
+)
+```
+
 ### PyTree Operations
 
 ```python
@@ -106,21 +142,23 @@ combined = pytree.tree_map(lambda x, y: x + y, data1, data2)
 
 ## Features
 
-- **🔥 torch.compile Optimized**: Built for maximum performance with PyTorch's JIT compiler
-- **🌳 Native PyTree Support**: Seamless integration with `torch.utils._pytree` for tree operations
-- **⚡ Zero-Copy Operations**: Efficient tensor sharing and manipulation without unnecessary copies
-- **🎯 Type Safety**: Full static typing support with IDE autocomplete and type checking
-- **📊 Batch Semantics**: Consistent batch/event dimension handling across all operations
-- **🔍 Shape Validation**: Automatic validation of tensor shapes and device consistency
-- **🏗️ Flexible Architecture**: Multiple container types for different use cases
-- **🧪 Comprehensive Testing**: Extensive test suite with compile compatibility verification
+- **torch.compile Optimized**: Compatible with PyTorch's JIT compiler
+- **PyTree Support**: Integration with `torch.utils._pytree` for tree operations
+- **Zero-Copy Operations**: Efficient tensor sharing and manipulation
+- **Type Safety**: Static typing support with IDE autocomplete and type checking
+- **Batch Semantics**: Consistent batch/event dimension handling
+- **Shape Validation**: Automatic validation of tensor shapes and device consistency
+- **Multiple Container Types**: Different container types for different use cases
+- **Probabilistic Support**: Distribution containers for probabilistic modeling
+- **Comprehensive Testing**: Extensive test suite with compile compatibility verification
+- **Memory Efficient**: Optimized memory usage with slots-based dataclasses
 
 ## API Overview
 
 ### Core Components
 
-- **`TensorContainer`**: Base class providing core tensor manipulation operations
-- **`TensorDict`**: Dictionary-like container for dynamic tensor collections
+- **`TensorContainer`**: Base class providing core tensor manipulation operations with batch/event dimension semantics
+- **`TensorDict`**: Dictionary-like container for dynamic tensor collections with nested structure support
 - **`TensorDataClass`**: DataClass-based container for static, typed tensor structures
 - **`TensorDistribution`**: Distribution wrapper for probabilistic tensor operations
 
@@ -130,25 +168,37 @@ combined = pytree.tree_map(lambda x, y: x + y, data1, data2)
 - **Event Dimensions**: Trailing dimensions beyond batch shape, can vary per tensor
 - **PyTree Integration**: All containers are registered PyTree nodes for seamless tree operations
 - **Device Consistency**: Automatic validation ensures all tensors reside on compatible devices
+- **Unsafe Construction**: Context manager for performance-critical scenarios with validation bypass
 
 ## torch.compile Compatibility
 
-Tensor Container is designed from the ground up for `torch.compile` compatibility:
+Tensor Container is designed for `torch.compile` compatibility:
 
 ```python
 @torch.compile
 def process_batch(data: TensorDict) -> TensorDict:
-    return data.apply(lambda x: torch.relu(x))
+    # PyTree operations compile efficiently
+    return TensorContainer._tree_map(lambda x: torch.relu(x), data)
 
-# Compiles efficiently with minimal graph breaks
+@torch.compile
+def sample_and_score(dist: TensorNormal, actions: torch.Tensor) -> torch.Tensor:
+    # Distribution operations are compile-safe
+    return dist.log_prob(actions)
+
+# All operations compile efficiently with minimal graph breaks
 compiled_result = process_batch(tensor_dict)
+log_probs = sample_and_score(normal_dist, action_tensor)
 ```
 
-Our testing framework includes comprehensive compile compatibility verification to ensure all operations work efficiently under JIT compilation.
+The testing framework includes compile compatibility verification to ensure operations work efficiently under JIT compilation, including:
+- Graph break detection and minimization
+- Recompilation tracking
+- Memory leak prevention
+- Performance benchmarking
 
 ## Contributing
 
-We welcome contributions! Tensor Container is a learning project for exploring PyTorch internals and tensor container implementations.
+Contributions are welcome! Tensor Container is a learning project for exploring PyTorch internals and tensor container implementations.
 
 ### Development Setup
 
@@ -168,13 +218,21 @@ pytest --strict-markers --cov=src
 # Run specific test modules
 pytest tests/tensor_dict/test_compile.py
 pytest tests/tensor_dataclass/
+pytest tests/tensor_distribution/
+
+# Run compile-specific tests
+pytest tests/tensor_dict/test_graph_breaks.py
+pytest tests/tensor_dict/test_recompilations.py
 ```
 
 ### Development Guidelines
 
 - All new features must maintain `torch.compile` compatibility
 - Comprehensive tests required, including compile compatibility verification
 - Follow existing code patterns and typing conventions
+- Distribution implementations must support KL divergence registration
+- Memory efficiency considerations for large-scale tensor operations
+- Unsafe construction patterns for performance-critical paths
 
 ### Contribution Process
 
@@ -184,13 +242,22 @@ pytest tests/tensor_dataclass/
 4. Ensure all tests pass and maintain coverage
 5. Submit a pull request with a clear description
 
+## Documentation
+
+The project includes documentation:
+
+- **[`docs/compatibility.md`](docs/compatibility.md)**: Python version compatibility guide and best practices
+- **[`docs/testing.md`](docs/testing.md)**: Testing philosophy, standards, and guidelines
+- **Source Code Documentation**: Extensive docstrings and type annotations throughout the codebase
+- **Test Coverage**: 643+ tests covering all major functionality with 86% code coverage
+
 ## License
 
 This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.
 
 ## Authors
 
-- **Tim Joseph** - *Creator and Lead Developer* - [mctigger](https://github.com/mctigger)
+- **Tim Joseph** - [mctigger](https://github.com/mctigger)
 
 ## Contact and Support
 

docs/compatibility.md

@@ -0,0 +1,75 @@
+# Python Compatibility Guide
+
+This document outlines compatibility considerations and solutions for supporting different Python versions in the tensorcontainer project.
+
+## Supported Python Versions
+
+**tensorcontainer requires Python 3.9 or higher.**
+
+## Compatibility Changes
+
+### Union Types and isinstance() Checks
+
+Union types can now be handled using `typing.get_args()`:
+
+```python
+from typing import Union, get_args
+import torch
+from tensorcontainer import TensorContainer
+
+TDCompatible = Union[torch.Tensor, TensorContainer]
+
+# Use get_args to extract types from Union:
+if isinstance(val, get_args(TDCompatible)):
+    pass
+```
+
+### Type Annotations
+
+**Note**: The `|` operator for union types was introduced in Python 3.10:
+
+```python
+# Python 3.10+ only:
+def func(x: int | str) -> None:
+    pass
+```
+
+For broader compatibility with Python 3.9+, use `Union` from typing:
+
+```python
+from typing import Union
+
+# Compatible with Python 3.9+:
+def func(x: Union[int, str]) -> None:
+    pass
+```
+
+## General Compatibility Tips
+
+### Import Compatibility
+
+Use `from __future__ import annotations` at the top of files to enable forward references and improve compatibility:
+
+```python
+from __future__ import annotations
+
+# This allows using string annotations that are evaluated later
+def func(x: 'SomeClass') -> 'SomeClass':
+    pass
+```
+
+### Version-Specific Features
+
+When using features that are only available in newer Python versions, use version checks:
+
+```python
+import sys
+
+if sys.version_info >= (3, 10):
+    # Use Python 3.10+ features
+    pass
+else:
+    # Fallback for Python 3.9
+    pass
+```
+

pyproject.toml

@@ -12,13 +12,12 @@ dependencies = [
     "torch"
 ]
 readme = "README.md"
-requires-python = ">=3.8"
+requires-python = ">=3.9"
 keywords = ["deep learning", "tensordict", "pytorch"]
 urls = {Homepage = "https://github.com/mctigger/tensor-container"}
 classifiers = [
     "License :: OSI Approved :: MIT License",
     "Programming Language :: Python :: 3",
-    "Programming Language :: Python :: 3.8",
     "Programming Language :: Python :: 3.9",
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
@@ -28,10 +27,9 @@ classifiers = [
 [project.optional-dependencies]
 dev = [
   "pytest>=7.2,<8.0",
-  # you can add extra test-related tools here:
   "pytest-cov",
-  # "hypothesis",
-  # "flake8",
+  "ruff"
+
 ]
 
 

src/tensorcontainer/__init__.py

@@ -2,4 +2,4 @@
 from .tensor_dict import TensorDict
 from .tensor_distribution import TensorDistribution
 
-__all__ = ["TensorDataClass", "TensorDict", "TensorDistribution"]
+__all__ = ["TensorDataClass", "TensorDict", "TensorDistribution"]

src/tensorcontainer/tensor_container.py

@@ -13,9 +13,9 @@
     Optional,
     Tuple,
     Type,
-    TypeAlias,
     Union,
 )
+from typing_extensions import TypeAlias
 
 import torch