Enhance TDLogger and Test Infrastructure

- Consolidated TDLogger functionality with new helper functions:
  - Added `contains_important_message()` to identify critical log messages
  - Implemented `is_setup_or_init()` for better message filtering
  - Created centralized message filtering and formatting utilities

- Improved test framework with standardized logging and capture:
  - Created `common/` directory for shared test utilities
  - Implemented `PrefixedTestCase` for consistent test logging
  - Added `log_diff.py` for log capture and verification
  - Enhanced test organization and discoverability

- Refactored test modules to use new infrastructure:
  - Updated test files to use `common/test_utils.py` and `common/log_diff.py`
  - Simplified log capture and comparison logic
  - Improved error handling and test output formatting

- Removed deprecated and unused test files
  - Cleaned up redundant test utilities
  - Standardized test module structure

- Version bump to 0.0.4 to reflect significant improvements in logging and testing

Author: unixwzrd

CHANGELOG.md

@@ -1,5 +1,41 @@
 # CHANGELOG
 
+## 2025-03-10 - 0.0.4
+
+### TDLogger Improvements
+- **Consolidated Duplicate Code**
+  - Added new helper functions to centralize common operations:
+    - `contains_important_message()` - Checks if a message contains important patterns that should always be logged
+    - `is_setup_or_init()` - Identifies setup and initialization functions for better filtering
+    - `format_message()` - Centralizes message formatting for consistency
+  - Reorganized code structure for better clarity and maintainability
+
+- **Optimized Message Filtering Logic**
+  - Improved the `should_skip_message()` function with clearer, more declarative logic
+  - Created constant `IMPORTANT_MESSAGE_PATTERNS` to centralize patterns that should always be logged
+  - Enhanced filtering to ensure critical messages are never skipped
+  - Simplified complex conditional logic for better readability and maintainability
+
+- **Improved Error Handling**
+  - Added robust try/except blocks in `log_message()` to catch and handle exceptions
+  - Added error reporting that directs error messages to stderr for easy debugging
+  - Ensured exceptions in the logger won't propagate to the main application
+  - Fixed edge cases that could cause unexpected behavior during logging
+
+- **Test Framework Enhancements**
+  - Improved `PrefixedTestCase` class with clear separation between test messages and TDLogger output
+  - Fixed formatting inconsistencies in the test output for better readability
+  - Updated expected output files to reflect new logger behavior
+  - Enhanced error handling in test utilities
+
+### Test Infrastructure Improvements
+- **Created Common Test Utilities Directory**
+  - Added `tests/common/` directory to house shared test infrastructure
+  - Implemented `test_utils.py` with the `PrefixedTestCase` class providing standardized logging capabilities
+  - Moved `log_diff.py` to common location for reuse across all test modules
+  - Standardized test setup and teardown procedures across all tests
+  - Improved test discoverability and organization
+
 ## 2025-03-08 - 0.0.3
 
 ### Logger Improvements

README.md

@@ -20,6 +20,7 @@ TorchDevice is a Python library that enables transparent code portability betwee
   - [Limitations](#limitations)
   - [Recent Updates](#recent-updates)
     - [March 2025 Updates](#march-2025-updates)
+      - [Improved Logging System](#improved-logging-system)
       - [Stream and Event Handling](#stream-and-event-handling)
       - [PyTorch Compatibility](#pytorch-compatibility)
       - [Code Quality](#code-quality)
@@ -194,6 +195,22 @@ Ensure these scripts are updated with any changes you made to `TorchDevice.py`.
 
 The TorchDevice library has undergone significant improvements to ensure seamless compatibility between CUDA and MPS devices:
 
+#### Improved Logging System
+
+- **TDLogger Module Enhancements**:
+  - Consolidated duplicate code with dedicated helper functions for message filtering and formatting
+  - Optimized message filtering logic with declarative pattern matching
+  - Improved error handling with robust exception capture
+  - Enhanced memory management using fixed-size collections for logging history
+  - Centralized important message patterns for consistent filtering
+  
+- **Test Framework Improvements**:
+  - Created standardized test infrastructure in the `common` directory
+  - Implemented `PrefixedTestCase` class for consistent test behavior and logging
+  - Added dedicated utilities for log capture and verification
+  - Improved log message formatting and separation between test output and redirected messages
+  - Enhanced test organization and discoverability
+
 #### Stream and Event Handling
 
 - **Enhanced Stream Support**: Implemented comprehensive CUDA stream functionality on MPS devices, including:

TODO.md

@@ -1,5 +1,41 @@
 # TorchDevice Project TODO List
 
+## TorchDevice.py Refactoring
+
+- **Modularize TorchDevice Core Functions**
+  - Split the large single-file implementation into logical modules:
+    - Core device detection and management
+    - CUDA function mocking
+    - Stream and Event handling
+    - Tensor operations interception
+    - PyTorch integration utilities
+  - Create proper module structure with appropriate imports
+  - Ensure backward compatibility during transition
+
+- **Fix Type Annotations and Linter Errors**
+  - Address all the linter errors in TorchDevice.py
+  - Add proper type annotations throughout the codebase
+  - Fix compatibility with static type checkers
+  - Maintain dynamic behavior while improving static analysis
+
+- **Implement CPU Override Functionality**
+  - Add mechanism to respect explicit CPU device requests
+  - Create policy for when to override vs. when to respect explicit device choices
+  - Add configuration option to control redirection behavior
+  - Update documenation to explain the new behavior
+
+- **Improve Class Structure**
+  - Review inheritance patterns for Stream and Event classes
+  - Ensure proper inheritance from PyTorch base classes
+  - Consolidate duplicate code in device handling
+  - Improve encapsulation of internal state
+
+- **Enhance Error Handling**
+  - Add more robust error handling for device operations
+  - Provide clearer error messages for device compatibility issues
+  - Add fallback mechanisms for unsupported operations
+  - Improve debugging information during redirection failures
+
 ## Logger Improvements
 
 - **Simplify the Frame Inspection Logic**
@@ -15,4 +51,49 @@
 - **Use Context Managers for Temporary State**
   - Identify operations that temporarily modify state
   - Implement context managers for these operations
-  - Ensure proper cleanup even in case of exceptions 
+  - Ensure proper cleanup even in case of exceptions
+
+- **Implement Lazy Formatting for Log Messages**
+  - Identify places where f-strings are used for messages that might be filtered out
+  - Replace these with format strings that are only evaluated if the message will be logged
+  - Verify that the changes don't affect the log output
+  - Measure the performance impact of the change
+
+- **Remove Unused Features and Code**
+  - Identify unused code or overly complex sections
+  - Simplify or remove these sections
+  - Verify that the changes don't affect the functionality
+
+## Device Handling Improvements
+
+- **Respect Explicit CPU Device Specifications**
+  - Modify TorchDevice to respect explicit CPU device requests
+  - Fix failing tests in test_cpu_mps_operations.py due to device redirection
+  - Ensure proper handling of operations involving both CPU and accelerated devices
+  - Add test cases for explicit CPU operations
+
+## Test Framework Improvements
+
+- **Refine Test Utilities into Proper Modules**
+  - Resolve import issues with common.log_diff and common.test_utils
+  - Restructure as proper Python modules with __init__.py files
+  - Update import paths in all test files
+  - Add proper documentation for all test utility functions
+
+- **Enhance Test Discoverability**
+  - Improve test organization for better discoverability
+  - Add test categorization (e.g., unit tests, integration tests)
+  - Create a more flexible test runner with better filtering options
+
+## Documentation
+
+- **Create Developer Documentation**
+  - Document the architecture and design decisions
+  - Add detailed explanations of key components
+  - Provide examples for common development tasks
+  - Include troubleshooting guides for common issues
+
+- **Improve API Documentation**
+  - Document all public interfaces
+  - Add examples for all major features
+  - Include type annotations for better IDE support 

TorchDevice/TorchDevice.py

@@ -67,31 +67,32 @@ def tensor_to_replacement(self, *args, **kwargs):
                 (hasattr(candidate, '__module__') and candidate.__module__ != 'torch')
             ):
             device_arg = candidate
+            # Replace args with default device if requested device not available
+            if isinstance(candidate, str) and candidate != default_device:
+                args = (default_device,) + args[1:]
+            elif isinstance(candidate, _ORIGINAL_TORCH_DEVICE_TYPE) and candidate.type != default_device:
+                args = (default_device,) + args[1:]
     elif 'device' in kwargs:
         device_arg = kwargs['device']
 
         # Handle device redirection in kwargs
         if isinstance(device_arg, str):
             device_type = device_arg.split(':')[0] if ':' in device_arg else device_arg
-            redirected_type = TorchDevice._redirect_device_type(device_type)
-            
-            if redirected_type != device_type:
+            if device_type != default_device:
                 if ':' in device_arg:
                     index = device_arg.split(':')[1]
-                    kwargs['device'] = f"{redirected_type}:{index}"
+                    kwargs['device'] = f"{default_device}:{index}"
                 else:
-                    kwargs['device'] = redirected_type
-                log_message(f"Redirecting tensor.to() from '{device_type}' to '{redirected_type}'.", "tensor.to")
+                    kwargs['device'] = default_device
+                log_message(f"Redirecting tensor.to() from '{device_type}' to '{default_device}'.", "tensor.to")
         elif hasattr(device_arg, 'type'):
             device_type = device_arg.type
-            redirected_type = TorchDevice._redirect_device_type(device_type)
-            
-            if redirected_type != device_type:
+            if device_type != default_device:
                 index = getattr(device_arg, 'index', 0)
                 if index is None:
                     index = 0
-                kwargs['device'] = TorchDevice.torch_device_replacement(f"{redirected_type}:{index}")
-                log_message(f"Redirecting tensor.to() from '{device_type}' to '{redirected_type}'.", "tensor.to")
+                kwargs['device'] = TorchDevice.torch_device_replacement(f"{default_device}:{index}")
+                log_message(f"Redirecting tensor.to() from '{device_type}' to '{default_device}'.", "tensor.to")
 
     # Log warning if target device doesn't match default
     if device_arg is not None:
@@ -103,6 +104,8 @@ def tensor_to_replacement(self, *args, **kwargs):
             target_device = str(device_arg)
         if target_device != default_device:
             log_message(f"tensor.to() called with device {device_arg} which does not match the default device {default_device}.", "tensor.to")
+    
+    # Call the original to() with the redirected device
     return _original_tensor_to(self, *args, **kwargs)
 
 def module_to_replacement(self, *args, **kwargs):