dmn_diagnostics.py

#!/usr/bin/env python3
"""
Complete DMN Diagnostics - Debug Tensor Shape Issues
====================================================

Diagnoses the DMN tensor shape inconsistencies in Complete BIOMIND training.
"""

import torch
import torch.nn as nn
import numpy as np
from typing import Dict, List, Optional, Tuple, Any
import time

# Import DMN components
from src.rcca.recursive_dmn import RecursiveDMNThinkingCycle, CounterfactualReasoningEngine

def run_dmn_diagnostics():
    """Run comprehensive DMN diagnostics"""
    print("[SEARCH] COMPLETE DMN DIAGNOSTICS")
    print("=" * 60)

    device = 'cpu'
    workspace_dim = 256

    # Test 1: Initialize DMN
    print("\n1. Testing DMN Initialization...")
    try:
        dmn = RecursiveDMNThinkingCycle(workspace_dim=workspace_dim)
        print("[OK] DMN initialized successfully")
    except Exception as e:
        print(f"[FAIL] DMN initialization failed: {e}")
        return

    # Test 2: Test different input shapes
    print("\n2. Testing Input Shape Handling...")

    test_inputs = [
        ("[256]", torch.randn(256)),
        ("[1, 256]", torch.randn(1, 256)),
        ("[4, 256]", torch.randn(4, 256)),
    ]

    for shape_name, test_input in test_inputs:
        print(f"\n   Testing input shape {shape_name}: {test_input.shape}")
        try:
            result = dmn.process_multi_step_reasoning(
                workspace_state=test_input,
                max_iterations=2,
                context="diagnostic_test"
            )

            final_state = result['final_reasoning_state']
            print(f"   [OK] Output shape: {final_state.shape}")
            print(f"   [OK] Output type: {type(final_state)}")
            print(f"   [OK] Output dim: {final_state.dim()}")

            # Check if shape is consistent
            if final_state.dim() == 1 and final_state.size(0) == workspace_dim:
                print("   [OK] Correct shape [256]")
            elif final_state.dim() == 2 and final_state.size(0) == 1 and final_state.size(1) == workspace_dim:
                print("   [OK] Correct shape [1, 256]")
            else:
                print(f"   [WARN] Unexpected shape: expected [256] or [1, 256], got {final_state.shape}")

        except Exception as e:
            print(f"   [FAIL] Failed: {e}")

    # Test 3: Test CounterfactualReasoningEngine directly
    print("\n3. Testing CounterfactualReasoningEngine...")

    try:
        cf_engine = CounterfactualReasoningEngine(workspace_dim=workspace_dim)

        # Test with different input shapes
        for shape_name, test_input in test_inputs[:2]:  # Only test [256] and [1, 256]
            print(f"\n   Testing CF Engine with {shape_name} input: {test_input.shape}")

            try:
                counterfactuals = cf_engine.generate_counterfactuals(
                    current_state=test_input,
                    narrative_context=test_input,
                    num_alternatives=3
                )

                print(f"   [OK] Generated {len(counterfactuals)} counterfactuals")

                for i, cf in enumerate(counterfactuals):
                    print(f"     CF {i}: content shape {cf.content.shape}, prob {cf.probability:.3f}")

                    # Check content shape
                    if cf.content.dim() == 1 and cf.content.size(0) == workspace_dim:
                        print("       [OK] Content shape [256]")
                    elif cf.content.dim() == 2 and cf.content.size(0) == 1 and cf.content.size(1) == workspace_dim:
                        print("       [OK] Content shape [1, 256]")
                    else:
                        print(f"       [WARN] Unexpected content shape: {cf.content.shape}")

            except Exception as e:
                print(f"   [FAIL] CF Engine failed: {e}")
                import traceback
                traceback.print_exc()

    except Exception as e:
        print(f"[FAIL] CounterfactualReasoningEngine initialization failed: {e}")

    # Test 4: Test network forward passes
    print("\n4. Testing Network Forward Passes...")

    try:
        cf_engine = CounterfactualReasoningEngine(workspace_dim=workspace_dim)

        test_tensor_1d = torch.randn(256)
        test_tensor_2d = torch.randn(1, 256)

        print(f"\n   Testing counterfactual_generator:")
        print(f"     Input [256]: {test_tensor_1d.shape} -> Output: {cf_engine.counterfactual_generator(test_tensor_1d).shape}")
        print(f"     Input [1, 256]: {test_tensor_2d.shape} -> Output: {cf_engine.counterfactual_generator(test_tensor_2d).shape}")

        print(f"\n   Testing probability_estimator:")
        prob_input_1d = torch.cat([test_tensor_1d, test_tensor_1d], dim=-1)  # [512]
        prob_input_2d = torch.cat([test_tensor_2d, test_tensor_2d], dim=-1)  # [1, 512]
        print(f"     Input [512]: {prob_input_1d.shape} -> Output: {cf_engine.probability_estimator(prob_input_1d).shape}")
        print(f"     Input [1, 512]: {prob_input_2d.shape} -> Output: {cf_engine.probability_estimator(prob_input_2d).shape}")

        print(f"\n   Testing self_state_predictor:")
        print(f"     Input [256]: {test_tensor_1d.shape} -> Output: {cf_engine.self_state_predictor(test_tensor_1d).shape}")
        print(f"     Input [1, 256]: {test_tensor_2d.shape} -> Output: {cf_engine.self_state_predictor(test_tensor_2d).shape}")

    except Exception as e:
        print(f"[FAIL] Network testing failed: {e}")
        import traceback
        traceback.print_exc()

    # Test 5: Test training scenario
    print("\n5. Testing Training Scenario...")

    try:
        # Simulate what happens in training
        batch_size = 4
        dmn_workspace = torch.randn(batch_size, 256)  # [4, 256]

        print(f"   Simulating batch processing with workspace shape: {dmn_workspace.shape}")

        dmn_features_list = []
        for i in range(batch_size):
            single_workspace = dmn_workspace[i:i+1]  # [1, 256]
            print(f"     Sample {i}: input shape {single_workspace.shape}")

            result = dmn.process_multi_step_reasoning(
                workspace_state=single_workspace,
                max_iterations=2,
                context=f"training_sample_{i}"
            )

            single_dmn_features = result['final_reasoning_state']
            print(f"     Sample {i}: output shape {single_dmn_features.shape}")

            # Apply the same validation as in training
            if single_dmn_features.dim() != 2 or single_dmn_features.size(1) != 256:
                print(f"     [WARN] Would reshape: {single_dmn_features.shape} -> [1, 256]")
                single_dmn_features = single_dmn_features.view(1, -1)[:, :256]
                print(f"     Reshaped to: {single_dmn_features.shape}")

            dmn_features_list.append(single_dmn_features)

        # Concatenate
        dmn_features = torch.cat(dmn_features_list, dim=0)
        print(f"   [OK] Final batch shape: {dmn_features.shape}")

        if dmn_features.shape == (batch_size, 256):
            print("   [OK] Correct final shape [4, 256]")
        else:
            print(f"   [WARN] Unexpected final shape: {dmn_features.shape}")

    except Exception as e:
        print(f"[FAIL] Training scenario test failed: {e}")
        import traceback
        traceback.print_exc()

    print("\n" + "=" * 60)
    print("[SEARCH] DIAGNOSTICS COMPLETE")
    print("=" * 60)

if __name__ == "__main__":
    run_dmn_diagnostics()</content>
<parameter name="filePath">c:\Users\morrossl\Documents\Private\eon\dmn_diagnostics.py