{"nbformat":4,"nbformat_minor":0,"metadata":{"colab":{"provenance":[],"collapsed_sections":["JAzPQnPRAV82","k5pJjdk3qRUr","FYHGvq016CIq","rB99d0u3CWOM","FtkQmzOpAP1p","W0U7A75M7wVQ","U68DMM5murm_","X5K51SXfMUbq","o4bkrLRFLPxO","Sf_zpjymdzNE","sYBO53Qekcxk"],"machine_shape":"hm","gpuType":"G4","authorship_tag":"ABX9TyPGajZneGEoyBPKriHRblIz"},"kernelspec":{"name":"python3","display_name":"Python 3"},"language_info":{"name":"python"},"accelerator":"GPU","widgets":{"application/vnd.jupyter.widget-state+json":{"a4d0fdf9703149c59898264e8eaf8789":{"model_module":"@jupyter-widgets/controls","model_name":"HBoxModel","model_module_version":"1.5.0","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"HBoxModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"HBoxView","box_style":"","children":["IPY_MODEL_07c1879a8b7049728078d1ca13a84710","IPY_MODEL_c27e770348a148409907dd7e794345ec","IPY_MODEL_0f94add3cab0471c922358e0d9c7f773"],"layout":"IPY_MODEL_72707cd9fdac4f3e962c8bfbe8ee43ae"}},"07c1879a8b7049728078d1ca13a84710":{"model_module":"@jupyter-widgets/controls","model_name":"HTMLModel","model_module_version":"1.5.0","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"HTMLModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"HTMLView","description":"","description_tooltip":null,"layout":"IPY_MODEL_bf10d23e46734d5f8f4c2bb61eab28ff","placeholder":"​","style":"IPY_MODEL_e6c10cf2e2ec4ba882522a811f5635f0","value":"Epochs: 100%"}},"c27e770348a148409907dd7e794345ec":{"model_module":"@jupyter-widgets/controls","model_name":"FloatProgressModel","model_module_version":"1.5.0","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"FloatProgressModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"ProgressView","bar_style":"success","description":"","description_tooltip":null,"layout":"IPY_MODEL_a2fd785ff8774099b265b56442ef3952","max":100,"min":0,"orientation":"horizontal","style":"IPY_MODEL_aa81112795f3486688f9d5c1015cf9af","value":100}},"0f94add3cab0471c922358e0d9c7f773":{"model_module":"@jupyter-widgets/controls","model_name":"HTMLModel","model_module_version":"1.5.0","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"HTMLModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"HTMLView","description":"","description_tooltip":null,"layout":"IPY_MODEL_7357538e5c964935aad818fbb0d1e29e","placeholder":"​","style":"IPY_MODEL_7ce103ccbf2145f8a9ba452694f8f314","value":" 100/100 [12:17&lt;00:00,  7.15s/it]"}},"72707cd9fdac4f3e962c8bfbe8ee43ae":{"model_module":"@jupyter-widgets/base","model_name":"LayoutModel","model_module_version":"1.2.0","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"bf10d23e46734d5f8f4c2bb61eab28ff":{"model_module":"@jupyter-widgets/base","model_name":"LayoutModel","model_module_version":"1.2.0","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"e6c10cf2e2ec4ba882522a811f5635f0":{"model_module":"@jupyter-widgets/controls","model_name":"DescriptionStyleModel","model_module_version":"1.5.0","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"DescriptionStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"StyleView","description_width":""}},"a2fd785ff8774099b265b56442ef3952":{"model_module":"@jupyter-widgets/base","model_name":"LayoutModel","model_module_version":"1.2.0","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"aa81112795f3486688f9d5c1015cf9af":{"model_module":"@jupyter-widgets/controls","model_name":"ProgressStyleModel","model_module_version":"1.5.0","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"ProgressStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"StyleView","bar_color":null,"description_width":""}},"7357538e5c964935aad818fbb0d1e29e":{"model_module":"@jupyter-widgets/base","model_name":"LayoutModel","model_module_version":"1.2.0","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"7ce103ccbf2145f8a9ba452694f8f314":{"model_module":"@jupyter-widgets/controls","model_name":"DescriptionStyleModel","model_module_version":"1.5.0","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"DescriptionStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"StyleView","description_width":""}},"66f0a0dd6a3b49afaaff8e55ddc8f232":{"model_module":"@jupyter-widgets/controls","model_name":"HBoxModel","model_module_version":"1.5.0","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"HBoxModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"HBoxView","box_style":"","children":["IPY_MODEL_6d12d020ec7349e1a1937ad42ae60992","IPY_MODEL_ab00fed336b84fd6bd79d2608beb0252","IPY_MODEL_bc211f72ad9d41859ba156da0b1f0a74"],"layout":"IPY_MODEL_d4a9013c24e74aebac09f2ee62a63148"}},"6d12d020ec7349e1a1937ad42ae60992":{"model_module":"@jupyter-widgets/controls","model_name":"HTMLModel","model_module_version":"1.5.0","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"HTMLModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"HTMLView","description":"","description_tooltip":null,"layout":"IPY_MODEL_663d2b99117244bd9ef4b6d9f0eafdb8","placeholder":"​","style":"IPY_MODEL_b0ae5117b9c64f71a3f5698d1546acb5","value":"Epochs: 100%"}},"ab00fed336b84fd6bd79d2608beb0252":{"model_module":"@jupyter-widgets/controls","model_name":"FloatProgressModel","model_module_version":"1.5.0","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"FloatProgressModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"ProgressView","bar_style":"success","description":"","description_tooltip":null,"layout":"IPY_MODEL_e7d75665dc104599af3afdc2011a61db","max":100,"min":0,"orientation":"horizontal","style":"IPY_MODEL_41ea6b42289f458bb56492ec1298f4f7","value":100}},"bc211f72ad9d41859ba156da0b1f0a74":{"model_module":"@jupyter-widgets/controls","model_name":"HTMLModel","model_module_version":"1.5.0","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"HTMLModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"HTMLView","description":"","description_tooltip":null,"layout":"IPY_MODEL_833456446d0c4e18996f40f68d2277da","placeholder":"​","style":"IPY_MODEL_e7808d06d7ad473a97990cdc41f9515e","value":" 100/100 [24:28&lt;00:00, 14.46s/it]"}},"d4a9013c24e74aebac09f2ee62a63148":{"model_module":"@jupyter-widgets/base","model_name":"LayoutModel","model_module_version":"1.2.0","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"663d2b99117244bd9ef4b6d9f0eafdb8":{"model_module":"@jupyter-widgets/base","model_name":"LayoutModel","model_module_version":"1.2.0","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"b0ae5117b9c64f71a3f5698d1546acb5":{"model_module":"@jupyter-widgets/controls","model_name":"DescriptionStyleModel","model_module_version":"1.5.0","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"DescriptionStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"StyleView","description_width":""}},"e7d75665dc104599af3afdc2011a61db":{"model_module":"@jupyter-widgets/base","model_name":"LayoutModel","model_module_version":"1.2.0","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"41ea6b42289f458bb56492ec1298f4f7":{"model_module":"@jupyter-widgets/controls","model_name":"ProgressStyleModel","model_module_version":"1.5.0","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"ProgressStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"StyleView","bar_color":null,"description_width":""}},"833456446d0c4e18996f40f68d2277da":{"model_module":"@jupyter-widgets/base","model_name":"LayoutModel","model_module_version":"1.2.0","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"e7808d06d7ad473a97990cdc41f9515e":{"model_module":"@jupyter-widgets/controls","model_name":"DescriptionStyleModel","model_module_version":"1.5.0","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"DescriptionStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"StyleView","description_width":""}},"8045de5b58de4308a4bf4d1e60023113":{"model_module":"@jupyter-widgets/controls","model_name":"HBoxModel","model_module_version":"1.5.0","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"HBoxModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"HBoxView","box_style":"","children":["IPY_MODEL_38fcc80f6c9145c1b3b75659a01f5d97","IPY_MODEL_5df6e33ca7104310a3eff8ba2eaf42d9","IPY_MODEL_424c6cf8e56947939bd29d98c104c573"],"layout":"IPY_MODEL_3531e2c63e4f42928c4e953442cba8d2"}},"38fcc80f6c9145c1b3b75659a01f5d97":{"model_module":"@jupyter-widgets/controls","model_name":"HTMLModel","model_module_version":"1.5.0","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"HTMLModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"HTMLView","description":"","description_tooltip":null,"layout":"IPY_MODEL_3fa201a8613a4f4caa09a3cb1d042fdf","placeholder":"​","style":"IPY_MODEL_dedb555f60e44516a1b5b80896a9634f","value":"Epochs: 100%"}},"5df6e33ca7104310a3eff8ba2eaf42d9":{"model_module":"@jupyter-widgets/controls","model_name":"FloatProgressModel","model_module_version":"1.5.0","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"FloatProgressModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"ProgressView","bar_style":"success","description":"","description_tooltip":null,"layout":"IPY_MODEL_1bb21ffdb1ba469ab6f86df3071f9dc5","max":100,"min":0,"orientation":"horizontal","style":"IPY_MODEL_963b762376c147048fd22d081017c320","value":100}},"424c6cf8e56947939bd29d98c104c573":{"model_module":"@jupyter-widgets/controls","model_name":"HTMLModel","model_module_version":"1.5.0","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"HTMLModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"HTMLView","description":"","description_tooltip":null,"layout":"IPY_MODEL_7e1e5cc2f9694b4e88067f96f7cc5f9f","placeholder":"​","style":"IPY_MODEL_210fc9a76302468ebacef0c0e65275b6","value":" 100/100 [23:55&lt;00:00, 13.87s/it]"}},"3531e2c63e4f42928c4e953442cba8d2":{"model_module":"@jupyter-widgets/base","model_name":"LayoutModel","model_module_version":"1.2.0","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"3fa201a8613a4f4caa09a3cb1d042fdf":{"model_module":"@jupyter-widgets/base","model_name":"LayoutModel","model_module_version":"1.2.0","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"dedb555f60e44516a1b5b80896a9634f":{"model_module":"@jupyter-widgets/controls","model_name":"DescriptionStyleModel","model_module_version":"1.5.0","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"DescriptionStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"StyleView","description_width":""}},"1bb21ffdb1ba469ab6f86df3071f9dc5":{"model_module":"@jupyter-widgets/base","model_name":"LayoutModel","model_module_version":"1.2.0","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"963b762376c147048fd22d081017c320":{"model_module":"@jupyter-widgets/controls","model_name":"ProgressStyleModel","model_module_version":"1.5.0","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"ProgressStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"StyleView","bar_color":null,"description_width":""}},"7e1e5cc2f9694b4e88067f96f7cc5f9f":{"model_module":"@jupyter-widgets/base","model_name":"LayoutModel","model_module_version":"1.2.0","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"210fc9a76302468ebacef0c0e65275b6":{"model_module":"@jupyter-widgets/controls","model_name":"DescriptionStyleModel","model_module_version":"1.5.0","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"DescriptionStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"StyleView","description_width":""}},"cba1d1d952cc48a99831fc3979d8ed0d":{"model_module":"@jupyter-widgets/controls","model_name":"HBoxModel","model_module_version":"1.5.0","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"HBoxModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"HBoxView","box_style":"","children":["IPY_MODEL_d1d1752b19484f5d82ee9caba836c53e","IPY_MODEL_bfb14422f4e7444fa83296f998bb27f9","IPY_MODEL_70e1b310d0f24624b224ea4585205e98"],"layout":"IPY_MODEL_e204ac6670454441bb2bb313ae4c82c0"}},"d1d1752b19484f5d82ee9caba836c53e":{"model_module":"@jupyter-widgets/controls","model_name":"HTMLModel","model_module_version":"1.5.0","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"HTMLModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"HTMLView","description":"","description_tooltip":null,"layout":"IPY_MODEL_e1f6b47538a54273a2e8230ac642d413","placeholder":"​","style":"IPY_MODEL_d46bfa89a3324ad6b563e2a82ffd152d","value":"Epochs: 100%"}},"bfb14422f4e7444fa83296f998bb27f9":{"model_module":"@jupyter-widgets/controls","model_name":"FloatProgressModel","model_module_version":"1.5.0","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"FloatProgressModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"ProgressView","bar_style":"success","description":"","description_tooltip":null,"layout":"IPY_MODEL_5bc38b424a504d84b304ee1e90537191","max":100,"min":0,"orientation":"horizontal","style":"IPY_MODEL_606bc83980ab41b5b3d3ac6eaa931d09","value":100}},"70e1b310d0f24624b224ea4585205e98":{"model_module":"@jupyter-widgets/controls","model_name":"HTMLModel","model_module_version":"1.5.0","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"HTMLModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"HTMLView","description":"","description_tooltip":null,"layout":"IPY_MODEL_1dd023b412e64d1293faa213af2a8631","placeholder":"​","style":"IPY_MODEL_5ff7475e37d44589b3fdfb342f9769c7","value":" 100/100 [50:40&lt;00:00, 29.49s/it]"}},"e204ac6670454441bb2bb313ae4c82c0":{"model_module":"@jupyter-widgets/base","model_name":"LayoutModel","model_module_version":"1.2.0","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"e1f6b47538a54273a2e8230ac642d413":{"model_module":"@jupyter-widgets/base","model_name":"LayoutModel","model_module_version":"1.2.0","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"d46bfa89a3324ad6b563e2a82ffd152d":{"model_module":"@jupyter-widgets/controls","model_name":"DescriptionStyleModel","model_module_version":"1.5.0","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"DescriptionStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"StyleView","description_width":""}},"5bc38b424a504d84b304ee1e90537191":{"model_module":"@jupyter-widgets/base","model_name":"LayoutModel","model_module_version":"1.2.0","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"606bc83980ab41b5b3d3ac6eaa931d09":{"model_module":"@jupyter-widgets/controls","model_name":"ProgressStyleModel","model_module_version":"1.5.0","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"ProgressStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"StyleView","bar_color":null,"description_width":""}},"1dd023b412e64d1293faa213af2a8631":{"model_module":"@jupyter-widgets/base","model_name":"LayoutModel","model_module_version":"1.2.0","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"5ff7475e37d44589b3fdfb342f9769c7":{"model_module":"@jupyter-widgets/controls","model_name":"DescriptionStyleModel","model_module_version":"1.5.0","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"DescriptionStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"StyleView","description_width":""}},"7ba7e8d6232b4b99991f0e44369e88ab":{"model_module":"@jupyter-widgets/controls","model_name":"HBoxModel","model_module_version":"1.5.0","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"HBoxModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"HBoxView","box_style":"","children":["IPY_MODEL_6697a72129524bd0bcd63850b874367c","IPY_MODEL_ff64e3fbea544b56bae47a7a2cd9d54c","IPY_MODEL_3c2531ab76d240ea82c66ba3cdd333d4"],"layout":"IPY_MODEL_bea929296b3648e79b1d201c19f0d172"}},"6697a72129524bd0bcd63850b874367c":{"model_module":"@jupyter-widgets/controls","model_name":"HTMLModel","model_module_version":"1.5.0","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"HTMLModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"HTMLView","description":"","description_tooltip":null,"layout":"IPY_MODEL_57192c34c42848cdbfdcb47f8eb6f23f","placeholder":"​","style":"IPY_MODEL_83d893f9e61f471d8f69fc092b262be6","value":"Epochs: 100%"}},"ff64e3fbea544b56bae47a7a2cd9d54c":{"model_module":"@jupyter-widgets/controls","model_name":"FloatProgressModel","model_module_version":"1.5.0","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"FloatProgressModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"ProgressView","bar_style":"success","description":"","description_tooltip":null,"layout":"IPY_MODEL_d73657ec98804890aea6306f5be33bd5","max":300,"min":0,"orientation":"horizontal","style":"IPY_MODEL_0eeacd094fab458ba80cf4273fc78bfb","value":300}},"3c2531ab76d240ea82c66ba3cdd333d4":{"model_module":"@jupyter-widgets/controls","model_name":"HTMLModel","model_module_version":"1.5.0","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"HTMLModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"HTMLView","description":"","description_tooltip":null,"layout":"IPY_MODEL_8bf14e45224943c38577cb55b46e8933","placeholder":"​","style":"IPY_MODEL_0fa1fe6d22a744ef81acf99dc3ed7918","value":" 300/300 [1:20:21&lt;00:00, 16.04s/it]"}},"bea929296b3648e79b1d201c19f0d172":{"model_module":"@jupyter-widgets/base","model_name":"LayoutModel","model_module_version":"1.2.0","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"57192c34c42848cdbfdcb47f8eb6f23f":{"model_module":"@jupyter-widgets/base","model_name":"LayoutModel","model_module_version":"1.2.0","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"83d893f9e61f471d8f69fc092b262be6":{"model_module":"@jupyter-widgets/controls","model_name":"DescriptionStyleModel","model_module_version":"1.5.0","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"DescriptionStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"StyleView","description_width":""}},"d73657ec98804890aea6306f5be33bd5":{"model_module":"@jupyter-widgets/base","model_name":"LayoutModel","model_module_version":"1.2.0","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"0eeacd094fab458ba80cf4273fc78bfb":{"model_module":"@jupyter-widgets/controls","model_name":"ProgressStyleModel","model_module_version":"1.5.0","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"ProgressStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"StyleView","bar_color":null,"description_width":""}},"8bf14e45224943c38577cb55b46e8933":{"model_module":"@jupyter-widgets/base","model_name":"LayoutModel","model_module_version":"1.2.0","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"0fa1fe6d22a744ef81acf99dc3ed7918":{"model_module":"@jupyter-widgets/controls","model_name":"DescriptionStyleModel","model_module_version":"1.5.0","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"DescriptionStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"StyleView","description_width":""}},"86d16df51c124161a2728f0003ebed64":{"model_module":"@jupyter-widgets/controls","model_name":"HBoxModel","model_module_version":"1.5.0","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"HBoxModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"HBoxView","box_style":"","children":["IPY_MODEL_b5b3f1a8eec049c78df39ea2e0a878d7","IPY_MODEL_80ab24bf7d1a4ad69b2618334aa90c11","IPY_MODEL_0b8d0ae380024fbb8a0fae99ac4ac102"],"layout":"IPY_MODEL_61cd81b77c3249a98f063dc6f9117df7"}},"b5b3f1a8eec049c78df39ea2e0a878d7":{"model_module":"@jupyter-widgets/controls","model_name":"HTMLModel","model_module_version":"1.5.0","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"HTMLModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"HTMLView","description":"","description_tooltip":null,"layout":"IPY_MODEL_d94819a2d40f4687994d35e75011b18c","placeholder":"​","style":"IPY_MODEL_6efd5fc8aa13488a815d2d74ee9fd2b0","value":"Epochs: 100%"}},"80ab24bf7d1a4ad69b2618334aa90c11":{"model_module":"@jupyter-widgets/controls","model_name":"FloatProgressModel","model_module_version":"1.5.0","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"FloatProgressModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"ProgressView","bar_style":"success","description":"","description_tooltip":null,"layout":"IPY_MODEL_85aa6598a0894f7b8e103c9a79dc897d","max":100,"min":0,"orientation":"horizontal","style":"IPY_MODEL_288bd07827f54042b5e8e7a342eb586a","value":100}},"0b8d0ae380024fbb8a0fae99ac4ac102":{"model_module":"@jupyter-widgets/controls","model_name":"HTMLModel","model_module_version":"1.5.0","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"HTMLModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"HTMLView","description":"","description_tooltip":null,"layout":"IPY_MODEL_bc430ece4e2645a28795a95a2e178010","placeholder":"​","style":"IPY_MODEL_bb75ffc809d0480e882295dc28ec3f5d","value":" 100/100 [57:04&lt;00:00, 34.11s/it]"}},"61cd81b77c3249a98f063dc6f9117df7":{"model_module":"@jupyter-widgets/base","model_name":"LayoutModel","model_module_version":"1.2.0","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"d94819a2d40f4687994d35e75011b18c":{"model_module":"@jupyter-widgets/base","model_name":"LayoutModel","model_module_version":"1.2.0","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"6efd5fc8aa13488a815d2d74ee9fd2b0":{"model_module":"@jupyter-widgets/controls","model_name":"DescriptionStyleModel","model_module_version":"1.5.0","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"DescriptionStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"StyleView","description_width":""}},"85aa6598a0894f7b8e103c9a79dc897d":{"model_module":"@jupyter-widgets/base","model_name":"LayoutModel","model_module_version":"1.2.0","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"288bd07827f54042b5e8e7a342eb586a":{"model_module":"@jupyter-widgets/controls","model_name":"ProgressStyleModel","model_module_version":"1.5.0","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"ProgressStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"StyleView","bar_color":null,"description_width":""}},"bc430ece4e2645a28795a95a2e178010":{"model_module":"@jupyter-widgets/base","model_name":"LayoutModel","model_module_version":"1.2.0","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"bb75ffc809d0480e882295dc28ec3f5d":{"model_module":"@jupyter-widgets/controls","model_name":"DescriptionStyleModel","model_module_version":"1.5.0","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"DescriptionStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"StyleView","description_width":""}},"8638148368d34ba78e9451e5d38ea54f":{"model_module":"@jupyter-widgets/controls","model_name":"HBoxModel","model_module_version":"1.5.0","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"HBoxModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"HBoxView","box_style":"","children":["IPY_MODEL_3fe3ced2159249f2abdd190983a48dd6","IPY_MODEL_1e30a8e7935241f9a58dede6e148656b","IPY_MODEL_a0153bdb2cde41d485969562898cd3d2"],"layout":"IPY_MODEL_5d7a4a42979c4622ad3c62594e48c786"}},"3fe3ced2159249f2abdd190983a48dd6":{"model_module":"@jupyter-widgets/controls","model_name":"HTMLModel","model_module_version":"1.5.0","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"HTMLModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"HTMLView","description":"","description_tooltip":null,"layout":"IPY_MODEL_f629365478a84dca84b8f64287eba2dd","placeholder":"​","style":"IPY_MODEL_54d2555ec2f94980bff8fd90a6d705fb","value":"Training:   0%"}},"1e30a8e7935241f9a58dede6e148656b":{"model_module":"@jupyter-widgets/controls","model_name":"FloatProgressModel","model_module_version":"1.5.0","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"FloatProgressModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"ProgressView","bar_style":"danger","description":"","description_tooltip":null,"layout":"IPY_MODEL_d67d42e594774904b7cd4367478da24c","max":200,"min":0,"orientation":"horizontal","style":"IPY_MODEL_02e36a82c7224d1ab00b1f54f5c78900","value":0}},"a0153bdb2cde41d485969562898cd3d2":{"model_module":"@jupyter-widgets/controls","model_name":"HTMLModel","model_module_version":"1.5.0","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"HTMLModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"HTMLView","description":"","description_tooltip":null,"layout":"IPY_MODEL_614125e4c7174d0983d0106216228730","placeholder":"​","style":"IPY_MODEL_9ea9af4b1a3c4255800e4260fde64577","value":" 0/200 [00:12&lt;?, ?ep/s]"}},"5d7a4a42979c4622ad3c62594e48c786":{"model_module":"@jupyter-widgets/base","model_name":"LayoutModel","model_module_version":"1.2.0","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"f629365478a84dca84b8f64287eba2dd":{"model_module":"@jupyter-widgets/base","model_name":"LayoutModel","model_module_version":"1.2.0","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"54d2555ec2f94980bff8fd90a6d705fb":{"model_module":"@jupyter-widgets/controls","model_name":"DescriptionStyleModel","model_module_version":"1.5.0","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"DescriptionStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"StyleView","description_width":""}},"d67d42e594774904b7cd4367478da24c":{"model_module":"@jupyter-widgets/base","model_name":"LayoutModel","model_module_version":"1.2.0","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"02e36a82c7224d1ab00b1f54f5c78900":{"model_module":"@jupyter-widgets/controls","model_name":"ProgressStyleModel","model_module_version":"1.5.0","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"ProgressStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"StyleView","bar_color":null,"description_width":""}},"614125e4c7174d0983d0106216228730":{"model_module":"@jupyter-widgets/base","model_name":"LayoutModel","model_module_version":"1.2.0","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"9ea9af4b1a3c4255800e4260fde64577":{"model_module":"@jupyter-widgets/controls","model_name":"DescriptionStyleModel","model_module_version":"1.5.0","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"DescriptionStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"StyleView","description_width":""}},"74ac9e73987f44b6bb162f0f9ba9c846":{"model_module":"@jupyter-widgets/controls","model_name":"HBoxModel","model_module_version":"1.5.0","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"HBoxModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"HBoxView","box_style":"","children":["IPY_MODEL_f41e6f87ec564b1c972a55f84ba7238a","IPY_MODEL_fd296886a4ae41818fc8f64787521637","IPY_MODEL_abc21c48f62b435a984e971dc69c8cb2"],"layout":"IPY_MODEL_5201c3b55b9d42138b84be227d368237"}},"f41e6f87ec564b1c972a55f84ba7238a":{"model_module":"@jupyter-widgets/controls","model_name":"HTMLModel","model_module_version":"1.5.0","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"HTMLModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"HTMLView","description":"","description_tooltip":null,"layout":"IPY_MODEL_4f7527e7353449319bd4370b6cd3184b","placeholder":"​","style":"IPY_MODEL_6718ec41897c4c01ae79ec4a6a5ecd4d","value":"Training:  43%"}},"fd296886a4ae41818fc8f64787521637":{"model_module":"@jupyter-widgets/controls","model_name":"FloatProgressModel","model_module_version":"1.5.0","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"FloatProgressModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"ProgressView","bar_style":"danger","description":"","description_tooltip":null,"layout":"IPY_MODEL_5ea7f24370ac4900809d0b9e52bf8017","max":200,"min":0,"orientation":"horizontal","style":"IPY_MODEL_f473b773987d4b848b1760bb55fd45ec","value":86}},"abc21c48f62b435a984e971dc69c8cb2":{"model_module":"@jupyter-widgets/controls","model_name":"HTMLModel","model_module_version":"1.5.0","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"HTMLModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"HTMLView","description":"","description_tooltip":null,"layout":"IPY_MODEL_f21f834c60a44e148fc08227a306dc09","placeholder":"​","style":"IPY_MODEL_f9408474f42248e79577eb763e72029a","value":" 86/200 [10:42&lt;09:54,  5.22s/ep, L=16.58 tr=0.054 te=0.057 best=0.058 3.3s]"}},"5201c3b55b9d42138b84be227d368237":{"model_module":"@jupyter-widgets/base","model_name":"LayoutModel","model_module_version":"1.2.0","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"4f7527e7353449319bd4370b6cd3184b":{"model_module":"@jupyter-widgets/base","model_name":"LayoutModel","model_module_version":"1.2.0","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"6718ec41897c4c01ae79ec4a6a5ecd4d":{"model_module":"@jupyter-widgets/controls","model_name":"DescriptionStyleModel","model_module_version":"1.5.0","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"DescriptionStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"StyleView","description_width":""}},"5ea7f24370ac4900809d0b9e52bf8017":{"model_module":"@jupyter-widgets/base","model_name":"LayoutModel","model_module_version":"1.2.0","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"f473b773987d4b848b1760bb55fd45ec":{"model_module":"@jupyter-widgets/controls","model_name":"ProgressStyleModel","model_module_version":"1.5.0","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"ProgressStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"StyleView","bar_color":null,"description_width":""}},"f21f834c60a44e148fc08227a306dc09":{"model_module":"@jupyter-widgets/base","model_name":"LayoutModel","model_module_version":"1.2.0","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"f9408474f42248e79577eb763e72029a":{"model_module":"@jupyter-widgets/controls","model_name":"DescriptionStyleModel","model_module_version":"1.5.0","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"DescriptionStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"StyleView","description_width":""}}}}},"cells":[{"cell_type":"code","source":["!pip uninstall geofractal geolip-core geometricvocab -y\n","!pip install \"git+https://github.com/AbstractEyes/geolip-core.git\""],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"DIS6kNZfAYXd","executionInfo":{"status":"ok","timestamp":1775261465956,"user_tz":420,"elapsed":15527,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"64eeb7b7-b967-435e-d2b2-0832bc198a6f"},"execution_count":1,"outputs":[{"output_type":"stream","name":"stdout","text":["\u001b[33mWARNING: Skipping geofractal as it is not installed.\u001b[0m\u001b[33m\n","\u001b[0m\u001b[33mWARNING: Skipping geolip-core as it is not installed.\u001b[0m\u001b[33m\n","\u001b[0m\u001b[33mWARNING: Skipping geometricvocab as it is not installed.\u001b[0m\u001b[33m\n","\u001b[0mCollecting git+https://github.com/AbstractEyes/geolip-core.git\n","  Cloning https://github.com/AbstractEyes/geolip-core.git to /tmp/pip-req-build-rc0ffh_a\n","  Running command git clone --filter=blob:none --quiet https://github.com/AbstractEyes/geolip-core.git /tmp/pip-req-build-rc0ffh_a\n","  Resolved https://github.com/AbstractEyes/geolip-core.git to commit 1ebad67db95faa13acb7ff5eff460aa841757a37\n","  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n","  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n","  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n","Collecting geofractal@ git+https://github.com/AbstractEyes/geofractal.git (from geolip-core==0.2.0)\n","  Cloning https://github.com/AbstractEyes/geofractal.git to /tmp/pip-install-5ujbswkn/geofractal_ade002c1d9624eb698edfc5b29fe5a41\n","  Running command git clone --filter=blob:none --quiet https://github.com/AbstractEyes/geofractal.git /tmp/pip-install-5ujbswkn/geofractal_ade002c1d9624eb698edfc5b29fe5a41\n","  Resolved https://github.com/AbstractEyes/geofractal.git to commit 323bb0591430c2d20e634174d0cc7294a778c01b\n","  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n","  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n","  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n","Requirement already satisfied: torch>=2.0 in /usr/local/lib/python3.12/dist-packages (from geolip-core==0.2.0) (2.10.0+cu128)\n","Requirement already satisfied: filelock in /usr/local/lib/python3.12/dist-packages (from torch>=2.0->geolip-core==0.2.0) (3.25.2)\n","Requirement already satisfied: typing-extensions>=4.10.0 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0->geolip-core==0.2.0) (4.15.0)\n","Requirement already satisfied: setuptools in /usr/local/lib/python3.12/dist-packages (from torch>=2.0->geolip-core==0.2.0) (75.2.0)\n","Requirement already satisfied: sympy>=1.13.3 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0->geolip-core==0.2.0) (1.14.0)\n","Requirement already satisfied: networkx>=2.5.1 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0->geolip-core==0.2.0) (3.6.1)\n","Requirement already satisfied: jinja2 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0->geolip-core==0.2.0) (3.1.6)\n","Requirement already satisfied: fsspec>=0.8.5 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0->geolip-core==0.2.0) (2025.3.0)\n","Requirement already satisfied: cuda-bindings==12.9.4 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0->geolip-core==0.2.0) (12.9.4)\n","Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.8.93 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0->geolip-core==0.2.0) (12.8.93)\n","Requirement already satisfied: nvidia-cuda-runtime-cu12==12.8.90 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0->geolip-core==0.2.0) (12.8.90)\n","Requirement already satisfied: nvidia-cuda-cupti-cu12==12.8.90 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0->geolip-core==0.2.0) (12.8.90)\n","Requirement already satisfied: nvidia-cudnn-cu12==9.10.2.21 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0->geolip-core==0.2.0) (9.10.2.21)\n","Requirement already satisfied: nvidia-cublas-cu12==12.8.4.1 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0->geolip-core==0.2.0) (12.8.4.1)\n","Requirement already satisfied: nvidia-cufft-cu12==11.3.3.83 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0->geolip-core==0.2.0) (11.3.3.83)\n","Requirement already satisfied: nvidia-curand-cu12==10.3.9.90 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0->geolip-core==0.2.0) (10.3.9.90)\n","Requirement already satisfied: nvidia-cusolver-cu12==11.7.3.90 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0->geolip-core==0.2.0) (11.7.3.90)\n","Requirement already satisfied: nvidia-cusparse-cu12==12.5.8.93 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0->geolip-core==0.2.0) (12.5.8.93)\n","Requirement already satisfied: nvidia-cusparselt-cu12==0.7.1 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0->geolip-core==0.2.0) (0.7.1)\n","Requirement already satisfied: nvidia-nccl-cu12==2.27.5 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0->geolip-core==0.2.0) (2.27.5)\n","Requirement already satisfied: nvidia-nvshmem-cu12==3.4.5 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0->geolip-core==0.2.0) (3.4.5)\n","Requirement already satisfied: nvidia-nvtx-cu12==12.8.90 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0->geolip-core==0.2.0) (12.8.90)\n","Requirement already satisfied: nvidia-nvjitlink-cu12==12.8.93 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0->geolip-core==0.2.0) (12.8.93)\n","Requirement already satisfied: nvidia-cufile-cu12==1.13.1.3 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0->geolip-core==0.2.0) (1.13.1.3)\n","Requirement already satisfied: triton==3.6.0 in /usr/local/lib/python3.12/dist-packages (from torch>=2.0->geolip-core==0.2.0) (3.6.0)\n","Requirement already satisfied: cuda-pathfinder~=1.1 in /usr/local/lib/python3.12/dist-packages (from cuda-bindings==12.9.4->torch>=2.0->geolip-core==0.2.0) (1.5.0)\n","Collecting geometricvocab@ git+https://github.com/AbstractEyes/lattice_vocabulary.git (from geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0)\n","  Cloning https://github.com/AbstractEyes/lattice_vocabulary.git to /tmp/pip-install-5ujbswkn/geometricvocab_362ebdca644a440d946edb2e2432d3c0\n","  Running command git clone --filter=blob:none --quiet https://github.com/AbstractEyes/lattice_vocabulary.git /tmp/pip-install-5ujbswkn/geometricvocab_362ebdca644a440d946edb2e2432d3c0\n","  Resolved https://github.com/AbstractEyes/lattice_vocabulary.git to commit bb3d81dbc95a848f33815983059fdbe425ca2d1a\n","  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n","  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n","  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n","Collecting wide_compiler@ git+https://github.com/AbstractEyes/pytorch-parallel-compiler.git (from geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0)\n","  Cloning https://github.com/AbstractEyes/pytorch-parallel-compiler.git to /tmp/pip-install-5ujbswkn/wide-compiler_62bce6c2cf3140db951e72e4c1e4b395\n","  Running command git clone --filter=blob:none --quiet https://github.com/AbstractEyes/pytorch-parallel-compiler.git /tmp/pip-install-5ujbswkn/wide-compiler_62bce6c2cf3140db951e72e4c1e4b395\n","  Resolved https://github.com/AbstractEyes/pytorch-parallel-compiler.git to commit 70bcba4f2cdddb26909df59856eb6afead753826\n","  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n","  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n","  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n","Requirement already satisfied: numpy>=1.23 in /usr/local/lib/python3.12/dist-packages (from geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (2.0.2)\n","Requirement already satisfied: datasets>=2.15 in /usr/local/lib/python3.12/dist-packages (from geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (4.0.0)\n","Requirement already satisfied: pyarrow>=15.0.0 in /usr/local/lib/python3.12/dist-packages (from datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (18.1.0)\n","Requirement already satisfied: dill<0.3.9,>=0.3.0 in /usr/local/lib/python3.12/dist-packages (from datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (0.3.8)\n","Requirement already satisfied: pandas in /usr/local/lib/python3.12/dist-packages (from datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (2.2.2)\n","Requirement already satisfied: requests>=2.32.2 in /usr/local/lib/python3.12/dist-packages (from datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (2.32.4)\n","Requirement already satisfied: tqdm>=4.66.3 in /usr/local/lib/python3.12/dist-packages (from datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (4.67.3)\n","Requirement already satisfied: xxhash in /usr/local/lib/python3.12/dist-packages (from datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (3.6.0)\n","Requirement already satisfied: multiprocess<0.70.17 in /usr/local/lib/python3.12/dist-packages (from datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (0.70.16)\n","Requirement already satisfied: huggingface-hub>=0.24.0 in /usr/local/lib/python3.12/dist-packages (from datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (1.8.0)\n","Requirement already satisfied: packaging in /usr/local/lib/python3.12/dist-packages (from datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (26.0)\n","Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.12/dist-packages (from datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (6.0.3)\n","Requirement already satisfied: mpmath<1.4,>=1.1.0 in /usr/local/lib/python3.12/dist-packages (from sympy>=1.13.3->torch>=2.0->geolip-core==0.2.0) (1.3.0)\n","Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.12/dist-packages (from jinja2->torch>=2.0->geolip-core==0.2.0) (3.0.3)\n","Requirement already satisfied: aiohttp!=4.0.0a0,!=4.0.0a1 in /usr/local/lib/python3.12/dist-packages (from fsspec[http]<=2025.3.0,>=2023.1.0->datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (3.13.4)\n","Requirement already satisfied: hf-xet<2.0.0,>=1.4.2 in /usr/local/lib/python3.12/dist-packages (from huggingface-hub>=0.24.0->datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (1.4.2)\n","Requirement already satisfied: httpx<1,>=0.23.0 in /usr/local/lib/python3.12/dist-packages (from huggingface-hub>=0.24.0->datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (0.28.1)\n","Requirement already satisfied: typer in /usr/local/lib/python3.12/dist-packages (from huggingface-hub>=0.24.0->datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (0.24.1)\n","Requirement already satisfied: charset_normalizer<4,>=2 in /usr/local/lib/python3.12/dist-packages (from requests>=2.32.2->datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (3.4.6)\n","Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.12/dist-packages (from requests>=2.32.2->datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (3.11)\n","Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.12/dist-packages (from requests>=2.32.2->datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (2.5.0)\n","Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.12/dist-packages (from requests>=2.32.2->datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (2026.2.25)\n","Requirement already satisfied: python-dateutil>=2.8.2 in /usr/local/lib/python3.12/dist-packages (from pandas->datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (2.9.0.post0)\n","Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.12/dist-packages (from pandas->datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (2025.2)\n","Requirement already satisfied: tzdata>=2022.7 in /usr/local/lib/python3.12/dist-packages (from pandas->datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (2025.3)\n","Requirement already satisfied: aiohappyeyeballs>=2.5.0 in /usr/local/lib/python3.12/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<=2025.3.0,>=2023.1.0->datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (2.6.1)\n","Requirement already satisfied: aiosignal>=1.4.0 in /usr/local/lib/python3.12/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<=2025.3.0,>=2023.1.0->datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (1.4.0)\n","Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.12/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<=2025.3.0,>=2023.1.0->datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (26.1.0)\n","Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.12/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<=2025.3.0,>=2023.1.0->datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (1.8.0)\n","Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.12/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<=2025.3.0,>=2023.1.0->datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (6.7.1)\n","Requirement already satisfied: propcache>=0.2.0 in /usr/local/lib/python3.12/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<=2025.3.0,>=2023.1.0->datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (0.4.1)\n","Requirement already satisfied: yarl<2.0,>=1.17.0 in /usr/local/lib/python3.12/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<=2025.3.0,>=2023.1.0->datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (1.23.0)\n","Requirement already satisfied: anyio in /usr/local/lib/python3.12/dist-packages (from httpx<1,>=0.23.0->huggingface-hub>=0.24.0->datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (4.13.0)\n","Requirement already satisfied: httpcore==1.* in /usr/local/lib/python3.12/dist-packages (from httpx<1,>=0.23.0->huggingface-hub>=0.24.0->datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (1.0.9)\n","Requirement already satisfied: h11>=0.16 in /usr/local/lib/python3.12/dist-packages (from httpcore==1.*->httpx<1,>=0.23.0->huggingface-hub>=0.24.0->datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (0.16.0)\n","Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.12/dist-packages (from python-dateutil>=2.8.2->pandas->datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (1.17.0)\n","Requirement already satisfied: click>=8.2.1 in /usr/local/lib/python3.12/dist-packages (from typer->huggingface-hub>=0.24.0->datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (8.3.1)\n","Requirement already satisfied: shellingham>=1.3.0 in /usr/local/lib/python3.12/dist-packages (from typer->huggingface-hub>=0.24.0->datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (1.5.4)\n","Requirement already satisfied: rich>=12.3.0 in /usr/local/lib/python3.12/dist-packages (from typer->huggingface-hub>=0.24.0->datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (13.9.4)\n","Requirement already satisfied: annotated-doc>=0.0.2 in /usr/local/lib/python3.12/dist-packages (from typer->huggingface-hub>=0.24.0->datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (0.0.4)\n","Requirement already satisfied: markdown-it-py>=2.2.0 in /usr/local/lib/python3.12/dist-packages (from rich>=12.3.0->typer->huggingface-hub>=0.24.0->datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (4.0.0)\n","Requirement already satisfied: pygments<3.0.0,>=2.13.0 in /usr/local/lib/python3.12/dist-packages (from rich>=12.3.0->typer->huggingface-hub>=0.24.0->datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (2.20.0)\n","Requirement already satisfied: mdurl~=0.1 in /usr/local/lib/python3.12/dist-packages (from markdown-it-py>=2.2.0->rich>=12.3.0->typer->huggingface-hub>=0.24.0->datasets>=2.15->geofractal@ git+https://github.com/AbstractEyes/geofractal.git->geolip-core==0.2.0) (0.1.2)\n","Building wheels for collected packages: geolip-core, geofractal, geometricvocab, wide_compiler\n","  Building wheel for geolip-core (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n","  Created wheel for geolip-core: filename=geolip_core-0.2.0-py3-none-any.whl size=138775 sha256=502d2c6b9822e019ee100bf07bb68aa66720bdd7542e0c68eaa14fe8d37d110a\n","  Stored in directory: /tmp/pip-ephem-wheel-cache-7_let9lk/wheels/72/5a/6a/cd41c6a9ab529f550d5e9a9a8ad6565e92c91620cc98d2135f\n","  Building wheel for geofractal (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n","  Created wheel for geofractal: filename=geofractal-1.2.0-py3-none-any.whl size=1100089 sha256=fc22e8090b2c3d5899ea14cb9b81c9f48231b8b97defb09cadf200360b1e3851\n","  Stored in directory: /tmp/pip-ephem-wheel-cache-7_let9lk/wheels/9d/42/04/523291e555da30d5f0c6528bf0709ac3828b2f5d85ab27dec8\n","  Building wheel for geometricvocab (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n","  Created wheel for geometricvocab: filename=geometricvocab-0.1.2-py3-none-any.whl size=1021226 sha256=24a3d8945315a589ef2bb8121cee166c1055e5cef4dae902e3f47016aaaa0998\n","  Stored in directory: /tmp/pip-ephem-wheel-cache-7_let9lk/wheels/56/ab/b9/e048fe5adb608dbe5936e03c6cbd813ec68b2fee4185d8a531\n","  Building wheel for wide_compiler (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n","  Created wheel for wide_compiler: filename=wide_compiler-0.7.0-py3-none-any.whl size=173287 sha256=13bbd2a09277ebb02fed20a0b5ff3d6255ce528a10e608dbc6e23a657a4ce0b2\n","  Stored in directory: /tmp/pip-ephem-wheel-cache-7_let9lk/wheels/b9/85/69/a8412e42739570934577bb1e071fd1e313abc6268369f19e90\n","Successfully built geolip-core geofractal geometricvocab wide_compiler\n","Installing collected packages: wide_compiler, geometricvocab, geofractal, geolip-core\n","Successfully installed geofractal-1.2.0 geolip-core-0.2.0 geometricvocab-0.1.2 wide_compiler-0.7.0\n"]}]},{"cell_type":"markdown","source":["# transformer smoke test"],"metadata":{"id":"JAzPQnPRAV82"}},{"cell_type":"code","source":["\"\"\"\n","Geometric Transformer — Component Profiling\n","\n","Times each stage of the forward pass to find bottlenecks.\n","Run in Colab after placing geometric_transformer.py in geolip_core.\n","\"\"\"\n","\n","import torch\n","import torch.nn as nn\n","import torch.nn.functional as F\n","import time\n","\n","device = torch.device('cuda')\n","print(f\"GPU: {torch.cuda.get_device_name()}\")\n","print(f\"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB\")\n","\n","from geolip_core.pipeline.components.geometric_transformer import (\n","    GeometricTransformer, GeometricTransformerLayer,\n","    ManifoldProjection, ConstellationObserver, PositionGeometricContext,\n","    ContentAttention, GeometricAttention, CayleyOrthogonal,\n","    QuaternionCompose, FiLMLayer, TorchComponent, BaseTower,\n",")\n","\n","# Also profile the SVD Input stage\n","try:\n","    from geolip_core.core.input.svd import SVDObserver\n","    _HAS_SVD = True\n","except ImportError:\n","    _HAS_SVD = False\n","\n","\n","def sync_time():\n","    torch.cuda.synchronize()\n","    return time.perf_counter()\n","\n","\n","def profile_component(name, fn, warmup=5, repeats=50):\n","    \"\"\"Time a callable with CUDA sync.\"\"\"\n","    for _ in range(warmup):\n","        fn()\n","    torch.cuda.synchronize()\n","\n","    times = []\n","    for _ in range(repeats):\n","        t0 = sync_time()\n","        fn()\n","        t1 = sync_time()\n","        times.append((t1 - t0) * 1000)  # ms\n","\n","    avg = sum(times) / len(times)\n","    std = (sum((t - avg)**2 for t in times) / len(times)) ** 0.5\n","    mn = min(times)\n","    print(f\"  {name:<40s}  {avg:>8.3f} ± {std:>6.3f} ms  (min {mn:.3f})\")\n","    return avg\n","\n","\n","def main():\n","    B, L, D = 128, 65, 256  # CIFAR-100: batch=128, 64 patches + CLS, d=256\n","    n_heads = 8\n","    n_anchors = 16\n","    manifold_dim = 128\n","    n_comp = 4\n","    d_comp = 16\n","    context_dim = 64\n","    quat_dim = 32\n","    conv_channels = 64\n","    svd_rank = 16\n","\n","    x = torch.randn(B, L, D, device=device)\n","    images = torch.randn(B, 3, 32, 32, device=device)\n","\n","    print(f\"\\nProfiling: B={B}, L={L}, D={D}\")\n","    print(f\"{'='*70}\")\n","\n","    # ── 1. SVD Input Stage ──\n","    print(\"\\n── SVD Input Stage ──\")\n","\n","    conv_frontend = nn.Sequential(\n","        nn.Conv2d(3, conv_channels, 3, padding=1, bias=False),\n","        nn.BatchNorm2d(conv_channels), nn.GELU(),\n","        nn.Conv2d(conv_channels, conv_channels, 3, padding=1, bias=False),\n","        nn.BatchNorm2d(conv_channels), nn.GELU(),\n","    ).to(device).eval()\n","\n","    profile_component(\"conv_frontend (3→64→64)\", lambda: conv_frontend(images))\n","\n","    feat = conv_frontend(images)\n","\n","    if _HAS_SVD:\n","        svd_obs = SVDObserver(conv_channels, svd_rank).to(device).eval()\n","        profile_component(\"SVDObserver (rank=16)\", lambda: svd_obs(feat))\n","    else:\n","        # Inline SVD timing\n","        to_svd = nn.Conv2d(conv_channels, svd_rank, 1, bias=False).to(device).eval()\n","        h = to_svd(feat)\n","        h_flat = h.permute(0, 2, 3, 1).reshape(B, 32*32, svd_rank)\n","\n","        def do_svd():\n","            with torch.no_grad():\n","                gram = torch.bmm(h_flat.transpose(1, 2), h_flat)\n","                torch.linalg.eigh(gram)\n","\n","        profile_component(\"1x1 conv (64→16)\", lambda: to_svd(feat))\n","        profile_component(\"gram_eigh_svd (B=128, 1024×16)\", do_svd)\n","\n","    patch_proj = nn.Conv2d(conv_channels, D, kernel_size=4, stride=4, bias=False).to(device).eval()\n","    profile_component(\"patch_proj (conv2d stride=4)\", lambda: patch_proj(feat))\n","\n","    svd_film_gamma = nn.Linear(2*svd_rank+2, D).to(device).eval()\n","    svd_feats = torch.randn(B, 2*svd_rank+2, device=device)\n","    tokens = torch.randn(B, L-1, D, device=device)\n","    profile_component(\"SVD FiLM (broadcast mul+add)\",\n","        lambda: svd_film_gamma(svd_feats).unsqueeze(1) * tokens)\n","\n","    # ── 2. Per-component in one transformer layer ──\n","    print(\"\\n── Transformer Layer Components ──\")\n","\n","    proj = ManifoldProjection('p', D, manifold_dim).to(device).eval()\n","    profile_component(\"ManifoldProjection (256→128, norm, L2)\", lambda: proj(x))\n","\n","    emb = proj(x)\n","    emb_flat = emb.reshape(B * L, -1)\n","\n","    obs = ConstellationObserver(\n","        dim=manifold_dim, n_anchors=n_anchors,\n","        n_comp=n_comp, d_comp=d_comp).to(device).eval()\n","    profile_component(f\"ConstellationObserver ({B*L}×{manifold_dim}, A={n_anchors})\",\n","        lambda: obs.observe(emb_flat))\n","\n","    obs_out = obs.observe(emb_flat)\n","    pw_dim = obs.curation.patchwork.output_dim\n","\n","    ctx_mod = PositionGeometricContext(\n","        'c', n_anchors, pw_dim, manifold_dim, context_dim).to(device).eval()\n","    profile_component(\"PositionGeometricContext → FiLM ctx\", lambda: ctx_mod(obs_out))\n","\n","    geo_ctx = ctx_mod(obs_out).reshape(B, L, -1)\n","\n","    content = ContentAttention('ca', D, n_heads).to(device).eval()\n","    profile_component(\"ContentAttention (standard MHA)\", lambda: content(x))\n","\n","    geo_attn = GeometricAttention('ga', D, n_heads, context_dim).to(device).eval()\n","    profile_component(\"GeometricAttention (FiLM Q,K + MHA + FiLM FFN)\",\n","        lambda: geo_attn(x, geo_ctx))\n","\n","    cayley = CayleyOrthogonal('cy', D).to(device).eval()\n","    profile_component(f\"CayleyOrthogonal (solve {D}×{D})\", lambda: cayley(x))\n","\n","    quat = QuaternionCompose('q', D, quat_dim).to(device).eval()\n","    a_out = content(x)\n","    b_out = geo_attn(x, geo_ctx)\n","    profile_component(\"QuaternionCompose (4-arm Hamilton)\",\n","        lambda: quat(a_out, b_out, a_out - b_out, a_out * b_out))\n","\n","    decode = nn.Sequential(\n","        nn.Linear(quat_dim * 4, D), nn.GELU(), nn.LayerNorm(D)).to(device).eval()\n","    composed = quat(a_out, b_out, a_out - b_out, a_out * b_out)\n","    profile_component(\"decode (Linear→GELU→LN)\", lambda: decode(composed))\n","\n","    gate = nn.Sequential(nn.Linear(D * 2, D), nn.Sigmoid()).to(device).eval()\n","    decoded = decode(composed)\n","    profile_component(\"gate (Linear→Sigmoid)\", lambda: gate(torch.cat([x, decoded], dim=-1)))\n","\n","    # ── 3. Full layer ──\n","    print(\"\\n── Full Layer ──\")\n","\n","    layer = GeometricTransformerLayer(\n","        'L0', D, n_heads, n_anchors, manifold_dim,\n","        n_comp, d_comp, context_dim, quat_dim).to(device).eval()\n","    if hasattr(layer, 'network_to'):\n","        layer.network_to(device=device, strict=False)\n","\n","    layer_time = profile_component(\"GeometricTransformerLayer (full)\", lambda: layer(x))\n","\n","    # ── 4. Full model ──\n","    print(\"\\n── Full Model (4 layers) ──\")\n","\n","    model = GeometricTransformer(\n","        'geo', d_model=D, n_heads=n_heads, n_layers=4,\n","        n_anchors=n_anchors, manifold_dim=manifold_dim,\n","        n_comp=n_comp, d_comp=d_comp, context_dim=context_dim,\n","        quat_dim=quat_dim).to(device).eval()\n","    if hasattr(model, 'network_to'):\n","        model.network_to(device=device, strict=False)\n","\n","    model_time = profile_component(\"GeometricTransformer (4 layers)\", lambda: model(x))\n","\n","    # With geo_state return\n","    profile_component(\"GeometricTransformer (4L + geo_state)\",\n","        lambda: model(x, return_geo_state=True))\n","\n","    # ── 5. Backward — run separately (needs grad) ──\n","    print(\"\\n── Backward Pass ──\")\n","    print(\"  (run separately — needs grad context)\")\n","\n","    # ── 6. Cayley bottleneck analysis ──\n","    print(\"\\n── Cayley Solve Scaling ──\")\n","    for d in [64, 128, 256, 512]:\n","        cy = CayleyOrthogonal('c', d).to(device).eval()\n","        x_d = torch.randn(B, L, d, device=device)\n","        profile_component(f\"CayleyOrthogonal d={d}\", lambda: cy(x_d), repeats=30)\n","\n","    # ── 7. Observer scaling ──\n","    print(\"\\n── ConstellationObserver Scaling ──\")\n","    for na in [8, 16, 32, 64]:\n","        o = ConstellationObserver(\n","            dim=manifold_dim, n_anchors=na,\n","            n_comp=4, d_comp=16).to(device).eval()\n","        profile_component(f\"Observer A={na}\", lambda: o.observe(emb_flat), repeats=30)\n","\n","    # ── Summary ──\n","    print(f\"\\n{'='*70}\")\n","    print(f\"  Single layer:  {layer_time:.1f} ms\")\n","    print(f\"  4-layer model: {model_time:.1f} ms\")\n","    print(f\"  Per-epoch est: {model_time * (50000 / B) / 1000:.1f} s  \"\n","          f\"({model_time * (50000 / B) / 60000:.1f} min)\")\n","    print(f\"  100 epochs:    {model_time * (50000 / B) * 100 / 60000:.1f} min\")\n","    print(f\"{'='*70}\")\n","\n","\n","if __name__ == '__main__':\n","    with torch.no_grad():\n","        main()\n","\n","    # Backward profiling needs grad context\n","    print(\"\\n── Backward Pass ──\")\n","    B, L, D = 128, 65, 256\n","    model = GeometricTransformer(\n","        'geo_bwd', d_model=D, n_heads=8, n_layers=4,\n","        n_anchors=16, manifold_dim=128, n_comp=4, d_comp=16,\n","        context_dim=64, quat_dim=32).to(device)\n","    if hasattr(model, 'network_to'):\n","        model.network_to(device=device, strict=False)\n","    model.train()\n","    x_grad = torch.randn(B, L, D, device=device, requires_grad=True)\n","\n","    def fwd_bwd():\n","        out = model(x_grad)\n","        loss = out.sum()\n","        loss.backward()\n","        model.zero_grad()\n","        if x_grad.grad is not None:\n","            x_grad.grad.zero_()\n","\n","    profile_component(\"forward + backward (4 layers)\", fwd_bwd)\n","    print(\"=\" * 70)"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"uNcr4fzSFzmV","executionInfo":{"status":"ok","timestamp":1774687338894,"user_tz":420,"elapsed":4782,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"c580ab89-ef38-44e4-b351-5eff464bfee8"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["GPU: NVIDIA RTX PRO 6000 Blackwell Server Edition\n","VRAM: 102.0 GB\n","\n","Profiling: B=128, L=65, D=256\n","======================================================================\n","\n","── SVD Input Stage ──\n","  conv_frontend (3→64→64)                      0.240 ±  0.002 ms  (min 0.237)\n","  SVDObserver (rank=16)                        0.469 ±  0.008 ms  (min 0.457)\n","  patch_proj (conv2d stride=4)                 0.068 ±  0.002 ms  (min 0.065)\n","  SVD FiLM (broadcast mul+add)                 0.041 ±  0.001 ms  (min 0.041)\n","\n","── Transformer Layer Components ──\n","  ManifoldProjection (256→128, norm, L2)       0.064 ±  0.001 ms  (min 0.062)\n","  ConstellationObserver (8320×128, A=16)       0.391 ±  0.025 ms  (min 0.369)\n","  PositionGeometricContext → FiLM ctx          0.108 ±  0.003 ms  (min 0.104)\n","  ContentAttention (standard MHA)              0.546 ±  0.003 ms  (min 0.542)\n","  GeometricAttention (FiLM Q,K + MHA + FiLM FFN)     0.770 ±  0.004 ms  (min 0.763)\n","  CayleyOrthogonal (solve 256×256)             0.613 ±  0.019 ms  (min 0.606)\n","  QuaternionCompose (4-arm Hamilton)           0.231 ±  0.004 ms  (min 0.226)\n","  decode (Linear→GELU→LN)                      0.055 ±  0.002 ms  (min 0.054)\n","  gate (Linear→Sigmoid)                        0.085 ±  0.001 ms  (min 0.084)\n","\n","── Full Layer ──\n"]},{"output_type":"stream","name":"stderr","text":["/tmp/ipykernel_60296/3447736175.py:169: UserWarning: Router 'L0': Using torch.to() bypasses router hardware safeguards. Use network_to() for full cohesion, or initialize with strict=False to silence this warning.\n","  n_comp, d_comp, context_dim, quat_dim).to(device).eval()\n"]},{"output_type":"stream","name":"stdout","text":["  GeometricTransformerLayer (full)             2.900 ±  0.022 ms  (min 2.868)\n","\n","── Full Model (4 layers) ──\n"]},{"output_type":"stream","name":"stderr","text":["/tmp/ipykernel_60296/3447736175.py:182: UserWarning: Router 'geo': Using torch.to() bypasses router hardware safeguards. Use network_to() for full cohesion, or initialize with strict=False to silence this warning.\n","  quat_dim=quat_dim).to(device).eval()\n"]},{"output_type":"stream","name":"stdout","text":["  GeometricTransformer (4 layers)             13.354 ±  0.040 ms  (min 13.311)\n","  GeometricTransformer (4L + geo_state)       13.361 ±  0.035 ms  (min 13.313)\n","\n","── Backward Pass ──\n","  (run separately — needs grad context)\n","\n","── Cayley Solve Scaling ──\n","  CayleyOrthogonal d=64                        0.208 ±  0.003 ms  (min 0.204)\n","  CayleyOrthogonal d=128                       0.324 ±  0.003 ms  (min 0.317)\n","  CayleyOrthogonal d=256                       0.611 ±  0.007 ms  (min 0.605)\n","  CayleyOrthogonal d=512                       1.374 ±  0.004 ms  (min 1.365)\n","\n","── ConstellationObserver Scaling ──\n","  Observer A=8                                 0.398 ±  0.030 ms  (min 0.372)\n","  Observer A=16                                0.396 ±  0.032 ms  (min 0.375)\n","  Observer A=32                                0.413 ±  0.031 ms  (min 0.394)\n","  Observer A=64                                0.415 ±  0.029 ms  (min 0.394)\n","\n","======================================================================\n","  Single layer:  2.9 ms\n","  4-layer model: 13.4 ms\n","  Per-epoch est: 5.2 s  (0.1 min)\n","  100 epochs:    8.7 min\n","======================================================================\n","\n","── Backward Pass ──\n"]},{"output_type":"stream","name":"stderr","text":["/tmp/ipykernel_60296/3447736175.py:231: UserWarning: Router 'geo_bwd': Using torch.to() bypasses router hardware safeguards. Use network_to() for full cohesion, or initialize with strict=False to silence this warning.\n","  context_dim=64, quat_dim=32).to(device)\n"]},{"output_type":"stream","name":"stdout","text":["  forward + backward (4 layers)               46.642 ±  0.048 ms  (min 46.550)\n","======================================================================\n"]}]},{"cell_type":"code","source":["\"\"\"\n","Geometric Transformer — CIFAR-100 Training\n","\n","Patches 32×32 images into tokens, feeds through the geometric\n","transformer with constellation-routed dual-stream attention.\n","\n","Patch strategy: 4×4 patches = 64 tokens per image\n","Each patch: (3, 4, 4) → flatten → project to d_model\n","CLS token prepended, classification from CLS output.\n","\n","!pip install geolip-core torchvision tqdm\n","\"\"\"\n","\n","import torch\n","import torch.nn as nn\n","import torch.nn.functional as F\n","import numpy as np\n","import time, json\n","from pathlib import Path\n","from tqdm.auto import tqdm\n","\n","device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n","print(f\"Device: {device}\")\n","if device.type == 'cuda':\n","    print(f\"  GPU: {torch.cuda.get_device_name()}\")\n","    print(f\"  VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB\")\n","\n","# Import the geometric transformer from geolip_core\n","from geolip_core.pipeline.components.geometric_transformer import (\n","    GeometricTransformer, TorchComponent, BaseTower\n",")\n","torch.set_float32_matmul_precision('high')\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# CONFIG\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","CONFIG = {\n","    # Model\n","    'd_model': 256,\n","    'n_heads': 8,\n","    'n_layers': 4,\n","    'n_anchors': 16,\n","    'manifold_dim': 128,\n","    'n_comp': 4,\n","    'd_comp': 16,\n","    'context_dim': 64,\n","    'quat_dim': 32,\n","    'dropout': 0.1,\n","\n","    # Input stage\n","    'patch_size': 4,       # 32/4 = 8×8 = 64 patches\n","    'img_size': 32,\n","    'in_channels': 3,\n","    'conv_channels': 64,   # conv frontend output channels\n","    'svd_rank': 32,        # SVD projection rank (≤32 for sub-ms)\n","\n","    # Training\n","    'epochs': 100,\n","    'batch_size': 1024,\n","    'lr': 3e-4,\n","    'weight_decay': 0.05,\n","    'warmup_epochs': 5,\n","    'label_smoothing': 0.1,\n","    'num_workers': 8,\n","\n","    # Data\n","    'num_classes': 100,\n","}\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# INPUT STAGE — conv frontend + SVD structural observation → tokens\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","try:\n","    from geolip_core.core.input.svd import SVDObserver\n","    _HAS_SVD = True\n","except ImportError:\n","    _HAS_SVD = False\n","\n","    class SVDObserver(nn.Module):\n","        \"\"\"Fallback SVDObserver matching geolip_core.core.input.svd interface.\"\"\"\n","        def __init__(self, in_channels, svd_rank=24):\n","            super().__init__()\n","            self.svd_rank = svd_rank\n","            self.to_svd = nn.Conv2d(in_channels, svd_rank, 1, bias=False)\n","            self.register_buffer('ema_s', torch.ones(svd_rank))\n","            self.register_buffer('ema_vh_flat', torch.eye(svd_rank).reshape(-1))\n","            self.ema_momentum = 0.99\n","\n","        def extract_features(self, S, Vh):\n","            B, k = S.shape\n","            S_safe = S.clamp(min=1e-6)\n","            s_norm = S_safe / (S_safe.sum(dim=-1, keepdim=True) + 1e-8)\n","            vh_diag = Vh.diagonal(dim1=-2, dim2=-1)\n","            vh_offdiag = (Vh.pow(2).sum((-2, -1)) - vh_diag.pow(2).sum(-1)).unsqueeze(-1).clamp(min=0)\n","            s_ent = -(s_norm * torch.log(s_norm.clamp(min=1e-8))).sum(-1, keepdim=True)\n","            out = torch.cat([s_norm, vh_diag, vh_offdiag, s_ent], dim=-1)\n","            return torch.where(torch.isfinite(out), out, torch.zeros_like(out))\n","\n","        def compute_novelty(self, S):\n","            return S - self.ema_s.clone().unsqueeze(0)\n","\n","        def forward(self, x):\n","            B, C, H, W = x.shape\n","            h = self.to_svd(x)\n","            h_flat = h.permute(0, 2, 3, 1).reshape(B, H * W, self.svd_rank)\n","            with torch.amp.autocast('cuda', enabled=False):\n","                with torch.no_grad():\n","                    gram = torch.bmm(h_flat.float().transpose(1, 2), h_flat.float())\n","                    evals, evecs = torch.linalg.eigh(gram)\n","                    evals = evals.flip(-1).clamp(min=1e-12)\n","                    S = evals.sqrt()\n","                    Vh = evecs.flip(-1).transpose(-2, -1)\n","                    S = torch.where(torch.isfinite(S), S, torch.ones_like(S))\n","                    Vh = torch.where(torch.isfinite(Vh), Vh, torch.zeros_like(Vh))\n","            features = self.extract_features(S, Vh)\n","            novelty = self.compute_novelty(S)\n","            return S, Vh, features, novelty\n","\n","        @torch.no_grad()\n","        def update_ema(self, S, Vh):\n","            m = self.ema_momentum\n","            self.ema_s.mul_(m).add_(S.detach().mean(0), alpha=1-m)\n","            self.ema_vh_flat.mul_(m).add_(Vh.detach().mean(0).reshape(-1), alpha=1-m)\n","\n","        @property\n","        def feature_dim(self):\n","            return 2 * self.svd_rank + 2\n","\n","\n","class ConvSVDPatchEmbedding(TorchComponent):\n","    \"\"\"Input stage: conv frontend → SVDObserver → patch tokens.\n","\n","    Pipeline:\n","        Image (B, 3, 32, 32)\n","          → Conv frontend (3→conv_channels, 2 layers, preserves spatial)\n","          → SVDObserver: 1×1 conv → svd_rank, gram_eigh_svd decomposition\n","              → S (singular values = energy distribution)\n","              → Vh (rotation = structural orientation)\n","              → features (2k+2 compact summary)\n","              → novelty (EMA deviation)\n","          → Patch projection: Conv2d(conv_channels, d_model, stride=patch_size)\n","              → (B, d_model, H/p, W/p) → reshape to (B, N, d_model) tokens\n","          → SVD context FiLM: modulate tokens with global structural info\n","          → CLS token + position embeddings\n","\n","    The constellation now triangulates against tokens that carry\n","    decomposed spatial structure, not random pixel projections.\n","\n","    Args:\n","        name: component identity\n","        img_size: input spatial size\n","        patch_size: patch spatial size (tokens = (img_size/patch_size)^2)\n","        in_channels: input image channels\n","        conv_channels: conv frontend output channels\n","        d_model: token embedding dimension\n","        svd_rank: SVD projection rank (≤32 for sub-ms)\n","    \"\"\"\n","    def __init__(self, name, img_size=32, patch_size=4, in_channels=3,\n","                 conv_channels=64, d_model=256, svd_rank=16):\n","        super().__init__(name)\n","        self.patch_size = patch_size\n","        self.n_patches = (img_size // patch_size) ** 2\n","        self.d_model = d_model\n","        self.svd_rank = svd_rank\n","\n","        # ── Conv frontend: extract spatial features ──\n","        self.conv_frontend = nn.Sequential(\n","            nn.Conv2d(in_channels, conv_channels, 3, padding=1, bias=False),\n","            nn.BatchNorm2d(conv_channels),\n","            nn.GELU(),\n","            nn.Conv2d(conv_channels, conv_channels, 3, padding=1, bias=False),\n","            nn.BatchNorm2d(conv_channels),\n","            nn.GELU(),\n","        )\n","\n","        # ── SVD structural observation ──\n","        self.svd_observer = SVDObserver(conv_channels, svd_rank)\n","\n","        # ── Patch projection: conv features → tokens ──\n","        # Conv2d with stride=patch_size is equivalent to unfold + project\n","        # but more efficient and lets the kernel learn spatial combinations\n","        self.patch_proj = nn.Conv2d(\n","            conv_channels, d_model, kernel_size=patch_size,\n","            stride=patch_size, bias=False)\n","        self.patch_norm = nn.LayerNorm(d_model)\n","\n","        # ── SVD context → FiLM on tokens ──\n","        # Global structural info modulates every token\n","        svd_feat_dim = self.svd_observer.feature_dim  # 2*svd_rank + 2\n","        self.svd_to_gamma = nn.Linear(svd_feat_dim, d_model)\n","        self.svd_to_beta = nn.Linear(svd_feat_dim, d_model)\n","        # Identity-init: γ=1, β=0 at start\n","        nn.init.zeros_(self.svd_to_gamma.weight); nn.init.ones_(self.svd_to_gamma.bias)\n","        nn.init.zeros_(self.svd_to_beta.weight); nn.init.zeros_(self.svd_to_beta.bias)\n","\n","        # ── CLS token + position embeddings ──\n","        self.cls_token = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)\n","        self.pos_embed = nn.Parameter(\n","            torch.randn(1, self.n_patches + 1, d_model) * 0.02)\n","\n","    def forward(self, x):\n","        \"\"\"\n","        Args:\n","            x: (B, C, H, W) input images\n","\n","        Returns:\n","            tokens: (B, N+1, d_model) token sequence with CLS\n","            svd_state: dict with SVD intermediates for diagnostics/loss\n","        \"\"\"\n","        B = x.shape[0]\n","\n","        # ── Conv frontend: raw pixels → spatial features ──\n","        feat = self.conv_frontend(x)  # (B, conv_channels, H, W)\n","\n","        # ── SVD: structural decomposition of feature map ──\n","        S, Vh, svd_features, novelty = self.svd_observer(feat)\n","\n","        # ── Patch projection: feature map → tokens ──\n","        tokens = self.patch_proj(feat)  # (B, d_model, H/p, W/p)\n","        tokens = tokens.flatten(2).transpose(1, 2)  # (B, N, d_model)\n","        tokens = self.patch_norm(tokens)\n","\n","        # ── FiLM: modulate tokens with global SVD context ──\n","        # svd_features is (B, 2k+2) — broadcast over all positions\n","        gamma = self.svd_to_gamma(svd_features).unsqueeze(1)  # (B, 1, d_model)\n","        beta = self.svd_to_beta(svd_features).unsqueeze(1)    # (B, 1, d_model)\n","        tokens = gamma * tokens + beta\n","\n","        # ── CLS token + position embeddings ──\n","        cls = self.cls_token.expand(B, -1, -1)\n","        tokens = torch.cat([cls, tokens], dim=1)  # (B, N+1, d_model)\n","        tokens = tokens + self.pos_embed\n","\n","        svd_state = {\n","            'singular_values': S,\n","            'Vh': Vh,\n","            'svd_features': svd_features,\n","            'novelty': novelty,\n","        }\n","\n","        # Update SVD EMA during training\n","        if self.training:\n","            self.svd_observer.update_ema(S, Vh)\n","\n","        return tokens, svd_state\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# CLASSIFICATION MODEL — patch embed + geometric transformer + head\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","class GeoViTClassifier(BaseTower):\n","    \"\"\"Geometric Vision Transformer for classification.\n","\n","    A BaseTower that composes:\n","        'patch_embed'  → ConvSVDPatchEmbedding (Input stage: conv + SVD + patch)\n","        'transformer'  → GeometricTransformer (processing tower)\n","        'head'         → classification head (Distinction stage)\n","\n","    Access:\n","        model['patch_embed']   → ConvSVDPatchEmbedding\n","        model['transformer']   → GeometricTransformer\n","        model['head']          → classifier MLP\n","    \"\"\"\n","    def __init__(self, name, config):\n","        super().__init__(name)\n","        self.config = config\n","\n","        # Input stage: conv frontend + SVD observation → patch tokens\n","        self.attach('patch_embed', ConvSVDPatchEmbedding(\n","            'patch_embed',\n","            img_size=config['img_size'],\n","            patch_size=config['patch_size'],\n","            in_channels=config['in_channels'],\n","            conv_channels=config['conv_channels'],\n","            d_model=config['d_model'],\n","            svd_rank=config['svd_rank'],\n","        ))\n","\n","        # Processing tower: geometric transformer\n","        self.attach('transformer', GeometricTransformer(\n","            'geo_cifar',\n","            d_model=config['d_model'],\n","            n_heads=config['n_heads'],\n","            n_layers=config['n_layers'],\n","            n_anchors=config['n_anchors'],\n","            manifold_dim=config['manifold_dim'],\n","            n_comp=config['n_comp'],\n","            d_comp=config['d_comp'],\n","            context_dim=config['context_dim'],\n","            quat_dim=config['quat_dim'],\n","            dropout=config['dropout'],\n","        ))\n","\n","        # Distinction stage: CLS token → class logits\n","        self.attach('head', nn.Sequential(\n","            nn.LayerNorm(config['d_model']),\n","            nn.Linear(config['d_model'], config['d_model']),\n","            nn.GELU(),\n","            nn.Dropout(config['dropout']),\n","            nn.Linear(config['d_model'], config['num_classes']),\n","        ))\n","\n","    def forward(self, x, return_geo_state=False):\n","        \"\"\"\n","        x: (B, 3, 32, 32) → (B, num_classes) logits\n","        \"\"\"\n","        tokens, svd_state = self['patch_embed'](x)  # (B, N+1, d_model), dict\n","\n","        if return_geo_state:\n","            features, geo_states = self['transformer'](tokens, return_geo_state=True)\n","        else:\n","            features = self['transformer'](tokens)\n","\n","        # CLS token is at position 0\n","        cls_out = features[:, 0]  # (B, d_model)\n","        logits = self['head'](cls_out)  # (B, num_classes)\n","\n","        if return_geo_state:\n","            return logits, geo_states, svd_state\n","        return logits\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# DATA\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","def get_dataloaders(config):\n","    import torchvision\n","    import torchvision.transforms as T\n","\n","    train_transform = T.Compose([\n","        T.RandomCrop(32, padding=4),\n","        T.RandomHorizontalFlip(),\n","        T.ToTensor(),\n","        T.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),\n","    ])\n","\n","    test_transform = T.Compose([\n","        T.ToTensor(),\n","        T.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),\n","    ])\n","\n","    train_ds = torchvision.datasets.CIFAR100(\n","        root='./data', train=True, download=True, transform=train_transform)\n","    test_ds = torchvision.datasets.CIFAR100(\n","        root='./data', train=False, download=True, transform=test_transform)\n","\n","    train_loader = torch.utils.data.DataLoader(\n","        train_ds, batch_size=config['batch_size'], shuffle=True,\n","        num_workers=config['num_workers'], pin_memory=True, drop_last=True)\n","    test_loader = torch.utils.data.DataLoader(\n","        test_ds, batch_size=config['batch_size'], shuffle=False,\n","        num_workers=config['num_workers'], pin_memory=True)\n","\n","    return train_loader, test_loader\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# TRAINING\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","def train_epoch(model, loader, optimizer, scheduler, epoch, config):\n","    model.train()\n","    total_loss = 0\n","    correct = 0\n","    total = 0\n","    criterion = nn.CrossEntropyLoss(label_smoothing=config['label_smoothing'])\n","\n","    for images, labels in loader:\n","        images = images.to(device)\n","        labels = labels.to(device)\n","\n","        logits = model(images)\n","        loss = criterion(logits, labels)\n","\n","        optimizer.zero_grad()\n","        loss.backward()\n","        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n","        optimizer.step()\n","        if scheduler is not None:\n","            scheduler.step()\n","\n","        total_loss += loss.item() * images.size(0)\n","        correct += (logits.argmax(1) == labels).sum().item()\n","        total += images.size(0)\n","\n","    return total_loss / total, correct / total\n","\n","\n","@torch.no_grad()\n","def evaluate(model, loader):\n","    model.eval()\n","    correct = 0\n","    total = 0\n","\n","    for images, labels in loader:\n","        images = images.to(device)\n","        labels = labels.to(device)\n","\n","        logits = model(images)\n","        correct += (logits.argmax(1) == labels).sum().item()\n","        total += images.size(0)\n","\n","    return correct / total\n","\n","\n","def main():\n","    config = CONFIG.copy()\n","\n","    print(\"=\" * 60)\n","    print(\"  Geometric Transformer — CIFAR-100\")\n","    print(f\"  Input: conv({config['in_channels']}→{config['conv_channels']}) + \"\n","          f\"SVD(rank={config['svd_rank']}) + \"\n","          f\"{config['patch_size']}×{config['patch_size']} patches = \"\n","          f\"{(config['img_size']//config['patch_size'])**2} tokens + CLS\")\n","    print(f\"  Model: d={config['d_model']}, heads={config['n_heads']}, \"\n","          f\"layers={config['n_layers']}, anchors={config['n_anchors']}\")\n","    print(\"=\" * 60)\n","\n","    # Data\n","    print(\"\\nLoading CIFAR-100...\")\n","    train_loader, test_loader = get_dataloaders(config)\n","    print(f\"  Train: {len(train_loader.dataset):,} | Test: {len(test_loader.dataset):,}\")\n","\n","    # Model\n","    model = GeoViTClassifier('geo_vit_cifar100', config)\n","    if hasattr(model, 'network_to'):\n","        model.network_to(device=device, strict=False)\n","        model.compile()\n","    else:\n","        model = model.to(device)\n","\n","    n_params = sum(p.numel() for p in model.parameters())\n","    n_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)\n","    print(f\"\\n  Total params:     {n_params:,}\")\n","    print(f\"  Trainable params: {n_trainable:,}\")\n","\n","    # Print breakdown\n","    for name, module in model.named_children():\n","        n = sum(p.numel() for p in module.parameters())\n","        if n > 0:\n","            print(f\"    {name:<20s}: {n:,}\")\n","\n","    # Optimizer + scheduler\n","    optimizer = torch.optim.AdamW(\n","        model.parameters(), lr=config['lr'], weight_decay=config['weight_decay'])\n","\n","    total_steps = config['epochs'] * len(train_loader)\n","    warmup_steps = config['warmup_epochs'] * len(train_loader)\n","\n","    def lr_lambda(step):\n","        if step < warmup_steps:\n","            return step / warmup_steps\n","        progress = (step - warmup_steps) / (total_steps - warmup_steps)\n","        return 0.5 * (1 + np.cos(np.pi * progress))\n","\n","    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)\n","\n","    # Training loop\n","    print(f\"\\n{'━'*60}\")\n","    print(f\"  Training for {config['epochs']} epochs\")\n","    print(f\"  Warmup: {config['warmup_epochs']} epochs, \"\n","          f\"LR: {config['lr']}, WD: {config['weight_decay']}\")\n","    print(f\"{'━'*60}\\n\")\n","\n","    best_acc = 0\n","    save_dir = Path('geo_cifar10'); save_dir.mkdir(exist_ok=True)\n","\n","    for epoch in tqdm(range(config['epochs']), desc=\"Epochs\"):\n","        t0 = time.time()\n","\n","        train_loss, train_acc = train_epoch(\n","            model, train_loader, optimizer, scheduler, epoch, config)\n","\n","        test_acc = evaluate(model, test_loader)\n","        elapsed = time.time() - t0\n","\n","        if test_acc > best_acc:\n","            best_acc = test_acc\n","            torch.save({\n","                'state_dict': {k: v.cpu() for k, v in model.state_dict().items()},\n","                'epoch': epoch,\n","                'test_acc': test_acc,\n","                'config': config,\n","            }, save_dir / 'best.pt')\n","\n","        if epoch % 5 == 0 or epoch == config['epochs'] - 1:\n","            lr = optimizer.param_groups[0]['lr']\n","            tqdm.write(\n","                f\"  E{epoch:>3d}  loss={train_loss:.4f}  \"\n","                f\"train={train_acc:.4f}  test={test_acc:.4f}  \"\n","                f\"best={best_acc:.4f}  lr={lr:.6f}  {elapsed:.1f}s\")\n","\n","    # Final summary\n","    print(f\"\\n{'═'*60}\")\n","    print(f\"  CIFAR-10 RESULTS\")\n","    print(f\"{'═'*60}\")\n","    print(f\"  Best test accuracy: {best_acc:.4f} ({best_acc*100:.2f}%)\")\n","    print(f\"  Parameters: {n_params:,}\")\n","    print(f\"  Checkpoint: {save_dir}/best.pt\")\n","\n","    # Geometric state inspection on a test batch\n","    print(f\"\\n  Geometric state inspection:\")\n","    model.eval()\n","    images, labels = next(iter(test_loader))\n","    images = images[:4].to(device)\n","    with torch.no_grad():\n","        logits, geo_states, svd_state = model(images, return_geo_state=True)\n","\n","    # SVD Input stage diagnostics\n","    S = svd_state['singular_values']\n","    s_norm = S / (S.sum(dim=-1, keepdim=True) + 1e-8)\n","    s_ent = -(s_norm * torch.log(s_norm.clamp(min=1e-8))).sum(-1)\n","    tqdm.write(\n","        f\"    SVD Input: S_top3={S[0, :3].tolist()}  \"\n","        f\"entropy={s_ent.mean().item():.3f}  \"\n","        f\"novelty_norm={svd_state['novelty'].norm(dim=-1).mean().item():.3f}\")\n","\n","    for i, gs in enumerate(geo_states):\n","        content = gs['content']\n","        geometric = gs['geometric']\n","\n","        # Measure stream divergence\n","        content_norm = content.norm(dim=-1).mean().item()\n","        geo_norm = geometric.norm(dim=-1).mean().item()\n","        agreement = F.cosine_similarity(\n","            content.reshape(-1, content.shape[-1]),\n","            geometric.reshape(-1, geometric.shape[-1]),\n","            dim=-1).mean().item()\n","\n","        # Anchor utilization\n","        tri = gs['triangulation']\n","        nearest_counts = torch.bincount(\n","            gs['nearest'].reshape(-1),\n","            minlength=tri.shape[-1]).float()\n","        anchor_entropy = -(nearest_counts / nearest_counts.sum() *\n","            torch.log(nearest_counts.clamp(min=1e-8) / nearest_counts.sum())).sum().item()\n","\n","        tqdm.write(\n","            f\"    Layer {i}: ‖content‖={content_norm:.3f}  \"\n","            f\"‖geo‖={geo_norm:.3f}  agreement={agreement:.3f}  \"\n","            f\"anchor_H={anchor_entropy:.3f}\")\n","\n","    print(f\"\\nDone.\")\n","\n","\n","if __name__ == '__main__':\n","    main()"],"metadata":{"colab":{"base_uri":"https://localhost:8080/","referenced_widgets":["a4d0fdf9703149c59898264e8eaf8789","07c1879a8b7049728078d1ca13a84710","c27e770348a148409907dd7e794345ec","0f94add3cab0471c922358e0d9c7f773","72707cd9fdac4f3e962c8bfbe8ee43ae","bf10d23e46734d5f8f4c2bb61eab28ff","e6c10cf2e2ec4ba882522a811f5635f0","a2fd785ff8774099b265b56442ef3952","aa81112795f3486688f9d5c1015cf9af","7357538e5c964935aad818fbb0d1e29e","7ce103ccbf2145f8a9ba452694f8f314"]},"id":"CMt8bAYRT44m","executionInfo":{"status":"ok","timestamp":1774688094511,"user_tz":420,"elapsed":741011,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"a23de2c8-5070-493f-dc08-476e6bec4750"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["Device: cuda\n","  GPU: NVIDIA RTX PRO 6000 Blackwell Server Edition\n","  VRAM: 102.0 GB\n","============================================================\n","  Geometric Transformer — CIFAR-100\n","  Input: conv(3→64) + SVD(rank=32) + 4×4 patches = 64 tokens + CLS\n","  Model: d=256, heads=8, layers=4, anchors=16\n","============================================================\n","\n","Loading CIFAR-100...\n","  Train: 50,000 | Test: 10,000\n","\n","  Total params:     8,837,092\n","  Trainable params: 8,837,092\n","    components          : 8,837,092\n","\n","━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\n","  Training for 100 epochs\n","  Warmup: 5 epochs, LR: 0.0003, WD: 0.05\n","━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\n","\n"]},{"output_type":"display_data","data":{"text/plain":["Epochs:   0%|          | 0/100 [00:00<?, ?it/s]"],"application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"a4d0fdf9703149c59898264e8eaf8789"}},"metadata":{}},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n","/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n","/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["  E  0  loss=4.5623  train=0.0286  test=0.0554  best=0.0554  lr=0.000060  27.5s\n","  E  5  loss=3.5923  train=0.1899  test=0.1988  best=0.1988  lr=0.000300  7.1s\n","  E 10  loss=2.9821  train=0.3333  test=0.3318  best=0.3318  lr=0.000297  7.1s\n","  E 15  loss=2.6232  train=0.4257  test=0.3976  best=0.3976  lr=0.000290  7.1s\n","  E 20  loss=2.3535  train=0.5002  test=0.4393  best=0.4393  lr=0.000279  7.1s\n","  E 25  loss=2.1309  train=0.5680  test=0.4630  best=0.4630  lr=0.000265  7.1s\n","  E 30  loss=1.8731  train=0.6495  test=0.4775  best=0.4775  lr=0.000248  7.5s\n","  E 35  loss=1.6509  train=0.7278  test=0.4795  best=0.4834  lr=0.000228  7.1s\n","  E 40  loss=1.4667  train=0.7928  test=0.4830  best=0.4891  lr=0.000206  7.1s\n","  E 45  loss=1.3058  train=0.8528  test=0.4818  best=0.4891  lr=0.000182  7.1s\n","  E 50  loss=1.1886  train=0.8974  test=0.4877  best=0.4891  lr=0.000157  7.1s\n","  E 55  loss=1.0898  train=0.9340  test=0.4880  best=0.4905  lr=0.000133  7.1s\n","  E 60  loss=1.0265  train=0.9560  test=0.4916  best=0.4916  lr=0.000108  7.1s\n","  E 65  loss=0.9735  train=0.9724  test=0.4894  best=0.4916  lr=0.000085  7.1s\n","  E 70  loss=0.9357  train=0.9828  test=0.4961  best=0.4961  lr=0.000064  7.1s\n","  E 75  loss=0.9068  train=0.9900  test=0.5000  best=0.5011  lr=0.000045  7.1s\n","  E 80  loss=0.8885  train=0.9935  test=0.4995  best=0.5011  lr=0.000029  7.1s\n","  E 85  loss=0.8775  train=0.9954  test=0.5018  best=0.5022  lr=0.000016  7.1s\n","  E 90  loss=0.8686  train=0.9967  test=0.5024  best=0.5028  lr=0.000007  7.1s\n","  E 95  loss=0.8650  train=0.9974  test=0.5024  best=0.5028  lr=0.000001  7.2s\n","  E 99  loss=0.8659  train=0.9972  test=0.5019  best=0.5028  lr=0.000000  7.2s\n","\n","════════════════════════════════════════════════════════════\n","  CIFAR-10 RESULTS\n","════════════════════════════════════════════════════════════\n","  Best test accuracy: 0.5028 (50.28%)\n","  Parameters: 8,837,092\n","  Checkpoint: geo_cifar10/best.pt\n","\n","  Geometric state inspection:\n","    SVD Input: S_top3=[25.29771614074707, 15.218098640441895, 12.499809265136719]  entropy=3.178  novelty_norm=28.339\n","    Layer 0: ‖content‖=15.470  ‖geo‖=15.478  agreement=0.709  anchor_H=2.051\n","    Layer 1: ‖content‖=15.492  ‖geo‖=15.475  agreement=0.741  anchor_H=2.387\n","    Layer 2: ‖content‖=15.539  ‖geo‖=15.476  agreement=0.719  anchor_H=2.066\n","    Layer 3: ‖content‖=15.578  ‖geo‖=15.506  agreement=0.644  anchor_H=2.405\n","\n","Done.\n"]}]},{"cell_type":"markdown","source":["# cifar100 bigger"],"metadata":{"id":"k5pJjdk3qRUr"}},{"cell_type":"code","source":["\"\"\"\n","Geometric Transformer — CIFAR-100 Training\n","\n","Patches 32×32 images into tokens, feeds through the geometric\n","transformer with constellation-routed dual-stream attention.\n","\n","Patch strategy: 4×4 patches = 64 tokens per image\n","Each patch: (3, 4, 4) → flatten → project to d_model\n","CLS token prepended, classification from CLS output.\n","\n","!pip install geolip-core torchvision tqdm\n","\"\"\"\n","\n","import torch\n","import torch.nn as nn\n","import torch.nn.functional as F\n","import numpy as np\n","import time, json\n","from pathlib import Path\n","from tqdm.auto import tqdm\n","\n","device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n","print(f\"Device: {device}\")\n","if device.type == 'cuda':\n","    print(f\"  GPU: {torch.cuda.get_device_name()}\")\n","    print(f\"  VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB\")\n","\n","# Import the geometric transformer from geolip_core\n","from geolip_core.pipeline.components.geometric_transformer import (\n","    GeometricTransformer, TorchComponent, BaseTower\n",")\n","torch.set_float32_matmul_precision('high')\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# CONFIG\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","CONFIG = {\n","    # Model\n","    'd_model': 256,\n","    'n_heads': 8,\n","    'n_layers': 8,\n","    'n_anchors': 128,\n","    'manifold_dim': 128,\n","    'n_comp': 4,\n","    'd_comp': 16,\n","    'context_dim': 64,\n","    'quat_dim': 32,\n","    'dropout': 0.1,\n","\n","    # Input stage\n","    'patch_size': 4,       # 32/4 = 8×8 = 64 patches\n","    'img_size': 32,\n","    'in_channels': 3,\n","    'conv_channels': 64,   # conv frontend output channels\n","    'svd_rank': 16,        # SVD projection rank (≤32 for sub-ms)\n","\n","    # Training\n","    'epochs': 100,\n","    'batch_size': 1024,\n","    'lr': 1e-3,\n","    'weight_decay': 0.05,\n","    'warmup_epochs': 5,\n","    'label_smoothing': 0.1,\n","    'num_workers': 8,\n","\n","    # Data\n","    'num_classes': 100,\n","}\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# INPUT STAGE — conv frontend + SVD structural observation → tokens\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","try:\n","    from geolip_core.core.input.svd import SVDObserver\n","    _HAS_SVD = True\n","except ImportError:\n","    _HAS_SVD = False\n","\n","    class SVDObserver(nn.Module):\n","        \"\"\"Fallback SVDObserver matching geolip_core.core.input.svd interface.\"\"\"\n","        def __init__(self, in_channels, svd_rank=24):\n","            super().__init__()\n","            self.svd_rank = svd_rank\n","            self.to_svd = nn.Conv2d(in_channels, svd_rank, 1, bias=False)\n","            self.register_buffer('ema_s', torch.ones(svd_rank))\n","            self.register_buffer('ema_vh_flat', torch.eye(svd_rank).reshape(-1))\n","            self.ema_momentum = 0.99\n","\n","        def extract_features(self, S, Vh):\n","            B, k = S.shape\n","            S_safe = S.clamp(min=1e-6)\n","            s_norm = S_safe / (S_safe.sum(dim=-1, keepdim=True) + 1e-8)\n","            vh_diag = Vh.diagonal(dim1=-2, dim2=-1)\n","            vh_offdiag = (Vh.pow(2).sum((-2, -1)) - vh_diag.pow(2).sum(-1)).unsqueeze(-1).clamp(min=0)\n","            s_ent = -(s_norm * torch.log(s_norm.clamp(min=1e-8))).sum(-1, keepdim=True)\n","            out = torch.cat([s_norm, vh_diag, vh_offdiag, s_ent], dim=-1)\n","            return torch.where(torch.isfinite(out), out, torch.zeros_like(out))\n","\n","        def compute_novelty(self, S):\n","            return S - self.ema_s.clone().unsqueeze(0)\n","\n","        def forward(self, x):\n","            B, C, H, W = x.shape\n","            h = self.to_svd(x)\n","            h_flat = h.permute(0, 2, 3, 1).reshape(B, H * W, self.svd_rank)\n","            with torch.amp.autocast('cuda', enabled=False):\n","                with torch.no_grad():\n","                    gram = torch.bmm(h_flat.float().transpose(1, 2), h_flat.float())\n","                    evals, evecs = torch.linalg.eigh(gram)\n","                    evals = evals.flip(-1).clamp(min=1e-12)\n","                    S = evals.sqrt()\n","                    Vh = evecs.flip(-1).transpose(-2, -1)\n","                    S = torch.where(torch.isfinite(S), S, torch.ones_like(S))\n","                    Vh = torch.where(torch.isfinite(Vh), Vh, torch.zeros_like(Vh))\n","            features = self.extract_features(S, Vh)\n","            novelty = self.compute_novelty(S)\n","            return S, Vh, features, novelty\n","\n","        @torch.no_grad()\n","        def update_ema(self, S, Vh):\n","            m = self.ema_momentum\n","            self.ema_s.mul_(m).add_(S.detach().mean(0), alpha=1-m)\n","            self.ema_vh_flat.mul_(m).add_(Vh.detach().mean(0).reshape(-1), alpha=1-m)\n","\n","        @property\n","        def feature_dim(self):\n","            return 2 * self.svd_rank + 2\n","\n","\n","class ConvSVDPatchEmbedding(TorchComponent):\n","    \"\"\"Input stage: conv frontend → SVDObserver → patch tokens.\n","\n","    Pipeline:\n","        Image (B, 3, 32, 32)\n","          → Conv frontend (3→conv_channels, 2 layers, preserves spatial)\n","          → SVDObserver: 1×1 conv → svd_rank, gram_eigh_svd decomposition\n","              → S (singular values = energy distribution)\n","              → Vh (rotation = structural orientation)\n","              → features (2k+2 compact summary)\n","              → novelty (EMA deviation)\n","          → Patch projection: Conv2d(conv_channels, d_model, stride=patch_size)\n","              → (B, d_model, H/p, W/p) → reshape to (B, N, d_model) tokens\n","          → SVD context FiLM: modulate tokens with global structural info\n","          → CLS token + position embeddings\n","\n","    The constellation now triangulates against tokens that carry\n","    decomposed spatial structure, not random pixel projections.\n","\n","    Args:\n","        name: component identity\n","        img_size: input spatial size\n","        patch_size: patch spatial size (tokens = (img_size/patch_size)^2)\n","        in_channels: input image channels\n","        conv_channels: conv frontend output channels\n","        d_model: token embedding dimension\n","        svd_rank: SVD projection rank (≤32 for sub-ms)\n","    \"\"\"\n","    def __init__(self, name, img_size=32, patch_size=4, in_channels=3,\n","                 conv_channels=64, d_model=256, svd_rank=16):\n","        super().__init__(name)\n","        self.patch_size = patch_size\n","        self.n_patches = (img_size // patch_size) ** 2\n","        self.d_model = d_model\n","        self.svd_rank = svd_rank\n","\n","        # ── Conv frontend: extract spatial features ──\n","        self.conv_frontend = nn.Sequential(\n","            nn.Conv2d(in_channels, conv_channels, 3, padding=1, bias=False),\n","            nn.BatchNorm2d(conv_channels),\n","            nn.GELU(),\n","            nn.Conv2d(conv_channels, conv_channels, 3, padding=1, bias=False),\n","            nn.BatchNorm2d(conv_channels),\n","            nn.GELU(),\n","        )\n","\n","        # ── SVD structural observation ──\n","        self.svd_observer = SVDObserver(conv_channels, svd_rank)\n","\n","        # ── Patch projection: conv features → tokens ──\n","        # Conv2d with stride=patch_size is equivalent to unfold + project\n","        # but more efficient and lets the kernel learn spatial combinations\n","        self.patch_proj = nn.Conv2d(\n","            conv_channels, d_model, kernel_size=patch_size,\n","            stride=patch_size, bias=False)\n","        self.patch_norm = nn.LayerNorm(d_model)\n","\n","        # ── SVD context → FiLM on tokens ──\n","        # Global structural info modulates every token\n","        svd_feat_dim = self.svd_observer.feature_dim  # 2*svd_rank + 2\n","        self.svd_to_gamma = nn.Linear(svd_feat_dim, d_model)\n","        self.svd_to_beta = nn.Linear(svd_feat_dim, d_model)\n","        # Identity-init: γ=1, β=0 at start\n","        nn.init.zeros_(self.svd_to_gamma.weight); nn.init.ones_(self.svd_to_gamma.bias)\n","        nn.init.zeros_(self.svd_to_beta.weight); nn.init.zeros_(self.svd_to_beta.bias)\n","\n","        # ── CLS token + position embeddings ──\n","        self.cls_token = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)\n","        self.pos_embed = nn.Parameter(\n","            torch.randn(1, self.n_patches + 1, d_model) * 0.02)\n","\n","    def forward(self, x):\n","        \"\"\"\n","        Args:\n","            x: (B, C, H, W) input images\n","\n","        Returns:\n","            tokens: (B, N+1, d_model) token sequence with CLS\n","            svd_state: dict with SVD intermediates for diagnostics/loss\n","        \"\"\"\n","        B = x.shape[0]\n","\n","        # ── Conv frontend: raw pixels → spatial features ──\n","        feat = self.conv_frontend(x)  # (B, conv_channels, H, W)\n","\n","        # ── SVD: structural decomposition of feature map ──\n","        S, Vh, svd_features, novelty = self.svd_observer(feat)\n","\n","        # ── Patch projection: feature map → tokens ──\n","        tokens = self.patch_proj(feat)  # (B, d_model, H/p, W/p)\n","        tokens = tokens.flatten(2).transpose(1, 2)  # (B, N, d_model)\n","        tokens = self.patch_norm(tokens)\n","\n","        # ── FiLM: modulate tokens with global SVD context ──\n","        # svd_features is (B, 2k+2) — broadcast over all positions\n","        gamma = self.svd_to_gamma(svd_features).unsqueeze(1)  # (B, 1, d_model)\n","        beta = self.svd_to_beta(svd_features).unsqueeze(1)    # (B, 1, d_model)\n","        tokens = gamma * tokens + beta\n","\n","        # ── CLS token + position embeddings ──\n","        cls = self.cls_token.expand(B, -1, -1)\n","        tokens = torch.cat([cls, tokens], dim=1)  # (B, N+1, d_model)\n","        tokens = tokens + self.pos_embed\n","\n","        svd_state = {\n","            'singular_values': S,\n","            'Vh': Vh,\n","            'svd_features': svd_features,\n","            'novelty': novelty,\n","        }\n","\n","        # Update SVD EMA during training\n","        if self.training:\n","            self.svd_observer.update_ema(S, Vh)\n","\n","        return tokens, svd_state\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# CLASSIFICATION MODEL — patch embed + geometric transformer + head\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","class GeoViTClassifier(BaseTower):\n","    \"\"\"Geometric Vision Transformer for classification.\n","\n","    A BaseTower that composes:\n","        'patch_embed'  → ConvSVDPatchEmbedding (Input stage: conv + SVD + patch)\n","        'transformer'  → GeometricTransformer (processing tower)\n","        'head'         → classification head (Distinction stage)\n","\n","    Access:\n","        model['patch_embed']   → ConvSVDPatchEmbedding\n","        model['transformer']   → GeometricTransformer\n","        model['head']          → classifier MLP\n","    \"\"\"\n","    def __init__(self, name, config):\n","        super().__init__(name)\n","        self.config = config\n","\n","        # Input stage: conv frontend + SVD observation → patch tokens\n","        self.attach('patch_embed', ConvSVDPatchEmbedding(\n","            'patch_embed',\n","            img_size=config['img_size'],\n","            patch_size=config['patch_size'],\n","            in_channels=config['in_channels'],\n","            conv_channels=config['conv_channels'],\n","            d_model=config['d_model'],\n","            svd_rank=config['svd_rank'],\n","        ))\n","\n","        # Processing tower: geometric transformer\n","        self.attach('transformer', GeometricTransformer(\n","            'geo_cifar',\n","            d_model=config['d_model'],\n","            n_heads=config['n_heads'],\n","            n_layers=config['n_layers'],\n","            n_anchors=config['n_anchors'],\n","            manifold_dim=config['manifold_dim'],\n","            n_comp=config['n_comp'],\n","            d_comp=config['d_comp'],\n","            context_dim=config['context_dim'],\n","            quat_dim=config['quat_dim'],\n","            dropout=config['dropout'],\n","        ))\n","\n","        # Distinction stage: CLS token → class logits\n","        self.attach('head', nn.Sequential(\n","            nn.LayerNorm(config['d_model']),\n","            nn.Linear(config['d_model'], config['d_model']),\n","            nn.GELU(),\n","            nn.Dropout(config['dropout']),\n","            nn.Linear(config['d_model'], config['num_classes']),\n","        ))\n","\n","    def forward(self, x, return_geo_state=False):\n","        \"\"\"\n","        x: (B, 3, 32, 32) → (B, num_classes) logits\n","        \"\"\"\n","        tokens, svd_state = self['patch_embed'](x)  # (B, N+1, d_model), dict\n","\n","        if return_geo_state:\n","            features, geo_states = self['transformer'](tokens, return_geo_state=True)\n","        else:\n","            features = self['transformer'](tokens)\n","\n","        # CLS token is at position 0\n","        cls_out = features[:, 0]  # (B, d_model)\n","        logits = self['head'](cls_out)  # (B, num_classes)\n","\n","        if return_geo_state:\n","            return logits, geo_states, svd_state\n","        return logits\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# DATA\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","def get_dataloaders(config):\n","    import torchvision\n","    import torchvision.transforms as T\n","\n","    train_transform = T.Compose([\n","        T.RandomCrop(32, padding=4),\n","        T.RandomHorizontalFlip(),\n","        T.ToTensor(),\n","        T.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),\n","    ])\n","\n","    test_transform = T.Compose([\n","        T.ToTensor(),\n","        T.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),\n","    ])\n","\n","    train_ds = torchvision.datasets.CIFAR100(\n","        root='./data', train=True, download=True, transform=train_transform)\n","    test_ds = torchvision.datasets.CIFAR100(\n","        root='./data', train=False, download=True, transform=test_transform)\n","\n","    train_loader = torch.utils.data.DataLoader(\n","        train_ds, batch_size=config['batch_size'], shuffle=True,\n","        num_workers=config['num_workers'], pin_memory=True, drop_last=True)\n","    test_loader = torch.utils.data.DataLoader(\n","        test_ds, batch_size=config['batch_size'], shuffle=False,\n","        num_workers=config['num_workers'], pin_memory=True)\n","\n","    return train_loader, test_loader\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# TRAINING\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","def train_epoch(model, loader, optimizer, scheduler, epoch, config):\n","    model.train()\n","    total_loss = 0\n","    correct = 0\n","    total = 0\n","    criterion = nn.CrossEntropyLoss(label_smoothing=config['label_smoothing'])\n","\n","    for images, labels in loader:\n","        images = images.to(device)\n","        labels = labels.to(device)\n","\n","        logits = model(images)\n","        loss = criterion(logits, labels)\n","\n","        optimizer.zero_grad()\n","        loss.backward()\n","        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n","        optimizer.step()\n","        if scheduler is not None:\n","            scheduler.step()\n","\n","        total_loss += loss.item() * images.size(0)\n","        correct += (logits.argmax(1) == labels).sum().item()\n","        total += images.size(0)\n","\n","    return total_loss / total, correct / total\n","\n","\n","@torch.no_grad()\n","def evaluate(model, loader):\n","    model.eval()\n","    correct = 0\n","    total = 0\n","\n","    for images, labels in loader:\n","        images = images.to(device)\n","        labels = labels.to(device)\n","\n","        logits = model(images)\n","        correct += (logits.argmax(1) == labels).sum().item()\n","        total += images.size(0)\n","\n","    return correct / total\n","\n","\n","def main():\n","    config = CONFIG.copy()\n","\n","    print(\"=\" * 60)\n","    print(\"  Geometric Transformer — CIFAR-100\")\n","    print(f\"  Input: conv({config['in_channels']}→{config['conv_channels']}) + \"\n","          f\"SVD(rank={config['svd_rank']}) + \"\n","          f\"{config['patch_size']}×{config['patch_size']} patches = \"\n","          f\"{(config['img_size']//config['patch_size'])**2} tokens + CLS\")\n","    print(f\"  Model: d={config['d_model']}, heads={config['n_heads']}, \"\n","          f\"layers={config['n_layers']}, anchors={config['n_anchors']}\")\n","    print(\"=\" * 60)\n","\n","    # Data\n","    print(\"\\nLoading CIFAR-100...\")\n","    train_loader, test_loader = get_dataloaders(config)\n","    print(f\"  Train: {len(train_loader.dataset):,} | Test: {len(test_loader.dataset):,}\")\n","\n","    # Model\n","    model = GeoViTClassifier('geo_vit_cifar100', config)\n","    if hasattr(model, 'network_to'):\n","        model.network_to(device=device, strict=False)\n","        model.compile()\n","    else:\n","        model = model.to(device)\n","\n","    n_params = sum(p.numel() for p in model.parameters())\n","    n_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)\n","    print(f\"\\n  Total params:     {n_params:,}\")\n","    print(f\"  Trainable params: {n_trainable:,}\")\n","\n","    # Print breakdown\n","    for name, module in model.named_children():\n","        n = sum(p.numel() for p in module.parameters())\n","        if n > 0:\n","            print(f\"    {name:<20s}: {n:,}\")\n","\n","    # Optimizer + scheduler\n","    optimizer = torch.optim.AdamW(\n","        model.parameters(), lr=config['lr'], weight_decay=config['weight_decay'])\n","\n","    total_steps = config['epochs'] * len(train_loader)\n","    warmup_steps = config['warmup_epochs'] * len(train_loader)\n","\n","    def lr_lambda(step):\n","        if step < warmup_steps:\n","            return step / warmup_steps\n","        progress = (step - warmup_steps) / (total_steps - warmup_steps)\n","        return 0.5 * (1 + np.cos(np.pi * progress))\n","\n","    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)\n","\n","    # Training loop\n","    print(f\"\\n{'━'*60}\")\n","    print(f\"  Training for {config['epochs']} epochs\")\n","    print(f\"  Warmup: {config['warmup_epochs']} epochs, \"\n","          f\"LR: {config['lr']}, WD: {config['weight_decay']}\")\n","    print(f\"{'━'*60}\\n\")\n","\n","    best_acc = 0\n","    save_dir = Path('geo_cifar10'); save_dir.mkdir(exist_ok=True)\n","\n","    for epoch in tqdm(range(config['epochs']), desc=\"Epochs\"):\n","        t0 = time.time()\n","\n","        train_loss, train_acc = train_epoch(\n","            model, train_loader, optimizer, scheduler, epoch, config)\n","\n","        test_acc = evaluate(model, test_loader)\n","        elapsed = time.time() - t0\n","\n","        if test_acc > best_acc:\n","            best_acc = test_acc\n","            torch.save({\n","                'state_dict': {k: v.cpu() for k, v in model.state_dict().items()},\n","                'epoch': epoch,\n","                'test_acc': test_acc,\n","                'config': config,\n","            }, save_dir / 'best.pt')\n","\n","        if epoch % 5 == 0 or epoch == config['epochs'] - 1:\n","            lr = optimizer.param_groups[0]['lr']\n","            tqdm.write(\n","                f\"  E{epoch:>3d}  loss={train_loss:.4f}  \"\n","                f\"train={train_acc:.4f}  test={test_acc:.4f}  \"\n","                f\"best={best_acc:.4f}  lr={lr:.6f}  {elapsed:.1f}s\")\n","\n","    # Final summary\n","    print(f\"\\n{'═'*60}\")\n","    print(f\"  CIFAR-10 RESULTS\")\n","    print(f\"{'═'*60}\")\n","    print(f\"  Best test accuracy: {best_acc:.4f} ({best_acc*100:.2f}%)\")\n","    print(f\"  Parameters: {n_params:,}\")\n","    print(f\"  Checkpoint: {save_dir}/best.pt\")\n","\n","    # Geometric state inspection on a test batch\n","    print(f\"\\n  Geometric state inspection:\")\n","    model.eval()\n","    images, labels = next(iter(test_loader))\n","    images = images[:4].to(device)\n","    with torch.no_grad():\n","        logits, geo_states, svd_state = model(images, return_geo_state=True)\n","\n","    # SVD Input stage diagnostics\n","    S = svd_state['singular_values']\n","    s_norm = S / (S.sum(dim=-1, keepdim=True) + 1e-8)\n","    s_ent = -(s_norm * torch.log(s_norm.clamp(min=1e-8))).sum(-1)\n","    tqdm.write(\n","        f\"    SVD Input: S_top3={S[0, :3].tolist()}  \"\n","        f\"entropy={s_ent.mean().item():.3f}  \"\n","        f\"novelty_norm={svd_state['novelty'].norm(dim=-1).mean().item():.3f}\")\n","\n","    for i, gs in enumerate(geo_states):\n","        content = gs['content']\n","        geometric = gs['geometric']\n","\n","        # Measure stream divergence\n","        content_norm = content.norm(dim=-1).mean().item()\n","        geo_norm = geometric.norm(dim=-1).mean().item()\n","        agreement = F.cosine_similarity(\n","            content.reshape(-1, content.shape[-1]),\n","            geometric.reshape(-1, geometric.shape[-1]),\n","            dim=-1).mean().item()\n","\n","        # Anchor utilization\n","        tri = gs['triangulation']\n","        nearest_counts = torch.bincount(\n","            gs['nearest'].reshape(-1),\n","            minlength=tri.shape[-1]).float()\n","        anchor_entropy = -(nearest_counts / nearest_counts.sum() *\n","            torch.log(nearest_counts.clamp(min=1e-8) / nearest_counts.sum())).sum().item()\n","\n","        tqdm.write(\n","            f\"    Layer {i}: ‖content‖={content_norm:.3f}  \"\n","            f\"‖geo‖={geo_norm:.3f}  agreement={agreement:.3f}  \"\n","            f\"anchor_H={anchor_entropy:.3f}\")\n","\n","    print(f\"\\nDone.\")\n","\n","\n","if __name__ == '__main__':\n","    main()"],"metadata":{"colab":{"base_uri":"https://localhost:8080/","referenced_widgets":["66f0a0dd6a3b49afaaff8e55ddc8f232","6d12d020ec7349e1a1937ad42ae60992","ab00fed336b84fd6bd79d2608beb0252","bc211f72ad9d41859ba156da0b1f0a74","d4a9013c24e74aebac09f2ee62a63148","663d2b99117244bd9ef4b6d9f0eafdb8","b0ae5117b9c64f71a3f5698d1546acb5","e7d75665dc104599af3afdc2011a61db","41ea6b42289f458bb56492ec1298f4f7","833456446d0c4e18996f40f68d2277da","e7808d06d7ad473a97990cdc41f9515e"]},"id":"7hMyNR2fqTKO","executionInfo":{"status":"ok","timestamp":1774690435131,"user_tz":420,"elapsed":1471598,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"ce03c8c1-69e5-4591-a217-5224f38c7ff1"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["Device: cuda\n","  GPU: NVIDIA RTX PRO 6000 Blackwell Server Edition\n","  VRAM: 102.0 GB\n","============================================================\n","  Geometric Transformer — CIFAR-100\n","  Input: conv(3→64) + SVD(rank=16) + 4×4 patches = 64 tokens + CLS\n","  Model: d=256, heads=8, layers=8, anchors=128\n","============================================================\n","\n","Loading CIFAR-100...\n","  Train: 50,000 | Test: 10,000\n","\n","  Total params:     17,615,780\n","  Trainable params: 17,615,780\n","    components          : 17,615,780\n","\n","━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\n","  Training for 100 epochs\n","  Warmup: 5 epochs, LR: 0.001, WD: 0.05\n","━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\n","\n"]},{"output_type":"display_data","data":{"text/plain":["Epochs:   0%|          | 0/100 [00:00<?, ?it/s]"],"application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"66f0a0dd6a3b49afaaff8e55ddc8f232"}},"metadata":{}},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n","/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n","/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["  E  0  loss=4.5471  train=0.0254  test=0.0496  best=0.0496  lr=0.000200  44.9s\n","  E  5  loss=3.5745  train=0.1833  test=0.1987  best=0.1987  lr=0.001000  14.2s\n","  E 10  loss=2.8993  train=0.3470  test=0.3395  best=0.3395  lr=0.000990  14.2s\n","  E 15  loss=2.5386  train=0.4458  test=0.4130  best=0.4130  lr=0.000967  14.2s\n","  E 20  loss=2.2387  train=0.5281  test=0.4710  best=0.4710  lr=0.000932  14.2s\n","  E 25  loss=1.9880  train=0.6039  test=0.5010  best=0.5017  lr=0.000884  14.2s\n","  E 30  loss=1.7633  train=0.6771  test=0.5126  best=0.5232  lr=0.000826  14.2s\n","  E 35  loss=1.5343  train=0.7559  test=0.5276  best=0.5315  lr=0.000759  14.3s\n","  E 40  loss=1.3194  train=0.8335  test=0.5394  best=0.5410  lr=0.000686  14.3s\n","  E 45  loss=1.1748  train=0.8867  test=0.5435  best=0.5486  lr=0.000607  14.3s\n","  E 50  loss=1.0529  train=0.9313  test=0.5524  best=0.5524  lr=0.000525  14.3s\n","  E 55  loss=0.9796  train=0.9542  test=0.5483  best=0.5569  lr=0.000442  14.2s\n","  E 60  loss=0.9222  train=0.9720  test=0.5594  best=0.5674  lr=0.000361  14.2s\n","  E 65  loss=0.8787  train=0.9836  test=0.5679  best=0.5679  lr=0.000284  14.2s\n","  E 70  loss=0.8515  train=0.9894  test=0.5709  best=0.5740  lr=0.000213  14.2s\n","  E 75  loss=0.8275  train=0.9953  test=0.5730  best=0.5783  lr=0.000149  14.2s\n","  E 80  loss=0.8117  train=0.9977  test=0.5805  best=0.5805  lr=0.000095  14.2s\n","  E 85  loss=0.8031  train=0.9991  test=0.5829  best=0.5867  lr=0.000053  14.2s\n","  E 90  loss=0.7999  train=0.9994  test=0.5860  best=0.5875  lr=0.000022  14.2s\n","  E 95  loss=0.7987  train=0.9996  test=0.5867  best=0.5876  lr=0.000004  14.2s\n","  E 99  loss=0.7984  train=0.9997  test=0.5854  best=0.5880  lr=0.000000  14.2s\n","\n","════════════════════════════════════════════════════════════\n","  CIFAR-10 RESULTS\n","════════════════════════════════════════════════════════════\n","  Best test accuracy: 0.5880 (58.80%)\n","  Parameters: 17,615,780\n","  Checkpoint: geo_cifar10/best.pt\n","\n","  Geometric state inspection:\n","    SVD Input: S_top3=[12.66508960723877, 8.487641334533691, 8.121267318725586]  entropy=2.651  novelty_norm=17.160\n","    Layer 0: ‖content‖=14.504  ‖geo‖=14.273  agreement=0.604  anchor_H=2.735\n","    Layer 1: ‖content‖=14.550  ‖geo‖=14.266  agreement=0.552  anchor_H=3.158\n","    Layer 2: ‖content‖=14.595  ‖geo‖=14.274  agreement=0.545  anchor_H=3.505\n","    Layer 3: ‖content‖=14.639  ‖geo‖=14.257  agreement=0.513  anchor_H=2.754\n","    Layer 4: ‖content‖=14.698  ‖geo‖=14.340  agreement=0.529  anchor_H=3.863\n","    Layer 5: ‖content‖=14.681  ‖geo‖=14.346  agreement=0.509  anchor_H=2.304\n","    Layer 6: ‖content‖=14.691  ‖geo‖=14.328  agreement=0.521  anchor_H=3.762\n","    Layer 7: ‖content‖=14.673  ‖geo‖=14.240  agreement=0.528  anchor_H=3.897\n","\n","Done.\n"]}]},{"cell_type":"code","source":["\"\"\"\n","Geometric Transformer — CIFAR-100 Training with Full Geometric Analysis\n","\n","TensorBoard logging of every geometric feature element:\n","  - CV (coefficient of variation) per layer — the pentachoron band metric\n","  - Stream agreement/divergence per layer\n","  - Anchor utilization, entropy, spread\n","  - Patchwork activation statistics\n","  - Bridge vs assignment consistency\n","  - Triangulation distance distributions\n","  - SVD spectrum, entropy, novelty\n","  - Quaternion arm norms and composition statistics\n","  - Cayley rotation ‖R-I‖ per layer\n","  - FiLM gamma/beta deviation from identity\n","  - Gate activation statistics\n","  - Gradient norms per component type\n","  - Weight norms per component type\n","\n","!pip install geolip-core torchvision tqdm tensorboard\n","\"\"\"\n","\n","import torch\n","import torch.nn as nn\n","import torch.nn.functional as F\n","import numpy as np\n","import time, json, math\n","from pathlib import Path\n","from tqdm.auto import tqdm\n","from torch.utils.tensorboard import SummaryWriter\n","\n","device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n","print(f\"Device: {device}\")\n","if device.type == 'cuda':\n","    print(f\"  GPU: {torch.cuda.get_device_name()}\")\n","    print(f\"  VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB\")\n","\n","from geolip_core.pipeline.components.geometric_transformer import (\n","    GeometricTransformer, GeometricTransformerLayer,\n","    CayleyOrthogonal, QuaternionCompose, FiLMLayer,\n","    ContentAttention, GeometricAttention,\n","    TorchComponent, BaseTower,\n",")\n","torch.set_float32_matmul_precision('high')\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# CONFIG\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","CONFIG = {\n","    # Model\n","    'd_model': 256,\n","    'n_heads': 8,\n","    'n_layers': 8,\n","    'n_anchors': 128,\n","    'manifold_dim': 128,\n","    'n_comp': 4,\n","    'd_comp': 16,\n","    'context_dim': 64,\n","    'quat_dim': 32,\n","    'dropout': 0.1,\n","\n","    # Input stage\n","    'patch_size': 4,\n","    'img_size': 32,\n","    'in_channels': 3,\n","    'conv_channels': 64,\n","    'svd_rank': 16,\n","\n","    # Training\n","    'epochs': 100,\n","    'batch_size': 1024,\n","    'lr': 1e-3,\n","    'weight_decay': 0.05,\n","    'warmup_epochs': 5,\n","    'label_smoothing': 0.1,\n","    'num_workers': 8,\n","\n","    # Data\n","    'num_classes': 100,\n","\n","    # Logging\n","    'log_geo_every': 5,      # full geometric analysis every N epochs\n","    'log_grads_every': 10,   # gradient norms every N epochs\n","    'log_dir': 'runs/geo_cifar100',\n","}\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# INPUT STAGE\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","try:\n","    from geolip_core.core.input.svd import SVDObserver\n","    _HAS_SVD = True\n","except ImportError:\n","    _HAS_SVD = False\n","\n","    class SVDObserver(nn.Module):\n","        \"\"\"Fallback SVDObserver.\"\"\"\n","        def __init__(self, in_channels, svd_rank=24):\n","            super().__init__()\n","            self.svd_rank = svd_rank\n","            self.to_svd = nn.Conv2d(in_channels, svd_rank, 1, bias=False)\n","            self.register_buffer('ema_s', torch.ones(svd_rank))\n","            self.register_buffer('ema_vh_flat', torch.eye(svd_rank).reshape(-1))\n","            self.ema_momentum = 0.99\n","\n","        def extract_features(self, S, Vh):\n","            B, k = S.shape\n","            S_safe = S.clamp(min=1e-6)\n","            s_norm = S_safe / (S_safe.sum(dim=-1, keepdim=True) + 1e-8)\n","            vh_diag = Vh.diagonal(dim1=-2, dim2=-1)\n","            vh_offdiag = (Vh.pow(2).sum((-2, -1)) - vh_diag.pow(2).sum(-1)).unsqueeze(-1).clamp(min=0)\n","            s_ent = -(s_norm * torch.log(s_norm.clamp(min=1e-8))).sum(-1, keepdim=True)\n","            out = torch.cat([s_norm, vh_diag, vh_offdiag, s_ent], dim=-1)\n","            return torch.where(torch.isfinite(out), out, torch.zeros_like(out))\n","\n","        def compute_novelty(self, S):\n","            return S - self.ema_s.clone().unsqueeze(0)\n","\n","        def forward(self, x):\n","            B, C, H, W = x.shape\n","            h = self.to_svd(x)\n","            h_flat = h.permute(0, 2, 3, 1).reshape(B, H * W, self.svd_rank)\n","            with torch.amp.autocast('cuda', enabled=False):\n","                with torch.no_grad():\n","                    gram = torch.bmm(h_flat.float().transpose(1, 2), h_flat.float())\n","                    evals, evecs = torch.linalg.eigh(gram)\n","                    evals = evals.flip(-1).clamp(min=1e-12)\n","                    S = evals.sqrt()\n","                    Vh = evecs.flip(-1).transpose(-2, -1)\n","                    S = torch.where(torch.isfinite(S), S, torch.ones_like(S))\n","                    Vh = torch.where(torch.isfinite(Vh), Vh, torch.zeros_like(Vh))\n","            features = self.extract_features(S, Vh)\n","            novelty = self.compute_novelty(S)\n","            return S, Vh, features, novelty\n","\n","        @torch.no_grad()\n","        def update_ema(self, S, Vh):\n","            m = self.ema_momentum\n","            self.ema_s.mul_(m).add_(S.detach().mean(0), alpha=1-m)\n","            self.ema_vh_flat.mul_(m).add_(Vh.detach().mean(0).reshape(-1), alpha=1-m)\n","\n","        @property\n","        def feature_dim(self):\n","            return 2 * self.svd_rank + 2\n","\n","\n","class ConvSVDPatchEmbedding(TorchComponent):\n","    \"\"\"Input stage: conv frontend → SVDObserver → patch tokens.\"\"\"\n","    def __init__(self, name, img_size=32, patch_size=4, in_channels=3,\n","                 conv_channels=64, d_model=256, svd_rank=16):\n","        super().__init__(name)\n","        self.patch_size = patch_size\n","        self.n_patches = (img_size // patch_size) ** 2\n","        self.d_model = d_model\n","        self.svd_rank = svd_rank\n","\n","        self.conv_frontend = nn.Sequential(\n","            nn.Conv2d(in_channels, conv_channels, 3, padding=1, bias=False),\n","            nn.BatchNorm2d(conv_channels), nn.GELU(),\n","            nn.Conv2d(conv_channels, conv_channels, 3, padding=1, bias=False),\n","            nn.BatchNorm2d(conv_channels), nn.GELU(),\n","        )\n","        self.svd_observer = SVDObserver(conv_channels, svd_rank)\n","        self.patch_proj = nn.Conv2d(\n","            conv_channels, d_model, kernel_size=patch_size,\n","            stride=patch_size, bias=False)\n","        self.patch_norm = nn.LayerNorm(d_model)\n","\n","        svd_feat_dim = self.svd_observer.feature_dim\n","        self.svd_to_gamma = nn.Linear(svd_feat_dim, d_model)\n","        self.svd_to_beta = nn.Linear(svd_feat_dim, d_model)\n","        nn.init.zeros_(self.svd_to_gamma.weight); nn.init.ones_(self.svd_to_gamma.bias)\n","        nn.init.zeros_(self.svd_to_beta.weight); nn.init.zeros_(self.svd_to_beta.bias)\n","\n","        self.cls_token = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)\n","        self.pos_embed = nn.Parameter(\n","            torch.randn(1, self.n_patches + 1, d_model) * 0.02)\n","\n","    def forward(self, x):\n","        B = x.shape[0]\n","        feat = self.conv_frontend(x)\n","        S, Vh, svd_features, novelty = self.svd_observer(feat)\n","        tokens = self.patch_proj(feat)\n","        tokens = tokens.flatten(2).transpose(1, 2)\n","        tokens = self.patch_norm(tokens)\n","        gamma = self.svd_to_gamma(svd_features).unsqueeze(1)\n","        beta = self.svd_to_beta(svd_features).unsqueeze(1)\n","        tokens = gamma * tokens + beta\n","        cls = self.cls_token.expand(B, -1, -1)\n","        tokens = torch.cat([cls, tokens], dim=1)\n","        tokens = tokens + self.pos_embed\n","        svd_state = {\n","            'singular_values': S, 'Vh': Vh,\n","            'svd_features': svd_features, 'novelty': novelty,\n","        }\n","        if self.training:\n","            self.svd_observer.update_ema(S, Vh)\n","        return tokens, svd_state\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# CLASSIFIER\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","class GeoViTClassifier(BaseTower):\n","    \"\"\"Geometric Vision Transformer for classification.\"\"\"\n","    def __init__(self, name, config):\n","        super().__init__(name)\n","        self.config = config\n","\n","        self.attach('patch_embed', ConvSVDPatchEmbedding(\n","            'patch_embed', img_size=config['img_size'],\n","            patch_size=config['patch_size'], in_channels=config['in_channels'],\n","            conv_channels=config['conv_channels'], d_model=config['d_model'],\n","            svd_rank=config['svd_rank'],\n","        ))\n","        self.attach('transformer', GeometricTransformer(\n","            'geo_cifar', d_model=config['d_model'], n_heads=config['n_heads'],\n","            n_layers=config['n_layers'], n_anchors=config['n_anchors'],\n","            manifold_dim=config['manifold_dim'], n_comp=config['n_comp'],\n","            d_comp=config['d_comp'], context_dim=config['context_dim'],\n","            quat_dim=config['quat_dim'], dropout=config['dropout'],\n","        ))\n","        self.attach('head', nn.Sequential(\n","            nn.LayerNorm(config['d_model']),\n","            nn.Linear(config['d_model'], config['d_model']),\n","            nn.GELU(), nn.Dropout(config['dropout']),\n","            nn.Linear(config['d_model'], config['num_classes']),\n","        ))\n","\n","    def forward(self, x, return_geo_state=False):\n","        tokens, svd_state = self['patch_embed'](x)\n","        if return_geo_state:\n","            features, geo_states = self['transformer'](tokens, return_geo_state=True)\n","        else:\n","            features = self['transformer'](tokens)\n","        cls_out = features[:, 0]\n","        logits = self['head'](cls_out)\n","        if return_geo_state:\n","            return logits, geo_states, svd_state\n","        return logits\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# GEOMETRIC ANALYSIS BATTERY\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","@torch.no_grad()\n","def compute_cv(points):\n","    \"\"\"Coefficient of variation on S^(d-1).\n","\n","    CV = std(pairwise_cosine_distances) / mean(pairwise_cosine_distances)\n","\n","    The pentachoron band: CV ∈ [0.20, 0.23] is the universal attractor.\n","    Below 0.20 = over-uniform (collapsed to regular simplex).\n","    Above 0.23 = under-structured (approaching random).\n","    Hard geometric floor ~0.11-0.12.\n","\n","    Args:\n","        points: (N, D) on S^(d-1)\n","\n","    Returns:\n","        cv: scalar\n","        mean_dist: scalar\n","        std_dist: scalar\n","    \"\"\"\n","    points = F.normalize(points.float(), dim=-1)\n","    cos_sim = points @ points.T\n","    # Extract upper triangle (excluding diagonal)\n","    n = points.shape[0]\n","    idx = torch.triu_indices(n, n, offset=1, device=points.device)\n","    pairwise_cos = cos_sim[idx[0], idx[1]]\n","    pairwise_dist = 1.0 - pairwise_cos  # cosine distance\n","\n","    mean_d = pairwise_dist.mean()\n","    std_d = pairwise_dist.std()\n","    cv = (std_d / (mean_d + 1e-8)).item()\n","    return cv, mean_d.item(), std_d.item()\n","\n","\n","@torch.no_grad()\n","def log_geometric_analysis(model, writer, epoch, test_loader, device, config):\n","    \"\"\"Full geometric analysis battery. Logs everything to TensorBoard.\"\"\"\n","    model.eval()\n","\n","    # Get a batch for analysis\n","    images, labels = next(iter(test_loader))\n","    images = images[:min(64, images.shape[0])].to(device)\n","    labels = labels[:min(64, labels.shape[0])].to(device)\n","\n","    logits, geo_states, svd_state = model(images, return_geo_state=True)\n","\n","    n_layers = len(geo_states)\n","    pred = logits.argmax(1)\n","    batch_acc = (pred == labels).float().mean().item()\n","    writer.add_scalar('analysis/batch_accuracy', batch_acc, epoch)\n","\n","    # ─── SVD Input Stage ───\n","    S = svd_state['singular_values']  # (B, k)\n","    s_norm = S / (S.sum(dim=-1, keepdim=True) + 1e-8)\n","    s_ent = -(s_norm * torch.log(s_norm.clamp(min=1e-8))).sum(-1)\n","    novelty = svd_state['novelty']\n","\n","    writer.add_scalar('svd/entropy_mean', s_ent.mean().item(), epoch)\n","    writer.add_scalar('svd/entropy_std', s_ent.std().item(), epoch)\n","    writer.add_scalar('svd/novelty_norm', novelty.norm(dim=-1).mean().item(), epoch)\n","    writer.add_scalar('svd/top1_ratio', (S[:, 0] / (S.sum(-1) + 1e-8)).mean().item(), epoch)\n","    writer.add_scalar('svd/condition_number',\n","        (S[:, 0] / (S[:, -1].clamp(min=1e-8))).mean().item(), epoch)\n","\n","    # Per-singular-value tracking\n","    for k in range(min(S.shape[1], 5)):\n","        writer.add_scalar(f'svd/S_{k}', S[:, k].mean().item(), epoch)\n","\n","    # SVD FiLM deviation from identity\n","    pe = model['patch_embed']\n","    gamma_w = pe.svd_to_gamma.weight.data\n","    gamma_b = pe.svd_to_gamma.bias.data\n","    beta_w = pe.svd_to_beta.weight.data\n","    beta_b = pe.svd_to_beta.bias.data\n","    writer.add_scalar('svd_film/gamma_weight_norm', gamma_w.norm().item(), epoch)\n","    writer.add_scalar('svd_film/gamma_bias_dev_from_1', (gamma_b - 1.0).abs().mean().item(), epoch)\n","    writer.add_scalar('svd_film/beta_weight_norm', beta_w.norm().item(), epoch)\n","    writer.add_scalar('svd_film/beta_bias_norm', beta_b.abs().mean().item(), epoch)\n","\n","    # ─── Per-Layer Geometric Analysis ───\n","    for i, gs in enumerate(geo_states):\n","        prefix = f'layer_{i}'\n","\n","        # === CV — the pentachoron band metric ===\n","        emb = gs['embedding']  # (B, L, manifold_dim)\n","        # CV of anchor positions\n","        transformer = model['transformer']\n","        # Access the observer's constellation anchors\n","        layer_module = None\n","        for name, mod in transformer.named_modules():\n","            if hasattr(mod, 'association') and hasattr(mod.association, 'constellation'):\n","                # Find the right layer's observer\n","                if f'layer_{i}' in name or (i == 0 and 'layer_0' in name):\n","                    anchors = F.normalize(mod.association.constellation.anchors, dim=-1)\n","                    cv_anchors, mean_d_anchors, std_d_anchors = compute_cv(anchors)\n","                    writer.add_scalar(f'{prefix}/cv_anchors', cv_anchors, epoch)\n","                    writer.add_scalar(f'{prefix}/anchor_mean_dist', mean_d_anchors, epoch)\n","                    writer.add_scalar(f'{prefix}/anchor_std_dist', std_d_anchors, epoch)\n","                    break\n","\n","        # CV of embeddings (sample from batch)\n","        emb_flat = emb.reshape(-1, emb.shape[-1])\n","        # Sample to keep computation tractable\n","        n_sample = min(512, emb_flat.shape[0])\n","        idx = torch.randperm(emb_flat.shape[0], device=device)[:n_sample]\n","        cv_emb, mean_d_emb, std_d_emb = compute_cv(emb_flat[idx])\n","        writer.add_scalar(f'{prefix}/cv_embeddings', cv_emb, epoch)\n","        writer.add_scalar(f'{prefix}/embedding_mean_dist', mean_d_emb, epoch)\n","        writer.add_scalar(f'{prefix}/embedding_std_dist', std_d_emb, epoch)\n","\n","        # === Stream Agreement ===\n","        content = gs['content']    # (B, L, D)\n","        geometric = gs['geometric']  # (B, L, D)\n","        agreement = F.cosine_similarity(\n","            content.reshape(-1, content.shape[-1]),\n","            geometric.reshape(-1, geometric.shape[-1]),\n","            dim=-1)\n","        writer.add_scalar(f'{prefix}/stream_agreement_mean', agreement.mean().item(), epoch)\n","        writer.add_scalar(f'{prefix}/stream_agreement_std', agreement.std().item(), epoch)\n","        writer.add_scalar(f'{prefix}/stream_agreement_min', agreement.min().item(), epoch)\n","\n","        # Stream norms\n","        writer.add_scalar(f'{prefix}/content_norm', content.norm(dim=-1).mean().item(), epoch)\n","        writer.add_scalar(f'{prefix}/geometric_norm', geometric.norm(dim=-1).mean().item(), epoch)\n","        writer.add_scalar(f'{prefix}/norm_ratio',\n","            (geometric.norm(dim=-1) / (content.norm(dim=-1) + 1e-8)).mean().item(), epoch)\n","\n","        # === Disagreement arm analysis ===\n","        disagree = content - geometric\n","        agree = content * geometric\n","        writer.add_scalar(f'{prefix}/disagree_norm', disagree.norm(dim=-1).mean().item(), epoch)\n","        writer.add_scalar(f'{prefix}/agree_norm', agree.norm(dim=-1).mean().item(), epoch)\n","        writer.add_scalar(f'{prefix}/disagree_to_content_ratio',\n","            (disagree.norm(dim=-1) / (content.norm(dim=-1) + 1e-8)).mean().item(), epoch)\n","\n","        # === Anchor Utilization ===\n","        tri = gs['triangulation']      # (B, L, A)\n","        assignment = gs['assignment']  # (B, L, A)\n","        nearest = gs['nearest']        # (B, L)\n","        n_anchors = tri.shape[-1]\n","\n","        nearest_flat = nearest.reshape(-1)\n","        counts = torch.bincount(nearest_flat, minlength=n_anchors).float()\n","        total_assignments = counts.sum()\n","\n","        # Entropy\n","        probs = counts / (total_assignments + 1e-8)\n","        anchor_entropy = -(probs * torch.log(probs.clamp(min=1e-8))).sum().item()\n","        max_entropy = math.log(n_anchors)\n","        writer.add_scalar(f'{prefix}/anchor_entropy', anchor_entropy, epoch)\n","        writer.add_scalar(f'{prefix}/anchor_entropy_normalized',\n","            anchor_entropy / (max_entropy + 1e-8), epoch)\n","\n","        # Active anchors (used at least once)\n","        active = (counts > 0).sum().item()\n","        writer.add_scalar(f'{prefix}/anchors_active', active, epoch)\n","        writer.add_scalar(f'{prefix}/anchors_active_frac', active / n_anchors, epoch)\n","\n","        # Anchor dominance (top-1 fraction)\n","        top1_frac = counts.max().item() / (total_assignments + 1e-8)\n","        writer.add_scalar(f'{prefix}/anchor_top1_frac', top1_frac, epoch)\n","\n","        # Dead anchors (never assigned)\n","        dead = (counts == 0).sum().item()\n","        writer.add_scalar(f'{prefix}/anchors_dead', dead, epoch)\n","\n","        # === Triangulation Statistics ===\n","        writer.add_scalar(f'{prefix}/tri_mean', tri.mean().item(), epoch)\n","        writer.add_scalar(f'{prefix}/tri_std', tri.std().item(), epoch)\n","        writer.add_scalar(f'{prefix}/tri_min', tri.min().item(), epoch)\n","        writer.add_scalar(f'{prefix}/tri_max', tri.max().item(), epoch)\n","\n","        # === Soft Assignment Statistics ===\n","        # Assignment entropy per position (how peaked is the assignment?)\n","        assign_ent = -(assignment * torch.log(assignment.clamp(min=1e-8))).sum(-1)\n","        writer.add_scalar(f'{prefix}/assignment_entropy_mean', assign_ent.mean().item(), epoch)\n","        writer.add_scalar(f'{prefix}/assignment_entropy_std', assign_ent.std().item(), epoch)\n","        # Max assignment confidence\n","        writer.add_scalar(f'{prefix}/assignment_max_prob', assignment.max(dim=-1).values.mean().item(), epoch)\n","\n","        # === Patchwork Statistics ===\n","        pw = gs['patchwork']  # (B, L, pw_dim)\n","        writer.add_scalar(f'{prefix}/patchwork_norm', pw.norm(dim=-1).mean().item(), epoch)\n","        writer.add_scalar(f'{prefix}/patchwork_std', pw.std().item(), epoch)\n","        # Sparsity: fraction of near-zero activations\n","        pw_sparsity = (pw.abs() < 0.01).float().mean().item()\n","        writer.add_scalar(f'{prefix}/patchwork_sparsity', pw_sparsity, epoch)\n","\n","        # === Bridge Consistency ===\n","        bridge = gs['bridge']  # (B, L, A)\n","        cos_to = gs['cos_to_anchors']  # (B, L, A)\n","        # Bridge tries to predict assignment from patchwork\n","        bridge_soft = F.softmax(bridge, dim=-1)\n","        bridge_assign_kl = F.kl_div(\n","            bridge_soft.log().reshape(-1, n_anchors),\n","            assignment.reshape(-1, n_anchors),\n","            reduction='batchmean', log_target=False)\n","        writer.add_scalar(f'{prefix}/bridge_assignment_kl', bridge_assign_kl.item(), epoch)\n","\n","        # === Quaternion Composition ===\n","        composed = gs['composed']  # (B, L, 4*quat_dim)\n","        writer.add_scalar(f'{prefix}/composed_norm', composed.norm(dim=-1).mean().item(), epoch)\n","        writer.add_scalar(f'{prefix}/composed_std', composed.std().item(), epoch)\n","\n","        # === Geo Context (compressed FiLM vector) ===\n","        geo_ctx = gs['geo_ctx']  # (B, L, context_dim)\n","        writer.add_scalar(f'{prefix}/geo_ctx_norm', geo_ctx.norm(dim=-1).mean().item(), epoch)\n","        writer.add_scalar(f'{prefix}/geo_ctx_std', geo_ctx.std().item(), epoch)\n","\n","    # ─── Cayley Rotation Analysis ───\n","    for name, mod in model.named_modules():\n","        if isinstance(mod, CayleyOrthogonal):\n","            R = mod.get_rotation()\n","            I = torch.eye(R.shape[0], device=R.device)\n","            r_dist = (R - I).norm().item()\n","            # Clean name for tensorboard\n","            clean_name = name.replace('.', '_')\n","            writer.add_scalar(f'cayley/{clean_name}_R_minus_I', r_dist, epoch)\n","            writer.add_scalar(f'cayley/{clean_name}_det', torch.det(R).item(), epoch)\n","\n","    # ─── FiLM Layer Analysis ───\n","    film_idx = 0\n","    for name, mod in model.named_modules():\n","        if isinstance(mod, FiLMLayer):\n","            g_w = mod.to_gamma.weight.data\n","            g_b = mod.to_gamma.bias.data\n","            b_w = mod.to_beta.weight.data\n","            b_b = mod.to_beta.bias.data\n","            writer.add_scalar(f'film/{film_idx}_gamma_dev',\n","                (g_b - 1.0).abs().mean().item(), epoch)\n","            writer.add_scalar(f'film/{film_idx}_beta_dev',\n","                b_b.abs().mean().item(), epoch)\n","            writer.add_scalar(f'film/{film_idx}_gamma_w_norm',\n","                g_w.norm().item(), epoch)\n","            writer.add_scalar(f'film/{film_idx}_beta_w_norm',\n","                b_w.norm().item(), epoch)\n","            film_idx += 1\n","\n","    # ─── Cross-Layer CV Trajectory ───\n","    cv_trajectory = []\n","    for i, gs in enumerate(geo_states):\n","        emb = gs['embedding']\n","        emb_flat = emb.reshape(-1, emb.shape[-1])\n","        n_sample = min(512, emb_flat.shape[0])\n","        idx = torch.randperm(emb_flat.shape[0], device=device)[:n_sample]\n","        cv, _, _ = compute_cv(emb_flat[idx])\n","        cv_trajectory.append(cv)\n","    writer.add_scalar('cv/trajectory_mean', np.mean(cv_trajectory), epoch)\n","    writer.add_scalar('cv/trajectory_std', np.std(cv_trajectory), epoch)\n","    writer.add_scalar('cv/trajectory_min', np.min(cv_trajectory), epoch)\n","    writer.add_scalar('cv/trajectory_max', np.max(cv_trajectory), epoch)\n","    # In pentachoron band?\n","    in_band = sum(1 for cv in cv_trajectory if 0.20 <= cv <= 0.23)\n","    writer.add_scalar('cv/layers_in_pentachoron_band', in_band, epoch)\n","    writer.add_scalar('cv/layers_in_band_frac', in_band / len(cv_trajectory), epoch)\n","\n","    return {\n","        'batch_acc': batch_acc,\n","        'cv_trajectory': cv_trajectory,\n","    }\n","\n","\n","@torch.no_grad()\n","def log_gradient_norms(model, writer, epoch):\n","    \"\"\"Log gradient norms per component type.\"\"\"\n","    type_grads = {}\n","    for name, param in model.named_parameters():\n","        if param.grad is not None:\n","            grad_norm = param.grad.norm().item()\n","            # Classify by component\n","            if 'projection' in name:\n","                key = 'manifold_proj'\n","            elif 'observer' in name or 'constellation' in name or 'anchor' in name:\n","                key = 'constellation'\n","            elif 'context' in name:\n","                key = 'geo_context'\n","            elif 'content' in name:\n","                key = 'content_attn'\n","            elif 'geometric' in name and 'film' not in name:\n","                key = 'geo_attn'\n","            elif 'film' in name:\n","                key = 'film'\n","            elif 'rotation' in name or 'cayley' in name or 'A_upper' in name:\n","                key = 'cayley'\n","            elif 'compose' in name or 'quat' in name or 'proj_w' in name:\n","                key = 'quaternion'\n","            elif 'decode' in name:\n","                key = 'decode'\n","            elif 'gate' in name:\n","                key = 'gate'\n","            elif 'conv' in name or 'patch' in name:\n","                key = 'input_stage'\n","            elif 'head' in name:\n","                key = 'head'\n","            elif 'svd' in name:\n","                key = 'svd'\n","            else:\n","                key = 'other'\n","\n","            if key not in type_grads:\n","                type_grads[key] = []\n","            type_grads[key].append(grad_norm)\n","\n","    for key, norms in type_grads.items():\n","        writer.add_scalar(f'grad_norm/{key}_mean', np.mean(norms), epoch)\n","        writer.add_scalar(f'grad_norm/{key}_max', np.max(norms), epoch)\n","\n","    # Total gradient norm\n","    total = sum(p.grad.norm().item() ** 2\n","                for p in model.parameters() if p.grad is not None) ** 0.5\n","    writer.add_scalar('grad_norm/total', total, epoch)\n","\n","\n","@torch.no_grad()\n","def log_weight_norms(model, writer, epoch):\n","    \"\"\"Log weight norms per component type.\"\"\"\n","    for name, param in model.named_parameters():\n","        if 'A_upper' in name:\n","            clean = name.replace('.', '_')\n","            writer.add_scalar(f'weights/{clean}_norm', param.norm().item(), epoch)\n","            writer.add_scalar(f'weights/{clean}_max', param.abs().max().item(), epoch)\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# DATA\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","def get_dataloaders(config):\n","    import torchvision\n","    import torchvision.transforms as T\n","\n","    train_transform = T.Compose([\n","        T.RandomCrop(32, padding=4),\n","        T.RandomHorizontalFlip(),\n","        T.ToTensor(),\n","        T.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),\n","    ])\n","    test_transform = T.Compose([\n","        T.ToTensor(),\n","        T.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),\n","    ])\n","\n","    train_ds = torchvision.datasets.CIFAR100(\n","        root='./data', train=True, download=True, transform=train_transform)\n","    test_ds = torchvision.datasets.CIFAR100(\n","        root='./data', train=False, download=True, transform=test_transform)\n","\n","    train_loader = torch.utils.data.DataLoader(\n","        train_ds, batch_size=config['batch_size'], shuffle=True,\n","        num_workers=config['num_workers'], pin_memory=True, drop_last=True)\n","    test_loader = torch.utils.data.DataLoader(\n","        test_ds, batch_size=config['batch_size'], shuffle=False,\n","        num_workers=config['num_workers'], pin_memory=True)\n","\n","    return train_loader, test_loader\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# TRAINING\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","def train_epoch(model, loader, optimizer, scheduler, epoch, config, writer):\n","    model.train()\n","    total_loss = 0\n","    correct = 0\n","    total = 0\n","    criterion = nn.CrossEntropyLoss(label_smoothing=config['label_smoothing'])\n","\n","    for batch_idx, (images, labels) in enumerate(loader):\n","        images = images.to(device)\n","        labels = labels.to(device)\n","\n","        logits = model(images)\n","        loss = criterion(logits, labels)\n","\n","        optimizer.zero_grad()\n","        loss.backward()\n","\n","        # Log gradient norms periodically\n","        if epoch % config['log_grads_every'] == 0 and batch_idx == 0:\n","            log_gradient_norms(model, writer, epoch)\n","\n","        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n","        optimizer.step()\n","        if scheduler is not None:\n","            scheduler.step()\n","\n","        total_loss += loss.item() * images.size(0)\n","        correct += (logits.argmax(1) == labels).sum().item()\n","        total += images.size(0)\n","\n","    return total_loss / total, correct / total\n","\n","\n","@torch.no_grad()\n","def evaluate(model, loader):\n","    model.eval()\n","    correct = 0\n","    total = 0\n","    for images, labels in loader:\n","        images = images.to(device)\n","        labels = labels.to(device)\n","        logits = model(images)\n","        correct += (logits.argmax(1) == labels).sum().item()\n","        total += images.size(0)\n","    return correct / total\n","\n","\n","def main():\n","    config = CONFIG.copy()\n","\n","    print(\"=\" * 60)\n","    print(\"  Geometric Transformer — CIFAR-100 (Full Analysis)\")\n","    print(f\"  Input: conv({config['in_channels']}→{config['conv_channels']}) + \"\n","          f\"SVD(rank={config['svd_rank']}) + \"\n","          f\"{config['patch_size']}×{config['patch_size']} patches = \"\n","          f\"{(config['img_size']//config['patch_size'])**2} tokens + CLS\")\n","    print(f\"  Model: d={config['d_model']}, heads={config['n_heads']}, \"\n","          f\"layers={config['n_layers']}, anchors={config['n_anchors']}\")\n","    print(\"=\" * 60)\n","\n","    # TensorBoard\n","    writer = SummaryWriter(config['log_dir'])\n","    writer.add_text('config', json.dumps(config, indent=2))\n","\n","    # Data\n","    print(\"\\nLoading CIFAR-100...\")\n","    train_loader, test_loader = get_dataloaders(config)\n","    print(f\"  Train: {len(train_loader.dataset):,} | Test: {len(test_loader.dataset):,}\")\n","\n","    # Model\n","    model = GeoViTClassifier('geo_vit_cifar100', config)\n","    if hasattr(model, 'network_to'):\n","        model.network_to(device=device, strict=False)\n","    else:\n","        model = model.to(device)\n","\n","    n_params = sum(p.numel() for p in model.parameters())\n","    n_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)\n","    print(f\"\\n  Total params:     {n_params:,}\")\n","    print(f\"  Trainable params: {n_trainable:,}\")\n","\n","    for name, module in model.named_children():\n","        n = sum(p.numel() for p in module.parameters())\n","        if n > 0:\n","            print(f\"    {name:<20s}: {n:,}\")\n","\n","    writer.add_scalar('model/total_params', n_params, 0)\n","\n","    # Optimizer + scheduler\n","    optimizer = torch.optim.AdamW(\n","        model.parameters(), lr=config['lr'], weight_decay=config['weight_decay'])\n","\n","    total_steps = config['epochs'] * len(train_loader)\n","    warmup_steps = config['warmup_epochs'] * len(train_loader)\n","\n","    def lr_lambda(step):\n","        if step < warmup_steps:\n","            return step / warmup_steps\n","        progress = (step - warmup_steps) / (total_steps - warmup_steps)\n","        return 0.5 * (1 + np.cos(np.pi * progress))\n","\n","    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)\n","\n","    print(f\"\\n{'━'*60}\")\n","    print(f\"  Training for {config['epochs']} epochs\")\n","    print(f\"  Warmup: {config['warmup_epochs']} epochs, \"\n","          f\"LR: {config['lr']}, WD: {config['weight_decay']}\")\n","    print(f\"  TensorBoard: {config['log_dir']}\")\n","    print(f\"  Geo analysis every {config['log_geo_every']} epochs\")\n","    print(f\"{'━'*60}\\n\")\n","\n","    best_acc = 0\n","    save_dir = Path('geo_cifar100'); save_dir.mkdir(exist_ok=True)\n","\n","    for epoch in tqdm(range(config['epochs']), desc=\"Epochs\"):\n","        t0 = time.time()\n","\n","        train_loss, train_acc = train_epoch(\n","            model, train_loader, optimizer, scheduler, epoch, config, writer)\n","\n","        test_acc = evaluate(model, test_loader)\n","        elapsed = time.time() - t0\n","\n","        # Core metrics every epoch\n","        lr = optimizer.param_groups[0]['lr']\n","        writer.add_scalar('train/loss', train_loss, epoch)\n","        writer.add_scalar('train/accuracy', train_acc, epoch)\n","        writer.add_scalar('test/accuracy', test_acc, epoch)\n","        writer.add_scalar('train/lr', lr, epoch)\n","        writer.add_scalar('train/epoch_time', elapsed, epoch)\n","\n","        # Generalization gap\n","        writer.add_scalar('gap/train_test', train_acc - test_acc, epoch)\n","        writer.add_scalar('gap/overfit_ratio',\n","            train_acc / (test_acc + 1e-8), epoch)\n","\n","        # Weight norms\n","        log_weight_norms(model, writer, epoch)\n","\n","        if test_acc > best_acc:\n","            best_acc = test_acc\n","            torch.save({\n","                'state_dict': {k: v.cpu() for k, v in model.state_dict().items()},\n","                'epoch': epoch,\n","                'test_acc': test_acc,\n","                'config': config,\n","            }, save_dir / 'best.pt')\n","\n","        # Full geometric analysis periodically\n","        if epoch % config['log_geo_every'] == 0 or epoch == config['epochs'] - 1:\n","            geo_info = log_geometric_analysis(\n","                model, writer, epoch, test_loader, device, config)\n","\n","            cv_str = ', '.join(f'{cv:.3f}' for cv in geo_info['cv_trajectory'])\n","            tqdm.write(\n","                f\"  E{epoch:>3d}  loss={train_loss:.4f}  \"\n","                f\"train={train_acc:.4f}  test={test_acc:.4f}  \"\n","                f\"best={best_acc:.4f}  lr={lr:.6f}  {elapsed:.1f}s\"\n","                f\"\\n        CV=[{cv_str}]\")\n","        elif epoch % 5 == 0:\n","            tqdm.write(\n","                f\"  E{epoch:>3d}  loss={train_loss:.4f}  \"\n","                f\"train={train_acc:.4f}  test={test_acc:.4f}  \"\n","                f\"best={best_acc:.4f}  lr={lr:.6f}  {elapsed:.1f}s\")\n","\n","    # Final summary\n","    print(f\"\\n{'═'*60}\")\n","    print(f\"  CIFAR-100 RESULTS\")\n","    print(f\"{'═'*60}\")\n","    print(f\"  Best test accuracy: {best_acc:.4f} ({best_acc*100:.2f}%)\")\n","    print(f\"  Parameters: {n_params:,}\")\n","    print(f\"  Checkpoint: {save_dir}/best.pt\")\n","    print(f\"  TensorBoard: {config['log_dir']}\")\n","\n","    # Final geometric state\n","    print(f\"\\n  Final geometric state:\")\n","    geo_info = log_geometric_analysis(\n","        model, writer, config['epochs'], test_loader, device, config)\n","\n","    writer.close()\n","    print(f\"\\nDone.\")\n","\n","\n","if __name__ == '__main__':\n","    main()"],"metadata":{"id":"UoYH0xhPzA_C"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["\"\"\"\n","Geometric Transformer — Standalone Analysis Battery\n","\n","Loads a trained checkpoint and runs comprehensive geometric analysis.\n","No training, no TensorBoard dependency — prints everything to console\n","and saves a JSON report.\n","\n","Usage:\n","    python analyze_geo_transformer.py                          # defaults\n","    python analyze_geo_transformer.py --checkpoint path/to/best.pt\n","    python analyze_geo_transformer.py --n_samples 256 --per_class\n","\n","Analysis categories:\n","    1. CV (coefficient of variation) — pentachoron band tracking\n","    2. Constellation geometry — anchor spread, utilization, dead anchors\n","    3. Stream dynamics — agreement, divergence, arm norms\n","    4. Triangulation — distance distributions, assignment peakedness\n","    5. Patchwork — activation patterns, sparsity, bridge consistency\n","    6. SVD Input — spectrum, entropy, condition number, novelty\n","    7. Cayley rotations — ‖R-I‖, determinant, parameter norms\n","    8. FiLM layers — deviation from identity\n","    9. Quaternion composition — arm balance, output statistics\n","    10. Gate — activation distribution, bypass fraction\n","    11. Per-class analysis — which classes the geometry handles well/poorly\n","    12. Cross-layer trajectories — how features evolve through depth\n","\n","!pip install geolip-core torchvision\n","\"\"\"\n","\n","import torch\n","import torch.nn as nn\n","import torch.nn.functional as F\n","import numpy as np\n","import json, math, argparse\n","from pathlib import Path\n","from collections import defaultdict\n","\n","device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n","\n","from geolip_core.pipeline.components.geometric_transformer import (\n","    GeometricTransformer, GeometricTransformerLayer,\n","    CayleyOrthogonal, QuaternionCompose, FiLMLayer,\n","    ContentAttention, GeometricAttention,\n","    ManifoldProjection, PositionGeometricContext,\n","    TorchComponent, BaseTower,\n",")\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# INPUT STAGE (must match training)\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","try:\n","    from geolip_core.core.input.svd import SVDObserver\n","    _HAS_SVD = True\n","except ImportError:\n","    _HAS_SVD = False\n","\n","    class SVDObserver(nn.Module):\n","        def __init__(self, in_channels, svd_rank=24):\n","            super().__init__()\n","            self.svd_rank = svd_rank\n","            self.to_svd = nn.Conv2d(in_channels, svd_rank, 1, bias=False)\n","            self.register_buffer('ema_s', torch.ones(svd_rank))\n","            self.register_buffer('ema_vh_flat', torch.eye(svd_rank).reshape(-1))\n","            self.ema_momentum = 0.99\n","\n","        def extract_features(self, S, Vh):\n","            B, k = S.shape\n","            S_safe = S.clamp(min=1e-6)\n","            s_norm = S_safe / (S_safe.sum(dim=-1, keepdim=True) + 1e-8)\n","            vh_diag = Vh.diagonal(dim1=-2, dim2=-1)\n","            vh_offdiag = (Vh.pow(2).sum((-2, -1)) - vh_diag.pow(2).sum(-1)).unsqueeze(-1).clamp(min=0)\n","            s_ent = -(s_norm * torch.log(s_norm.clamp(min=1e-8))).sum(-1, keepdim=True)\n","            out = torch.cat([s_norm, vh_diag, vh_offdiag, s_ent], dim=-1)\n","            return torch.where(torch.isfinite(out), out, torch.zeros_like(out))\n","\n","        def compute_novelty(self, S):\n","            return S - self.ema_s.clone().unsqueeze(0)\n","\n","        def forward(self, x):\n","            B, C, H, W = x.shape\n","            h = self.to_svd(x)\n","            h_flat = h.permute(0, 2, 3, 1).reshape(B, H * W, self.svd_rank)\n","            with torch.amp.autocast('cuda', enabled=False):\n","                with torch.no_grad():\n","                    gram = torch.bmm(h_flat.float().transpose(1, 2), h_flat.float())\n","                    evals, evecs = torch.linalg.eigh(gram)\n","                    evals = evals.flip(-1).clamp(min=1e-12)\n","                    S = evals.sqrt()\n","                    Vh = evecs.flip(-1).transpose(-2, -1)\n","                    S = torch.where(torch.isfinite(S), S, torch.ones_like(S))\n","                    Vh = torch.where(torch.isfinite(Vh), Vh, torch.zeros_like(Vh))\n","            return S, Vh, self.extract_features(S, Vh), self.compute_novelty(S)\n","\n","        @property\n","        def feature_dim(self):\n","            return 2 * self.svd_rank + 2\n","\n","\n","class ConvSVDPatchEmbedding(TorchComponent):\n","    def __init__(self, name, img_size=32, patch_size=4, in_channels=3,\n","                 conv_channels=64, d_model=256, svd_rank=16):\n","        super().__init__(name)\n","        self.patch_size = patch_size\n","        self.n_patches = (img_size // patch_size) ** 2\n","        self.d_model = d_model\n","        self.svd_rank = svd_rank\n","        self.conv_frontend = nn.Sequential(\n","            nn.Conv2d(in_channels, conv_channels, 3, padding=1, bias=False),\n","            nn.BatchNorm2d(conv_channels), nn.GELU(),\n","            nn.Conv2d(conv_channels, conv_channels, 3, padding=1, bias=False),\n","            nn.BatchNorm2d(conv_channels), nn.GELU(),\n","        )\n","        self.svd_observer = SVDObserver(conv_channels, svd_rank)\n","        self.patch_proj = nn.Conv2d(conv_channels, d_model, kernel_size=patch_size,\n","                                     stride=patch_size, bias=False)\n","        self.patch_norm = nn.LayerNorm(d_model)\n","        svd_feat_dim = self.svd_observer.feature_dim\n","        self.svd_to_gamma = nn.Linear(svd_feat_dim, d_model)\n","        self.svd_to_beta = nn.Linear(svd_feat_dim, d_model)\n","        nn.init.zeros_(self.svd_to_gamma.weight); nn.init.ones_(self.svd_to_gamma.bias)\n","        nn.init.zeros_(self.svd_to_beta.weight); nn.init.zeros_(self.svd_to_beta.bias)\n","        self.cls_token = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)\n","        self.pos_embed = nn.Parameter(torch.randn(1, self.n_patches + 1, d_model) * 0.02)\n","\n","    def forward(self, x):\n","        B = x.shape[0]\n","        feat = self.conv_frontend(x)\n","        S, Vh, svd_features, novelty = self.svd_observer(feat)\n","        tokens = self.patch_proj(feat)\n","        tokens = tokens.flatten(2).transpose(1, 2)\n","        tokens = self.patch_norm(tokens)\n","        gamma = self.svd_to_gamma(svd_features).unsqueeze(1)\n","        beta = self.svd_to_beta(svd_features).unsqueeze(1)\n","        tokens = gamma * tokens + beta\n","        cls = self.cls_token.expand(B, -1, -1)\n","        tokens = torch.cat([cls, tokens], dim=1)\n","        tokens = tokens + self.pos_embed\n","        return tokens, {\n","            'singular_values': S, 'Vh': Vh,\n","            'svd_features': svd_features, 'novelty': novelty,\n","        }\n","\n","\n","class GeoViTClassifier(BaseTower):\n","    def __init__(self, name, config):\n","        super().__init__(name)\n","        self.config = config\n","        self.attach('patch_embed', ConvSVDPatchEmbedding(\n","            'patch_embed', img_size=config['img_size'],\n","            patch_size=config['patch_size'], in_channels=config['in_channels'],\n","            conv_channels=config['conv_channels'], d_model=config['d_model'],\n","            svd_rank=config['svd_rank'],\n","        ))\n","        self.attach('transformer', GeometricTransformer(\n","            'geo_cifar', d_model=config['d_model'], n_heads=config['n_heads'],\n","            n_layers=config['n_layers'], n_anchors=config['n_anchors'],\n","            manifold_dim=config['manifold_dim'], n_comp=config['n_comp'],\n","            d_comp=config['d_comp'], context_dim=config['context_dim'],\n","            quat_dim=config['quat_dim'], dropout=config['dropout'],\n","        ))\n","        self.attach('head', nn.Sequential(\n","            nn.LayerNorm(config['d_model']),\n","            nn.Linear(config['d_model'], config['d_model']),\n","            nn.GELU(), nn.Dropout(config['dropout']),\n","            nn.Linear(config['d_model'], config['num_classes']),\n","        ))\n","\n","    def forward(self, x, return_geo_state=False):\n","        tokens, svd_state = self['patch_embed'](x)\n","        if return_geo_state:\n","            features, geo_states = self['transformer'](tokens, return_geo_state=True)\n","        else:\n","            features = self['transformer'](tokens)\n","        cls_out = features[:, 0]\n","        logits = self['head'](cls_out)\n","        if return_geo_state:\n","            return logits, geo_states, svd_state\n","        return logits\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# GEOMETRIC PRIMITIVES\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","@torch.no_grad()\n","def compute_cv(points, max_points=1024):\n","    \"\"\"Coefficient of variation of pairwise cosine distances on S^(d-1).\"\"\"\n","    points = F.normalize(points.float(), dim=-1)\n","    if points.shape[0] > max_points:\n","        idx = torch.randperm(points.shape[0], device=points.device)[:max_points]\n","        points = points[idx]\n","    cos_sim = points @ points.T\n","    n = points.shape[0]\n","    idx = torch.triu_indices(n, n, offset=1, device=points.device)\n","    pairwise_dist = 1.0 - cos_sim[idx[0], idx[1]]\n","    mean_d = pairwise_dist.mean()\n","    std_d = pairwise_dist.std()\n","    return {\n","        'cv': (std_d / (mean_d + 1e-8)).item(),\n","        'mean_dist': mean_d.item(),\n","        'std_dist': std_d.item(),\n","        'min_dist': pairwise_dist.min().item(),\n","        'max_dist': pairwise_dist.max().item(),\n","        'n_points': points.shape[0],\n","    }\n","\n","\n","@torch.no_grad()\n","def compute_anchor_geometry(anchors):\n","    \"\"\"Full geometric analysis of anchor positions on S^(d-1).\"\"\"\n","    anchors = F.normalize(anchors.float(), dim=-1)\n","    n, d = anchors.shape\n","    cos_sim = anchors @ anchors.T\n","\n","    idx = torch.triu_indices(n, n, offset=1, device=anchors.device)\n","    pairwise_cos = cos_sim[idx[0], idx[1]]\n","    pairwise_dist = 1.0 - pairwise_cos\n","\n","    # Nearest neighbor distances\n","    cos_sim_no_diag = cos_sim.clone()\n","    cos_sim_no_diag.fill_diagonal_(-2.0)\n","    nn_cos = cos_sim_no_diag.max(dim=1).values\n","    nn_dist = 1.0 - nn_cos\n","\n","    cv_info = compute_cv(anchors)\n","\n","    return {\n","        **cv_info,\n","        'n_anchors': n,\n","        'dim': d,\n","        'pairwise_cos_mean': pairwise_cos.mean().item(),\n","        'pairwise_cos_std': pairwise_cos.std().item(),\n","        'nn_dist_mean': nn_dist.mean().item(),\n","        'nn_dist_std': nn_dist.std().item(),\n","        'nn_dist_min': nn_dist.min().item(),\n","        'nn_dist_max': nn_dist.max().item(),\n","        'max_cos_similarity': pairwise_cos.max().item(),\n","        'min_cos_similarity': pairwise_cos.min().item(),\n","    }\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# ANALYSIS BATTERY\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","@torch.no_grad()\n","def run_analysis(model, test_loader, device, n_batches=4, per_class=False):\n","    \"\"\"Run the complete geometric analysis battery.\"\"\"\n","    model.eval()\n","    report = {}\n","\n","    # ─── Collect predictions and geo states across multiple batches ───\n","    all_logits = []\n","    all_labels = []\n","    all_geo_states = None\n","    all_svd_state = None\n","\n","    for batch_idx, (images, labels) in enumerate(test_loader):\n","        if batch_idx >= n_batches:\n","            break\n","        images = images.to(device)\n","        labels = labels.to(device)\n","\n","        logits, geo_states, svd_state = model(images, return_geo_state=True)\n","        all_logits.append(logits.cpu())\n","        all_labels.append(labels.cpu())\n","\n","        # Keep last batch geo_states for detailed analysis\n","        if batch_idx == 0:\n","            all_geo_states = geo_states\n","            all_svd_state = svd_state\n","\n","    all_logits = torch.cat(all_logits, dim=0)\n","    all_labels = torch.cat(all_labels, dim=0)\n","    n_total = all_labels.shape[0]\n","    pred = all_logits.argmax(1)\n","    overall_acc = (pred == all_labels).float().mean().item()\n","\n","    report['overall'] = {\n","        'accuracy': overall_acc,\n","        'n_samples': n_total,\n","        'n_classes': all_logits.shape[1],\n","    }\n","\n","    n_layers = len(all_geo_states)\n","\n","    # ─── 1. SVD Input Stage ───\n","    print(\"\\n\" + \"=\" * 70)\n","    print(\"  1. SVD INPUT STAGE\")\n","    print(\"=\" * 70)\n","\n","    S = all_svd_state['singular_values']\n","    s_norm = S / (S.sum(dim=-1, keepdim=True) + 1e-8)\n","    s_ent = -(s_norm * torch.log(s_norm.clamp(min=1e-8))).sum(-1)\n","    novelty = all_svd_state['novelty']\n","\n","    svd_report = {\n","        'entropy_mean': s_ent.mean().item(),\n","        'entropy_std': s_ent.std().item(),\n","        'novelty_norm': novelty.norm(dim=-1).mean().item(),\n","        'top1_ratio': (S[:, 0] / (S.sum(-1) + 1e-8)).mean().item(),\n","        'condition_number': (S[:, 0] / (S[:, -1].clamp(min=1e-8))).mean().item(),\n","        'singular_values': [S[:, k].mean().item() for k in range(S.shape[1])],\n","    }\n","\n","    # SVD FiLM deviation\n","    pe = model['patch_embed']\n","    svd_report['film_gamma_bias_dev'] = (pe.svd_to_gamma.bias.data - 1.0).abs().mean().item()\n","    svd_report['film_beta_bias_norm'] = pe.svd_to_beta.bias.data.abs().mean().item()\n","    svd_report['film_gamma_weight_norm'] = pe.svd_to_gamma.weight.data.norm().item()\n","    svd_report['film_beta_weight_norm'] = pe.svd_to_beta.weight.data.norm().item()\n","\n","    report['svd'] = svd_report\n","\n","    print(f\"  Entropy:          {svd_report['entropy_mean']:.3f} ± {svd_report['entropy_std']:.3f}\")\n","    print(f\"  Novelty norm:     {svd_report['novelty_norm']:.3f}\")\n","    print(f\"  Top-1 ratio:      {svd_report['top1_ratio']:.3f}\")\n","    print(f\"  Condition number: {svd_report['condition_number']:.1f}\")\n","    print(f\"  Singular values:  {', '.join(f'{v:.2f}' for v in svd_report['singular_values'][:5])}...\")\n","    print(f\"  FiLM γ dev:       {svd_report['film_gamma_bias_dev']:.4f}\")\n","    print(f\"  FiLM β norm:      {svd_report['film_beta_bias_norm']:.4f}\")\n","\n","    # ─── 2. Per-Layer Geometric Analysis ───\n","    report['layers'] = {}\n","\n","    for i in range(n_layers):\n","        gs = all_geo_states[i]\n","        lr = {}\n","\n","        print(f\"\\n{'─'*70}\")\n","        print(f\"  LAYER {i}\")\n","        print(f\"{'─'*70}\")\n","\n","        # === 2a. CV — Pentachoron Band ===\n","        emb = gs['embedding']\n","        emb_flat = emb.reshape(-1, emb.shape[-1])\n","        cv_emb = compute_cv(emb_flat)\n","        lr['cv_embeddings'] = cv_emb\n","\n","        # Anchor CV\n","        anchor_cv = None\n","        for name, mod in model.named_modules():\n","            if (hasattr(mod, 'association') and hasattr(mod.association, 'constellation')\n","                    and f'layer_{i}' in name):\n","                anchors = mod.association.constellation.anchors.data\n","                anchor_cv = compute_anchor_geometry(anchors)\n","                lr['anchor_geometry'] = anchor_cv\n","                break\n","\n","        in_band = 0.20 <= cv_emb['cv'] <= 0.23\n","        print(f\"  CV embeddings:    {cv_emb['cv']:.4f}  {'✓ IN BAND' if in_band else '✗ outside'}\")\n","        print(f\"    mean_dist={cv_emb['mean_dist']:.4f}  std={cv_emb['std_dist']:.4f}\")\n","        if anchor_cv:\n","            a_in_band = 0.20 <= anchor_cv['cv'] <= 0.23\n","            print(f\"  CV anchors:       {anchor_cv['cv']:.4f}  {'✓ IN BAND' if a_in_band else '✗ outside'}\")\n","            print(f\"    nn_dist: mean={anchor_cv['nn_dist_mean']:.4f}  \"\n","                  f\"min={anchor_cv['nn_dist_min']:.4f}  max={anchor_cv['nn_dist_max']:.4f}\")\n","            print(f\"    max_cos_sim:    {anchor_cv['max_cos_similarity']:.4f}\")\n","\n","        # === 2b. Stream Dynamics ===\n","        content = gs['content']\n","        geometric = gs['geometric']\n","\n","        agreement = F.cosine_similarity(\n","            content.reshape(-1, content.shape[-1]),\n","            geometric.reshape(-1, geometric.shape[-1]), dim=-1)\n","\n","        disagree = content - geometric\n","        agree = content * geometric\n","\n","        stream = {\n","            'agreement_mean': agreement.mean().item(),\n","            'agreement_std': agreement.std().item(),\n","            'agreement_min': agreement.min().item(),\n","            'agreement_max': agreement.max().item(),\n","            'content_norm': content.norm(dim=-1).mean().item(),\n","            'geometric_norm': geometric.norm(dim=-1).mean().item(),\n","            'norm_ratio': (geometric.norm(dim=-1) / (content.norm(dim=-1) + 1e-8)).mean().item(),\n","            'disagree_norm': disagree.norm(dim=-1).mean().item(),\n","            'agree_norm': agree.norm(dim=-1).mean().item(),\n","            'disagree_to_content': (disagree.norm(dim=-1) / (content.norm(dim=-1) + 1e-8)).mean().item(),\n","        }\n","        lr['stream'] = stream\n","\n","        print(f\"\\n  Stream agreement: {stream['agreement_mean']:.4f} ± {stream['agreement_std']:.4f}\")\n","        print(f\"    range: [{stream['agreement_min']:.4f}, {stream['agreement_max']:.4f}]\")\n","        print(f\"  ‖content‖:        {stream['content_norm']:.3f}\")\n","        print(f\"  ‖geometric‖:      {stream['geometric_norm']:.3f}\")\n","        print(f\"  ‖disagree‖:       {stream['disagree_norm']:.3f}  \"\n","              f\"(ratio: {stream['disagree_to_content']:.3f})\")\n","        print(f\"  ‖agree‖:          {stream['agree_norm']:.3f}\")\n","\n","        # === 2c. Anchor Utilization ===\n","        tri = gs['triangulation']\n","        assignment = gs['assignment']\n","        nearest = gs['nearest']\n","        n_anchors = tri.shape[-1]\n","\n","        nearest_flat = nearest.reshape(-1)\n","        counts = torch.bincount(nearest_flat, minlength=n_anchors).float()\n","        total_assign = counts.sum()\n","\n","        probs = counts / (total_assign + 1e-8)\n","        entropy = -(probs * torch.log(probs.clamp(min=1e-8))).sum().item()\n","        max_entropy = math.log(n_anchors)\n","\n","        active = (counts > 0).sum().item()\n","        dead = (counts == 0).sum().item()\n","        top1_frac = counts.max().item() / (total_assign + 1e-8)\n","\n","        # Top-5 and bottom-5 anchors\n","        sorted_counts, sorted_idx = counts.sort(descending=True)\n","\n","        anchor_util = {\n","            'entropy': entropy,\n","            'entropy_normalized': entropy / (max_entropy + 1e-8),\n","            'max_entropy': max_entropy,\n","            'active': int(active),\n","            'dead': int(dead),\n","            'total_anchors': n_anchors,\n","            'top1_frac': top1_frac,\n","            'top5_counts': sorted_counts[:5].tolist(),\n","            'top5_indices': sorted_idx[:5].tolist(),\n","            'bottom5_counts': sorted_counts[-5:].tolist(),\n","        }\n","        lr['anchor_utilization'] = anchor_util\n","\n","        print(f\"\\n  Anchors: {active}/{n_anchors} active, {dead} dead\")\n","        print(f\"  Entropy: {entropy:.3f} / {max_entropy:.3f} \"\n","              f\"(normalized: {anchor_util['entropy_normalized']:.3f})\")\n","        print(f\"  Top-1 dominance: {top1_frac:.4f}\")\n","        print(f\"  Top-5 counts: {sorted_counts[:5].int().tolist()}\")\n","\n","        # === 2d. Triangulation ===\n","        tri_stats = {\n","            'mean': tri.mean().item(),\n","            'std': tri.std().item(),\n","            'min': tri.min().item(),\n","            'max': tri.max().item(),\n","            'median': tri.median().item(),\n","        }\n","        lr['triangulation'] = tri_stats\n","\n","        # === 2e. Soft Assignment ===\n","        assign_ent = -(assignment * torch.log(assignment.clamp(min=1e-8))).sum(-1)\n","        max_prob = assignment.max(dim=-1).values\n","\n","        assign_stats = {\n","            'entropy_mean': assign_ent.mean().item(),\n","            'entropy_std': assign_ent.std().item(),\n","            'max_prob_mean': max_prob.mean().item(),\n","            'max_prob_std': max_prob.std().item(),\n","        }\n","        lr['assignment'] = assign_stats\n","\n","        print(f\"\\n  Assignment entropy: {assign_stats['entropy_mean']:.3f} ± {assign_stats['entropy_std']:.3f}\")\n","        print(f\"  Max prob:           {assign_stats['max_prob_mean']:.4f} ± {assign_stats['max_prob_std']:.4f}\")\n","\n","        # === 2f. Patchwork ===\n","        pw = gs['patchwork']\n","        pw_stats = {\n","            'norm': pw.norm(dim=-1).mean().item(),\n","            'std': pw.std().item(),\n","            'sparsity': (pw.abs() < 0.01).float().mean().item(),\n","            'dead_frac': (pw.abs() < 1e-6).float().mean().item(),\n","        }\n","        lr['patchwork'] = pw_stats\n","\n","        print(f\"\\n  Patchwork norm:     {pw_stats['norm']:.3f}\")\n","        print(f\"  Patchwork sparsity: {pw_stats['sparsity']:.3f}  \"\n","              f\"(dead: {pw_stats['dead_frac']:.4f})\")\n","\n","        # === 2g. Bridge Consistency ===\n","        bridge = gs['bridge']\n","        bridge_soft = F.softmax(bridge, dim=-1)\n","        bridge_kl = F.kl_div(\n","            bridge_soft.log().reshape(-1, n_anchors),\n","            assignment.reshape(-1, n_anchors),\n","            reduction='batchmean', log_target=False).item()\n","        lr['bridge_kl'] = bridge_kl\n","        print(f\"  Bridge ↔ assignment KL: {bridge_kl:.4f}\")\n","\n","        # === 2h. Quaternion Composition ===\n","        composed = gs['composed']\n","        comp_stats = {\n","            'norm': composed.norm(dim=-1).mean().item(),\n","            'std': composed.std().item(),\n","        }\n","        lr['composed'] = comp_stats\n","\n","        # === 2i. Geo Context ===\n","        geo_ctx = gs['geo_ctx']\n","        ctx_stats = {\n","            'norm': geo_ctx.norm(dim=-1).mean().item(),\n","            'std': geo_ctx.std().item(),\n","        }\n","        lr['geo_ctx'] = ctx_stats\n","\n","        print(f\"\\n  Geo context norm:   {ctx_stats['norm']:.3f}\")\n","        print(f\"  Composed norm:      {comp_stats['norm']:.3f}\")\n","\n","        report['layers'][f'layer_{i}'] = lr\n","\n","    # ─── 3. Cayley Rotations ───\n","    print(f\"\\n{'='*70}\")\n","    print(\"  3. CAYLEY ROTATIONS\")\n","    print(\"=\" * 70)\n","\n","    cayley_report = {}\n","    for name, mod in model.named_modules():\n","        if isinstance(mod, CayleyOrthogonal):\n","            R = mod.get_rotation()\n","            I = torch.eye(R.shape[0], device=R.device)\n","            r_dist = (R - I).norm().item()\n","            det = torch.det(R).item()\n","            a_norm = mod.A_upper.data.norm().item()\n","            a_max = mod.A_upper.data.abs().max().item()\n","\n","            clean = name.replace('.', '_')\n","            cayley_report[clean] = {\n","                'R_minus_I': r_dist,\n","                'det': det,\n","                'A_upper_norm': a_norm,\n","                'A_upper_max': a_max,\n","                'dim': mod.dim,\n","            }\n","            print(f\"  {name}\")\n","            print(f\"    ‖R-I‖={r_dist:.4f}  det={det:.6f}  \"\n","                  f\"‖A‖={a_norm:.4f}  max|A|={a_max:.6f}\")\n","\n","    report['cayley'] = cayley_report\n","\n","    # ─── 4. FiLM Layers ───\n","    print(f\"\\n{'='*70}\")\n","    print(\"  4. FiLM LAYERS\")\n","    print(\"=\" * 70)\n","\n","    film_report = {}\n","    film_idx = 0\n","    for name, mod in model.named_modules():\n","        if isinstance(mod, FiLMLayer):\n","            g_b = mod.to_gamma.bias.data\n","            b_b = mod.to_beta.bias.data\n","            g_dev = (g_b - 1.0).abs().mean().item()\n","            b_dev = b_b.abs().mean().item()\n","            g_w_norm = mod.to_gamma.weight.data.norm().item()\n","            b_w_norm = mod.to_beta.weight.data.norm().item()\n","\n","            film_report[f'film_{film_idx}'] = {\n","                'name': name,\n","                'gamma_dev': g_dev,\n","                'beta_dev': b_dev,\n","                'gamma_w_norm': g_w_norm,\n","                'beta_w_norm': b_w_norm,\n","            }\n","\n","            status = \"ACTIVE\" if g_dev > 0.05 or b_dev > 0.05 else \"near-identity\"\n","            print(f\"  {film_idx:>2d} [{name[-40:]:>40s}]  \"\n","                  f\"γ_dev={g_dev:.4f}  β_dev={b_dev:.4f}  {status}\")\n","            film_idx += 1\n","\n","    report['film'] = film_report\n","\n","    # ─── 5. Cross-Layer Trajectories ───\n","    print(f\"\\n{'='*70}\")\n","    print(\"  5. CROSS-LAYER TRAJECTORIES\")\n","    print(\"=\" * 70)\n","\n","    trajectories = {\n","        'cv': [], 'agreement': [], 'disagree_norm': [],\n","        'anchor_entropy': [], 'patchwork_norm': [],\n","        'bridge_kl': [], 'composed_norm': [],\n","    }\n","\n","    for i in range(n_layers):\n","        lr = report['layers'][f'layer_{i}']\n","        trajectories['cv'].append(lr['cv_embeddings']['cv'])\n","        trajectories['agreement'].append(lr['stream']['agreement_mean'])\n","        trajectories['disagree_norm'].append(lr['stream']['disagree_norm'])\n","        trajectories['anchor_entropy'].append(lr['anchor_utilization']['entropy_normalized'])\n","        trajectories['patchwork_norm'].append(lr['patchwork']['norm'])\n","        trajectories['bridge_kl'].append(lr['bridge_kl'])\n","        trajectories['composed_norm'].append(lr['composed']['norm'])\n","\n","    report['trajectories'] = trajectories\n","\n","    for key, vals in trajectories.items():\n","        vals_str = '  '.join(f'{v:.4f}' for v in vals)\n","        trend = vals[-1] - vals[0]\n","        trend_str = f\"+{trend:.4f}\" if trend > 0 else f\"{trend:.4f}\"\n","        print(f\"  {key:<20s}: [{vals_str}]  Δ={trend_str}\")\n","\n","    # Pentachoron band analysis\n","    cv_vals = trajectories['cv']\n","    in_band = [0.20 <= cv <= 0.23 for cv in cv_vals]\n","    print(f\"\\n  Pentachoron band (0.20-0.23):\")\n","    print(f\"    Layers in band: {sum(in_band)}/{len(in_band)}\")\n","    print(f\"    CV range: [{min(cv_vals):.4f}, {max(cv_vals):.4f}]\")\n","    print(f\"    CV mean:  {np.mean(cv_vals):.4f}\")\n","\n","    # ─── 6. Per-Class Analysis ───\n","    if per_class:\n","        print(f\"\\n{'='*70}\")\n","        print(\"  6. PER-CLASS ANALYSIS\")\n","        print(\"=\" * 70)\n","\n","        n_classes = all_logits.shape[1]\n","        class_stats = {}\n","        class_accs = []\n","\n","        for c in range(n_classes):\n","            mask = all_labels == c\n","            if mask.sum() == 0:\n","                continue\n","            c_pred = pred[mask]\n","            c_labels = all_labels[mask]\n","            c_acc = (c_pred == c_labels).float().mean().item()\n","            c_conf = F.softmax(all_logits[mask], dim=-1)[:, c].mean().item()\n","            c_entropy = -(F.softmax(all_logits[mask], dim=-1) *\n","                         F.log_softmax(all_logits[mask], dim=-1)).sum(-1).mean().item()\n","\n","            class_stats[int(c)] = {\n","                'accuracy': c_acc,\n","                'confidence': c_conf,\n","                'entropy': c_entropy,\n","                'n_samples': int(mask.sum().item()),\n","            }\n","            class_accs.append(c_acc)\n","\n","        report['per_class'] = class_stats\n","\n","        class_accs_arr = np.array(class_accs)\n","        print(f\"  Class accuracy distribution:\")\n","        print(f\"    Mean:   {class_accs_arr.mean():.4f}\")\n","        print(f\"    Std:    {class_accs_arr.std():.4f}\")\n","        print(f\"    Min:    {class_accs_arr.min():.4f}\")\n","        print(f\"    Max:    {class_accs_arr.max():.4f}\")\n","        print(f\"    P25:    {np.percentile(class_accs_arr, 25):.4f}\")\n","        print(f\"    P75:    {np.percentile(class_accs_arr, 75):.4f}\")\n","\n","        # Best and worst classes\n","        sorted_classes = sorted(class_stats.items(), key=lambda x: x[1]['accuracy'])\n","        print(f\"\\n  Bottom 5 classes:\")\n","        for c, stats in sorted_classes[:5]:\n","            print(f\"    Class {c:>3d}: acc={stats['accuracy']:.3f}  \"\n","                  f\"conf={stats['confidence']:.3f}  ent={stats['entropy']:.3f}\")\n","        print(f\"\\n  Top 5 classes:\")\n","        for c, stats in sorted_classes[-5:]:\n","            print(f\"    Class {c:>3d}: acc={stats['accuracy']:.3f}  \"\n","                  f\"conf={stats['confidence']:.3f}  ent={stats['entropy']:.3f}\")\n","\n","    # ─── 7. Model Summary ───\n","    print(f\"\\n{'='*70}\")\n","    print(\"  7. MODEL SUMMARY\")\n","    print(\"=\" * 70)\n","\n","    n_params = sum(p.numel() for p in model.parameters())\n","    n_cayley = sum(1 for _ in (m for m in model.modules() if isinstance(m, CayleyOrthogonal)))\n","    n_film = sum(1 for _ in (m for m in model.modules() if isinstance(m, FiLMLayer)))\n","    n_observers = sum(1 for name, m in model.named_modules()\n","                      if hasattr(m, 'association') and hasattr(m.association, 'constellation'))\n","\n","    print(f\"  Parameters:     {n_params:,}\")\n","    print(f\"  Layers:         {n_layers}\")\n","    print(f\"  Cayley modules: {n_cayley}\")\n","    print(f\"  FiLM modules:   {n_film}\")\n","    print(f\"  Observers:      {n_observers}\")\n","    print(f\"  Overall acc:    {overall_acc:.4f}\")\n","\n","    report['model'] = {\n","        'n_params': n_params,\n","        'n_layers': n_layers,\n","        'n_cayley': n_cayley,\n","        'n_film': n_film,\n","        'n_observers': n_observers,\n","        'accuracy': overall_acc,\n","    }\n","\n","    return report\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# DATA\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","def get_test_loader(config):\n","    import torchvision\n","    import torchvision.transforms as T\n","\n","    test_transform = T.Compose([\n","        T.ToTensor(),\n","        T.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),\n","    ])\n","    test_ds = torchvision.datasets.CIFAR100(\n","        root='./data', train=False, download=True, transform=test_transform)\n","    return torch.utils.data.DataLoader(\n","        test_ds, batch_size=config.get('batch_size', 256), shuffle=False,\n","        num_workers=4, pin_memory=True)\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# MAIN\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","def main():\n","    parser = argparse.ArgumentParser(description='Geometric Transformer Analysis')\n","    parser.add_argument('--checkpoint', type=str, default='geo_cifar100/best.pt',\n","                        help='Path to checkpoint')\n","    parser.add_argument('--n_batches', type=int, default=4,\n","                        help='Number of test batches to analyze')\n","    parser.add_argument('--per_class', action='store_true',\n","                        help='Run per-class analysis')\n","    parser.add_argument('--output', type=str, default='geo_analysis.json',\n","                        help='Output JSON report path')\n","    args, _ = parser.parse_known_args()\n","\n","    print(f\"Device: {device}\")\n","    if device.type == 'cuda':\n","        print(f\"  GPU: {torch.cuda.get_device_name()}\")\n","\n","    # Load checkpoint\n","    print(f\"\\nLoading checkpoint: {args.checkpoint}\")\n","    ckpt = torch.load(args.checkpoint, map_location='cpu', weights_only=False)\n","    config = ckpt['config']\n","    print(f\"  Epoch: {ckpt['epoch']}\")\n","    print(f\"  Test acc: {ckpt['test_acc']:.4f}\")\n","    print(f\"  Config: d={config['d_model']}, layers={config['n_layers']}, \"\n","          f\"anchors={config['n_anchors']}\")\n","\n","    # Build model\n","    model = GeoViTClassifier('analysis', config)\n","    model.load_state_dict(ckpt['state_dict'], strict=False)\n","    if hasattr(model, 'network_to'):\n","        model.network_to(device=device, strict=False)\n","    else:\n","        model = model.to(device)\n","    model.eval()\n","\n","    n_params = sum(p.numel() for p in model.parameters())\n","    print(f\"  Parameters: {n_params:,}\")\n","\n","    # Data\n","    test_loader = get_test_loader(config)\n","\n","    # Run analysis\n","    print(f\"\\n{'━'*70}\")\n","    print(f\"  GEOMETRIC ANALYSIS BATTERY\")\n","    print(f\"{'━'*70}\")\n","\n","    report = run_analysis(\n","        model, test_loader, device,\n","        n_batches=args.n_batches,\n","        per_class=args.per_class,\n","    )\n","\n","    # Save report\n","    output_path = Path(args.output)\n","    # Convert non-serializable types\n","    def clean_for_json(obj):\n","        if isinstance(obj, torch.Tensor):\n","            if obj.dim() == 0:\n","                return obj.item()\n","            return obj.tolist()\n","        if isinstance(obj, (np.float32, np.float64)):\n","            return float(obj)\n","        if isinstance(obj, (np.int32, np.int64)):\n","            return int(obj)\n","        if isinstance(obj, np.ndarray):\n","            return obj.tolist()\n","        if isinstance(obj, dict):\n","            return {k: clean_for_json(v) for k, v in obj.items()}\n","        if isinstance(obj, list):\n","            return [clean_for_json(v) for v in obj]\n","        return obj\n","\n","    class TensorEncoder(json.JSONEncoder):\n","        def default(self, obj):\n","            if isinstance(obj, torch.Tensor):\n","                return obj.cpu().tolist() if obj.dim() > 0 else obj.item()\n","            if isinstance(obj, (np.float32, np.float64, np.floating)):\n","                return float(obj)\n","            if isinstance(obj, (np.int32, np.int64, np.integer)):\n","                return int(obj)\n","            if isinstance(obj, np.ndarray):\n","                return obj.tolist()\n","            return super().default(obj)\n","\n","    with open(output_path, 'w') as f:\n","        json.dump(clean_for_json(report), f, indent=2, cls=TensorEncoder)\n","    print(f\"\\nReport saved: {output_path}\")\n","\n","    print(f\"\\n{'═'*70}\")\n","    print(f\"  ANALYSIS COMPLETE\")\n","    print(f\"{'═'*70}\")\n","\n","\n","if __name__ == '__main__':\n","    main()"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"DvXZRyMgz_H5","executionInfo":{"status":"ok","timestamp":1774691488733,"user_tz":420,"elapsed":1455,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"3f78d3ad-9abc-4ea3-a4ee-e84aab04100a"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["Device: cuda\n","  GPU: NVIDIA RTX PRO 6000 Blackwell Server Edition\n","\n","Loading checkpoint: geo_cifar10/best.pt\n","  Epoch: 97\n","  Test acc: 0.5880\n","  Config: d=256, layers=8, anchors=128\n","  Parameters: 17,615,780\n","\n","━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\n","  GEOMETRIC ANALYSIS BATTERY\n","━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\n","\n","======================================================================\n","  1. SVD INPUT STAGE\n","======================================================================\n","  Entropy:          2.614 ± 0.048\n","  Novelty norm:     13.374\n","  Top-1 ratio:      0.157\n","  Condition number: 7.9\n","  Singular values:  16.12, 11.66, 10.05, 8.84, 7.82...\n","  FiLM γ dev:       0.1060\n","  FiLM β norm:      0.0082\n","\n","──────────────────────────────────────────────────────────────────────\n","  LAYER 0\n","──────────────────────────────────────────────────────────────────────\n","  CV embeddings:    0.5818  ✗ outside\n","    mean_dist=0.9453  std=0.5499\n","  CV anchors:       0.3309  ✗ outside\n","    nn_dist: mean=0.5348  min=0.3366  max=0.8810\n","    max_cos_sim:    0.6634\n","\n","  Stream agreement: 0.5710 ± 0.1056\n","    range: [-0.1664, 0.8103]\n","  ‖content‖:        14.496\n","  ‖geometric‖:      14.276\n","  ‖disagree‖:       13.239  (ratio: 0.913)\n","  ‖agree‖:          19.366\n","\n","  Anchors: 108/128 active, 20 dead\n","  Entropy: 3.238 / 4.852 (normalized: 0.667)\n","  Top-1 dominance: 0.1317\n","  Top-5 counts: [8767, 8647, 5682, 4655, 3884]\n","\n","  Assignment entropy: 3.689 ± 0.266\n","  Max prob:           0.0889 ± 0.0356\n","\n","  Patchwork norm:     6.757\n","  Patchwork sparsity: 0.030  (dead: 0.0000)\n","  Bridge ↔ assignment KL: 1.2455\n","\n","  Geo context norm:   5.601\n","  Composed norm:      5.657\n","\n","──────────────────────────────────────────────────────────────────────\n","  LAYER 1\n","──────────────────────────────────────────────────────────────────────\n","  CV embeddings:    0.5212  ✗ outside\n","    mean_dist=0.9869  std=0.5144\n","  CV anchors:       0.2926  ✗ outside\n","    nn_dist: mean=0.5817  min=0.3401  max=0.9386\n","    max_cos_sim:    0.6599\n","\n","  Stream agreement: 0.4951 ± 0.0908\n","    range: [0.0567, 0.7982]\n","  ‖content‖:        14.552\n","  ‖geometric‖:      14.259\n","  ‖disagree‖:       14.419  (ratio: 0.991)\n","  ‖agree‖:          20.954\n","\n","  Anchors: 122/128 active, 6 dead\n","  Entropy: 3.374 / 4.852 (normalized: 0.695)\n","  Top-1 dominance: 0.1648\n","  Top-5 counts: [10972, 9111, 5272, 3543, 2881]\n","\n","  Assignment entropy: 3.623 ± 0.436\n","  Max prob:           0.0926 ± 0.0401\n","\n","  Patchwork norm:     6.537\n","  Patchwork sparsity: 0.049  (dead: 0.0000)\n","  Bridge ↔ assignment KL: 1.2699\n","\n","  Geo context norm:   5.387\n","  Composed norm:      5.657\n","\n","──────────────────────────────────────────────────────────────────────\n","  LAYER 2\n","──────────────────────────────────────────────────────────────────────\n","  CV embeddings:    0.3559  ✗ outside\n","    mean_dist=0.9558  std=0.3402\n","  CV anchors:       0.0728  ✗ outside\n","    nn_dist: mean=0.8668  min=0.6892  max=0.9660\n","    max_cos_sim:    0.3108\n","\n","  Stream agreement: 0.5040 ± 0.0911\n","    range: [0.0057, 0.8542]\n","  ‖content‖:        14.611\n","  ‖geometric‖:      14.265\n","  ‖disagree‖:       14.323  (ratio: 0.980)\n","  ‖agree‖:          22.637\n","\n","  Anchors: 128/128 active, 0 dead\n","  Entropy: 3.707 / 4.852 (normalized: 0.764)\n","  Top-1 dominance: 0.1378\n","  Top-5 counts: [9174, 6442, 4740, 4082, 3106]\n","\n","  Assignment entropy: 4.032 ± 0.326\n","  Max prob:           0.0872 ± 0.0369\n","\n","  Patchwork norm:     6.993\n","  Patchwork sparsity: 0.010  (dead: 0.0000)\n","  Bridge ↔ assignment KL: 0.8795\n","\n","  Geo context norm:   4.169\n","  Composed norm:      5.657\n","\n","──────────────────────────────────────────────────────────────────────\n","  LAYER 3\n","──────────────────────────────────────────────────────────────────────\n","  CV embeddings:    0.5173  ✗ outside\n","    mean_dist=0.9940  std=0.5142\n","  CV anchors:       0.2275  ✓ IN BAND\n","    nn_dist: mean=0.6495  min=0.3554  max=0.9487\n","    max_cos_sim:    0.6446\n","\n","  Stream agreement: 0.4752 ± 0.1013\n","    range: [0.0166, 0.8098]\n","  ‖content‖:        14.639\n","  ‖geometric‖:      14.253\n","  ‖disagree‖:       14.733  (ratio: 1.006)\n","  ‖agree‖:          22.272\n","\n","  Anchors: 126/128 active, 2 dead\n","  Entropy: 2.992 / 4.852 (normalized: 0.617)\n","  Top-1 dominance: 0.3002\n","  Top-5 counts: [19983, 8341, 5374, 4567, 2886]\n","\n","  Assignment entropy: 3.734 ± 0.331\n","  Max prob:           0.0992 ± 0.0362\n","\n","  Patchwork norm:     6.941\n","  Patchwork sparsity: 0.024  (dead: 0.0000)\n","  Bridge ↔ assignment KL: 1.2729\n","\n","  Geo context norm:   5.199\n","  Composed norm:      5.657\n","\n","──────────────────────────────────────────────────────────────────────\n","  LAYER 4\n","──────────────────────────────────────────────────────────────────────\n","  CV embeddings:    0.2149  ✓ IN BAND\n","    mean_dist=0.9653  std=0.2075\n","  CV anchors:       0.0274  ✗ outside\n","    nn_dist: mean=0.9506  min=0.9262  max=0.9663\n","    max_cos_sim:    0.0738\n","\n","  Stream agreement: 0.4917 ± 0.0985\n","    range: [0.0370, 0.7965]\n","  ‖content‖:        14.675\n","  ‖geometric‖:      14.321\n","  ‖disagree‖:       14.551  (ratio: 0.992)\n","  ‖agree‖:          24.259\n","\n","  Anchors: 128/128 active, 0 dead\n","  Entropy: 4.351 / 4.852 (normalized: 0.897)\n","  Top-1 dominance: 0.0932\n","  Top-5 counts: [6205, 2616, 2100, 1967, 1804]\n","\n","  Assignment entropy: 4.444 ± 0.070\n","  Max prob:           0.0586 ± 0.0222\n","\n","  Patchwork norm:     7.212\n","  Patchwork sparsity: 0.013  (dead: 0.0000)\n","  Bridge ↔ assignment KL: 0.5672\n","\n","  Geo context norm:   3.356\n","  Composed norm:      5.657\n","\n","──────────────────────────────────────────────────────────────────────\n","  LAYER 5\n","──────────────────────────────────────────────────────────────────────\n","  CV embeddings:    0.5516  ✗ outside\n","    mean_dist=0.9760  std=0.5384\n","  CV anchors:       0.2100  ✓ IN BAND\n","    nn_dist: mean=0.6753  min=0.3498  max=0.9591\n","    max_cos_sim:    0.6502\n","\n","  Stream agreement: 0.5019 ± 0.1150\n","    range: [-0.0193, 0.8089]\n","  ‖content‖:        14.668\n","  ‖geometric‖:      14.326\n","  ‖disagree‖:       14.376  (ratio: 0.980)\n","  ‖agree‖:          24.178\n","\n","  Anchors: 127/128 active, 1 dead\n","  Entropy: 2.929 / 4.852 (normalized: 0.604)\n","  Top-1 dominance: 0.2254\n","  Top-5 counts: [15004, 11686, 10564, 4106, 2471]\n","\n","  Assignment entropy: 3.591 ± 0.378\n","  Max prob:           0.1064 ± 0.0357\n","\n","  Patchwork norm:     6.922\n","  Patchwork sparsity: 0.016  (dead: 0.0000)\n","  Bridge ↔ assignment KL: 1.4438\n","\n","  Geo context norm:   4.382\n","  Composed norm:      5.657\n","\n","──────────────────────────────────────────────────────────────────────\n","  LAYER 6\n","──────────────────────────────────────────────────────────────────────\n","  CV embeddings:    0.1839  ✗ outside\n","    mean_dist=0.9620  std=0.1769\n","  CV anchors:       0.0286  ✗ outside\n","    nn_dist: mean=0.9472  min=0.9222  max=0.9663\n","    max_cos_sim:    0.0778\n","\n","  Stream agreement: 0.4953 ± 0.1109\n","    range: [-0.0383, 0.8394]\n","  ‖content‖:        14.707\n","  ‖geometric‖:      14.320\n","  ‖disagree‖:       14.498  (ratio: 0.986)\n","  ‖agree‖:          22.482\n","\n","  Anchors: 128/128 active, 0 dead\n","  Entropy: 4.508 / 4.852 (normalized: 0.929)\n","  Top-1 dominance: 0.0456\n","  Top-5 counts: [3038, 2576, 1633, 1623, 1600]\n","\n","  Assignment entropy: 4.441 ± 0.072\n","  Max prob:           0.0576 ± 0.0217\n","\n","  Patchwork norm:     7.264\n","  Patchwork sparsity: 0.010  (dead: 0.0000)\n","  Bridge ↔ assignment KL: 0.5440\n","\n","  Geo context norm:   4.090\n","  Composed norm:      5.657\n","\n","──────────────────────────────────────────────────────────────────────\n","  LAYER 7\n","──────────────────────────────────────────────────────────────────────\n","  CV embeddings:    0.1524  ✗ outside\n","    mean_dist=0.9659  std=0.1472\n","  CV anchors:       0.0267  ✗ outside\n","    nn_dist: mean=0.9553  min=0.9330  max=0.9697\n","    max_cos_sim:    0.0670\n","\n","  Stream agreement: 0.5149 ± 0.1081\n","    range: [0.0578, 0.8161]\n","  ‖content‖:        14.676\n","  ‖geometric‖:      14.262\n","  ‖disagree‖:       14.168  (ratio: 0.965)\n","  ‖agree‖:          18.667\n","\n","  Anchors: 128/128 active, 0 dead\n","  Entropy: 4.571 / 4.852 (normalized: 0.942)\n","  Top-1 dominance: 0.0294\n","  Top-5 counts: [1958, 1840, 1752, 1626, 1469]\n","\n","  Assignment entropy: 4.458 ± 0.060\n","  Max prob:           0.0574 ± 0.0213\n","\n","  Patchwork norm:     7.199\n","  Patchwork sparsity: 0.009  (dead: 0.0000)\n","  Bridge ↔ assignment KL: 0.5030\n","\n","  Geo context norm:   4.999\n","  Composed norm:      5.657\n","\n","======================================================================\n","  3. CAYLEY ROTATIONS\n","======================================================================\n","  components.transformer.components.layer_0.components.rotation\n","    ‖R-I‖=7.9747  det=1.000004  ‖A‖=3.1221  max|A|=0.079771\n","  components.transformer.components.layer_1.components.rotation\n","    ‖R-I‖=9.0687  det=1.000003  ‖A‖=3.6030  max|A|=0.081425\n","  components.transformer.components.layer_2.components.rotation\n","    ‖R-I‖=9.0947  det=0.999997  ‖A‖=3.5666  max|A|=0.086583\n","  components.transformer.components.layer_3.components.rotation\n","    ‖R-I‖=9.4239  det=0.999994  ‖A‖=3.7388  max|A|=0.127077\n","  components.transformer.components.layer_4.components.rotation\n","    ‖R-I‖=9.8966  det=0.999998  ‖A‖=3.9441  max|A|=0.094444\n","  components.transformer.components.layer_5.components.rotation\n","    ‖R-I‖=10.0317  det=0.999999  ‖A‖=4.0110  max|A|=0.099184\n","  components.transformer.components.layer_6.components.rotation\n","    ‖R-I‖=10.6463  det=1.000002  ‖A‖=4.3021  max|A|=0.111726\n","  components.transformer.components.layer_7.components.rotation\n","    ‖R-I‖=10.8984  det=1.000011  ‖A‖=4.4340  max|A|=0.135460\n","  components.transformer.components.cross_rot_0\n","    ‖R-I‖=7.7315  det=1.000002  ‖A‖=2.9510  max|A|=0.076292\n","  components.transformer.components.cross_rot_1\n","    ‖R-I‖=8.0402  det=1.000008  ‖A‖=3.0686  max|A|=0.084121\n","  components.transformer.components.cross_rot_2\n","    ‖R-I‖=8.2310  det=1.000011  ‖A‖=3.1494  max|A|=0.081548\n","  components.transformer.components.cross_rot_3\n","    ‖R-I‖=8.6737  det=1.000002  ‖A‖=3.3424  max|A|=0.096540\n","  components.transformer.components.cross_rot_4\n","    ‖R-I‖=8.9007  det=1.000002  ‖A‖=3.4457  max|A|=0.082012\n","  components.transformer.components.cross_rot_5\n","    ‖R-I‖=9.1983  det=0.999993  ‖A‖=3.5820  max|A|=0.089175\n","  components.transformer.components.cross_rot_6\n","    ‖R-I‖=9.4788  det=0.999981  ‖A‖=3.7370  max|A|=0.096476\n","\n","======================================================================\n","  4. FiLM LAYERS\n","======================================================================\n","   0 [ents.layer_0.components.geometric.film_q]  γ_dev=0.1582  β_dev=0.0256  ACTIVE\n","   1 [ents.layer_0.components.geometric.film_k]  γ_dev=0.1479  β_dev=0.0000  ACTIVE\n","   2 [ts.layer_0.components.geometric.film_ffn]  γ_dev=0.1384  β_dev=0.0034  ACTIVE\n","   3 [ents.layer_1.components.geometric.film_q]  γ_dev=0.1362  β_dev=0.0373  ACTIVE\n","   4 [ents.layer_1.components.geometric.film_k]  γ_dev=0.1296  β_dev=0.0000  ACTIVE\n","   5 [ts.layer_1.components.geometric.film_ffn]  γ_dev=0.1445  β_dev=0.0046  ACTIVE\n","   6 [ents.layer_2.components.geometric.film_q]  γ_dev=0.1242  β_dev=0.0351  ACTIVE\n","   7 [ents.layer_2.components.geometric.film_k]  γ_dev=0.1213  β_dev=0.0000  ACTIVE\n","   8 [ts.layer_2.components.geometric.film_ffn]  γ_dev=0.1347  β_dev=0.0071  ACTIVE\n","   9 [ents.layer_3.components.geometric.film_q]  γ_dev=0.1326  β_dev=0.0384  ACTIVE\n","  10 [ents.layer_3.components.geometric.film_k]  γ_dev=0.1226  β_dev=0.0000  ACTIVE\n","  11 [ts.layer_3.components.geometric.film_ffn]  γ_dev=0.1335  β_dev=0.0052  ACTIVE\n","  12 [ents.layer_4.components.geometric.film_q]  γ_dev=0.1279  β_dev=0.0392  ACTIVE\n","  13 [ents.layer_4.components.geometric.film_k]  γ_dev=0.1237  β_dev=0.0000  ACTIVE\n","  14 [ts.layer_4.components.geometric.film_ffn]  γ_dev=0.1190  β_dev=0.0081  ACTIVE\n","  15 [ents.layer_5.components.geometric.film_q]  γ_dev=0.1255  β_dev=0.0410  ACTIVE\n","  16 [ents.layer_5.components.geometric.film_k]  γ_dev=0.1180  β_dev=0.0000  ACTIVE\n","  17 [ts.layer_5.components.geometric.film_ffn]  γ_dev=0.1185  β_dev=0.0065  ACTIVE\n","  18 [ents.layer_6.components.geometric.film_q]  γ_dev=0.1333  β_dev=0.0394  ACTIVE\n","  19 [ents.layer_6.components.geometric.film_k]  γ_dev=0.1217  β_dev=0.0000  ACTIVE\n","  20 [ts.layer_6.components.geometric.film_ffn]  γ_dev=0.1019  β_dev=0.0063  ACTIVE\n","  21 [ents.layer_7.components.geometric.film_q]  γ_dev=0.1063  β_dev=0.0355  ACTIVE\n","  22 [ents.layer_7.components.geometric.film_k]  γ_dev=0.1033  β_dev=0.0000  ACTIVE\n","  23 [ts.layer_7.components.geometric.film_ffn]  γ_dev=0.0848  β_dev=0.0060  ACTIVE\n","\n","======================================================================\n","  5. CROSS-LAYER TRAJECTORIES\n","======================================================================\n","  cv                  : [0.5818  0.5212  0.3559  0.5173  0.2149  0.5516  0.1839  0.1524]  Δ=-0.4294\n","  agreement           : [0.5710  0.4951  0.5040  0.4752  0.4917  0.5019  0.4953  0.5149]  Δ=-0.0561\n","  disagree_norm       : [13.2387  14.4192  14.3227  14.7334  14.5513  14.3758  14.4981  14.1681]  Δ=+0.9294\n","  anchor_entropy      : [0.6674  0.6954  0.7640  0.6166  0.8967  0.6036  0.9291  0.9421]  Δ=+0.2746\n","  patchwork_norm      : [6.7572  6.5370  6.9934  6.9409  7.2119  6.9216  7.2638  7.1994]  Δ=+0.4422\n","  bridge_kl           : [1.2455  1.2699  0.8795  1.2729  0.5672  1.4438  0.5440  0.5030]  Δ=-0.7425\n","  composed_norm       : [5.6569  5.6569  5.6569  5.6569  5.6569  5.6569  5.6569  5.6569]  Δ=0.0000\n","\n","  Pentachoron band (0.20-0.23):\n","    Layers in band: 1/8\n","    CV range: [0.1524, 0.5818]\n","    CV mean:  0.3849\n","\n","======================================================================\n","  7. MODEL SUMMARY\n","======================================================================\n","  Parameters:     17,615,780\n","  Layers:         8\n","  Cayley modules: 15\n","  FiLM modules:   24\n","  Observers:      8\n","  Overall acc:    0.5759\n","\n","Report saved: geo_analysis.json\n","\n","══════════════════════════════════════════════════════════════════════\n","  ANALYSIS COMPLETE\n","══════════════════════════════════════════════════════════════════════\n"]}]},{"cell_type":"code","source":["\"\"\"\n","Geometric Transformer — GeoLIP Pipeline Integration\n","=====================================================\n","Dual-stream transformer with constellation-routed attention,\n","quaternion composition, and per-layer Cayley alignment.\n","\n","Uses REAL geolip_core components:\n","    core.associate.constellation  — ConstellationObserver (anchors + triangulation + patchwork)\n","    core.curate.gate              — AnchorGate (CM determinant validity)\n","    core.align.procrustes         — CayleyOrthogonal rotation in SO(d)\n","    pipeline.observer             — TorchComponent / BaseTower interfaces\n","\n","NEW components (transformer-specific):\n","    ManifoldProjection            — Input stage: hidden_state → S^(d-1)\n","    PositionGeometricContext      — Curation: constellation output → FiLM context\n","    FiLMLayer                     — Feature-wise Linear Modulation (proven in Ryan Spearman)\n","    GeometricAttention            — Attention with FiLM on Q,K from curated constellation\n","    QuaternionCompose             — Hamilton product of dual-stream outputs (proven)\n","    CayleyOrthogonal              — SO(d) rotation via Cayley map (proven)\n","    DualStreamBlock               — Content + geometric streams, aligned + composed\n","    GeometricTransformerLayer     — Full layer: project → observe → attend → compose\n","    GeometricTransformer          — Stack of layers with cross-layer rotation\n","\n","Architecture per layer:\n","    1. ManifoldProjection:  h_i → emb_i on S^(d-1) per position\n","    2. ConstellationObserver: emb_i → {triangulation, assignment, patchwork, bridge}\n","    3. PositionGeometricContext: constellation output → (B, L, context_dim)\n","    4. Stream A (content):  standard self-attention\n","    5. Stream B (geometric): attention with FiLM(Q,K | geo_ctx), V unmodulated\n","    6. CayleyOrthogonal:    align B → A basis\n","    7. QuaternionCompose:   w=content, i=aligned_geo, j=disagree, k=agree\n","    8. Gated residual\n","\n","Design principles from Ryan Spearman (ρ=0.309, 76/84 wins):\n","    - FiLM on Q,K ONLY — geometry routes attention, V stays pure\n","    - FiLM on individual arms BEFORE composition, not after\n","    - Quaternion algebra as structural regularizer (non-commutative coupling)\n","    - Disagreement arm (j) carries the transferable signal\n","    - CayleyOrthogonal guarantees pure rotation (det=1 always)\n","    - Never global average pool — per-position geometric context\n","\n","Usage:\n","    from geometric_transformer import GeometricTransformer\n","\n","    model = GeometricTransformer('geo_xfmr', d_model=512, n_layers=4)\n","    out = model(hidden_states)\n","\n","    # Or as a head on frozen ESM-2:\n","    model = GeometricTransformer('esm2_geo', d_model=1280, n_layers=6)\n","    out = model(esm2_hidden_states)\n","\n","Dependencies:\n","    pip install geolip-core  (includes constellation, patchwork, gate, observer interfaces)\n","\"\"\"\n","\n","import math\n","import torch\n","import torch.nn as nn\n","import torch.nn.functional as F\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# GEOLIP IMPORTS — real components, not reimplementations\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","try:\n","    from geolip_core.core.associate.constellation import (\n","        ConstellationObserver, ConstellationAssociation, ConstellationCuration,\n","        Constellation, init_anchors_repulsion,\n","    )\n","    from geolip_core.core.curate.gate import AnchorGate\n","    from geolip_core.pipeline.observer import (\n","        TorchComponent, BaseTower, Input, Curation, Distinction,\n","    )\n","    _HAS_GEOLIP = True\n","except ImportError:\n","    _HAS_GEOLIP = False\n","\n","    # ── Fallback stubs ──\n","    class TorchComponent(nn.Module):\n","        def __init__(self, name=None, **kwargs):\n","            super().__init__()\n","            self._component_name = name or self.__class__.__name__\n","\n","    class BaseTower(nn.Module):\n","        def __init__(self, name=None, **kwargs):\n","            super().__init__()\n","            self._tower_name = name or self.__class__.__name__\n","            self._components = nn.ModuleDict()\n","            self._cache = {}\n","\n","        def attach(self, name, module):\n","            if isinstance(module, nn.Module):\n","                self._components[name] = module\n","            return self\n","\n","        def has(self, name):\n","            return name in self._components\n","\n","        def __getitem__(self, key):\n","            return self._components[key]\n","\n","        def cache_set(self, key, value):\n","            self._cache[key] = value\n","\n","        def cache_get(self, key, default=None):\n","            return self._cache.get(key, default)\n","\n","        def cache_clear(self):\n","            self._cache.clear()\n","\n","    Input = TorchComponent\n","    Curation = TorchComponent\n","    Distinction = TorchComponent\n","\n","    class Constellation(nn.Module):\n","        \"\"\"Learned anchors on S^(d-1). Triangulates input embeddings.\"\"\"\n","        def __init__(self, n_anchors, dim, anchor_drop=0.0, anchor_init='repulsion'):\n","            super().__init__()\n","            self.n_anchors = n_anchors\n","            self.dim = dim\n","            self.anchor_drop = anchor_drop\n","            anchors = torch.randn(n_anchors, dim)\n","            # Repulsion-initialized\n","            anchors = F.normalize(anchors, dim=-1)\n","            for _ in range(200):\n","                sim = anchors @ anchors.T\n","                sim.fill_diagonal_(-2.0)\n","                anchors = F.normalize(anchors - 0.05 * anchors[sim.argmax(dim=1)], dim=-1)\n","            self.anchors = nn.Parameter(anchors)\n","\n","        def triangulate(self, emb, training=False):\n","            anchors = F.normalize(self.anchors, dim=-1)\n","            cos = emb @ anchors.T\n","            tri = 1.0 - cos\n","            _, nearest = cos.max(dim=-1)\n","            return tri, nearest\n","\n","        def forward(self, emb, training=False):\n","            return self.triangulate(emb, training)\n","\n","    class ConstellationAssociation(TorchComponent):\n","        \"\"\"Association through constellation anchors.\"\"\"\n","        def __init__(self, dim=256, n_anchors=32, anchor_drop=0.0,\n","                     anchor_init='repulsion', assign_temp=0.1, **kwargs):\n","            super().__init__(**kwargs)\n","            self.assign_temp = assign_temp\n","            self.constellation = Constellation(n_anchors, dim, anchor_drop, anchor_init)\n","\n","        @property\n","        def frame_dim(self):\n","            return self.constellation.n_anchors\n","\n","        def associate(self, emb, **context):\n","            anchors_n = F.normalize(self.constellation.anchors, dim=-1)\n","            cos = emb @ anchors_n.T\n","            tri = 1.0 - cos\n","            _, nearest = cos.max(dim=-1)\n","            soft_assign = F.softmax(cos / self.assign_temp, dim=-1)\n","            mag = context.get('mag', None)\n","            distances_weighted = tri * mag if mag is not None else tri\n","            return {\n","                'distances': tri, 'distances_weighted': distances_weighted,\n","                'cos_to_anchors': cos, 'assignment': soft_assign,\n","                'nearest': nearest,\n","            }\n","\n","        def forward(self, emb, **context):\n","            return self.associate(emb, **context)\n","\n","    class Patchwork(nn.Module):\n","        \"\"\"Round-robin patchwork compartments.\"\"\"\n","        def __init__(self, n_anchors, n_comp=8, d_comp=32, activation='gelu'):\n","            super().__init__()\n","            self.n_comp = n_comp\n","            anchors_per = max(1, n_anchors // n_comp)\n","            self.compartments = nn.ModuleList([\n","                nn.Sequential(nn.Linear(anchors_per, d_comp), nn.GELU(), nn.Linear(d_comp, d_comp))\n","                for _ in range(n_comp)\n","            ])\n","            self.output_dim = n_comp * d_comp\n","            self.anchors_per = anchors_per\n","\n","        def forward(self, distances):\n","            parts = []\n","            for i, comp in enumerate(self.compartments):\n","                start = i * self.anchors_per\n","                end = start + self.anchors_per\n","                chunk = distances[..., start:end]\n","                if chunk.shape[-1] < self.anchors_per:\n","                    chunk = F.pad(chunk, (0, self.anchors_per - chunk.shape[-1]))\n","                parts.append(comp(chunk))\n","            return torch.cat(parts, dim=-1)\n","\n","    class ConstellationCuration(Curation):\n","        \"\"\"Curation through patchwork compartments + bridge.\"\"\"\n","        def __init__(self, n_anchors=32, dim=256, n_comp=8, d_comp=32,\n","                     activation='gelu', **kwargs):\n","            super().__init__(**kwargs)\n","            self.dim = dim\n","            self.n_anchors = n_anchors\n","            self.patchwork = Patchwork(n_anchors, n_comp, d_comp, activation)\n","            pw_dim = self.patchwork.output_dim\n","            self.bridge = nn.Linear(pw_dim, n_anchors)\n","            self._feature_dim = n_anchors + pw_dim + dim\n","\n","        @property\n","        def feature_dim(self):\n","            return self._feature_dim\n","\n","        def curate_full(self, association_output, emb=None, **context):\n","            distances = association_output['distances_weighted']\n","            assignment = association_output['assignment']\n","            pw = self.patchwork(distances)\n","            bridge = self.bridge(pw)\n","            parts = [assignment, pw]\n","            if emb is not None:\n","                parts.append(emb)\n","            features = torch.cat(parts, dim=-1)\n","            return {'patchwork': pw, 'bridge': bridge, 'features': features}\n","\n","        def forward(self, association_output, emb=None, **context):\n","            return self.curate_full(association_output, emb=emb, **context)['features']\n","\n","    class ConstellationObserver(nn.Module):\n","        \"\"\"Composed association + curation.\"\"\"\n","        def __init__(self, dim=256, n_anchors=32, n_comp=8, d_comp=32,\n","                     anchor_drop=0.0, anchor_init='repulsion',\n","                     activation='gelu', assign_temp=0.1):\n","            super().__init__()\n","            self.association = ConstellationAssociation(\n","                dim=dim, n_anchors=n_anchors, anchor_drop=anchor_drop,\n","                anchor_init=anchor_init, assign_temp=assign_temp)\n","            self.curation = ConstellationCuration(\n","                n_anchors=n_anchors, dim=dim, n_comp=n_comp,\n","                d_comp=d_comp, activation=activation)\n","\n","        @property\n","        def constellation(self):\n","            return self.association.constellation\n","\n","        @property\n","        def patchwork(self):\n","            return self.curation.patchwork\n","\n","        @property\n","        def feature_dim(self):\n","            return self.curation.feature_dim\n","\n","        def observe(self, emb, **context):\n","            a_out = self.association(emb, **context)\n","            c_out = self.curation.curate_full(a_out, emb=emb, **context)\n","            return {\n","                'embedding': emb, 'features': c_out['features'],\n","                'triangulation': a_out['distances'],\n","                'cos_to_anchors': a_out['cos_to_anchors'],\n","                'nearest': a_out['nearest'],\n","                'assignment': a_out['assignment'],\n","                'patchwork': c_out['patchwork'], 'bridge': c_out['bridge'],\n","            }\n","\n","        def forward(self, emb, **context):\n","            return self.observe(emb, **context)\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# PROVEN COMPONENTS — from Ryan Spearman (unchanged, tested)\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","class FiLMLayer(TorchComponent):\n","    \"\"\"Feature-wise Linear Modulation. Proven in Ryan Spearman.\n","\n","    Produces γ * x + β from geometric context.\n","    Identity-initialized: γ=1, β=0 at init.\n","    \"\"\"\n","    def __init__(self, name, feature_dim, context_dim):\n","        super().__init__(name)\n","        self.to_gamma = nn.Linear(context_dim, feature_dim)\n","        self.to_beta = nn.Linear(context_dim, feature_dim)\n","        nn.init.zeros_(self.to_gamma.weight); nn.init.ones_(self.to_gamma.bias)\n","        nn.init.zeros_(self.to_beta.weight); nn.init.zeros_(self.to_beta.bias)\n","\n","    def forward(self, x, ctx):\n","        \"\"\"x: (B, L, D), ctx: (B, L, C) → (B, L, D)\"\"\"\n","        return self.to_gamma(ctx) * x + self.to_beta(ctx)\n","\n","\n","class CayleyOrthogonal(TorchComponent):\n","    \"\"\"Guaranteed SO(d) rotation via Cayley map. Proven in Procrustes alignment.\n","\n","    Q = (I - A)(I + A)^(-1) where A is skew-symmetric.\n","    det(Q) = 1 always. ‖R-I‖ ≈ 4.1 at convergence in SO(256).\n","\n","    Always computes fresh — no caching, no graph-lifetime issues,\n","    clean trace for CompileRouter.\n","    \"\"\"\n","    def __init__(self, name, dim):\n","        super().__init__(name)\n","        self.dim = dim\n","        self.A_upper = nn.Parameter(torch.zeros(dim * (dim - 1) // 2) * 0.01)\n","        # Static index tensors as buffers for device tracking\n","        idx = torch.triu_indices(dim, dim, offset=1)\n","        self.register_buffer('_triu_row', idx[0], persistent=False)\n","        self.register_buffer('_triu_col', idx[1], persistent=False)\n","        self.register_buffer('_eye', torch.eye(dim), persistent=False)\n","\n","    def get_rotation(self):\n","        \"\"\"Build SO(d) rotation from skew-symmetric parameters.\"\"\"\n","        d = self.dim\n","        A = torch.zeros(d, d, device=self.A_upper.device, dtype=self.A_upper.dtype)\n","        A[self._triu_row, self._triu_col] = self.A_upper\n","        A = A - A.T\n","        return torch.linalg.solve(self._eye + A, self._eye - A)\n","\n","    def forward(self, x):\n","        \"\"\"(..., dim) → (..., dim) rotated.\"\"\"\n","        return x @ self.get_rotation().T\n","\n","\n","def quaternion_multiply(q1, q2):\n","    \"\"\"Hamilton product. q = (w, x, y, z) along dim=-2.\n","\n","    Supports batched: (..., 4, D) × (..., 4, D) → (..., 4, D)\n","    Or scalar:        (..., 4) × (..., 4) → (..., 4)\n","    \"\"\"\n","    w1, x1, y1, z1 = q1.unbind(-2) if q1.dim() >= 2 and q1.shape[-2] == 4 else q1.unbind(-1)\n","    w2, x2, y2, z2 = q2.unbind(-2) if q2.dim() >= 2 and q2.shape[-2] == 4 else q2.unbind(-1)\n","    stack_dim = -2 if q1.dim() >= 2 and q1.shape[-2] == 4 else -1\n","    return torch.stack([\n","        w1*w2 - x1*x2 - y1*y2 - z1*z2,\n","        w1*x2 + x1*w2 + y1*z2 - z1*y2,\n","        w1*y2 - x1*z2 + y1*w2 + z1*x2,\n","        w1*z2 + x1*y2 - y1*x2 + z1*w2,\n","    ], dim=stack_dim)\n","\n","\n","def quaternion_multiply_batched(q1, q2):\n","    \"\"\"Hamilton product on (B, 4, D) tensors. Fully vectorized, no loops.\n","\n","    Each of the 4 slices along dim=1 is one quaternion component.\n","    The D dimension is batched — all D quaternions multiplied in parallel.\n","    \"\"\"\n","    w1, x1, y1, z1 = q1[:, 0], q1[:, 1], q1[:, 2], q1[:, 3]\n","    w2, x2, y2, z2 = q2[:, 0], q2[:, 1], q2[:, 2], q2[:, 3]\n","    return torch.stack([\n","        w1*w2 - x1*x2 - y1*y2 - z1*z2,\n","        w1*x2 + x1*w2 + y1*z2 - z1*y2,\n","        w1*y2 - x1*z2 + y1*w2 + z1*x2,\n","        w1*z2 + x1*y2 - y1*x2 + z1*w2,\n","    ], dim=1)  # (B, 4, D)\n","\n","\n","class QuaternionCompose(TorchComponent):\n","    \"\"\"Four-arm Hamilton product composition. Proven in GeoQuat head.\n","\n","    The algebra forces cross-term interactions between arms.\n","    Arms cannot independently memorize — the non-commutative\n","    product couples their outputs as structural regularizer.\n","\n","    Fully vectorized: single batched Hamilton product, no Python loops.\n","    \"\"\"\n","    def __init__(self, name, input_dim, quat_dim=64):\n","        super().__init__(name)\n","        self.quat_dim = quat_dim\n","        self.proj_w = nn.Linear(input_dim, quat_dim)\n","        self.proj_i = nn.Linear(input_dim, quat_dim)\n","        self.proj_j = nn.Linear(input_dim, quat_dim)\n","        self.proj_k = nn.Linear(input_dim, quat_dim)\n","        self.rotation = nn.Parameter(torch.randn(1, 4, quat_dim) * 0.1)\n","\n","    @property\n","    def output_dim(self):\n","        return self.quat_dim * 4\n","\n","    def forward(self, arm_w, arm_i, arm_j, arm_k):\n","        \"\"\"Each arm: (B, L, D) → composed: (B, L, 4*quat_dim)\"\"\"\n","        shape = arm_w.shape[:-1]\n","        D = arm_w.shape[-1]\n","        flat = arm_w.dim() > 2\n","        if flat:\n","            arm_w = arm_w.reshape(-1, D); arm_i = arm_i.reshape(-1, D)\n","            arm_j = arm_j.reshape(-1, D); arm_k = arm_k.reshape(-1, D)\n","\n","        # q: (N, 4, quat_dim) — stack 4 projected arms as quaternion components\n","        q = torch.stack([self.proj_w(arm_w), self.proj_i(arm_i),\n","                         self.proj_j(arm_j), self.proj_k(arm_k)], dim=1)\n","        q = q / (q.norm(dim=1, keepdim=True) + 1e-8)\n","\n","        # r: (N, 4, quat_dim) — broadcast learned rotation\n","        r = self.rotation.expand(q.shape[0], -1, -1)\n","        r = r / (r.norm(dim=1, keepdim=True) + 1e-8)\n","\n","        # Single batched Hamilton product over all quat_dim simultaneously\n","        # (N, 4, quat_dim) × (N, 4, quat_dim) → (N, 4, quat_dim)\n","        composed = quaternion_multiply_batched(r, q)\n","\n","        # Flatten 4 × quat_dim → 4*quat_dim\n","        composed = composed.reshape(q.shape[0], -1)\n","\n","        if flat:\n","            composed = composed.reshape(*shape, -1)\n","        return composed\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# NEW COMPONENTS — transformer-specific, built for this architecture\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","class ManifoldProjection(TorchComponent):\n","    \"\"\"Input stage: project transformer hidden states to S^(d-1).\n","\n","    Per-position, per-layer projection from model space to the\n","    constellation's embedding space. L2-normalized to sit on the\n","    unit hypersphere.\n","\n","    This is the tap — it reads the representation without modifying it.\n","    \"\"\"\n","    def __init__(self, name, d_model, manifold_dim):\n","        super().__init__(name)\n","        self.proj = nn.Linear(d_model, manifold_dim)\n","        self.norm = nn.LayerNorm(manifold_dim)\n","\n","    def forward(self, hidden_states):\n","        \"\"\"(B, L, D) → (B, L, manifold_dim) on S^(manifold_dim - 1)\"\"\"\n","        h = self.norm(self.proj(hidden_states))\n","        return F.normalize(h, dim=-1)\n","\n","\n","class PositionGeometricContext(TorchComponent):\n","    \"\"\"Curation stage: constellation observation + geometric history → FiLM context.\n","\n","    Three streams:\n","        anchor:     cos_to_anchors + assignment + triangulation — WHERE on the manifold\n","        structural: patchwork + embedding — WHAT the local geometry looks like\n","        history:    geo_residual from previous layers — WHAT prior layers observed\n","\n","    First layer receives zero history. Each subsequent layer reads accumulated\n","    geometric context from all prior constellations.\n","    \"\"\"\n","    def __init__(self, name, n_anchors, pw_dim, manifold_dim, context_dim):\n","        super().__init__(name)\n","        self.context_dim = context_dim\n","        self.pw_dim = pw_dim\n","\n","        # Anchor features: cos + assignment + triangulation = 3 * n_anchors\n","        self.anchor_mlp = nn.Sequential(\n","            nn.Linear(n_anchors * 3, context_dim),\n","            nn.GELU(),\n","            nn.LayerNorm(context_dim),\n","        )\n","        # Structural features: patchwork + embedding\n","        self.struct_mlp = nn.Sequential(\n","            nn.Linear(pw_dim + manifold_dim, context_dim),\n","            nn.GELU(),\n","            nn.LayerNorm(context_dim),\n","        )\n","        # History features: accumulated geo_residual from prior layers\n","        self.history_mlp = nn.Sequential(\n","            nn.Linear(pw_dim, context_dim),\n","            nn.GELU(),\n","            nn.LayerNorm(context_dim),\n","        )\n","        # Fuse anchor + structural + history\n","        self.fuse = nn.Sequential(\n","            nn.Linear(context_dim * 3, context_dim),\n","            nn.GELU(),\n","            nn.LayerNorm(context_dim),\n","        )\n","\n","    def forward(self, obs_dict, geo_residual=None):\n","        \"\"\"\n","        Args:\n","            obs_dict: from ConstellationObserver.observe()\n","            geo_residual: (B*L, pw_dim) accumulated geometric context,\n","                          or None for first layer (uses zeros)\n","        Returns:\n","            (B*L, context_dim) geometric context for FiLM\n","        \"\"\"\n","        anchor_feats = torch.cat([\n","            obs_dict['cos_to_anchors'],\n","            obs_dict['assignment'],\n","            obs_dict['triangulation'],\n","        ], dim=-1)\n","\n","        struct_feats = torch.cat([\n","            obs_dict['patchwork'],\n","            obs_dict['embedding'],\n","        ], dim=-1)\n","\n","        a = self.anchor_mlp(anchor_feats)\n","        s = self.struct_mlp(struct_feats)\n","\n","        # History — zeros on first layer\n","        if geo_residual is not None:\n","            h = self.history_mlp(geo_residual)\n","        else:\n","            h = torch.zeros_like(a)\n","\n","        return self.fuse(torch.cat([a, s, h], dim=-1))\n","\n","\n","class GeometricAttention(TorchComponent):\n","    \"\"\"Attention with FiLM from curated constellation. Stream B.\n","\n","    FiLM modulates Q and K BEFORE attention — the constellation\n","    position controls WHERE attention flows. V stays unmodulated.\n","    FiLM between FFN layers conditions the nonlinearity.\n","\n","    Proven principle: context before composition, not after.\n","    \"\"\"\n","    def __init__(self, name, d_model, n_heads=8, context_dim=128, dropout=0.1):\n","        super().__init__(name)\n","        self.d_model = d_model\n","        self.n_heads = n_heads\n","        self.head_dim = d_model // n_heads\n","        self.scale = self.head_dim ** -0.5\n","\n","        self.w_q = nn.Linear(d_model, d_model)\n","        self.w_k = nn.Linear(d_model, d_model)\n","        self.w_v = nn.Linear(d_model, d_model)\n","        self.w_o = nn.Linear(d_model, d_model)\n","        self.dropout = nn.Dropout(dropout)\n","\n","        # FiLM on Q and K — geometry routes attention\n","        self.film_q = FiLMLayer(f'{name}_film_q', d_model, context_dim)\n","        self.film_k = FiLMLayer(f'{name}_film_k', d_model, context_dim)\n","\n","        self.norm = nn.LayerNorm(d_model)\n","\n","        # FFN with FiLM between layers\n","        self.ffn1 = nn.Linear(d_model, d_model * 4)\n","        self.film_ffn = FiLMLayer(f'{name}_film_ffn', d_model * 4, context_dim)\n","        self.ffn2 = nn.Linear(d_model * 4, d_model)\n","        self.ffn_drop = nn.Dropout(dropout)\n","        self.ffn_norm = nn.LayerNorm(d_model)\n","\n","    def forward(self, x, geo_ctx, attn_mask=None, key_padding_mask=None):\n","        \"\"\"\n","        x: (B, L, D), geo_ctx: (B, L, C) → (B, L, D)\n","        \"\"\"\n","        B, L, D = x.shape\n","        H, HD = self.n_heads, self.head_dim\n","\n","        Q = self.film_q(self.w_q(x), geo_ctx)\n","        K = self.film_k(self.w_k(x), geo_ctx)\n","        V = self.w_v(x)  # V unmodulated — content stays pure\n","\n","        Q = Q.view(B, L, H, HD).transpose(1, 2)\n","        K = K.view(B, L, H, HD).transpose(1, 2)\n","        V = V.view(B, L, H, HD).transpose(1, 2)\n","\n","        scores = (Q @ K.transpose(-2, -1)) * self.scale\n","        if attn_mask is not None:\n","            scores = scores + attn_mask\n","        if key_padding_mask is not None:\n","            scores = scores.masked_fill(\n","                key_padding_mask.unsqueeze(1).unsqueeze(2), float('-inf'))\n","        attn_out = (self.dropout(F.softmax(scores, dim=-1)) @ V)\n","        attn_out = attn_out.transpose(1, 2).reshape(B, L, D)\n","\n","        x = self.norm(x + self.w_o(attn_out))\n","\n","        # FFN with geometric FiLM between layers\n","        h = F.gelu(self.ffn1(x))\n","        h = self.film_ffn(h, geo_ctx)\n","        x = self.ffn_norm(x + self.ffn_drop(self.ffn2(h)))\n","\n","        return x\n","\n","\n","class ContentAttention(TorchComponent):\n","    \"\"\"Standard self-attention. Stream A. No geometric conditioning.\"\"\"\n","    def __init__(self, name, d_model, n_heads=8, dropout=0.1):\n","        super().__init__(name)\n","        self.attn = nn.MultiheadAttention(\n","            d_model, n_heads, dropout=dropout, batch_first=True)\n","        self.norm = nn.LayerNorm(d_model)\n","        self.ffn = nn.Sequential(\n","            nn.Linear(d_model, d_model * 4), nn.GELU(),\n","            nn.Linear(d_model * 4, d_model), nn.Dropout(dropout))\n","        self.ffn_norm = nn.LayerNorm(d_model)\n","\n","    def forward(self, x, attn_mask=None, key_padding_mask=None):\n","        a, _ = self.attn(x, x, x, attn_mask=attn_mask,\n","                         key_padding_mask=key_padding_mask)\n","        x = self.norm(x + a)\n","        x = self.ffn_norm(x + self.ffn(x))\n","        return x\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# LAYER — dual-stream with constellation routing\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","class GeometricTransformerLayer(BaseTower):\n","    \"\"\"One layer of the geometric transformer.\n","\n","    Pipeline per layer:\n","        1. ManifoldProjection: h_i → emb_i on S^(manifold_dim - 1)\n","        2. ConstellationObserver: emb_i → {triangulation, assignment, patchwork, ...}\n","        3. PositionGeometricContext: observation → FiLM context (B, L, context_dim)\n","        4. ContentAttention (Stream A): standard MHA\n","        5. GeometricAttention (Stream B): FiLM(Q,K | geo_ctx), V pure\n","        6. CayleyOrthogonal: align B basis → A basis\n","        7. QuaternionCompose: w=A, i=aligned_B, j=A-B, k=A*B\n","        8. Decode + gated residual\n","\n","    Access:\n","        layer['projection']  → ManifoldProjection\n","        layer['observer']    → ConstellationObserver\n","        layer['context']     → PositionGeometricContext\n","        layer['content']     → ContentAttention\n","        layer['geometric']   → GeometricAttention\n","        layer['rotation']    → CayleyOrthogonal\n","        layer['compose']     → QuaternionCompose\n","    \"\"\"\n","    def __init__(self, name, d_model, n_heads=8, n_anchors=32,\n","                 manifold_dim=256, n_comp=8, d_comp=32,\n","                 context_dim=128, quat_dim=64, dropout=0.1):\n","        super().__init__(name)\n","        self.d_model = d_model\n","\n","        # 1. Project to manifold\n","        self.attach('projection', ManifoldProjection(\n","            f'{name}_proj', d_model, manifold_dim))\n","\n","        # 2. Constellation observer (real association + curation)\n","        self.attach('observer', ConstellationObserver(\n","            dim=manifold_dim, n_anchors=n_anchors,\n","            n_comp=n_comp, d_comp=d_comp))\n","\n","        # 3. Fuse observation into FiLM context\n","        pw_dim = self['observer'].curation.patchwork.output_dim\n","        self.attach('context', PositionGeometricContext(\n","            f'{name}_ctx', n_anchors, pw_dim, manifold_dim, context_dim))\n","\n","        # 4. Stream A: content\n","        self.attach('content', ContentAttention(\n","            f'{name}_content', d_model, n_heads, dropout))\n","\n","        # 5. Stream B: geometric\n","        self.attach('geometric', GeometricAttention(\n","            f'{name}_geo', d_model, n_heads, context_dim, dropout))\n","\n","        # 6. Cayley rotation: align B → A\n","        self.attach('rotation', CayleyOrthogonal(f'{name}_cayley', d_model))\n","\n","        # 7. Quaternion composition\n","        self.attach('compose', QuaternionCompose(\n","            f'{name}_quat', d_model, quat_dim))\n","\n","        # 8. Decode + gate\n","        self.attach('decode', nn.Sequential(\n","            nn.Linear(quat_dim * 4, d_model), nn.GELU(), nn.LayerNorm(d_model)))\n","        self.attach('gate', nn.Sequential(\n","            nn.Linear(d_model * 2, d_model), nn.Sigmoid()))\n","\n","        # 9. Geometric residual stream — accumulates across layers\n","        #    geo_gate: learned scalar gate per dimension, init near-zero (sigmoid(-3) ≈ 0.05)\n","        #    geo_proj: project patchwork into residual space\n","        self._pw_dim = pw_dim\n","        self.attach('geo_proj', nn.Sequential(\n","            nn.Linear(pw_dim, pw_dim), nn.LayerNorm(pw_dim)))\n","        self.attach('geo_gate_proj', nn.Linear(pw_dim, pw_dim))\n","        # Init gate bias to -3 so sigmoid starts near zero (minimal contribution at init)\n","        nn.init.zeros_(self['geo_gate_proj'].weight)\n","        nn.init.constant_(self['geo_gate_proj'].bias, -3.0)\n","\n","    def forward(self, x, geo_residual=None, attn_mask=None, key_padding_mask=None):\n","        \"\"\"\n","        Args:\n","            x: (B, L, D) input hidden states\n","            geo_residual: (B, L, pw_dim) accumulated geometric context from prior layers,\n","                          or None for the first layer (initialized to zeros)\n","\n","        Returns:\n","            x_out: (B, L, D) transformed hidden states\n","            geo_residual_out: (B, L, pw_dim) updated geometric residual\n","            geo_state: dict with full geometric state (11 fields)\n","        \"\"\"\n","        B, L, D = x.shape\n","\n","        # 1. Project to manifold: per-position embedding on S^(d-1)\n","        emb = self['projection'](x)  # (B, L, manifold_dim)\n","\n","        # 2. Constellation observation: flatten to (B*L, manifold_dim) for observer\n","        emb_flat = emb.reshape(B * L, -1)\n","        obs = self['observer'].observe(emb_flat)\n","\n","        # 3. Build FiLM context WITH geometric history\n","        geo_res_flat = None\n","        if geo_residual is not None:\n","            geo_res_flat = geo_residual.reshape(B * L, -1)\n","\n","        geo_ctx_flat = self['context'](obs, geo_residual=geo_res_flat)\n","        geo_ctx = geo_ctx_flat.reshape(B, L, -1)  # (B, L, context_dim)\n","\n","        # 4. Stream A: content attention\n","        a_out = self['content'](x, attn_mask=attn_mask,\n","                                key_padding_mask=key_padding_mask)\n","\n","        # 5. Stream B: geometric attention\n","        b_out = self['geometric'](x, geo_ctx, attn_mask=attn_mask,\n","                                   key_padding_mask=key_padding_mask)\n","\n","        # 6. Cayley rotation: align B → A\n","        b_aligned = self['rotation'](b_out)\n","\n","        # 7. Quaternion composition\n","        composed = self['compose'](\n","            arm_w=a_out, arm_i=b_aligned,\n","            arm_j=a_out - b_aligned, arm_k=a_out * b_aligned)\n","\n","        # 8. Decode + gated residual\n","        decoded = self['decode'](composed)\n","        g = self['gate'](torch.cat([x, decoded], dim=-1))\n","        x_out = g * decoded + (1 - g) * x\n","\n","        # 9. Update geometric residual stream\n","        #    Patchwork features carry the constellation's interpretation.\n","        #    Gated accumulation lets each layer contribute without overwriting.\n","        pw_unflat = obs['patchwork'].reshape(B, L, -1)  # (B, L, pw_dim)\n","        geo_update = self['geo_proj'](pw_unflat)\n","        geo_gate = torch.sigmoid(self['geo_gate_proj'](pw_unflat))\n","\n","        if geo_residual is None:\n","            geo_residual_out = geo_gate * geo_update\n","        else:\n","            geo_residual_out = geo_residual + geo_gate * geo_update\n","\n","        # 10. Build full geometric state\n","        def unflatten(t):\n","            if t is None: return None\n","            if t.dim() == 1: return t.reshape(B, L)\n","            return t.reshape(B, L, *t.shape[1:])\n","\n","        geo_state = {\n","            'embedding':      emb,\n","            'geo_ctx':        geo_ctx,\n","            'triangulation':  unflatten(obs['triangulation']),\n","            'cos_to_anchors': unflatten(obs['cos_to_anchors']),\n","            'assignment':     unflatten(obs['assignment']),\n","            'nearest':        unflatten(obs['nearest']),\n","            'patchwork':      unflatten(obs['patchwork']),\n","            'bridge':         unflatten(obs['bridge']),\n","            'content':        a_out,\n","            'geometric':      b_out,\n","            'composed':       composed,\n","            'geo_residual':   geo_residual_out,\n","        }\n","\n","        return x_out, geo_residual_out, geo_state\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# FULL MODEL — stack of layers\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","class GeometricTransformer(BaseTower):\n","    \"\"\"Geometric Transformer — dual-stream with constellation routing.\n","\n","    Stack of GeometricTransformerLayers. Optional cross-layer Cayley\n","    rotation aligns each layer's output basis to the next layer's\n","    expected input.\n","\n","    Access:\n","        model['layer_0']      → first layer\n","        model['cross_rot_0']  → cross-layer rotation 0→1\n","        model['final_norm']   → output normalization\n","\n","    Args:\n","        name: tower identity\n","        d_model: transformer model dimension\n","        n_heads: attention heads per stream\n","        n_layers: number of geometric transformer layers\n","        n_anchors: constellation anchor points\n","        manifold_dim: dimension of S^(d-1) for constellation\n","        n_comp: patchwork compartments\n","        d_comp: hidden dim per compartment\n","        context_dim: FiLM conditioning dimension\n","        quat_dim: quaternion space dimension\n","        dropout: dropout rate\n","        cross_layer_rotation: add Cayley rotation between layers\n","        vocab_size: if set, adds embedding + output head\n","    \"\"\"\n","    def __init__(self, name, d_model=512, n_heads=8, n_layers=4,\n","                 n_anchors=32, manifold_dim=256, n_comp=8, d_comp=32,\n","                 context_dim=128, quat_dim=64, dropout=0.1,\n","                 cross_layer_rotation=True, vocab_size=None, max_seq_len=2048):\n","        super().__init__(name)\n","        self.d_model = d_model\n","        self.n_layers = n_layers\n","\n","        if vocab_size is not None:\n","            self.attach('embed', nn.Embedding(vocab_size, d_model))\n","            self.attach('pos_embed', nn.Embedding(max_seq_len, d_model))\n","            self.attach('head', nn.Linear(d_model, vocab_size, bias=False))\n","\n","        for i in range(n_layers):\n","            self.attach(f'layer_{i}', GeometricTransformerLayer(\n","                f'{name}_L{i}', d_model, n_heads, n_anchors,\n","                manifold_dim, n_comp, d_comp, context_dim, quat_dim, dropout))\n","\n","        if cross_layer_rotation and n_layers > 1:\n","            for i in range(n_layers - 1):\n","                self.attach(f'cross_rot_{i}', CayleyOrthogonal(\n","                    f'{name}_xrot_{i}', d_model))\n","\n","        self.attach('final_norm', nn.LayerNorm(d_model))\n","\n","        self._config = dict(\n","            d_model=d_model, n_heads=n_heads, n_layers=n_layers,\n","            n_anchors=n_anchors, manifold_dim=manifold_dim,\n","            n_comp=n_comp, d_comp=d_comp, context_dim=context_dim,\n","            quat_dim=quat_dim, dropout=dropout,\n","            cross_layer_rotation=cross_layer_rotation,\n","            vocab_size=vocab_size,\n","        )\n","\n","    @property\n","    def config(self):\n","        return self._config.copy()\n","\n","    def param_report(self):\n","        total = 0\n","        name = getattr(self, '_tower_name', getattr(self, 'name', self.__class__.__name__))\n","        print(f\"\\n  {name} — parameter report\")\n","        print(f\"  {'Component':<35s}  {'Params':>12s}\")\n","        print(f\"  {'─'*35}  {'─'*12}\")\n","        for cname, module in self.named_children():\n","            n = sum(p.numel() for p in module.parameters())\n","            total += n\n","            print(f\"  {cname:<35s}  {n:>12,}\")\n","        print(f\"  {'─'*35}  {'─'*12}\")\n","        print(f\"  {'TOTAL':<35s}  {total:>12,}\")\n","        return total\n","\n","    def forward(self, x, attn_mask=None, key_padding_mask=None,\n","                return_geo_state=False):\n","        \"\"\"\n","        Args:\n","            x: (B, L, D) hidden states or (B, L) token ids\n","            return_geo_state: if True, return per-layer geometric state dicts\n","\n","        Returns:\n","            out: (B, L, D) transformed hidden states (or logits if head attached)\n","            geo_states: list of per-layer geo_state dicts (if return_geo_state)\n","                Each dict contains: embedding, geo_ctx, triangulation,\n","                cos_to_anchors, assignment, nearest, patchwork, bridge,\n","                content, geometric, composed, geo_residual\n","        \"\"\"\n","        if self.has('embed') and x.dtype in (torch.long, torch.int32, torch.int64):\n","            pos = torch.arange(x.shape[1], device=x.device)\n","            x = self['embed'](x) + self['pos_embed'](pos)\n","\n","        geo_states = []\n","        has_xrot = self.has('cross_rot_0')\n","        geo_residual = None  # First layer establishes the stream\n","\n","        for i in range(self.n_layers):\n","            x, geo_residual, geo_state = self[f'layer_{i}'](\n","                x, geo_residual=geo_residual,\n","                attn_mask=attn_mask, key_padding_mask=key_padding_mask)\n","            if return_geo_state:\n","                geo_states.append(geo_state)\n","            if has_xrot and i < self.n_layers - 1:\n","                x = self[f'cross_rot_{i}'](x)\n","                # Note: geo_residual is NOT rotated by cross_rot.\n","                # It lives in its own space — the patchwork interpretation\n","                # space — which is independent of the hidden state basis.\n","                # This is intentional: the geometric residual carries\n","                # constellation observations, not hidden state features.\n","\n","        x = self['final_norm'](x)\n","        if self.has('head'):\n","            x = self['head'](x)\n","\n","        return (x, geo_states) if return_geo_state else x\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# FACTORIES\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","def geo_transformer_esm2(name='geo_esm2', n_layers=6, **kw):\n","    \"\"\"Pre-configured for ESM-2 650M (d=1280).\"\"\"\n","    return GeometricTransformer(name, d_model=1280, n_heads=16,\n","        n_layers=n_layers, n_anchors=32, manifold_dim=256,\n","        n_comp=8, d_comp=32, context_dim=128, quat_dim=64, **kw)\n","\n","def geo_transformer_small(name='geo_small', n_layers=4, **kw):\n","    \"\"\"Small config for prototyping.\"\"\"\n","    return GeometricTransformer(name, d_model=256, n_heads=8,\n","        n_layers=n_layers, n_anchors=16, manifold_dim=128,\n","        n_comp=4, d_comp=16, context_dim=64, quat_dim=32, **kw)\n","\n","def geo_transformer_vision(name='geo_vit', n_layers=4, **kw):\n","    \"\"\"For scatter/SVD vision pipeline (patches as tokens).\"\"\"\n","    return GeometricTransformer(name, d_model=384, n_heads=8,\n","        n_layers=n_layers, n_anchors=32, manifold_dim=128,\n","        n_comp=8, d_comp=16, context_dim=64, quat_dim=32, **kw)\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# SELF-TEST\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","if __name__ == '__main__':\n","    print(\"Geometric Transformer — Self-Test\")\n","    print(f\"  geolip_core available: {_HAS_GEOLIP}\")\n","    print(\"=\" * 60)\n","\n","    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n","\n","    model = geo_transformer_small('test', n_layers=2)\n","    if hasattr(model, 'network_to'):\n","        model.network_to(device=device, strict=False)\n","    else:\n","        model = model.to(device)\n","    total = model.param_report()\n","\n","    B, L, D = 2, 32, 256\n","    x = torch.randn(B, L, D, device=device)\n","\n","    out, geos = model(x, return_geo_state=True)\n","    assert out.shape == (B, L, D), f\"Expected ({B},{L},{D}), got {out.shape}\"\n","    assert len(geos) == 2\n","\n","    print(f\"\\n  Input:  ({B}, {L}, {D})\")\n","    print(f\"  Output: {out.shape}\")\n","    print(f\"  Geo states: {len(geos)} layers\")\n","    print(f\"  State keys: {sorted(geos[0].keys())}\")\n","    for k, v in geos[0].items():\n","        if v is not None:\n","            shape = v.shape if hasattr(v, 'shape') else type(v).__name__\n","            print(f\"    {k:<18s}: {shape}\")\n","\n","    # Verify geo_residual stream continuity\n","    assert 'geo_residual' in geos[0], \"geo_residual missing from layer 0\"\n","    assert 'geo_residual' in geos[1], \"geo_residual missing from layer 1\"\n","    gr0 = geos[0]['geo_residual']\n","    gr1 = geos[1]['geo_residual']\n","    print(f\"\\n  Geo residual stream:\")\n","    print(f\"    Layer 0: {gr0.shape}  norm={gr0.norm(dim=-1).mean():.4f}\")\n","    print(f\"    Layer 1: {gr1.shape}  norm={gr1.norm(dim=-1).mean():.4f}\")\n","    # Layer 1 should have accumulated from layer 0 — norms should differ\n","    print(f\"    Accumulated: {'YES' if gr1.norm() > gr0.norm() * 0.5 else 'NO (gate closed)'}\")\n","\n","    # Verify rotations\n","    for name, module in model.named_modules():\n","        if isinstance(module, CayleyOrthogonal):\n","            R = module.get_rotation()\n","            I = torch.eye(R.shape[0], device=R.device)\n","            print(f\"  {name}: ‖RRᵀ-I‖={((R@R.T)-I).norm():.8f}  det={torch.det(R):.4f}\")\n","\n","    # ESM-2 scale overhead\n","    print(f\"\\n  ESM-2 scale:\")\n","    esm = geo_transformer_esm2('esm2', n_layers=6)\n","    if hasattr(esm, 'network_to'):\n","        esm.network_to(device=device, strict=False)\n","    else:\n","        esm = esm.to(device)\n","    n = esm.param_report()\n","    print(f\"  Overhead on 650M base: {n/1e6:.1f}M ({n/650e6*100:.1f}%)\")\n","\n","    print(f\"\\n{'='*60}\")\n","    print(f\"  PASSED\")\n","    print(f\"{'='*60}\")"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"mGU9Fv2g5vLK","executionInfo":{"status":"ok","timestamp":1774692269048,"user_tz":420,"elapsed":935,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"1484436b-b755-4131-bc38-3b20f3f76157"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["Geometric Transformer — Self-Test\n","  geolip_core available: True\n","============================================================\n","\n","  test — parameter report\n","  Component                                  Params\n","  ───────────────────────────────────  ────────────\n","  components                              4,212,768\n","  stages                                          0\n","  ───────────────────────────────────  ────────────\n","  TOTAL                                   4,212,768\n","\n","  Input:  (2, 32, 256)\n","  Output: torch.Size([2, 32, 256])\n","  Geo states: 2 layers\n","  State keys: ['assignment', 'bridge', 'composed', 'content', 'cos_to_anchors', 'embedding', 'geo_ctx', 'geo_residual', 'geometric', 'nearest', 'patchwork', 'triangulation']\n","    embedding         : torch.Size([2, 32, 128])\n","    geo_ctx           : torch.Size([2, 32, 64])\n","    triangulation     : torch.Size([2, 32, 16])\n","    cos_to_anchors    : torch.Size([2, 32, 16])\n","    assignment        : torch.Size([2, 32, 16])\n","    nearest           : torch.Size([2, 32])\n","    patchwork         : torch.Size([2, 32, 64])\n","    bridge            : torch.Size([2, 32, 16])\n","    content           : torch.Size([2, 32, 256])\n","    geometric         : torch.Size([2, 32, 256])\n","    composed          : torch.Size([2, 32, 128])\n","    geo_residual      : torch.Size([2, 32, 64])\n","\n","  Geo residual stream:\n","    Layer 0: torch.Size([2, 32, 64])  norm=0.3794\n","    Layer 1: torch.Size([2, 32, 64])  norm=0.5534\n","    Accumulated: YES\n","  components.layer_0.components.rotation: ‖RRᵀ-I‖=0.00000000  det=1.0000\n","  components.layer_1.components.rotation: ‖RRᵀ-I‖=0.00000000  det=1.0000\n","  components.cross_rot_0: ‖RRᵀ-I‖=0.00000000  det=1.0000\n","\n","  ESM-2 scale:\n","\n","  esm2 — parameter report\n","  Component                                  Params\n","  ───────────────────────────────────  ────────────\n","  components                            284,599,104\n","  stages                                          0\n","  ───────────────────────────────────  ────────────\n","  TOTAL                                 284,599,104\n","  Overhead on 650M base: 284.6M (43.8%)\n","\n","============================================================\n","  PASSED\n","============================================================\n"]}]},{"cell_type":"markdown","source":["# v 1.1"],"metadata":{"id":"FYHGvq016CIq"}},{"cell_type":"code","source":["\"\"\"\n","Geometric Transformer — CIFAR-100 Training\n","\n","Patches 32×32 images into tokens, feeds through the geometric\n","transformer with constellation-routed dual-stream attention.\n","\n","Patch strategy: 4×4 patches = 64 tokens per image\n","Each patch: (3, 4, 4) → flatten → project to d_model\n","CLS token prepended, classification from CLS output.\n","\n","!pip install geolip-core torchvision tqdm\n","\"\"\"\n","\n","import torch\n","import torch.nn as nn\n","import torch.nn.functional as F\n","import numpy as np\n","import time, json\n","from pathlib import Path\n","from tqdm.auto import tqdm\n","\n","device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n","print(f\"Device: {device}\")\n","if device.type == 'cuda':\n","    print(f\"  GPU: {torch.cuda.get_device_name()}\")\n","    print(f\"  VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB\")\n","\n","# Import the geometric transformer from geolip_core\n","from geolip_core.pipeline.components.geometric_transformer import (\n","    GeometricTransformer, TorchComponent, BaseTower\n",")\n","torch.set_float32_matmul_precision('high')\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# CONFIG\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","CONFIG = {\n","    # Model\n","    'd_model': 256,\n","    'n_heads': 8,\n","    'n_layers': 8,\n","    'n_anchors': 128,\n","    'manifold_dim': 128,\n","    'n_comp': 4,\n","    'd_comp': 16,\n","    'context_dim': 64,\n","    'quat_dim': 32,\n","    'dropout': 0.1,\n","\n","    # Input stage\n","    'patch_size': 4,       # 32/4 = 8×8 = 64 patches\n","    'img_size': 32,\n","    'in_channels': 3,\n","    'conv_channels': 64,   # conv frontend output channels\n","    'svd_rank': 16,        # SVD projection rank (≤32 for sub-ms)\n","\n","    # Training\n","    'epochs': 100,\n","    'batch_size': 1024,\n","    'lr': 1e-3,\n","    'weight_decay': 0.05,\n","    'warmup_epochs': 5,\n","    'label_smoothing': 0.1,\n","    'num_workers': 8,\n","\n","    # Data\n","    'num_classes': 100,\n","}\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# INPUT STAGE — conv frontend + SVD structural observation → tokens\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","try:\n","    from geolip_core.core.input.svd import SVDObserver\n","    _HAS_SVD = True\n","except ImportError:\n","    _HAS_SVD = False\n","\n","    class SVDObserver(nn.Module):\n","        \"\"\"Fallback SVDObserver matching geolip_core.core.input.svd interface.\"\"\"\n","        def __init__(self, in_channels, svd_rank=24):\n","            super().__init__()\n","            self.svd_rank = svd_rank\n","            self.to_svd = nn.Conv2d(in_channels, svd_rank, 1, bias=False)\n","            self.register_buffer('ema_s', torch.ones(svd_rank))\n","            self.register_buffer('ema_vh_flat', torch.eye(svd_rank).reshape(-1))\n","            self.ema_momentum = 0.99\n","\n","        def extract_features(self, S, Vh):\n","            B, k = S.shape\n","            S_safe = S.clamp(min=1e-6)\n","            s_norm = S_safe / (S_safe.sum(dim=-1, keepdim=True) + 1e-8)\n","            vh_diag = Vh.diagonal(dim1=-2, dim2=-1)\n","            vh_offdiag = (Vh.pow(2).sum((-2, -1)) - vh_diag.pow(2).sum(-1)).unsqueeze(-1).clamp(min=0)\n","            s_ent = -(s_norm * torch.log(s_norm.clamp(min=1e-8))).sum(-1, keepdim=True)\n","            out = torch.cat([s_norm, vh_diag, vh_offdiag, s_ent], dim=-1)\n","            return torch.where(torch.isfinite(out), out, torch.zeros_like(out))\n","\n","        def compute_novelty(self, S):\n","            return S - self.ema_s.clone().unsqueeze(0)\n","\n","        def forward(self, x):\n","            B, C, H, W = x.shape\n","            h = self.to_svd(x)\n","            h_flat = h.permute(0, 2, 3, 1).reshape(B, H * W, self.svd_rank)\n","            with torch.amp.autocast('cuda', enabled=False):\n","                with torch.no_grad():\n","                    gram = torch.bmm(h_flat.float().transpose(1, 2), h_flat.float())\n","                    evals, evecs = torch.linalg.eigh(gram)\n","                    evals = evals.flip(-1).clamp(min=1e-12)\n","                    S = evals.sqrt()\n","                    Vh = evecs.flip(-1).transpose(-2, -1)\n","                    S = torch.where(torch.isfinite(S), S, torch.ones_like(S))\n","                    Vh = torch.where(torch.isfinite(Vh), Vh, torch.zeros_like(Vh))\n","            features = self.extract_features(S, Vh)\n","            novelty = self.compute_novelty(S)\n","            return S, Vh, features, novelty\n","\n","        @torch.no_grad()\n","        def update_ema(self, S, Vh):\n","            m = self.ema_momentum\n","            self.ema_s.mul_(m).add_(S.detach().mean(0), alpha=1-m)\n","            self.ema_vh_flat.mul_(m).add_(Vh.detach().mean(0).reshape(-1), alpha=1-m)\n","\n","        @property\n","        def feature_dim(self):\n","            return 2 * self.svd_rank + 2\n","\n","\n","class ConvSVDPatchEmbedding(TorchComponent):\n","    \"\"\"Input stage: conv frontend → SVDObserver → patch tokens.\n","\n","    Pipeline:\n","        Image (B, 3, 32, 32)\n","          → Conv frontend (3→conv_channels, 2 layers, preserves spatial)\n","          → SVDObserver: 1×1 conv → svd_rank, gram_eigh_svd decomposition\n","              → S (singular values = energy distribution)\n","              → Vh (rotation = structural orientation)\n","              → features (2k+2 compact summary)\n","              → novelty (EMA deviation)\n","          → Patch projection: Conv2d(conv_channels, d_model, stride=patch_size)\n","              → (B, d_model, H/p, W/p) → reshape to (B, N, d_model) tokens\n","          → SVD context FiLM: modulate tokens with global structural info\n","          → CLS token + position embeddings\n","\n","    The constellation now triangulates against tokens that carry\n","    decomposed spatial structure, not random pixel projections.\n","\n","    Args:\n","        name: component identity\n","        img_size: input spatial size\n","        patch_size: patch spatial size (tokens = (img_size/patch_size)^2)\n","        in_channels: input image channels\n","        conv_channels: conv frontend output channels\n","        d_model: token embedding dimension\n","        svd_rank: SVD projection rank (≤32 for sub-ms)\n","    \"\"\"\n","    def __init__(self, name, img_size=32, patch_size=4, in_channels=3,\n","                 conv_channels=64, d_model=256, svd_rank=16):\n","        super().__init__(name)\n","        self.patch_size = patch_size\n","        self.n_patches = (img_size // patch_size) ** 2\n","        self.d_model = d_model\n","        self.svd_rank = svd_rank\n","\n","        # ── Conv frontend: extract spatial features ──\n","        self.conv_frontend = nn.Sequential(\n","            nn.Conv2d(in_channels, conv_channels, 3, padding=1, bias=False),\n","            nn.BatchNorm2d(conv_channels),\n","            nn.GELU(),\n","            nn.Conv2d(conv_channels, conv_channels, 3, padding=1, bias=False),\n","            nn.BatchNorm2d(conv_channels),\n","            nn.GELU(),\n","        )\n","\n","        # ── SVD structural observation ──\n","        self.svd_observer = SVDObserver(conv_channels, svd_rank)\n","\n","        # ── Patch projection: conv features → tokens ──\n","        # Conv2d with stride=patch_size is equivalent to unfold + project\n","        # but more efficient and lets the kernel learn spatial combinations\n","        self.patch_proj = nn.Conv2d(\n","            conv_channels, d_model, kernel_size=patch_size,\n","            stride=patch_size, bias=False)\n","        self.patch_norm = nn.LayerNorm(d_model)\n","\n","        # ── SVD context → FiLM on tokens ──\n","        # Global structural info modulates every token\n","        svd_feat_dim = self.svd_observer.feature_dim  # 2*svd_rank + 2\n","        self.svd_to_gamma = nn.Linear(svd_feat_dim, d_model)\n","        self.svd_to_beta = nn.Linear(svd_feat_dim, d_model)\n","        # Identity-init: γ=1, β=0 at start\n","        nn.init.zeros_(self.svd_to_gamma.weight); nn.init.ones_(self.svd_to_gamma.bias)\n","        nn.init.zeros_(self.svd_to_beta.weight); nn.init.zeros_(self.svd_to_beta.bias)\n","\n","        # ── CLS token + position embeddings ──\n","        self.cls_token = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)\n","        self.pos_embed = nn.Parameter(\n","            torch.randn(1, self.n_patches + 1, d_model) * 0.02)\n","\n","    def forward(self, x):\n","        \"\"\"\n","        Args:\n","            x: (B, C, H, W) input images\n","\n","        Returns:\n","            tokens: (B, N+1, d_model) token sequence with CLS\n","            svd_state: dict with SVD intermediates for diagnostics/loss\n","        \"\"\"\n","        B = x.shape[0]\n","\n","        # ── Conv frontend: raw pixels → spatial features ──\n","        feat = self.conv_frontend(x)  # (B, conv_channels, H, W)\n","\n","        # ── SVD: structural decomposition of feature map ──\n","        S, Vh, svd_features, novelty = self.svd_observer(feat)\n","\n","        # ── Patch projection: feature map → tokens ──\n","        tokens = self.patch_proj(feat)  # (B, d_model, H/p, W/p)\n","        tokens = tokens.flatten(2).transpose(1, 2)  # (B, N, d_model)\n","        tokens = self.patch_norm(tokens)\n","\n","        # ── FiLM: modulate tokens with global SVD context ──\n","        # svd_features is (B, 2k+2) — broadcast over all positions\n","        gamma = self.svd_to_gamma(svd_features).unsqueeze(1)  # (B, 1, d_model)\n","        beta = self.svd_to_beta(svd_features).unsqueeze(1)    # (B, 1, d_model)\n","        tokens = gamma * tokens + beta\n","\n","        # ── CLS token + position embeddings ──\n","        cls = self.cls_token.expand(B, -1, -1)\n","        tokens = torch.cat([cls, tokens], dim=1)  # (B, N+1, d_model)\n","        tokens = tokens + self.pos_embed\n","\n","        svd_state = {\n","            'singular_values': S,\n","            'Vh': Vh,\n","            'svd_features': svd_features,\n","            'novelty': novelty,\n","        }\n","\n","        # Update SVD EMA during training\n","        if self.training:\n","            self.svd_observer.update_ema(S, Vh)\n","\n","        return tokens, svd_state\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# CLASSIFICATION MODEL — patch embed + geometric transformer + head\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","class GeoViTClassifier(BaseTower):\n","    \"\"\"Geometric Vision Transformer for classification.\n","\n","    A BaseTower that composes:\n","        'patch_embed'  → ConvSVDPatchEmbedding (Input stage: conv + SVD + patch)\n","        'transformer'  → GeometricTransformer (processing tower)\n","        'head'         → classification head (Distinction stage)\n","\n","    Access:\n","        model['patch_embed']   → ConvSVDPatchEmbedding\n","        model['transformer']   → GeometricTransformer\n","        model['head']          → classifier MLP\n","    \"\"\"\n","    def __init__(self, name, config):\n","        super().__init__(name)\n","        self.config = config\n","\n","        # Input stage: conv frontend + SVD observation → patch tokens\n","        self.attach('patch_embed', ConvSVDPatchEmbedding(\n","            'patch_embed',\n","            img_size=config['img_size'],\n","            patch_size=config['patch_size'],\n","            in_channels=config['in_channels'],\n","            conv_channels=config['conv_channels'],\n","            d_model=config['d_model'],\n","            svd_rank=config['svd_rank'],\n","        ))\n","\n","        # Processing tower: geometric transformer\n","        self.attach('transformer', GeometricTransformer(\n","            'geo_cifar',\n","            d_model=config['d_model'],\n","            n_heads=config['n_heads'],\n","            n_layers=config['n_layers'],\n","            n_anchors=config['n_anchors'],\n","            manifold_dim=config['manifold_dim'],\n","            n_comp=config['n_comp'],\n","            d_comp=config['d_comp'],\n","            context_dim=config['context_dim'],\n","            quat_dim=config['quat_dim'],\n","            dropout=config['dropout'],\n","        ))\n","\n","        # Distinction stage: CLS token → class logits\n","        self.attach('head', nn.Sequential(\n","            nn.LayerNorm(config['d_model']),\n","            nn.Linear(config['d_model'], config['d_model']),\n","            nn.GELU(),\n","            nn.Dropout(config['dropout']),\n","            nn.Linear(config['d_model'], config['num_classes']),\n","        ))\n","\n","    def forward(self, x, return_geo_state=False):\n","        \"\"\"\n","        x: (B, 3, 32, 32) → (B, num_classes) logits\n","        \"\"\"\n","        tokens, svd_state = self['patch_embed'](x)  # (B, N+1, d_model), dict\n","\n","        if return_geo_state:\n","            features, geo_states = self['transformer'](tokens, return_geo_state=True)\n","        else:\n","            features = self['transformer'](tokens)\n","\n","        # CLS token is at position 0\n","        cls_out = features[:, 0]  # (B, d_model)\n","        logits = self['head'](cls_out)  # (B, num_classes)\n","\n","        if return_geo_state:\n","            return logits, geo_states, svd_state\n","        return logits\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# DATA\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","def get_dataloaders(config):\n","    import torchvision\n","    import torchvision.transforms as T\n","\n","    train_transform = T.Compose([\n","        T.RandomCrop(32, padding=4),\n","        T.RandomHorizontalFlip(),\n","        T.ToTensor(),\n","        T.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),\n","    ])\n","\n","    test_transform = T.Compose([\n","        T.ToTensor(),\n","        T.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),\n","    ])\n","\n","    train_ds = torchvision.datasets.CIFAR100(\n","        root='./data', train=True, download=True, transform=train_transform)\n","    test_ds = torchvision.datasets.CIFAR100(\n","        root='./data', train=False, download=True, transform=test_transform)\n","\n","    train_loader = torch.utils.data.DataLoader(\n","        train_ds, batch_size=config['batch_size'], shuffle=True,\n","        num_workers=config['num_workers'], pin_memory=True, drop_last=True)\n","    test_loader = torch.utils.data.DataLoader(\n","        test_ds, batch_size=config['batch_size'], shuffle=False,\n","        num_workers=config['num_workers'], pin_memory=True)\n","\n","    return train_loader, test_loader\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# TRAINING\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","def train_epoch(model, loader, optimizer, scheduler, epoch, config):\n","    model.train()\n","    total_loss = 0\n","    correct = 0\n","    total = 0\n","    criterion = nn.CrossEntropyLoss(label_smoothing=config['label_smoothing'])\n","\n","    for images, labels in loader:\n","        images = images.to(device)\n","        labels = labels.to(device)\n","\n","        logits = model(images)\n","        loss = criterion(logits, labels)\n","\n","        optimizer.zero_grad()\n","        loss.backward()\n","        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n","        optimizer.step()\n","        if scheduler is not None:\n","            scheduler.step()\n","\n","        total_loss += loss.item() * images.size(0)\n","        correct += (logits.argmax(1) == labels).sum().item()\n","        total += images.size(0)\n","\n","    return total_loss / total, correct / total\n","\n","\n","@torch.no_grad()\n","def evaluate(model, loader):\n","    model.eval()\n","    correct = 0\n","    total = 0\n","\n","    for images, labels in loader:\n","        images = images.to(device)\n","        labels = labels.to(device)\n","\n","        logits = model(images)\n","        correct += (logits.argmax(1) == labels).sum().item()\n","        total += images.size(0)\n","\n","    return correct / total\n","\n","\n","def main():\n","    config = CONFIG.copy()\n","\n","    print(\"=\" * 60)\n","    print(\"  Geometric Transformer — CIFAR-100\")\n","    print(f\"  Input: conv({config['in_channels']}→{config['conv_channels']}) + \"\n","          f\"SVD(rank={config['svd_rank']}) + \"\n","          f\"{config['patch_size']}×{config['patch_size']} patches = \"\n","          f\"{(config['img_size']//config['patch_size'])**2} tokens + CLS\")\n","    print(f\"  Model: d={config['d_model']}, heads={config['n_heads']}, \"\n","          f\"layers={config['n_layers']}, anchors={config['n_anchors']}\")\n","    print(\"=\" * 60)\n","\n","    # Data\n","    print(\"\\nLoading CIFAR-100...\")\n","    train_loader, test_loader = get_dataloaders(config)\n","    print(f\"  Train: {len(train_loader.dataset):,} | Test: {len(test_loader.dataset):,}\")\n","\n","    # Model\n","    model = GeoViTClassifier('geo_vit_cifar100', config)\n","    if hasattr(model, 'network_to'):\n","        model.network_to(device=device, strict=False)\n","        model.compile()\n","    else:\n","        model = model.to(device)\n","\n","    n_params = sum(p.numel() for p in model.parameters())\n","    n_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)\n","    print(f\"\\n  Total params:     {n_params:,}\")\n","    print(f\"  Trainable params: {n_trainable:,}\")\n","\n","    # Print breakdown\n","    for name, module in model.named_children():\n","        n = sum(p.numel() for p in module.parameters())\n","        if n > 0:\n","            print(f\"    {name:<20s}: {n:,}\")\n","\n","    # Optimizer + scheduler\n","    optimizer = torch.optim.AdamW(\n","        model.parameters(), lr=config['lr'], weight_decay=config['weight_decay'])\n","\n","    total_steps = config['epochs'] * len(train_loader)\n","    warmup_steps = config['warmup_epochs'] * len(train_loader)\n","\n","    def lr_lambda(step):\n","        if step < warmup_steps:\n","            return step / warmup_steps\n","        progress = (step - warmup_steps) / (total_steps - warmup_steps)\n","        return 0.5 * (1 + np.cos(np.pi * progress))\n","\n","    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)\n","\n","    # Training loop\n","    print(f\"\\n{'━'*60}\")\n","    print(f\"  Training for {config['epochs']} epochs\")\n","    print(f\"  Warmup: {config['warmup_epochs']} epochs, \"\n","          f\"LR: {config['lr']}, WD: {config['weight_decay']}\")\n","    print(f\"{'━'*60}\\n\")\n","\n","    best_acc = 0\n","    save_dir = Path('geo_cifar100'); save_dir.mkdir(exist_ok=True)\n","\n","    for epoch in tqdm(range(config['epochs']), desc=\"Epochs\"):\n","        t0 = time.time()\n","\n","        train_loss, train_acc = train_epoch(\n","            model, train_loader, optimizer, scheduler, epoch, config)\n","\n","        test_acc = evaluate(model, test_loader)\n","        elapsed = time.time() - t0\n","\n","        if test_acc > best_acc:\n","            best_acc = test_acc\n","            torch.save({\n","                'state_dict': {k: v.cpu() for k, v in model.state_dict().items()},\n","                'epoch': epoch,\n","                'test_acc': test_acc,\n","                'config': config,\n","            }, save_dir / 'best.pt')\n","\n","        if epoch % 5 == 0 or epoch == config['epochs'] - 1:\n","            lr = optimizer.param_groups[0]['lr']\n","            tqdm.write(\n","                f\"  E{epoch:>3d}  loss={train_loss:.4f}  \"\n","                f\"train={train_acc:.4f}  test={test_acc:.4f}  \"\n","                f\"best={best_acc:.4f}  lr={lr:.6f}  {elapsed:.1f}s\")\n","\n","    # Final summary\n","    print(f\"\\n{'═'*60}\")\n","    print(f\"  CIFAR-10 RESULTS\")\n","    print(f\"{'═'*60}\")\n","    print(f\"  Best test accuracy: {best_acc:.4f} ({best_acc*100:.2f}%)\")\n","    print(f\"  Parameters: {n_params:,}\")\n","    print(f\"  Checkpoint: {save_dir}/best.pt\")\n","\n","    # Geometric state inspection on a test batch\n","    print(f\"\\n  Geometric state inspection:\")\n","    model.eval()\n","    images, labels = next(iter(test_loader))\n","    images = images[:4].to(device)\n","    with torch.no_grad():\n","        logits, geo_states, svd_state = model(images, return_geo_state=True)\n","\n","    # SVD Input stage diagnostics\n","    S = svd_state['singular_values']\n","    s_norm = S / (S.sum(dim=-1, keepdim=True) + 1e-8)\n","    s_ent = -(s_norm * torch.log(s_norm.clamp(min=1e-8))).sum(-1)\n","    tqdm.write(\n","        f\"    SVD Input: S_top3={S[0, :3].tolist()}  \"\n","        f\"entropy={s_ent.mean().item():.3f}  \"\n","        f\"novelty_norm={svd_state['novelty'].norm(dim=-1).mean().item():.3f}\")\n","\n","    for i, gs in enumerate(geo_states):\n","        content = gs['content']\n","        geometric = gs['geometric']\n","\n","        # Measure stream divergence\n","        content_norm = content.norm(dim=-1).mean().item()\n","        geo_norm = geometric.norm(dim=-1).mean().item()\n","        agreement = F.cosine_similarity(\n","            content.reshape(-1, content.shape[-1]),\n","            geometric.reshape(-1, geometric.shape[-1]),\n","            dim=-1).mean().item()\n","\n","        # Anchor utilization\n","        tri = gs['triangulation']\n","        nearest_counts = torch.bincount(\n","            gs['nearest'].reshape(-1),\n","            minlength=tri.shape[-1]).float()\n","        anchor_entropy = -(nearest_counts / nearest_counts.sum() *\n","            torch.log(nearest_counts.clamp(min=1e-8) / nearest_counts.sum())).sum().item()\n","\n","        tqdm.write(\n","            f\"    Layer {i}: ‖content‖={content_norm:.3f}  \"\n","            f\"‖geo‖={geo_norm:.3f}  agreement={agreement:.3f}  \"\n","            f\"anchor_H={anchor_entropy:.3f}\")\n","\n","    print(f\"\\nDone.\")\n","\n","\n","if __name__ == '__main__':\n","    main()"],"metadata":{"colab":{"base_uri":"https://localhost:8080/","referenced_widgets":["8045de5b58de4308a4bf4d1e60023113","38fcc80f6c9145c1b3b75659a01f5d97","5df6e33ca7104310a3eff8ba2eaf42d9","424c6cf8e56947939bd29d98c104c573","3531e2c63e4f42928c4e953442cba8d2","3fa201a8613a4f4caa09a3cb1d042fdf","dedb555f60e44516a1b5b80896a9634f","1bb21ffdb1ba469ab6f86df3071f9dc5","963b762376c147048fd22d081017c320","7e1e5cc2f9694b4e88067f96f7cc5f9f","210fc9a76302468ebacef0c0e65275b6"]},"id":"2ctwQjGJ6B0D","executionInfo":{"status":"ok","timestamp":1774693879924,"user_tz":420,"elapsed":1439949,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"286e2d96-fc6f-4e30-b8e4-64ae2dfdb4fe"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["Device: cuda\n","  GPU: NVIDIA RTX PRO 6000 Blackwell Server Edition\n","  VRAM: 102.0 GB\n","============================================================\n","  Geometric Transformer — CIFAR-100\n","  Input: conv(3→64) + SVD(rank=16) + 4×4 patches = 64 tokens + CLS\n","  Model: d=256, heads=8, layers=8, anchors=128\n","============================================================\n","\n","Loading CIFAR-100...\n","  Train: 50,000 | Test: 10,000\n","\n","  Total params:     17,750,436\n","  Trainable params: 17,750,436\n","    components          : 17,750,436\n","\n","━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\n","  Training for 100 epochs\n","  Warmup: 5 epochs, LR: 0.001, WD: 0.05\n","━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\n","\n"]},{"output_type":"display_data","data":{"text/plain":["Epochs:   0%|          | 0/100 [00:00<?, ?it/s]"],"application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"8045de5b58de4308a4bf4d1e60023113"}},"metadata":{}},{"output_type":"stream","name":"stdout","text":["  E  0  loss=4.5288  train=0.0247  test=0.0460  best=0.0460  lr=0.000200  56.3s\n","  E  5  loss=3.5206  train=0.1944  test=0.2092  best=0.2092  lr=0.001000  13.8s\n","  E 10  loss=2.9069  train=0.3468  test=0.3407  best=0.3407  lr=0.000990  13.8s\n","  E 15  loss=2.5525  train=0.4388  test=0.4082  best=0.4082  lr=0.000967  13.8s\n","  E 20  loss=2.2735  train=0.5202  test=0.4599  best=0.4599  lr=0.000932  13.8s\n","  E 25  loss=2.0192  train=0.5974  test=0.4872  best=0.4872  lr=0.000884  13.8s\n","  E 30  loss=1.7773  train=0.6736  test=0.5185  best=0.5185  lr=0.000826  13.8s\n","  E 35  loss=1.5384  train=0.7547  test=0.5283  best=0.5288  lr=0.000759  13.8s\n","  E 40  loss=1.3364  train=0.8273  test=0.5277  best=0.5343  lr=0.000686  14.3s\n","  E 45  loss=1.1758  train=0.8839  test=0.5309  best=0.5347  lr=0.000607  13.8s\n","  E 50  loss=1.0603  train=0.9267  test=0.5375  best=0.5401  lr=0.000525  13.8s\n","  E 55  loss=0.9841  train=0.9530  test=0.5437  best=0.5451  lr=0.000442  13.8s\n","  E 60  loss=0.9290  train=0.9682  test=0.5494  best=0.5532  lr=0.000361  13.9s\n","  E 65  loss=0.8827  train=0.9814  test=0.5558  best=0.5587  lr=0.000284  13.8s\n","  E 70  loss=0.8486  train=0.9897  test=0.5568  best=0.5600  lr=0.000213  13.8s\n","  E 75  loss=0.8284  train=0.9940  test=0.5576  best=0.5600  lr=0.000149  13.8s\n","  E 80  loss=0.8126  train=0.9975  test=0.5591  best=0.5657  lr=0.000095  13.8s\n","  E 85  loss=0.8036  train=0.9987  test=0.5698  best=0.5698  lr=0.000053  13.8s\n","  E 90  loss=0.7998  train=0.9993  test=0.5718  best=0.5731  lr=0.000022  13.8s\n","  E 95  loss=0.7987  train=0.9994  test=0.5728  best=0.5731  lr=0.000004  13.8s\n","  E 99  loss=0.7982  train=0.9996  test=0.5723  best=0.5731  lr=0.000000  13.8s\n","\n","════════════════════════════════════════════════════════════\n","  CIFAR-10 RESULTS\n","════════════════════════════════════════════════════════════\n","  Best test accuracy: 0.5731 (57.31%)\n","  Parameters: 17,750,436\n","  Checkpoint: geo_cifar100/best.pt\n","\n","  Geometric state inspection:\n","    SVD Input: S_top3=[12.710785865783691, 9.01830768585205, 8.61393928527832]  entropy=2.641  novelty_norm=23.703\n","    Layer 0: ‖content‖=14.489  ‖geo‖=14.332  agreement=0.606  anchor_H=2.793\n","    Layer 1: ‖content‖=14.595  ‖geo‖=14.304  agreement=0.416  anchor_H=3.257\n","    Layer 2: ‖content‖=14.625  ‖geo‖=14.323  agreement=0.497  anchor_H=2.767\n","    Layer 3: ‖content‖=14.627  ‖geo‖=14.289  agreement=0.526  anchor_H=3.270\n","    Layer 4: ‖content‖=14.653  ‖geo‖=14.258  agreement=0.509  anchor_H=3.126\n","    Layer 5: ‖content‖=14.620  ‖geo‖=14.301  agreement=0.528  anchor_H=3.183\n","    Layer 6: ‖content‖=14.660  ‖geo‖=14.259  agreement=0.549  anchor_H=3.534\n","    Layer 7: ‖content‖=14.649  ‖geo‖=14.273  agreement=0.532  anchor_H=3.065\n","\n","Done.\n"]}]},{"cell_type":"code","source":["\"\"\"\n","Geometric Transformer — Standalone Analysis Battery\n","\n","Loads a trained checkpoint and runs comprehensive geometric analysis.\n","No training, no TensorBoard dependency — prints everything to console\n","and saves a JSON report.\n","\n","Usage:\n","    python analyze_geo_transformer.py                          # defaults\n","    python analyze_geo_transformer.py --checkpoint path/to/best.pt\n","    python analyze_geo_transformer.py --n_samples 256 --per_class\n","\n","Analysis categories:\n","    1. CV (coefficient of variation) — pentachoron band tracking\n","    2. Constellation geometry — anchor spread, utilization, dead anchors\n","    3. Stream dynamics — agreement, divergence, arm norms\n","    4. Triangulation — distance distributions, assignment peakedness\n","    5. Patchwork — activation patterns, sparsity, bridge consistency\n","    6. SVD Input — spectrum, entropy, condition number, novelty\n","    7. Cayley rotations — ‖R-I‖, determinant, parameter norms\n","    8. FiLM layers — deviation from identity\n","    9. Quaternion composition — arm balance, output statistics\n","    10. Gate — activation distribution, bypass fraction\n","    11. Per-class analysis — which classes the geometry handles well/poorly\n","    12. Cross-layer trajectories — how features evolve through depth\n","\n","!pip install geolip-core torchvision\n","\"\"\"\n","\n","import torch\n","import torch.nn as nn\n","import torch.nn.functional as F\n","import numpy as np\n","import json, math, argparse\n","from pathlib import Path\n","from collections import defaultdict\n","\n","device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n","\n","from geolip_core.pipeline.components.geometric_transformer import (\n","    GeometricTransformer, GeometricTransformerLayer,\n","    CayleyOrthogonal, QuaternionCompose, FiLMLayer,\n","    ContentAttention, GeometricAttention,\n","    ManifoldProjection, PositionGeometricContext,\n","    TorchComponent, BaseTower,\n",")\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# INPUT STAGE (must match training)\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","try:\n","    from geolip_core.core.input.svd import SVDObserver\n","    _HAS_SVD = True\n","except ImportError:\n","    _HAS_SVD = False\n","\n","    class SVDObserver(nn.Module):\n","        def __init__(self, in_channels, svd_rank=24):\n","            super().__init__()\n","            self.svd_rank = svd_rank\n","            self.to_svd = nn.Conv2d(in_channels, svd_rank, 1, bias=False)\n","            self.register_buffer('ema_s', torch.ones(svd_rank))\n","            self.register_buffer('ema_vh_flat', torch.eye(svd_rank).reshape(-1))\n","            self.ema_momentum = 0.99\n","\n","        def extract_features(self, S, Vh):\n","            B, k = S.shape\n","            S_safe = S.clamp(min=1e-6)\n","            s_norm = S_safe / (S_safe.sum(dim=-1, keepdim=True) + 1e-8)\n","            vh_diag = Vh.diagonal(dim1=-2, dim2=-1)\n","            vh_offdiag = (Vh.pow(2).sum((-2, -1)) - vh_diag.pow(2).sum(-1)).unsqueeze(-1).clamp(min=0)\n","            s_ent = -(s_norm * torch.log(s_norm.clamp(min=1e-8))).sum(-1, keepdim=True)\n","            out = torch.cat([s_norm, vh_diag, vh_offdiag, s_ent], dim=-1)\n","            return torch.where(torch.isfinite(out), out, torch.zeros_like(out))\n","\n","        def compute_novelty(self, S):\n","            return S - self.ema_s.clone().unsqueeze(0)\n","\n","        def forward(self, x):\n","            B, C, H, W = x.shape\n","            h = self.to_svd(x)\n","            h_flat = h.permute(0, 2, 3, 1).reshape(B, H * W, self.svd_rank)\n","            with torch.amp.autocast('cuda', enabled=False):\n","                with torch.no_grad():\n","                    gram = torch.bmm(h_flat.float().transpose(1, 2), h_flat.float())\n","                    evals, evecs = torch.linalg.eigh(gram)\n","                    evals = evals.flip(-1).clamp(min=1e-12)\n","                    S = evals.sqrt()\n","                    Vh = evecs.flip(-1).transpose(-2, -1)\n","                    S = torch.where(torch.isfinite(S), S, torch.ones_like(S))\n","                    Vh = torch.where(torch.isfinite(Vh), Vh, torch.zeros_like(Vh))\n","            return S, Vh, self.extract_features(S, Vh), self.compute_novelty(S)\n","\n","        @property\n","        def feature_dim(self):\n","            return 2 * self.svd_rank + 2\n","\n","\n","class ConvSVDPatchEmbedding(TorchComponent):\n","    def __init__(self, name, img_size=32, patch_size=4, in_channels=3,\n","                 conv_channels=64, d_model=256, svd_rank=16):\n","        super().__init__(name)\n","        self.patch_size = patch_size\n","        self.n_patches = (img_size // patch_size) ** 2\n","        self.d_model = d_model\n","        self.svd_rank = svd_rank\n","        self.conv_frontend = nn.Sequential(\n","            nn.Conv2d(in_channels, conv_channels, 3, padding=1, bias=False),\n","            nn.BatchNorm2d(conv_channels), nn.GELU(),\n","            nn.Conv2d(conv_channels, conv_channels, 3, padding=1, bias=False),\n","            nn.BatchNorm2d(conv_channels), nn.GELU(),\n","        )\n","        self.svd_observer = SVDObserver(conv_channels, svd_rank)\n","        self.patch_proj = nn.Conv2d(conv_channels, d_model, kernel_size=patch_size,\n","                                     stride=patch_size, bias=False)\n","        self.patch_norm = nn.LayerNorm(d_model)\n","        svd_feat_dim = self.svd_observer.feature_dim\n","        self.svd_to_gamma = nn.Linear(svd_feat_dim, d_model)\n","        self.svd_to_beta = nn.Linear(svd_feat_dim, d_model)\n","        nn.init.zeros_(self.svd_to_gamma.weight); nn.init.ones_(self.svd_to_gamma.bias)\n","        nn.init.zeros_(self.svd_to_beta.weight); nn.init.zeros_(self.svd_to_beta.bias)\n","        self.cls_token = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)\n","        self.pos_embed = nn.Parameter(torch.randn(1, self.n_patches + 1, d_model) * 0.02)\n","\n","    def forward(self, x):\n","        B = x.shape[0]\n","        feat = self.conv_frontend(x)\n","        S, Vh, svd_features, novelty = self.svd_observer(feat)\n","        tokens = self.patch_proj(feat)\n","        tokens = tokens.flatten(2).transpose(1, 2)\n","        tokens = self.patch_norm(tokens)\n","        gamma = self.svd_to_gamma(svd_features).unsqueeze(1)\n","        beta = self.svd_to_beta(svd_features).unsqueeze(1)\n","        tokens = gamma * tokens + beta\n","        cls = self.cls_token.expand(B, -1, -1)\n","        tokens = torch.cat([cls, tokens], dim=1)\n","        tokens = tokens + self.pos_embed\n","        return tokens, {\n","            'singular_values': S, 'Vh': Vh,\n","            'svd_features': svd_features, 'novelty': novelty,\n","        }\n","\n","\n","class GeoViTClassifier(BaseTower):\n","    def __init__(self, name, config):\n","        super().__init__(name)\n","        self.config = config\n","        self.attach('patch_embed', ConvSVDPatchEmbedding(\n","            'patch_embed', img_size=config['img_size'],\n","            patch_size=config['patch_size'], in_channels=config['in_channels'],\n","            conv_channels=config['conv_channels'], d_model=config['d_model'],\n","            svd_rank=config['svd_rank'],\n","        ))\n","        self.attach('transformer', GeometricTransformer(\n","            'geo_cifar', d_model=config['d_model'], n_heads=config['n_heads'],\n","            n_layers=config['n_layers'], n_anchors=config['n_anchors'],\n","            manifold_dim=config['manifold_dim'], n_comp=config['n_comp'],\n","            d_comp=config['d_comp'], context_dim=config['context_dim'],\n","            quat_dim=config['quat_dim'], dropout=config['dropout'],\n","        ))\n","        self.attach('head', nn.Sequential(\n","            nn.LayerNorm(config['d_model']),\n","            nn.Linear(config['d_model'], config['d_model']),\n","            nn.GELU(), nn.Dropout(config['dropout']),\n","            nn.Linear(config['d_model'], config['num_classes']),\n","        ))\n","\n","    def forward(self, x, return_geo_state=False):\n","        tokens, svd_state = self['patch_embed'](x)\n","        if return_geo_state:\n","            features, geo_states = self['transformer'](tokens, return_geo_state=True)\n","        else:\n","            features = self['transformer'](tokens)\n","        cls_out = features[:, 0]\n","        logits = self['head'](cls_out)\n","        if return_geo_state:\n","            return logits, geo_states, svd_state\n","        return logits\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# GEOMETRIC PRIMITIVES\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","@torch.no_grad()\n","def compute_cv(points, max_points=1024):\n","    \"\"\"Coefficient of variation of pairwise cosine distances on S^(d-1).\"\"\"\n","    points = F.normalize(points.float(), dim=-1)\n","    if points.shape[0] > max_points:\n","        idx = torch.randperm(points.shape[0], device=points.device)[:max_points]\n","        points = points[idx]\n","    cos_sim = points @ points.T\n","    n = points.shape[0]\n","    idx = torch.triu_indices(n, n, offset=1, device=points.device)\n","    pairwise_dist = 1.0 - cos_sim[idx[0], idx[1]]\n","    mean_d = pairwise_dist.mean()\n","    std_d = pairwise_dist.std()\n","    return {\n","        'cv': (std_d / (mean_d + 1e-8)).item(),\n","        'mean_dist': mean_d.item(),\n","        'std_dist': std_d.item(),\n","        'min_dist': pairwise_dist.min().item(),\n","        'max_dist': pairwise_dist.max().item(),\n","        'n_points': points.shape[0],\n","    }\n","\n","\n","@torch.no_grad()\n","def compute_anchor_geometry(anchors):\n","    \"\"\"Full geometric analysis of anchor positions on S^(d-1).\"\"\"\n","    anchors = F.normalize(anchors.float(), dim=-1)\n","    n, d = anchors.shape\n","    cos_sim = anchors @ anchors.T\n","\n","    idx = torch.triu_indices(n, n, offset=1, device=anchors.device)\n","    pairwise_cos = cos_sim[idx[0], idx[1]]\n","    pairwise_dist = 1.0 - pairwise_cos\n","\n","    # Nearest neighbor distances\n","    cos_sim_no_diag = cos_sim.clone()\n","    cos_sim_no_diag.fill_diagonal_(-2.0)\n","    nn_cos = cos_sim_no_diag.max(dim=1).values\n","    nn_dist = 1.0 - nn_cos\n","\n","    cv_info = compute_cv(anchors)\n","\n","    return {\n","        **cv_info,\n","        'n_anchors': n,\n","        'dim': d,\n","        'pairwise_cos_mean': pairwise_cos.mean().item(),\n","        'pairwise_cos_std': pairwise_cos.std().item(),\n","        'nn_dist_mean': nn_dist.mean().item(),\n","        'nn_dist_std': nn_dist.std().item(),\n","        'nn_dist_min': nn_dist.min().item(),\n","        'nn_dist_max': nn_dist.max().item(),\n","        'max_cos_similarity': pairwise_cos.max().item(),\n","        'min_cos_similarity': pairwise_cos.min().item(),\n","    }\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# ANALYSIS BATTERY\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","@torch.no_grad()\n","def run_analysis(model, test_loader, device, n_batches=4, per_class=False):\n","    \"\"\"Run the complete geometric analysis battery.\"\"\"\n","    model.eval()\n","    report = {}\n","\n","    # ─── Collect predictions and geo states across multiple batches ───\n","    all_logits = []\n","    all_labels = []\n","    all_geo_states = None\n","    all_svd_state = None\n","\n","    for batch_idx, (images, labels) in enumerate(test_loader):\n","        if batch_idx >= n_batches:\n","            break\n","        images = images.to(device)\n","        labels = labels.to(device)\n","\n","        logits, geo_states, svd_state = model(images, return_geo_state=True)\n","        all_logits.append(logits.cpu())\n","        all_labels.append(labels.cpu())\n","\n","        # Keep last batch geo_states for detailed analysis\n","        if batch_idx == 0:\n","            all_geo_states = geo_states\n","            all_svd_state = svd_state\n","\n","    all_logits = torch.cat(all_logits, dim=0)\n","    all_labels = torch.cat(all_labels, dim=0)\n","    n_total = all_labels.shape[0]\n","    pred = all_logits.argmax(1)\n","    overall_acc = (pred == all_labels).float().mean().item()\n","\n","    report['overall'] = {\n","        'accuracy': overall_acc,\n","        'n_samples': n_total,\n","        'n_classes': all_logits.shape[1],\n","    }\n","\n","    n_layers = len(all_geo_states)\n","\n","    # ─── 1. SVD Input Stage ───\n","    print(\"\\n\" + \"=\" * 70)\n","    print(\"  1. SVD INPUT STAGE\")\n","    print(\"=\" * 70)\n","\n","    S = all_svd_state['singular_values']\n","    s_norm = S / (S.sum(dim=-1, keepdim=True) + 1e-8)\n","    s_ent = -(s_norm * torch.log(s_norm.clamp(min=1e-8))).sum(-1)\n","    novelty = all_svd_state['novelty']\n","\n","    svd_report = {\n","        'entropy_mean': s_ent.mean().item(),\n","        'entropy_std': s_ent.std().item(),\n","        'novelty_norm': novelty.norm(dim=-1).mean().item(),\n","        'top1_ratio': (S[:, 0] / (S.sum(-1) + 1e-8)).mean().item(),\n","        'condition_number': (S[:, 0] / (S[:, -1].clamp(min=1e-8))).mean().item(),\n","        'singular_values': [S[:, k].mean().item() for k in range(S.shape[1])],\n","    }\n","\n","    # SVD FiLM deviation\n","    pe = model['patch_embed']\n","    svd_report['film_gamma_bias_dev'] = (pe.svd_to_gamma.bias.data - 1.0).abs().mean().item()\n","    svd_report['film_beta_bias_norm'] = pe.svd_to_beta.bias.data.abs().mean().item()\n","    svd_report['film_gamma_weight_norm'] = pe.svd_to_gamma.weight.data.norm().item()\n","    svd_report['film_beta_weight_norm'] = pe.svd_to_beta.weight.data.norm().item()\n","\n","    report['svd'] = svd_report\n","\n","    print(f\"  Entropy:          {svd_report['entropy_mean']:.3f} ± {svd_report['entropy_std']:.3f}\")\n","    print(f\"  Novelty norm:     {svd_report['novelty_norm']:.3f}\")\n","    print(f\"  Top-1 ratio:      {svd_report['top1_ratio']:.3f}\")\n","    print(f\"  Condition number: {svd_report['condition_number']:.1f}\")\n","    print(f\"  Singular values:  {', '.join(f'{v:.2f}' for v in svd_report['singular_values'][:5])}...\")\n","    print(f\"  FiLM γ dev:       {svd_report['film_gamma_bias_dev']:.4f}\")\n","    print(f\"  FiLM β norm:      {svd_report['film_beta_bias_norm']:.4f}\")\n","\n","    # ─── 2. Per-Layer Geometric Analysis ───\n","    report['layers'] = {}\n","\n","    for i in range(n_layers):\n","        gs = all_geo_states[i]\n","        lr = {}\n","\n","        print(f\"\\n{'─'*70}\")\n","        print(f\"  LAYER {i}\")\n","        print(f\"{'─'*70}\")\n","\n","        # === 2a. CV — Pentachoron Band ===\n","        emb = gs['embedding']\n","        emb_flat = emb.reshape(-1, emb.shape[-1])\n","        cv_emb = compute_cv(emb_flat)\n","        lr['cv_embeddings'] = cv_emb\n","\n","        # Anchor CV\n","        anchor_cv = None\n","        for name, mod in model.named_modules():\n","            if (hasattr(mod, 'association') and hasattr(mod.association, 'constellation')\n","                    and f'layer_{i}' in name):\n","                anchors = mod.association.constellation.anchors.data\n","                anchor_cv = compute_anchor_geometry(anchors)\n","                lr['anchor_geometry'] = anchor_cv\n","                break\n","\n","        in_band = 0.20 <= cv_emb['cv'] <= 0.23\n","        print(f\"  CV embeddings:    {cv_emb['cv']:.4f}  {'✓ IN BAND' if in_band else '✗ outside'}\")\n","        print(f\"    mean_dist={cv_emb['mean_dist']:.4f}  std={cv_emb['std_dist']:.4f}\")\n","        if anchor_cv:\n","            a_in_band = 0.20 <= anchor_cv['cv'] <= 0.23\n","            print(f\"  CV anchors:       {anchor_cv['cv']:.4f}  {'✓ IN BAND' if a_in_band else '✗ outside'}\")\n","            print(f\"    nn_dist: mean={anchor_cv['nn_dist_mean']:.4f}  \"\n","                  f\"min={anchor_cv['nn_dist_min']:.4f}  max={anchor_cv['nn_dist_max']:.4f}\")\n","            print(f\"    max_cos_sim:    {anchor_cv['max_cos_similarity']:.4f}\")\n","\n","        # === 2b. Stream Dynamics ===\n","        content = gs['content']\n","        geometric = gs['geometric']\n","\n","        agreement = F.cosine_similarity(\n","            content.reshape(-1, content.shape[-1]),\n","            geometric.reshape(-1, geometric.shape[-1]), dim=-1)\n","\n","        disagree = content - geometric\n","        agree = content * geometric\n","\n","        stream = {\n","            'agreement_mean': agreement.mean().item(),\n","            'agreement_std': agreement.std().item(),\n","            'agreement_min': agreement.min().item(),\n","            'agreement_max': agreement.max().item(),\n","            'content_norm': content.norm(dim=-1).mean().item(),\n","            'geometric_norm': geometric.norm(dim=-1).mean().item(),\n","            'norm_ratio': (geometric.norm(dim=-1) / (content.norm(dim=-1) + 1e-8)).mean().item(),\n","            'disagree_norm': disagree.norm(dim=-1).mean().item(),\n","            'agree_norm': agree.norm(dim=-1).mean().item(),\n","            'disagree_to_content': (disagree.norm(dim=-1) / (content.norm(dim=-1) + 1e-8)).mean().item(),\n","        }\n","        lr['stream'] = stream\n","\n","        print(f\"\\n  Stream agreement: {stream['agreement_mean']:.4f} ± {stream['agreement_std']:.4f}\")\n","        print(f\"    range: [{stream['agreement_min']:.4f}, {stream['agreement_max']:.4f}]\")\n","        print(f\"  ‖content‖:        {stream['content_norm']:.3f}\")\n","        print(f\"  ‖geometric‖:      {stream['geometric_norm']:.3f}\")\n","        print(f\"  ‖disagree‖:       {stream['disagree_norm']:.3f}  \"\n","              f\"(ratio: {stream['disagree_to_content']:.3f})\")\n","        print(f\"  ‖agree‖:          {stream['agree_norm']:.3f}\")\n","\n","        # === 2c. Anchor Utilization ===\n","        tri = gs['triangulation']\n","        assignment = gs['assignment']\n","        nearest = gs['nearest']\n","        n_anchors = tri.shape[-1]\n","\n","        nearest_flat = nearest.reshape(-1)\n","        counts = torch.bincount(nearest_flat, minlength=n_anchors).float()\n","        total_assign = counts.sum()\n","\n","        probs = counts / (total_assign + 1e-8)\n","        entropy = -(probs * torch.log(probs.clamp(min=1e-8))).sum().item()\n","        max_entropy = math.log(n_anchors)\n","\n","        active = (counts > 0).sum().item()\n","        dead = (counts == 0).sum().item()\n","        top1_frac = counts.max().item() / (total_assign + 1e-8)\n","\n","        # Top-5 and bottom-5 anchors\n","        sorted_counts, sorted_idx = counts.sort(descending=True)\n","\n","        anchor_util = {\n","            'entropy': entropy,\n","            'entropy_normalized': entropy / (max_entropy + 1e-8),\n","            'max_entropy': max_entropy,\n","            'active': int(active),\n","            'dead': int(dead),\n","            'total_anchors': n_anchors,\n","            'top1_frac': top1_frac,\n","            'top5_counts': sorted_counts[:5].tolist(),\n","            'top5_indices': sorted_idx[:5].tolist(),\n","            'bottom5_counts': sorted_counts[-5:].tolist(),\n","        }\n","        lr['anchor_utilization'] = anchor_util\n","\n","        print(f\"\\n  Anchors: {active}/{n_anchors} active, {dead} dead\")\n","        print(f\"  Entropy: {entropy:.3f} / {max_entropy:.3f} \"\n","              f\"(normalized: {anchor_util['entropy_normalized']:.3f})\")\n","        print(f\"  Top-1 dominance: {top1_frac:.4f}\")\n","        print(f\"  Top-5 counts: {sorted_counts[:5].int().tolist()}\")\n","\n","        # === 2d. Triangulation ===\n","        tri_stats = {\n","            'mean': tri.mean().item(),\n","            'std': tri.std().item(),\n","            'min': tri.min().item(),\n","            'max': tri.max().item(),\n","            'median': tri.median().item(),\n","        }\n","        lr['triangulation'] = tri_stats\n","\n","        # === 2e. Soft Assignment ===\n","        assign_ent = -(assignment * torch.log(assignment.clamp(min=1e-8))).sum(-1)\n","        max_prob = assignment.max(dim=-1).values\n","\n","        assign_stats = {\n","            'entropy_mean': assign_ent.mean().item(),\n","            'entropy_std': assign_ent.std().item(),\n","            'max_prob_mean': max_prob.mean().item(),\n","            'max_prob_std': max_prob.std().item(),\n","        }\n","        lr['assignment'] = assign_stats\n","\n","        print(f\"\\n  Assignment entropy: {assign_stats['entropy_mean']:.3f} ± {assign_stats['entropy_std']:.3f}\")\n","        print(f\"  Max prob:           {assign_stats['max_prob_mean']:.4f} ± {assign_stats['max_prob_std']:.4f}\")\n","\n","        # === 2f. Patchwork ===\n","        pw = gs['patchwork']\n","        pw_stats = {\n","            'norm': pw.norm(dim=-1).mean().item(),\n","            'std': pw.std().item(),\n","            'sparsity': (pw.abs() < 0.01).float().mean().item(),\n","            'dead_frac': (pw.abs() < 1e-6).float().mean().item(),\n","        }\n","        lr['patchwork'] = pw_stats\n","\n","        print(f\"\\n  Patchwork norm:     {pw_stats['norm']:.3f}\")\n","        print(f\"  Patchwork sparsity: {pw_stats['sparsity']:.3f}  \"\n","              f\"(dead: {pw_stats['dead_frac']:.4f})\")\n","\n","        # === 2g. Bridge Consistency ===\n","        bridge = gs['bridge']\n","        bridge_soft = F.softmax(bridge, dim=-1)\n","        bridge_kl = F.kl_div(\n","            bridge_soft.log().reshape(-1, n_anchors),\n","            assignment.reshape(-1, n_anchors),\n","            reduction='batchmean', log_target=False).item()\n","        lr['bridge_kl'] = bridge_kl\n","        print(f\"  Bridge ↔ assignment KL: {bridge_kl:.4f}\")\n","\n","        # === 2h. Quaternion Composition ===\n","        composed = gs['composed']\n","        comp_stats = {\n","            'norm': composed.norm(dim=-1).mean().item(),\n","            'std': composed.std().item(),\n","        }\n","        lr['composed'] = comp_stats\n","\n","        # === 2i. Geo Context ===\n","        geo_ctx = gs['geo_ctx']\n","        ctx_stats = {\n","            'norm': geo_ctx.norm(dim=-1).mean().item(),\n","            'std': geo_ctx.std().item(),\n","        }\n","        lr['geo_ctx'] = ctx_stats\n","\n","        # === 2j. Geometric Residual Stream ===\n","        geo_res = gs.get('geo_residual')\n","        if geo_res is not None:\n","            # Per-position residual statistics\n","            res_norms = geo_res.norm(dim=-1)  # (B, L)\n","            res_stats = {\n","                'norm_mean': res_norms.mean().item(),\n","                'norm_std': res_norms.std().item(),\n","                'norm_min': res_norms.min().item(),\n","                'norm_max': res_norms.max().item(),\n","                'feature_std': geo_res.std().item(),\n","                'sparsity': (geo_res.abs() < 0.01).float().mean().item(),\n","            }\n","\n","            # Cross-position consistency: how similar is the residual across positions?\n","            # High consistency = global signal, low = position-specific\n","            geo_res_flat = geo_res.reshape(-1, geo_res.shape[-1])  # (B*L, pw_dim)\n","            if geo_res_flat.shape[0] > 1:\n","                n_sample = min(256, geo_res_flat.shape[0])\n","                idx = torch.randperm(geo_res_flat.shape[0], device=device)[:n_sample]\n","                sampled = F.normalize(geo_res_flat[idx], dim=-1)\n","                cos_mat = sampled @ sampled.T\n","                triu_idx = torch.triu_indices(n_sample, n_sample, offset=1, device=device)\n","                pos_consistency = cos_mat[triu_idx[0], triu_idx[1]].mean().item()\n","                res_stats['position_consistency'] = pos_consistency\n","            else:\n","                res_stats['position_consistency'] = 0.0\n","\n","            # Growth rate: how much did this layer add to the residual?\n","            # Compare norm to previous layer (tracked in trajectory)\n","            lr['geo_residual'] = res_stats\n","\n","            print(f\"\\n  Geo residual norm:  {res_stats['norm_mean']:.3f} ± {res_stats['norm_std']:.3f}\")\n","            print(f\"    range: [{res_stats['norm_min']:.3f}, {res_stats['norm_max']:.3f}]\")\n","            print(f\"    position consistency: {res_stats['position_consistency']:.3f}\")\n","            print(f\"    sparsity: {res_stats['sparsity']:.3f}\")\n","        else:\n","            lr['geo_residual'] = {'norm_mean': 0.0, 'position_consistency': 0.0}\n","            print(f\"\\n  Geo residual: NOT PRESENT (old model without residual stream)\")\n","\n","        print(f\"  Geo context norm:   {ctx_stats['norm']:.3f}\")\n","        print(f\"  Composed norm:      {comp_stats['norm']:.3f}\")\n","\n","        report['layers'][f'layer_{i}'] = lr\n","\n","    # ─── 3. Cayley Rotations ───\n","    print(f\"\\n{'='*70}\")\n","    print(\"  3. CAYLEY ROTATIONS\")\n","    print(\"=\" * 70)\n","\n","    cayley_report = {}\n","    for name, mod in model.named_modules():\n","        if isinstance(mod, CayleyOrthogonal):\n","            R = mod.get_rotation()\n","            I = torch.eye(R.shape[0], device=R.device)\n","            r_dist = (R - I).norm().item()\n","            det = torch.det(R).item()\n","            a_norm = mod.A_upper.data.norm().item()\n","            a_max = mod.A_upper.data.abs().max().item()\n","\n","            clean = name.replace('.', '_')\n","            cayley_report[clean] = {\n","                'R_minus_I': r_dist,\n","                'det': det,\n","                'A_upper_norm': a_norm,\n","                'A_upper_max': a_max,\n","                'dim': mod.dim,\n","            }\n","            print(f\"  {name}\")\n","            print(f\"    ‖R-I‖={r_dist:.4f}  det={det:.6f}  \"\n","                  f\"‖A‖={a_norm:.4f}  max|A|={a_max:.6f}\")\n","\n","    report['cayley'] = cayley_report\n","\n","    # ─── 4. FiLM Layers ───\n","    print(f\"\\n{'='*70}\")\n","    print(\"  4. FiLM LAYERS\")\n","    print(\"=\" * 70)\n","\n","    film_report = {}\n","    film_idx = 0\n","    for name, mod in model.named_modules():\n","        if isinstance(mod, FiLMLayer):\n","            g_b = mod.to_gamma.bias.data\n","            b_b = mod.to_beta.bias.data\n","            g_dev = (g_b - 1.0).abs().mean().item()\n","            b_dev = b_b.abs().mean().item()\n","            g_w_norm = mod.to_gamma.weight.data.norm().item()\n","            b_w_norm = mod.to_beta.weight.data.norm().item()\n","\n","            film_report[f'film_{film_idx}'] = {\n","                'name': name,\n","                'gamma_dev': g_dev,\n","                'beta_dev': b_dev,\n","                'gamma_w_norm': g_w_norm,\n","                'beta_w_norm': b_w_norm,\n","            }\n","\n","            status = \"ACTIVE\" if g_dev > 0.05 or b_dev > 0.05 else \"near-identity\"\n","            print(f\"  {film_idx:>2d} [{name[-40:]:>40s}]  \"\n","                  f\"γ_dev={g_dev:.4f}  β_dev={b_dev:.4f}  {status}\")\n","            film_idx += 1\n","\n","    report['film'] = film_report\n","\n","    # ─── 5. Cross-Layer Trajectories ───\n","    print(f\"\\n{'='*70}\")\n","    print(\"  5. CROSS-LAYER TRAJECTORIES\")\n","    print(\"=\" * 70)\n","\n","    trajectories = {\n","        'cv': [], 'agreement': [], 'disagree_norm': [],\n","        'anchor_entropy': [], 'patchwork_norm': [],\n","        'bridge_kl': [], 'composed_norm': [],\n","        'geo_res_norm': [], 'geo_res_consistency': [],\n","    }\n","\n","    for i in range(n_layers):\n","        lr = report['layers'][f'layer_{i}']\n","        trajectories['cv'].append(lr['cv_embeddings']['cv'])\n","        trajectories['agreement'].append(lr['stream']['agreement_mean'])\n","        trajectories['disagree_norm'].append(lr['stream']['disagree_norm'])\n","        trajectories['anchor_entropy'].append(lr['anchor_utilization']['entropy_normalized'])\n","        trajectories['patchwork_norm'].append(lr['patchwork']['norm'])\n","        trajectories['bridge_kl'].append(lr['bridge_kl'])\n","        trajectories['composed_norm'].append(lr['composed']['norm'])\n","        trajectories['geo_res_norm'].append(lr['geo_residual']['norm_mean'])\n","        trajectories['geo_res_consistency'].append(lr['geo_residual'].get('position_consistency', 0.0))\n","\n","    report['trajectories'] = trajectories\n","\n","    for key, vals in trajectories.items():\n","        vals_str = '  '.join(f'{v:.4f}' for v in vals)\n","        trend = vals[-1] - vals[0]\n","        trend_str = f\"+{trend:.4f}\" if trend > 0 else f\"{trend:.4f}\"\n","        print(f\"  {key:<24s}: [{vals_str}]  Δ={trend_str}\")\n","\n","    # Pentachoron band analysis\n","    cv_vals = trajectories['cv']\n","    in_band = [0.20 <= cv <= 0.23 for cv in cv_vals]\n","    print(f\"\\n  Pentachoron band (0.20-0.23):\")\n","    print(f\"    Layers in band: {sum(in_band)}/{len(in_band)}\")\n","    print(f\"    CV range: [{min(cv_vals):.4f}, {max(cv_vals):.4f}]\")\n","    print(f\"    CV mean:  {np.mean(cv_vals):.4f}\")\n","\n","    # Geometric residual accumulation analysis\n","    res_norms = trajectories['geo_res_norm']\n","    if any(n > 0 for n in res_norms):\n","        print(f\"\\n  Geometric residual stream:\")\n","        print(f\"    Accumulation: {res_norms[0]:.4f} → {res_norms[-1]:.4f}\")\n","        growth_rates = [res_norms[i+1] - res_norms[i] for i in range(len(res_norms)-1)]\n","        print(f\"    Growth per layer: [{', '.join(f'{g:+.4f}' for g in growth_rates)}]\")\n","        print(f\"    Growth mean: {np.mean(growth_rates):+.4f}\")\n","\n","        # Cooperation metric: do CV and bridge_kl improve as geo_residual grows?\n","        # Negative correlation = layers cooperating (more history → better geometry)\n","        if len(res_norms) >= 4:\n","            cv_corr = np.corrcoef(res_norms, cv_vals)[0, 1]\n","            bkl_corr = np.corrcoef(res_norms, trajectories['bridge_kl'])[0, 1]\n","            print(f\"\\n  Cooperation analysis:\")\n","            print(f\"    geo_res ↔ CV correlation:        {cv_corr:+.3f}  \"\n","                  f\"{'COOPERATING' if cv_corr < -0.3 else 'NEUTRAL' if abs(cv_corr) < 0.3 else 'ADVERSARIAL'}\")\n","            print(f\"    geo_res ↔ bridge_kl correlation:  {bkl_corr:+.3f}  \"\n","                  f\"{'COOPERATING' if bkl_corr < -0.3 else 'NEUTRAL' if abs(bkl_corr) < 0.3 else 'ADVERSARIAL'}\")\n","\n","            report['cooperation'] = {\n","                'geo_res_cv_corr': float(cv_corr),\n","                'geo_res_bridge_kl_corr': float(bkl_corr),\n","                'growth_rates': growth_rates,\n","            }\n","    else:\n","        print(f\"\\n  Geometric residual stream: NOT ACTIVE (old model)\")\n","\n","    # ─── 6. Per-Class Analysis ───\n","    if per_class:\n","        print(f\"\\n{'='*70}\")\n","        print(\"  6. PER-CLASS ANALYSIS\")\n","        print(\"=\" * 70)\n","\n","        n_classes = all_logits.shape[1]\n","        class_stats = {}\n","        class_accs = []\n","\n","        for c in range(n_classes):\n","            mask = all_labels == c\n","            if mask.sum() == 0:\n","                continue\n","            c_pred = pred[mask]\n","            c_labels = all_labels[mask]\n","            c_acc = (c_pred == c_labels).float().mean().item()\n","            c_conf = F.softmax(all_logits[mask], dim=-1)[:, c].mean().item()\n","            c_entropy = -(F.softmax(all_logits[mask], dim=-1) *\n","                         F.log_softmax(all_logits[mask], dim=-1)).sum(-1).mean().item()\n","\n","            class_stats[int(c)] = {\n","                'accuracy': c_acc,\n","                'confidence': c_conf,\n","                'entropy': c_entropy,\n","                'n_samples': int(mask.sum().item()),\n","            }\n","            class_accs.append(c_acc)\n","\n","        report['per_class'] = class_stats\n","\n","        class_accs_arr = np.array(class_accs)\n","        print(f\"  Class accuracy distribution:\")\n","        print(f\"    Mean:   {class_accs_arr.mean():.4f}\")\n","        print(f\"    Std:    {class_accs_arr.std():.4f}\")\n","        print(f\"    Min:    {class_accs_arr.min():.4f}\")\n","        print(f\"    Max:    {class_accs_arr.max():.4f}\")\n","        print(f\"    P25:    {np.percentile(class_accs_arr, 25):.4f}\")\n","        print(f\"    P75:    {np.percentile(class_accs_arr, 75):.4f}\")\n","\n","        # Best and worst classes\n","        sorted_classes = sorted(class_stats.items(), key=lambda x: x[1]['accuracy'])\n","        print(f\"\\n  Bottom 5 classes:\")\n","        for c, stats in sorted_classes[:5]:\n","            print(f\"    Class {c:>3d}: acc={stats['accuracy']:.3f}  \"\n","                  f\"conf={stats['confidence']:.3f}  ent={stats['entropy']:.3f}\")\n","        print(f\"\\n  Top 5 classes:\")\n","        for c, stats in sorted_classes[-5:]:\n","            print(f\"    Class {c:>3d}: acc={stats['accuracy']:.3f}  \"\n","                  f\"conf={stats['confidence']:.3f}  ent={stats['entropy']:.3f}\")\n","\n","    # ─── 7. Model Summary ───\n","    print(f\"\\n{'='*70}\")\n","    print(\"  7. MODEL SUMMARY\")\n","    print(\"=\" * 70)\n","\n","    n_params = sum(p.numel() for p in model.parameters())\n","    n_cayley = sum(1 for _ in (m for m in model.modules() if isinstance(m, CayleyOrthogonal)))\n","    n_film = sum(1 for _ in (m for m in model.modules() if isinstance(m, FiLMLayer)))\n","    n_observers = sum(1 for name, m in model.named_modules()\n","                      if hasattr(m, 'association') and hasattr(m.association, 'constellation'))\n","\n","    print(f\"  Parameters:     {n_params:,}\")\n","    print(f\"  Layers:         {n_layers}\")\n","    print(f\"  Cayley modules: {n_cayley}\")\n","    print(f\"  FiLM modules:   {n_film}\")\n","    print(f\"  Observers:      {n_observers}\")\n","    print(f\"  Overall acc:    {overall_acc:.4f}\")\n","\n","    report['model'] = {\n","        'n_params': n_params,\n","        'n_layers': n_layers,\n","        'n_cayley': n_cayley,\n","        'n_film': n_film,\n","        'n_observers': n_observers,\n","        'accuracy': overall_acc,\n","    }\n","\n","    return report\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# DATA\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","def get_test_loader(config):\n","    import torchvision\n","    import torchvision.transforms as T\n","\n","    test_transform = T.Compose([\n","        T.ToTensor(),\n","        T.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),\n","    ])\n","    test_ds = torchvision.datasets.CIFAR100(\n","        root='./data', train=False, download=True, transform=test_transform)\n","    return torch.utils.data.DataLoader(\n","        test_ds, batch_size=config.get('batch_size', 256), shuffle=False,\n","        num_workers=4, pin_memory=True)\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# MAIN\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","def main():\n","    parser = argparse.ArgumentParser(description='Geometric Transformer Analysis')\n","    parser.add_argument('--checkpoint', type=str, default='geo_cifar100/best.pt',\n","                        help='Path to checkpoint')\n","    parser.add_argument('--n_batches', type=int, default=4,\n","                        help='Number of test batches to analyze')\n","    parser.add_argument('--per_class', action='store_true',\n","                        help='Run per-class analysis')\n","    parser.add_argument('--output', type=str, default='geo_analysis.json',\n","                        help='Output JSON report path')\n","    args, _ = parser.parse_known_args()\n","\n","    print(f\"Device: {device}\")\n","    if device.type == 'cuda':\n","        print(f\"  GPU: {torch.cuda.get_device_name()}\")\n","\n","    # Load checkpoint\n","    print(f\"\\nLoading checkpoint: {args.checkpoint}\")\n","    ckpt = torch.load(args.checkpoint, map_location='cpu', weights_only=False)\n","    config = ckpt['config']\n","    print(f\"  Epoch: {ckpt['epoch']}\")\n","    print(f\"  Test acc: {ckpt['test_acc']:.4f}\")\n","    print(f\"  Config: d={config['d_model']}, layers={config['n_layers']}, \"\n","          f\"anchors={config['n_anchors']}\")\n","\n","    # Build model\n","    model = GeoViTClassifier('analysis', config)\n","    model.load_state_dict(ckpt['state_dict'], strict=False)\n","    if hasattr(model, 'network_to'):\n","        model.network_to(device=device, strict=False)\n","    else:\n","        model = model.to(device)\n","    model.eval()\n","\n","    n_params = sum(p.numel() for p in model.parameters())\n","    print(f\"  Parameters: {n_params:,}\")\n","\n","    # Data\n","    test_loader = get_test_loader(config)\n","\n","    # Run analysis\n","    print(f\"\\n{'━'*70}\")\n","    print(f\"  GEOMETRIC ANALYSIS BATTERY\")\n","    print(f\"{'━'*70}\")\n","\n","    report = run_analysis(\n","        model, test_loader, device,\n","        n_batches=args.n_batches,\n","        per_class=args.per_class,\n","    )\n","\n","    # Save report\n","    output_path = Path(args.output)\n","    # Convert non-serializable types\n","    def clean_for_json(obj):\n","        if isinstance(obj, torch.Tensor):\n","            if obj.dim() == 0:\n","                return obj.item()\n","            return obj.tolist()\n","        if isinstance(obj, (np.float32, np.float64)):\n","            return float(obj)\n","        if isinstance(obj, (np.int32, np.int64)):\n","            return int(obj)\n","        if isinstance(obj, np.ndarray):\n","            return obj.tolist()\n","        if isinstance(obj, dict):\n","            return {k: clean_for_json(v) for k, v in obj.items()}\n","        if isinstance(obj, list):\n","            return [clean_for_json(v) for v in obj]\n","        return obj\n","\n","    class TensorEncoder(json.JSONEncoder):\n","        def default(self, obj):\n","            if isinstance(obj, torch.Tensor):\n","                return obj.cpu().tolist() if obj.dim() > 0 else obj.item()\n","            if isinstance(obj, (np.float32, np.float64, np.floating)):\n","                return float(obj)\n","            if isinstance(obj, (np.int32, np.int64, np.integer)):\n","                return int(obj)\n","            if isinstance(obj, np.ndarray):\n","                return obj.tolist()\n","            return super().default(obj)\n","\n","    with open(output_path, 'w') as f:\n","        json.dump(clean_for_json(report), f, indent=2, cls=TensorEncoder)\n","    print(f\"\\nReport saved: {output_path}\")\n","\n","    print(f\"\\n{'═'*70}\")\n","    print(f\"  ANALYSIS COMPLETE\")\n","    print(f\"{'═'*70}\")\n","\n","\n","if __name__ == '__main__':\n","    main()"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"gM8PyZvMAHX4","executionInfo":{"status":"ok","timestamp":1774694301868,"user_tz":420,"elapsed":1424,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"f2dbe007-fb8d-4c52-e119-3166335ad947"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["Device: cuda\n","  GPU: NVIDIA RTX PRO 6000 Blackwell Server Edition\n","\n","Loading checkpoint: geo_cifar100/best.pt\n","  Epoch: 86\n","  Test acc: 0.5731\n","  Config: d=256, layers=8, anchors=128\n","  Parameters: 17,750,436\n","\n","━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\n","  GEOMETRIC ANALYSIS BATTERY\n","━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\n","\n","======================================================================\n","  1. SVD INPUT STAGE\n","======================================================================\n","  Entropy:          2.593 ± 0.057\n","  Novelty norm:     17.345\n","  Top-1 ratio:      0.175\n","  Condition number: 9.0\n","  Singular values:  18.48, 11.86, 10.05, 8.71, 7.68...\n","  FiLM γ dev:       0.1042\n","  FiLM β norm:      0.0086\n","\n","──────────────────────────────────────────────────────────────────────\n","  LAYER 0\n","──────────────────────────────────────────────────────────────────────\n","  CV embeddings:    0.4830  ✗ outside\n","    mean_dist=0.9527  std=0.4602\n","  CV anchors:       0.2260  ✓ IN BAND\n","    nn_dist: mean=0.6016  min=0.3577  max=0.8813\n","    max_cos_sim:    0.6423\n","\n","  Stream agreement: 0.5415 ± 0.1073\n","    range: [-0.2147, 0.8100]\n","  ‖content‖:        14.484\n","  ‖geometric‖:      14.356\n","  ‖disagree‖:       13.726  (ratio: 0.948)\n","  ‖agree‖:          19.217\n","\n","  Anchors: 113/128 active, 15 dead\n","  Entropy: 3.453 / 4.852 (normalized: 0.712)\n","  Top-1 dominance: 0.2117\n","  Top-5 counts: [14093, 4485, 3959, 2841, 2829]\n","\n","  Assignment entropy: 3.570 ± 0.336\n","  Max prob:           0.1194 ± 0.0496\n","\n","  Patchwork norm:     6.600\n","  Patchwork sparsity: 0.030  (dead: 0.0000)\n","  Bridge ↔ assignment KL: 1.4133\n","\n","  Geo residual norm:  1.341 ± 0.354\n","    range: [0.649, 1.792]\n","    position consistency: 0.562\n","    sparsity: 0.300\n","  Geo context norm:   5.578\n","  Composed norm:      5.657\n","\n","──────────────────────────────────────────────────────────────────────\n","  LAYER 1\n","──────────────────────────────────────────────────────────────────────\n","  CV embeddings:    0.4528  ✗ outside\n","    mean_dist=0.9836  std=0.4454\n","  CV anchors:       0.2439  ✗ outside\n","    nn_dist: mean=0.6545  min=0.4845  max=0.9270\n","    max_cos_sim:    0.5155\n","\n","  Stream agreement: 0.4463 ± 0.0837\n","    range: [0.0717, 0.7363]\n","  ‖content‖:        14.601\n","  ‖geometric‖:      14.297\n","  ‖disagree‖:       15.163  (ratio: 1.038)\n","  ‖agree‖:          18.566\n","\n","  Anchors: 124/128 active, 4 dead\n","  Entropy: 3.665 / 4.852 (normalized: 0.755)\n","  Top-1 dominance: 0.1127\n","  Top-5 counts: [7502, 5366, 5318, 4305, 4127]\n","\n","  Assignment entropy: 3.827 ± 0.293\n","  Max prob:           0.0827 ± 0.0281\n","\n","  Patchwork norm:     6.736\n","  Patchwork sparsity: 0.022  (dead: 0.0000)\n","  Bridge ↔ assignment KL: 1.1953\n","\n","  Geo residual norm:  1.626 ± 0.428\n","    range: [0.913, 2.549]\n","    position consistency: 0.400\n","    sparsity: 0.171\n","  Geo context norm:   4.755\n","  Composed norm:      5.657\n","\n","──────────────────────────────────────────────────────────────────────\n","  LAYER 2\n","──────────────────────────────────────────────────────────────────────\n","  CV embeddings:    0.5751  ✗ outside\n","    mean_dist=0.9567  std=0.5502\n","  CV anchors:       0.2579  ✗ outside\n","    nn_dist: mean=0.6218  min=0.3429  max=0.9133\n","    max_cos_sim:    0.6571\n","\n","  Stream agreement: 0.4693 ± 0.0906\n","    range: [0.0427, 0.7770]\n","  ‖content‖:        14.629\n","  ‖geometric‖:      14.312\n","  ‖disagree‖:       14.857  (ratio: 1.016)\n","  ‖agree‖:          20.458\n","\n","  Anchors: 125/128 active, 3 dead\n","  Entropy: 3.016 / 4.852 (normalized: 0.622)\n","  Top-1 dominance: 0.2611\n","  Top-5 counts: [17381, 7601, 5457, 5260, 3108]\n","\n","  Assignment entropy: 3.572 ± 0.309\n","  Max prob:           0.1141 ± 0.0396\n","\n","  Patchwork norm:     6.783\n","  Patchwork sparsity: 0.020  (dead: 0.0000)\n","  Bridge ↔ assignment KL: 1.4194\n","\n","  Geo residual norm:  1.916 ± 0.344\n","    range: [0.887, 2.643]\n","    position consistency: 0.362\n","    sparsity: 0.118\n","  Geo context norm:   4.908\n","  Composed norm:      5.657\n","\n","──────────────────────────────────────────────────────────────────────\n","  LAYER 3\n","──────────────────────────────────────────────────────────────────────\n","  CV embeddings:    0.4028  ✗ outside\n","    mean_dist=0.9844  std=0.3965\n","  CV anchors:       0.1056  ✗ outside\n","    nn_dist: mean=0.8128  min=0.5646  max=0.9530\n","    max_cos_sim:    0.4354\n","\n","  Stream agreement: 0.4895 ± 0.1002\n","    range: [0.0506, 0.7934]\n","  ‖content‖:        14.625\n","  ‖geometric‖:      14.309\n","  ‖disagree‖:       14.549  (ratio: 0.995)\n","  ‖agree‖:          23.034\n","\n","  Anchors: 128/128 active, 0 dead\n","  Entropy: 3.493 / 4.852 (normalized: 0.720)\n","  Top-1 dominance: 0.1334\n","  Top-5 counts: [8877, 6760, 5917, 3582, 3467]\n","\n","  Assignment entropy: 3.770 ± 0.422\n","  Max prob:           0.1129 ± 0.0521\n","\n","  Patchwork norm:     7.552\n","  Patchwork sparsity: 0.015  (dead: 0.0000)\n","  Bridge ↔ assignment KL: 1.2033\n","\n","  Geo residual norm:  1.929 ± 0.344\n","    range: [0.885, 2.649]\n","    position consistency: 0.367\n","    sparsity: 0.102\n","  Geo context norm:   3.470\n","  Composed norm:      5.657\n","\n","──────────────────────────────────────────────────────────────────────\n","  LAYER 4\n","──────────────────────────────────────────────────────────────────────\n","  CV embeddings:    0.4371  ✗ outside\n","    mean_dist=0.9817  std=0.4291\n","  CV anchors:       0.0469  ✗ outside\n","    nn_dist: mean=0.9012  min=0.8053  max=0.9538\n","    max_cos_sim:    0.1947\n","\n","  Stream agreement: 0.4925 ± 0.1056\n","    range: [-0.0263, 0.7960]\n","  ‖content‖:        14.666\n","  ‖geometric‖:      14.289\n","  ‖disagree‖:       14.510  (ratio: 0.989)\n","  ‖agree‖:          23.155\n","\n","  Anchors: 127/128 active, 1 dead\n","  Entropy: 3.319 / 4.852 (normalized: 0.684)\n","  Top-1 dominance: 0.2135\n","  Top-5 counts: [14211, 4698, 4655, 4486, 4250]\n","\n","  Assignment entropy: 4.062 ± 0.246\n","  Max prob:           0.1029 ± 0.0560\n","\n","  Patchwork norm:     7.345\n","  Patchwork sparsity: 0.022  (dead: 0.0000)\n","  Bridge ↔ assignment KL: 1.0032\n","\n","  Geo residual norm:  1.937 ± 0.340\n","    range: [0.860, 2.651]\n","    position consistency: 0.387\n","    sparsity: 0.099\n","  Geo context norm:   4.274\n","  Composed norm:      5.657\n","\n","──────────────────────────────────────────────────────────────────────\n","  LAYER 5\n","──────────────────────────────────────────────────────────────────────\n","  CV embeddings:    0.3482  ✗ outside\n","    mean_dist=0.9608  std=0.3346\n","  CV anchors:       0.0756  ✗ outside\n","    nn_dist: mean=0.8577  min=0.6880  max=0.9575\n","    max_cos_sim:    0.3120\n","\n","  Stream agreement: 0.4926 ± 0.1092\n","    range: [0.0378, 0.8113]\n","  ‖content‖:        14.624\n","  ‖geometric‖:      14.320\n","  ‖disagree‖:       14.497  (ratio: 0.991)\n","  ‖agree‖:          22.058\n","\n","  Anchors: 128/128 active, 0 dead\n","  Entropy: 3.709 / 4.852 (normalized: 0.764)\n","  Top-1 dominance: 0.1253\n","  Top-5 counts: [8340, 6734, 5079, 3750, 3704]\n","\n","  Assignment entropy: 4.004 ± 0.298\n","  Max prob:           0.0914 ± 0.0379\n","\n","  Patchwork norm:     7.478\n","  Patchwork sparsity: 0.025  (dead: 0.0000)\n","  Bridge ↔ assignment KL: 0.9496\n","\n","  Geo residual norm:  1.948 ± 0.334\n","    range: [0.870, 2.652]\n","    position consistency: 0.390\n","    sparsity: 0.097\n","  Geo context norm:   4.542\n","  Composed norm:      5.657\n","\n","──────────────────────────────────────────────────────────────────────\n","  LAYER 6\n","──────────────────────────────────────────────────────────────────────\n","  CV embeddings:    0.3469  ✗ outside\n","    mean_dist=0.9668  std=0.3354\n","  CV anchors:       0.0549  ✗ outside\n","    nn_dist: mean=0.8869  min=0.7636  max=0.9492\n","    max_cos_sim:    0.2364\n","\n","  Stream agreement: 0.5023 ± 0.1034\n","    range: [0.0404, 0.8000]\n","  ‖content‖:        14.680\n","  ‖geometric‖:      14.318\n","  ‖disagree‖:       14.392  (ratio: 0.980)\n","  ‖agree‖:          20.846\n","\n","  Anchors: 128/128 active, 0 dead\n","  Entropy: 3.865 / 4.852 (normalized: 0.797)\n","  Top-1 dominance: 0.0738\n","  Top-5 counts: [4913, 4575, 4549, 3796, 3783]\n","\n","  Assignment entropy: 4.112 ± 0.238\n","  Max prob:           0.0858 ± 0.0344\n","\n","  Patchwork norm:     7.283\n","  Patchwork sparsity: 0.027  (dead: 0.0000)\n","  Bridge ↔ assignment KL: 0.9126\n","\n","  Geo residual norm:  1.974 ± 0.325\n","    range: [0.877, 2.658]\n","    position consistency: 0.380\n","    sparsity: 0.075\n","  Geo context norm:   4.199\n","  Composed norm:      5.657\n","\n","──────────────────────────────────────────────────────────────────────\n","  LAYER 7\n","──────────────────────────────────────────────────────────────────────\n","  CV embeddings:    0.3158  ✗ outside\n","    mean_dist=0.9251  std=0.2922\n","  CV anchors:       0.0368  ✗ outside\n","    nn_dist: mean=0.9233  min=0.8506  max=0.9642\n","    max_cos_sim:    0.1494\n","\n","  Stream agreement: 0.5238 ± 0.1151\n","    range: [-0.0632, 0.8245]\n","  ‖content‖:        14.664\n","  ‖geometric‖:      14.302\n","  ‖disagree‖:       14.039  (ratio: 0.957)\n","  ‖agree‖:          18.541\n","\n","  Anchors: 128/128 active, 0 dead\n","  Entropy: 3.715 / 4.852 (normalized: 0.766)\n","  Top-1 dominance: 0.1567\n","  Top-5 counts: [10428, 6251, 6089, 5098, 3218]\n","\n","  Assignment entropy: 4.323 ± 0.174\n","  Max prob:           0.0723 ± 0.0314\n","\n","  Patchwork norm:     7.435\n","  Patchwork sparsity: 0.015  (dead: 0.0000)\n","  Bridge ↔ assignment KL: 0.6352\n","\n","  Geo residual norm:  2.015 ± 0.315\n","    range: [0.905, 2.672]\n","    position consistency: 0.432\n","    sparsity: 0.049\n","  Geo context norm:   6.111\n","  Composed norm:      5.657\n","\n","======================================================================\n","  3. CAYLEY ROTATIONS\n","======================================================================\n","  components.transformer.components.layer_0.components.rotation\n","    ‖R-I‖=7.5051  det=1.000003  ‖A‖=2.9140  max|A|=0.089733\n","  components.transformer.components.layer_1.components.rotation\n","    ‖R-I‖=9.1009  det=1.000018  ‖A‖=3.5962  max|A|=0.103736\n","  components.transformer.components.layer_2.components.rotation\n","    ‖R-I‖=9.6409  det=0.999993  ‖A‖=3.8321  max|A|=0.092645\n","  components.transformer.components.layer_3.components.rotation\n","    ‖R-I‖=10.1452  det=0.999992  ‖A‖=4.0722  max|A|=0.101473\n","  components.transformer.components.layer_4.components.rotation\n","    ‖R-I‖=10.1594  det=0.999997  ‖A‖=4.0807  max|A|=0.105992\n","  components.transformer.components.layer_5.components.rotation\n","    ‖R-I‖=10.2185  det=1.000003  ‖A‖=4.1135  max|A|=0.107747\n","  components.transformer.components.layer_6.components.rotation\n","    ‖R-I‖=11.1442  det=1.000001  ‖A‖=4.5629  max|A|=0.116672\n","  components.transformer.components.layer_7.components.rotation\n","    ‖R-I‖=10.6947  det=0.999996  ‖A‖=4.3455  max|A|=0.112253\n","  components.transformer.components.cross_rot_0\n","    ‖R-I‖=7.6806  det=0.999996  ‖A‖=2.9250  max|A|=0.081570\n","  components.transformer.components.cross_rot_1\n","    ‖R-I‖=8.2690  det=1.000010  ‖A‖=3.1725  max|A|=0.077519\n","  components.transformer.components.cross_rot_2\n","    ‖R-I‖=8.7539  det=1.000013  ‖A‖=3.3866  max|A|=0.085116\n","  components.transformer.components.cross_rot_3\n","    ‖R-I‖=9.0362  det=0.999996  ‖A‖=3.5113  max|A|=0.097553\n","  components.transformer.components.cross_rot_4\n","    ‖R-I‖=9.1991  det=1.000003  ‖A‖=3.5779  max|A|=0.092572\n","  components.transformer.components.cross_rot_5\n","    ‖R-I‖=9.4783  det=0.999989  ‖A‖=3.7214  max|A|=0.124684\n","  components.transformer.components.cross_rot_6\n","    ‖R-I‖=9.5608  det=0.999994  ‖A‖=3.7816  max|A|=0.092843\n","\n","======================================================================\n","  4. FiLM LAYERS\n","======================================================================\n","   0 [ents.layer_0.components.geometric.film_q]  γ_dev=0.1561  β_dev=0.0261  ACTIVE\n","   1 [ents.layer_0.components.geometric.film_k]  γ_dev=0.1476  β_dev=0.0000  ACTIVE\n","   2 [ts.layer_0.components.geometric.film_ffn]  γ_dev=0.1370  β_dev=0.0037  ACTIVE\n","   3 [ents.layer_1.components.geometric.film_q]  γ_dev=0.1341  β_dev=0.0372  ACTIVE\n","   4 [ents.layer_1.components.geometric.film_k]  γ_dev=0.1234  β_dev=0.0000  ACTIVE\n","   5 [ts.layer_1.components.geometric.film_ffn]  γ_dev=0.1248  β_dev=0.0048  ACTIVE\n","   6 [ents.layer_2.components.geometric.film_q]  γ_dev=0.1379  β_dev=0.0420  ACTIVE\n","   7 [ents.layer_2.components.geometric.film_k]  γ_dev=0.1246  β_dev=0.0000  ACTIVE\n","   8 [ts.layer_2.components.geometric.film_ffn]  γ_dev=0.1296  β_dev=0.0056  ACTIVE\n","   9 [ents.layer_3.components.geometric.film_q]  γ_dev=0.1278  β_dev=0.0385  ACTIVE\n","  10 [ents.layer_3.components.geometric.film_k]  γ_dev=0.1231  β_dev=0.0000  ACTIVE\n","  11 [ts.layer_3.components.geometric.film_ffn]  γ_dev=0.1280  β_dev=0.0081  ACTIVE\n","  12 [ents.layer_4.components.geometric.film_q]  γ_dev=0.1348  β_dev=0.0418  ACTIVE\n","  13 [ents.layer_4.components.geometric.film_k]  γ_dev=0.1213  β_dev=0.0000  ACTIVE\n","  14 [ts.layer_4.components.geometric.film_ffn]  γ_dev=0.1285  β_dev=0.0068  ACTIVE\n","  15 [ents.layer_5.components.geometric.film_q]  γ_dev=0.1277  β_dev=0.0412  ACTIVE\n","  16 [ents.layer_5.components.geometric.film_k]  γ_dev=0.1199  β_dev=0.0000  ACTIVE\n","  17 [ts.layer_5.components.geometric.film_ffn]  γ_dev=0.1332  β_dev=0.0064  ACTIVE\n","  18 [ents.layer_6.components.geometric.film_q]  γ_dev=0.1243  β_dev=0.0379  ACTIVE\n","  19 [ents.layer_6.components.geometric.film_k]  γ_dev=0.1206  β_dev=0.0000  ACTIVE\n","  20 [ts.layer_6.components.geometric.film_ffn]  γ_dev=0.1145  β_dev=0.0062  ACTIVE\n","  21 [ents.layer_7.components.geometric.film_q]  γ_dev=0.1186  β_dev=0.0405  ACTIVE\n","  22 [ents.layer_7.components.geometric.film_k]  γ_dev=0.1106  β_dev=0.0000  ACTIVE\n","  23 [ts.layer_7.components.geometric.film_ffn]  γ_dev=0.0836  β_dev=0.0051  ACTIVE\n","\n","======================================================================\n","  5. CROSS-LAYER TRAJECTORIES\n","======================================================================\n","  cv                      : [0.4830  0.4528  0.5751  0.4028  0.4371  0.3482  0.3469  0.3158]  Δ=-0.1672\n","  agreement               : [0.5415  0.4463  0.4693  0.4895  0.4925  0.4926  0.5023  0.5238]  Δ=-0.0177\n","  disagree_norm           : [13.7264  15.1626  14.8567  14.5495  14.5105  14.4970  14.3916  14.0393]  Δ=+0.3129\n","  anchor_entropy          : [0.7117  0.7554  0.6215  0.7199  0.6840  0.7645  0.7966  0.7656]  Δ=+0.0539\n","  patchwork_norm          : [6.6001  6.7358  6.7834  7.5518  7.3452  7.4780  7.2830  7.4353]  Δ=+0.8352\n","  bridge_kl               : [1.4133  1.1953  1.4194  1.2033  1.0032  0.9496  0.9126  0.6352]  Δ=-0.7782\n","  composed_norm           : [5.6569  5.6569  5.6569  5.6569  5.6569  5.6569  5.6569  5.6569]  Δ=0.0000\n","  geo_res_norm            : [1.3414  1.6262  1.9160  1.9291  1.9373  1.9483  1.9741  2.0149]  Δ=+0.6735\n","  geo_res_consistency     : [0.5621  0.3999  0.3621  0.3668  0.3866  0.3900  0.3795  0.4318]  Δ=-0.1303\n","\n","  Pentachoron band (0.20-0.23):\n","    Layers in band: 0/8\n","    CV range: [0.3158, 0.5751]\n","    CV mean:  0.4202\n","\n","  Geometric residual stream:\n","    Accumulation: 1.3414 → 2.0149\n","    Growth per layer: [+0.2848, +0.2898, +0.0131, +0.0082, +0.0110, +0.0258, +0.0408]\n","    Growth mean: +0.0962\n","\n","  Cooperation analysis:\n","    geo_res ↔ CV correlation:        -0.449  COOPERATING\n","    geo_res ↔ bridge_kl correlation:  -0.632  COOPERATING\n","\n","======================================================================\n","  7. MODEL SUMMARY\n","======================================================================\n","  Parameters:     17,750,436\n","  Layers:         8\n","  Cayley modules: 15\n","  FiLM modules:   24\n","  Observers:      8\n","  Overall acc:    0.5713\n","\n","Report saved: geo_analysis.json\n","\n","══════════════════════════════════════════════════════════════════════\n","  ANALYSIS COMPLETE\n","══════════════════════════════════════════════════════════════════════\n"]}]},{"cell_type":"markdown","source":["# transformer version 2"],"metadata":{"id":"rB99d0u3CWOM"}},{"cell_type":"code","source":["\"\"\"\n","Geometric Transformer — CIFAR-100 Training with CM-Validated Analysis\n","\n","Changes from previous version:\n","  - CM gate diagnostics per layer: active anchors, gate_mean, cm_positive_frac\n","  - CM quality in geometric residual analysis (replaces blind gate)\n","  - Geometric regularization losses (CV target + anchor spread) in training loop\n","  - Anchor diagnostics via model.anchor_diagnostics()\n","  - CM quality trajectory alongside CV and bridge KL for cooperation analysis\n","\n","TensorBoard logging of every geometric feature element:\n","  - CV (coefficient of variation) per layer — the pentachoron band metric\n","  - CM gate: active anchors, gate mean, cm_positive_frac, quality per position\n","  - Stream agreement/divergence per layer\n","  - Anchor utilization, entropy, spread\n","  - Patchwork activation statistics (from CM-validated triangulation)\n","  - Bridge vs assignment consistency\n","  - Triangulation distance distributions\n","  - SVD spectrum, entropy, novelty\n","  - Quaternion arm norms and composition statistics\n","  - Cayley rotation ‖R-I‖ per layer\n","  - FiLM gamma/beta deviation from identity\n","  - Gate activation statistics\n","  - Gradient norms per component type (including cm_gate)\n","  - Weight norms per component type\n","  - Geometric regularization: CV loss, spread loss per epoch\n","\n","!pip install geolip-core torchvision tqdm tensorboard\n","\"\"\"\n","\n","import torch\n","import torch.nn as nn\n","import torch.nn.functional as F\n","import numpy as np\n","import time, json, math\n","from pathlib import Path\n","from tqdm.auto import tqdm\n","from torch.utils.tensorboard import SummaryWriter\n","\n","device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n","print(f\"Device: {device}\")\n","if device.type == 'cuda':\n","    print(f\"  GPU: {torch.cuda.get_device_name()}\")\n","    print(f\"  VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB\")\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# IMPORT TRANSFORMER\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","# Try geolip_core installed package first, fall back to local file\n","try:\n","    from geolip_core.pipeline.components.geometric_transformer import (\n","        GeometricTransformer, GeometricTransformerLayer,\n","        CayleyOrthogonal, QuaternionCompose, FiLMLayer,\n","        ContentAttention, GeometricAttention, CMValidatedGate,\n","        TorchComponent, BaseTower,\n","        anchor_neighborhood_cm,\n","    )\n","    print(\"  Imported from geolip_core (installed)\")\n","except ImportError:\n","    try:\n","        from geometric_transformer import (\n","            GeometricTransformer, GeometricTransformerLayer,\n","            CayleyOrthogonal, QuaternionCompose, FiLMLayer,\n","            ContentAttention, GeometricAttention, CMValidatedGate,\n","            TorchComponent, BaseTower,\n","            anchor_neighborhood_cm,\n","        )\n","        print(\"  Imported from local geometric_transformer.py\")\n","    except ImportError:\n","        raise ImportError(\n","            \"Cannot find geometric_transformer. Place geometric_transformer.py \"\n","            \"in the working directory or install geolip-core.\")\n","\n","torch.set_float32_matmul_precision('high')\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# CONFIG\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","CONFIG = {\n","    # Model\n","    'd_model': 256,\n","    'n_heads': 4,\n","    'n_layers': 4,\n","    'n_anchors': 128,\n","    'manifold_dim': 128,\n","    'n_comp': 8,\n","    'd_comp': 16,\n","    'context_dim': 64,\n","    'quat_dim': 64,\n","    'dropout': 0.1,\n","    'cm_neighbors': 4,       # CM simplex neighbors\n","\n","    # Input stage\n","    'patch_size': 4,\n","    'img_size': 32,\n","    'in_channels': 3,\n","    'conv_channels': 64,\n","    'svd_rank': 24,\n","\n","    # Training\n","    'epochs': 300,\n","    'batch_size': 1024,\n","    'lr': 1e-3,\n","    'weight_decay': 0.05,\n","    'warmup_epochs': 5,\n","    'label_smoothing': 0.1,\n","    'num_workers': 8,\n","\n","    # Geometric regularization\n","    'cv_target': 0.20,      # pentachoron band center\n","    'cv_weight': 0.1,        # CV loss weight\n","    'spread_weight': 0.01,   # anchor spread loss weight\n","\n","    # Augmentation — tuned for CM gate training\n","    'cutmix_alpha': 1.0,     # CutMix beta distribution α (1.0 = uniform box sizes)\n","    'cutmix_prob': 0.5,      # probability of applying CutMix per batch\n","    'random_erasing_p': 0.25,  # probability of erasing per image\n","\n","    # Data\n","    'num_classes': 100,\n","\n","    # Logging\n","    'log_geo_every': 5,      # full geometric analysis every N epochs\n","    'log_grads_every': 10,   # gradient norms every N epochs\n","    'log_dir': 'runs/geo_cifar100_2',\n","}\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# INPUT STAGE\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","try:\n","    from geolip_core.core.input.svd import SVDObserver\n","    _HAS_SVD = True\n","except ImportError:\n","    _HAS_SVD = False\n","\n","    class SVDObserver(nn.Module):\n","        \"\"\"Fallback SVDObserver.\"\"\"\n","        def __init__(self, in_channels, svd_rank=24):\n","            super().__init__()\n","            self.svd_rank = svd_rank\n","            self.to_svd = nn.Conv2d(in_channels, svd_rank, 1, bias=False)\n","            self.register_buffer('ema_s', torch.ones(svd_rank))\n","            self.register_buffer('ema_vh_flat', torch.eye(svd_rank).reshape(-1))\n","            self.ema_momentum = 0.99\n","\n","        def extract_features(self, S, Vh):\n","            B, k = S.shape\n","            S_safe = S.clamp(min=1e-6)\n","            s_norm = S_safe / (S_safe.sum(dim=-1, keepdim=True) + 1e-8)\n","            vh_diag = Vh.diagonal(dim1=-2, dim2=-1)\n","            vh_offdiag = (Vh.pow(2).sum((-2, -1)) - vh_diag.pow(2).sum(-1)).unsqueeze(-1).clamp(min=0)\n","            s_ent = -(s_norm * torch.log(s_norm.clamp(min=1e-8))).sum(-1, keepdim=True)\n","            out = torch.cat([s_norm, vh_diag, vh_offdiag, s_ent], dim=-1)\n","            return torch.where(torch.isfinite(out), out, torch.zeros_like(out))\n","\n","        def compute_novelty(self, S):\n","            return S - self.ema_s.clone().unsqueeze(0)\n","\n","        def forward(self, x):\n","            B, C, H, W = x.shape\n","            h = self.to_svd(x)\n","            h_flat = h.permute(0, 2, 3, 1).reshape(B, H * W, self.svd_rank)\n","            with torch.amp.autocast('cuda', enabled=False):\n","                with torch.no_grad():\n","                    gram = torch.bmm(h_flat.float().transpose(1, 2), h_flat.float())\n","                    evals, evecs = torch.linalg.eigh(gram)\n","                    evals = evals.flip(-1).clamp(min=1e-12)\n","                    S = evals.sqrt()\n","                    Vh = evecs.flip(-1).transpose(-2, -1)\n","                    S = torch.where(torch.isfinite(S), S, torch.ones_like(S))\n","                    Vh = torch.where(torch.isfinite(Vh), Vh, torch.zeros_like(Vh))\n","            features = self.extract_features(S, Vh)\n","            novelty = self.compute_novelty(S)\n","            return S, Vh, features, novelty\n","\n","        @torch.no_grad()\n","        def update_ema(self, S, Vh):\n","            m = self.ema_momentum\n","            self.ema_s.mul_(m).add_(S.detach().mean(0), alpha=1-m)\n","            self.ema_vh_flat.mul_(m).add_(Vh.detach().mean(0).reshape(-1), alpha=1-m)\n","\n","        @property\n","        def feature_dim(self):\n","            return 2 * self.svd_rank + 2\n","\n","\n","class ConvSVDPatchEmbedding(TorchComponent):\n","    \"\"\"Input stage: conv frontend → SVDObserver → patch tokens.\"\"\"\n","    def __init__(self, name, img_size=32, patch_size=4, in_channels=3,\n","                 conv_channels=64, d_model=256, svd_rank=16):\n","        super().__init__(name)\n","        self.patch_size = patch_size\n","        self.n_patches = (img_size // patch_size) ** 2\n","        self.d_model = d_model\n","        self.svd_rank = svd_rank\n","\n","        self.conv_frontend = nn.Sequential(\n","            nn.Conv2d(in_channels, conv_channels, 3, padding=1, bias=False),\n","            nn.BatchNorm2d(conv_channels), nn.GELU(),\n","            nn.Conv2d(conv_channels, conv_channels, 3, padding=1, bias=False),\n","            nn.BatchNorm2d(conv_channels), nn.GELU(),\n","        )\n","        self.svd_observer = SVDObserver(conv_channels, svd_rank)\n","        self.patch_proj = nn.Conv2d(\n","            conv_channels, d_model, kernel_size=patch_size,\n","            stride=patch_size, bias=False)\n","        self.patch_norm = nn.LayerNorm(d_model)\n","\n","        svd_feat_dim = self.svd_observer.feature_dim\n","        self.svd_to_gamma = nn.Linear(svd_feat_dim, d_model)\n","        self.svd_to_beta = nn.Linear(svd_feat_dim, d_model)\n","        nn.init.zeros_(self.svd_to_gamma.weight); nn.init.ones_(self.svd_to_gamma.bias)\n","        nn.init.zeros_(self.svd_to_beta.weight); nn.init.zeros_(self.svd_to_beta.bias)\n","\n","        self.cls_token = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)\n","        self.pos_embed = nn.Parameter(\n","            torch.randn(1, self.n_patches + 1, d_model) * 0.02)\n","\n","    def forward(self, x):\n","        B = x.shape[0]\n","        feat = self.conv_frontend(x)\n","        S, Vh, svd_features, novelty = self.svd_observer(feat)\n","        tokens = self.patch_proj(feat)\n","        tokens = tokens.flatten(2).transpose(1, 2)\n","        tokens = self.patch_norm(tokens)\n","        gamma = self.svd_to_gamma(svd_features).unsqueeze(1)\n","        beta = self.svd_to_beta(svd_features).unsqueeze(1)\n","        tokens = gamma * tokens + beta\n","        cls = self.cls_token.expand(B, -1, -1)\n","        tokens = torch.cat([cls, tokens], dim=1)\n","        tokens = tokens + self.pos_embed\n","        svd_state = {\n","            'singular_values': S, 'Vh': Vh,\n","            'svd_features': svd_features, 'novelty': novelty,\n","        }\n","        if self.training:\n","            self.svd_observer.update_ema(S, Vh)\n","        return tokens, svd_state\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# CLASSIFIER ( uses GeometricTransformer with CM gates)\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","class GeoViTClassifier(BaseTower):\n","    \"\"\"Geometric Vision Transformer for classification.\n","\n","    Wraps ConvSVDPatchEmbedding + GeometricTransformer + task head.\n","    Exposes geometric_losses() for regularization during training.\n","    \"\"\"\n","    def __init__(self, name, config):\n","        super().__init__(name)\n","        self.config = config\n","\n","        self.attach('patch_embed', ConvSVDPatchEmbedding(\n","            'patch_embed', img_size=config['img_size'],\n","            patch_size=config['patch_size'], in_channels=config['in_channels'],\n","            conv_channels=config['conv_channels'], d_model=config['d_model'],\n","            svd_rank=config['svd_rank'],\n","        ))\n","        self.attach('transformer', GeometricTransformer(\n","            'geo_cifar', d_model=config['d_model'], n_heads=config['n_heads'],\n","            n_layers=config['n_layers'], n_anchors=config['n_anchors'],\n","            manifold_dim=config['manifold_dim'], n_comp=config['n_comp'],\n","            d_comp=config['d_comp'], context_dim=config['context_dim'],\n","            quat_dim=config['quat_dim'], dropout=config['dropout'],\n","            cm_neighbors=config.get('cm_neighbors', 3),\n","        ))\n","        self.attach('head', nn.Sequential(\n","            nn.LayerNorm(config['d_model']),\n","            nn.Linear(config['d_model'], config['d_model']),\n","            nn.GELU(), nn.Dropout(config['dropout']),\n","            nn.Linear(config['d_model'], config['num_classes']),\n","        ))\n","\n","    def forward(self, x, return_geo_state=False):\n","        tokens, svd_state = self['patch_embed'](x)\n","        if return_geo_state:\n","            features, geo_states = self['transformer'](tokens, return_geo_state=True)\n","        else:\n","            features = self['transformer'](tokens)\n","        cls_out = features[:, 0]\n","        logits = self['head'](cls_out)\n","        if return_geo_state:\n","            return logits, geo_states, svd_state\n","        return logits\n","\n","    def geometric_losses(self):\n","        \"\"\"Delegate to transformer's built-in geometric regularization.\"\"\"\n","        return self['transformer'].geometric_losses(\n","            cv_target=self.config.get('cv_target', 0.215),\n","            cv_weight=self.config.get('cv_weight', 0.1),\n","            spread_weight=self.config.get('spread_weight', 0.01),\n","        )\n","\n","    def anchor_diagnostics(self):\n","        \"\"\"Delegate to transformer's anchor diagnostics.\"\"\"\n","        return self['transformer'].anchor_diagnostics()\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# GEOMETRIC ANALYSIS BATTERY ( includes CM diagnostics)\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","@torch.no_grad()\n","def compute_cv(points):\n","    \"\"\"Coefficient of variation on S^(d-1).\n","    CV = std(pairwise_cosine_distances) / mean(pairwise_cosine_distances)\n","    Pentachoron band: CV ∈ [0.20, 0.23].\n","    \"\"\"\n","    points = F.normalize(points.float(), dim=-1)\n","    cos_sim = points @ points.T\n","    n = points.shape[0]\n","    idx = torch.triu_indices(n, n, offset=1, device=points.device)\n","    pairwise_dist = 1.0 - cos_sim[idx[0], idx[1]]\n","    mean_d = pairwise_dist.mean()\n","    std_d = pairwise_dist.std()\n","    cv = (std_d / (mean_d + 1e-8)).item()\n","    return cv, mean_d.item(), std_d.item()\n","\n","\n","@torch.no_grad()\n","def log_geometric_analysis(model, writer, epoch, test_loader, device, config):\n","    \"\"\"Full geometric analysis battery with CM diagnostics.\"\"\"\n","    model.eval()\n","\n","    images, labels = next(iter(test_loader))\n","    images = images[:min(64, images.shape[0])].to(device)\n","    labels = labels[:min(64, labels.shape[0])].to(device)\n","\n","    logits, geo_states, svd_state = model(images, return_geo_state=True)\n","\n","    n_layers = len(geo_states)\n","    pred = logits.argmax(1)\n","    batch_acc = (pred == labels).float().mean().item()\n","    writer.add_scalar('analysis/batch_accuracy', batch_acc, epoch)\n","\n","    # ─── SVD Input Stage ───\n","    S = svd_state['singular_values']\n","    s_norm = S / (S.sum(dim=-1, keepdim=True) + 1e-8)\n","    s_ent = -(s_norm * torch.log(s_norm.clamp(min=1e-8))).sum(-1)\n","    novelty = svd_state['novelty']\n","\n","    writer.add_scalar('svd/entropy_mean', s_ent.mean().item(), epoch)\n","    writer.add_scalar('svd/entropy_std', s_ent.std().item(), epoch)\n","    writer.add_scalar('svd/novelty_norm', novelty.norm(dim=-1).mean().item(), epoch)\n","    writer.add_scalar('svd/top1_ratio', (S[:, 0] / (S.sum(-1) + 1e-8)).mean().item(), epoch)\n","    writer.add_scalar('svd/condition_number',\n","        (S[:, 0] / (S[:, -1].clamp(min=1e-8))).mean().item(), epoch)\n","    for k in range(min(S.shape[1], 5)):\n","        writer.add_scalar(f'svd/S_{k}', S[:, k].mean().item(), epoch)\n","\n","    # SVD FiLM deviation\n","    pe = model['patch_embed']\n","    writer.add_scalar('svd_film/gamma_weight_norm', pe.svd_to_gamma.weight.data.norm().item(), epoch)\n","    writer.add_scalar('svd_film/gamma_bias_dev_from_1',\n","        (pe.svd_to_gamma.bias.data - 1.0).abs().mean().item(), epoch)\n","    writer.add_scalar('svd_film/beta_weight_norm', pe.svd_to_beta.weight.data.norm().item(), epoch)\n","    writer.add_scalar('svd_film/beta_bias_norm', pe.svd_to_beta.bias.data.abs().mean().item(), epoch)\n","\n","    # ─── Anchor Diagnostics (built-in) ───\n","    anchor_diag = model.anchor_diagnostics()\n","    for layer_name, d in anchor_diag.items():\n","        for k, v in d.items():\n","            writer.add_scalar(f'anchor_diag/{layer_name}_{k}', v, epoch)\n","\n","    # ─── Per-Layer Geometric Analysis ───\n","    for i, gs in enumerate(geo_states):\n","        prefix = f'layer_{i}'\n","\n","        # === CV — pentachoron band metric ===\n","        emb = gs['embedding']\n","        # Anchor CV\n","        transformer = model['transformer']\n","        layer = transformer[f'layer_{i}']\n","        anchors = F.normalize(\n","            layer['observer'].association.constellation.anchors, dim=-1)\n","        cv_anchors, mean_d_anchors, std_d_anchors = compute_cv(anchors)\n","        writer.add_scalar(f'{prefix}/cv_anchors', cv_anchors, epoch)\n","        writer.add_scalar(f'{prefix}/anchor_mean_dist', mean_d_anchors, epoch)\n","        writer.add_scalar(f'{prefix}/anchor_std_dist', std_d_anchors, epoch)\n","\n","        # Embedding CV\n","        emb_flat = emb.reshape(-1, emb.shape[-1])\n","        n_sample = min(512, emb_flat.shape[0])\n","        idx = torch.randperm(emb_flat.shape[0], device=device)[:n_sample]\n","        cv_emb, mean_d_emb, std_d_emb = compute_cv(emb_flat[idx])\n","        writer.add_scalar(f'{prefix}/cv_embeddings', cv_emb, epoch)\n","        writer.add_scalar(f'{prefix}/embedding_mean_dist', mean_d_emb, epoch)\n","\n","        # === CM Gate Diagnostics ===\n","        gate_info = gs.get('gate_info', {})\n","        gate_values = gs.get('gate_values')\n","        cm_quality = gs.get('cm_quality')\n","\n","        if gate_info:\n","            writer.add_scalar(f'{prefix}/cm_active_anchors',\n","                gate_info.get('active', 0), epoch)\n","            writer.add_scalar(f'{prefix}/cm_gate_mean',\n","                gate_info.get('gate_mean', 0), epoch)\n","            writer.add_scalar(f'{prefix}/cm_positive_frac',\n","                gate_info.get('cm_positive_frac', 0), epoch)\n","\n","        if gate_values is not None:\n","            gv = gate_values\n","            writer.add_scalar(f'{prefix}/gate_values_min', gv.min().item(), epoch)\n","            writer.add_scalar(f'{prefix}/gate_values_max', gv.max().item(), epoch)\n","            writer.add_scalar(f'{prefix}/gate_values_std', gv.std().item(), epoch)\n","            # Per-anchor gate mean (which anchors are consistently open/closed)\n","            gv_per_anchor = gv.mean(dim=0).mean(dim=0)  # average over B and L\n","            writer.add_scalar(f'{prefix}/gate_anchor_spread',\n","                gv_per_anchor.std().item(), epoch)\n","            # Fraction of positions with >50% anchors open\n","            if gv.dim() == 3:\n","                pos_open_frac = (gv.mean(dim=-1) > 0.5).float().mean().item()\n","            else:\n","                pos_open_frac = (gv > 0.5).float().mean().item()\n","            writer.add_scalar(f'{prefix}/gate_positions_open_frac', pos_open_frac, epoch)\n","\n","        if cm_quality is not None:\n","            writer.add_scalar(f'{prefix}/cm_quality_mean', cm_quality.mean().item(), epoch)\n","            writer.add_scalar(f'{prefix}/cm_quality_std', cm_quality.std().item(), epoch)\n","            writer.add_scalar(f'{prefix}/cm_quality_min', cm_quality.min().item(), epoch)\n","\n","        # === Stream Agreement ===\n","        content = gs['content']\n","        geometric = gs['geometric']\n","        agreement = F.cosine_similarity(\n","            content.reshape(-1, content.shape[-1]),\n","            geometric.reshape(-1, geometric.shape[-1]), dim=-1)\n","        writer.add_scalar(f'{prefix}/stream_agreement_mean', agreement.mean().item(), epoch)\n","        writer.add_scalar(f'{prefix}/stream_agreement_std', agreement.std().item(), epoch)\n","\n","        writer.add_scalar(f'{prefix}/content_norm', content.norm(dim=-1).mean().item(), epoch)\n","        writer.add_scalar(f'{prefix}/geometric_norm', geometric.norm(dim=-1).mean().item(), epoch)\n","\n","        # === Disagreement arm analysis ===\n","        disagree = content - geometric\n","        agree = content * geometric\n","        writer.add_scalar(f'{prefix}/disagree_norm', disagree.norm(dim=-1).mean().item(), epoch)\n","        writer.add_scalar(f'{prefix}/agree_norm', agree.norm(dim=-1).mean().item(), epoch)\n","\n","        # === Anchor Utilization ===\n","        tri = gs['triangulation']\n","        assignment = gs['assignment']\n","        nearest = gs['nearest']\n","        n_anchors = tri.shape[-1]\n","\n","        nearest_flat = nearest.reshape(-1)\n","        counts = torch.bincount(nearest_flat, minlength=n_anchors).float()\n","        total_assignments = counts.sum()\n","\n","        probs = counts / (total_assignments + 1e-8)\n","        anchor_entropy = -(probs * torch.log(probs.clamp(min=1e-8))).sum().item()\n","        max_entropy = math.log(n_anchors)\n","        writer.add_scalar(f'{prefix}/anchor_entropy_normalized',\n","            anchor_entropy / (max_entropy + 1e-8), epoch)\n","        active = (counts > 0).sum().item()\n","        writer.add_scalar(f'{prefix}/anchors_active', active, epoch)\n","        writer.add_scalar(f'{prefix}/anchors_active_frac', active / n_anchors, epoch)\n","        dead = (counts == 0).sum().item()\n","        writer.add_scalar(f'{prefix}/anchors_dead', dead, epoch)\n","\n","        # === Triangulation Statistics ===\n","        writer.add_scalar(f'{prefix}/tri_mean', tri.mean().item(), epoch)\n","        writer.add_scalar(f'{prefix}/tri_std', tri.std().item(), epoch)\n","\n","        # === Soft Assignment Statistics ===\n","        assign_ent = -(assignment * torch.log(assignment.clamp(min=1e-8))).sum(-1)\n","        writer.add_scalar(f'{prefix}/assignment_entropy_mean', assign_ent.mean().item(), epoch)\n","        writer.add_scalar(f'{prefix}/assignment_max_prob',\n","            assignment.max(dim=-1).values.mean().item(), epoch)\n","\n","        # === Patchwork Statistics (now from CM-validated triangulation) ===\n","        pw = gs['patchwork']\n","        writer.add_scalar(f'{prefix}/patchwork_norm', pw.norm(dim=-1).mean().item(), epoch)\n","        writer.add_scalar(f'{prefix}/patchwork_std', pw.std().item(), epoch)\n","        pw_sparsity = (pw.abs() < 0.01).float().mean().item()\n","        writer.add_scalar(f'{prefix}/patchwork_sparsity', pw_sparsity, epoch)\n","\n","        # === Bridge Consistency ===\n","        bridge = gs['bridge']\n","        bridge_soft = F.softmax(bridge, dim=-1)\n","        bridge_assign_kl = F.kl_div(\n","            bridge_soft.log().reshape(-1, n_anchors),\n","            assignment.reshape(-1, n_anchors),\n","            reduction='batchmean', log_target=False)\n","        writer.add_scalar(f'{prefix}/bridge_assignment_kl', bridge_assign_kl.item(), epoch)\n","\n","        # === Quaternion Composition ===\n","        composed = gs['composed']\n","        writer.add_scalar(f'{prefix}/composed_norm', composed.norm(dim=-1).mean().item(), epoch)\n","\n","        # === Geo Context ===\n","        geo_ctx = gs['geo_ctx']\n","        writer.add_scalar(f'{prefix}/geo_ctx_norm', geo_ctx.norm(dim=-1).mean().item(), epoch)\n","\n","        # === Geometric Residual Stream (CM-conditioned) ===\n","        geo_res = gs.get('geo_residual')\n","        if geo_res is not None:\n","            res_norms = geo_res.norm(dim=-1)\n","            writer.add_scalar(f'{prefix}/geo_res_norm', res_norms.mean().item(), epoch)\n","            writer.add_scalar(f'{prefix}/geo_res_std', geo_res.std().item(), epoch)\n","            writer.add_scalar(f'{prefix}/geo_res_sparsity',\n","                (geo_res.abs() < 0.01).float().mean().item(), epoch)\n","            # Cross-position consistency\n","            geo_res_flat = geo_res.reshape(-1, geo_res.shape[-1])\n","            n_s = min(256, geo_res_flat.shape[0])\n","            idx_s = torch.randperm(geo_res_flat.shape[0], device=geo_res.device)[:n_s]\n","            sampled = F.normalize(geo_res_flat[idx_s], dim=-1)\n","            cos_mat = sampled @ sampled.T\n","            triu = torch.triu_indices(n_s, n_s, offset=1, device=geo_res.device)\n","            writer.add_scalar(f'{prefix}/geo_res_consistency',\n","                cos_mat[triu[0], triu[1]].mean().item(), epoch)\n","\n","    # ─── Cayley Rotation Analysis ───\n","    for name, mod in model.named_modules():\n","        if isinstance(mod, CayleyOrthogonal):\n","            R = mod.get_rotation()\n","            I = torch.eye(R.shape[0], device=R.device)\n","            r_dist = (R - I).norm().item()\n","            clean_name = name.replace('.', '_')\n","            writer.add_scalar(f'cayley/{clean_name}_R_minus_I', r_dist, epoch)\n","\n","    # ─── FiLM Layer Analysis ───\n","    film_idx = 0\n","    for name, mod in model.named_modules():\n","        if isinstance(mod, FiLMLayer):\n","            g_b = mod.to_gamma.bias.data\n","            b_b = mod.to_beta.bias.data\n","            writer.add_scalar(f'film/{film_idx}_gamma_dev',\n","                (g_b - 1.0).abs().mean().item(), epoch)\n","            writer.add_scalar(f'film/{film_idx}_beta_dev',\n","                b_b.abs().mean().item(), epoch)\n","            film_idx += 1\n","\n","    # ─── Cross-Layer Trajectories ───\n","    cv_trajectory = []\n","    cm_quality_trajectory = []\n","    res_norms = []\n","    bridge_kls = []\n","\n","    for i, gs in enumerate(geo_states):\n","        # CV\n","        emb = gs['embedding']\n","        emb_flat = emb.reshape(-1, emb.shape[-1])\n","        n_sample = min(512, emb_flat.shape[0])\n","        idx = torch.randperm(emb_flat.shape[0], device=device)[:n_sample]\n","        cv, _, _ = compute_cv(emb_flat[idx])\n","        cv_trajectory.append(cv)\n","\n","        # CM quality\n","        cm_q = gs.get('cm_quality')\n","        if cm_q is not None:\n","            cm_quality_trajectory.append(cm_q.mean().item())\n","\n","        # Geo residual norms\n","        geo_res = gs.get('geo_residual')\n","        if geo_res is not None:\n","            res_norms.append(geo_res.norm(dim=-1).mean().item())\n","\n","        # Bridge KL\n","        n_anchors = gs['assignment'].shape[-1]\n","        bridge_soft = F.softmax(gs['bridge'], dim=-1)\n","        bkl = F.kl_div(\n","            bridge_soft.log().reshape(-1, n_anchors),\n","            gs['assignment'].reshape(-1, n_anchors),\n","            reduction='batchmean', log_target=False).item()\n","        bridge_kls.append(bkl)\n","\n","    # CV trajectory\n","    writer.add_scalar('cv/trajectory_mean', np.mean(cv_trajectory), epoch)\n","    writer.add_scalar('cv/trajectory_std', np.std(cv_trajectory), epoch)\n","    in_band = sum(1 for cv in cv_trajectory if 0.20 <= cv <= 0.23)\n","    writer.add_scalar('cv/layers_in_pentachoron_band', in_band, epoch)\n","    writer.add_scalar('cv/layers_in_band_frac', in_band / len(cv_trajectory), epoch)\n","\n","    # CM quality trajectory\n","    if cm_quality_trajectory:\n","        writer.add_scalar('cm/quality_trajectory_mean',\n","            np.mean(cm_quality_trajectory), epoch)\n","        writer.add_scalar('cm/quality_trajectory_std',\n","            np.std(cm_quality_trajectory), epoch)\n","        writer.add_scalar('cm/quality_min_layer',\n","            np.min(cm_quality_trajectory), epoch)\n","        writer.add_scalar('cm/quality_max_layer',\n","            np.max(cm_quality_trajectory), epoch)\n","\n","    # Geometric residual trajectory\n","    if res_norms:\n","        writer.add_scalar('geo_res/trajectory_start', res_norms[0], epoch)\n","        writer.add_scalar('geo_res/trajectory_end', res_norms[-1], epoch)\n","        writer.add_scalar('geo_res/accumulation_ratio',\n","            res_norms[-1] / (res_norms[0] + 1e-8), epoch)\n","        growth = [res_norms[j+1] - res_norms[j] for j in range(len(res_norms)-1)]\n","        writer.add_scalar('geo_res/growth_mean', np.mean(growth), epoch)\n","        writer.add_scalar('geo_res/growth_std', np.std(growth), epoch)\n","\n","        # Cooperation analysis (includes CM quality)\n","        if len(res_norms) >= 4:\n","            cv_corr = float(np.corrcoef(res_norms, cv_trajectory)[0, 1])\n","            bkl_corr = float(np.corrcoef(res_norms, bridge_kls)[0, 1])\n","            writer.add_scalar('cooperation/geo_res_vs_cv', cv_corr, epoch)\n","            writer.add_scalar('cooperation/geo_res_vs_bridge_kl', bkl_corr, epoch)\n","\n","            if len(cm_quality_trajectory) == len(res_norms):\n","                cm_corr = float(np.corrcoef(\n","                    res_norms, cm_quality_trajectory)[0, 1])\n","                writer.add_scalar('cooperation/geo_res_vs_cm_quality', cm_corr, epoch)\n","                # CM vs CV: do layers with better CM quality also have better CV?\n","                cm_cv_corr = float(np.corrcoef(\n","                    cm_quality_trajectory, cv_trajectory)[0, 1])\n","                writer.add_scalar('cooperation/cm_quality_vs_cv', cm_cv_corr, epoch)\n","\n","    return {\n","        'batch_acc': batch_acc,\n","        'cv_trajectory': cv_trajectory,\n","        'cm_quality_trajectory': cm_quality_trajectory,\n","        'res_norms': res_norms,\n","        'bridge_kls': bridge_kls,\n","    }\n","\n","\n","@torch.no_grad()\n","def log_gradient_norms(model, writer, epoch):\n","    \"\"\"Log gradient norms per component type (includes cm_gate).\"\"\"\n","    type_grads = {}\n","    for name, param in model.named_parameters():\n","        if param.grad is not None:\n","            grad_norm = param.grad.norm().item()\n","            if 'projection' in name and 'proj' in name:\n","                key = 'manifold_proj'\n","            elif 'cm_gate' in name:\n","                key = 'cm_gate'\n","            elif 'observer' in name or 'constellation' in name or 'anchor' in name:\n","                key = 'constellation'\n","            elif 'context' in name:\n","                key = 'geo_context'\n","            elif 'content' in name:\n","                key = 'content_attn'\n","            elif 'geometric' in name and 'film' not in name:\n","                key = 'geo_attn'\n","            elif 'film' in name:\n","                key = 'film'\n","            elif 'rotation' in name or 'cayley' in name or 'A_upper' in name:\n","                key = 'cayley'\n","            elif 'compose' in name or 'quat' in name or 'proj_w' in name:\n","                key = 'quaternion'\n","            elif 'decode' in name:\n","                key = 'decode'\n","            elif 'gate' in name:\n","                key = 'gate'\n","            elif 'conv' in name or 'patch' in name:\n","                key = 'input_stage'\n","            elif 'head' in name:\n","                key = 'head'\n","            elif 'svd' in name:\n","                key = 'svd'\n","            elif 'geo_proj' in name:\n","                key = 'geo_residual_proj'\n","            else:\n","                key = 'other'\n","\n","            if key not in type_grads:\n","                type_grads[key] = []\n","            type_grads[key].append(grad_norm)\n","\n","    for key, norms in type_grads.items():\n","        writer.add_scalar(f'grad_norm/{key}_mean', np.mean(norms), epoch)\n","        writer.add_scalar(f'grad_norm/{key}_max', np.max(norms), epoch)\n","\n","    total = sum(p.grad.norm().item() ** 2\n","                for p in model.parameters() if p.grad is not None) ** 0.5\n","    writer.add_scalar('grad_norm/total', total, epoch)\n","\n","\n","@torch.no_grad()\n","def log_weight_norms(model, writer, epoch):\n","    \"\"\"Log weight norms per component type.\"\"\"\n","    for name, param in model.named_parameters():\n","        if 'A_upper' in name:\n","            clean = name.replace('.', '_')\n","            writer.add_scalar(f'weights/{clean}_norm', param.norm().item(), epoch)\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# DATA\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","def get_dataloaders(config):\n","    import torchvision\n","    import torchvision.transforms as T\n","\n","    # Augmentation pipeline tuned for geometric transformer:\n","    #   TrivialAugmentWide: continuous severity spectrum of geometric + photometric\n","    #     transforms. Exercises CM gate across full quality range — mild distortion\n","    #     keeps CM high, severe distortion creates partially-degenerate simplices.\n","    #   RandomErasing: creates degenerate manifold projections (zero-volume CM simplices).\n","    #     Trains CM gate to close on corrupted regions.\n","    #   CutMix applied at batch level in train_epoch (not here).\n","    train_transform = T.Compose([\n","        T.RandomCrop(32, padding=4),\n","        T.RandomHorizontalFlip(),\n","        T.TrivialAugmentWide(),\n","        T.ToTensor(),\n","        T.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),\n","        T.RandomErasing(p=config.get('random_erasing_p', 0.25),\n","                        scale=(0.02, 0.33), ratio=(0.3, 3.3)),\n","    ])\n","    test_transform = T.Compose([\n","        T.ToTensor(),\n","        T.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),\n","    ])\n","\n","    train_ds = torchvision.datasets.CIFAR100(\n","        root='./data', train=True, download=True, transform=train_transform)\n","    test_ds = torchvision.datasets.CIFAR100(\n","        root='./data', train=False, download=True, transform=test_transform)\n","\n","    train_loader = torch.utils.data.DataLoader(\n","        train_ds, batch_size=config['batch_size'], shuffle=True,\n","        num_workers=config['num_workers'], pin_memory=True, drop_last=True)\n","    test_loader = torch.utils.data.DataLoader(\n","        test_ds, batch_size=config['batch_size'], shuffle=False,\n","        num_workers=config['num_workers'], pin_memory=True)\n","\n","    return train_loader, test_loader\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# CUTMIX — batch-level augmentation for CM gate boundary training\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","def cutmix_batch(images, labels, alpha=1.0):\n","    \"\"\"Apply CutMix to a batch. Returns mixed images + label pairs + lambda.\n","\n","    CutMix replaces a rectangular region of image A with image B.\n","    Positions inside each region have coherent geometry — valid CM simplices.\n","    The boundary between regions has mixed geometric context — the CM gate\n","    should learn to suppress these positions.\n","\n","    Args:\n","        images: (B, C, H, W) batch\n","        labels: (B,) integer labels\n","        alpha: Beta distribution parameter (1.0 = uniform box sizes)\n","\n","    Returns:\n","        images: (B, C, H, W) mixed batch (modified in-place)\n","        labels_a: (B,) labels for region A\n","        labels_b: (B,) labels for region B\n","        lam: float, fraction of image A remaining\n","    \"\"\"\n","    lam = np.random.beta(alpha, alpha)\n","    B = images.size(0)\n","    idx = torch.randperm(B, device=images.device)\n","\n","    H, W = images.shape[2], images.shape[3]\n","    cut_ratio = (1.0 - lam) ** 0.5\n","    cw = int(W * cut_ratio)\n","    ch = int(H * cut_ratio)\n","    cx = np.random.randint(W)\n","    cy = np.random.randint(H)\n","    x1 = max(cx - cw // 2, 0); x2 = min(cx + cw // 2, W)\n","    y1 = max(cy - ch // 2, 0); y2 = min(cy + ch // 2, H)\n","\n","    images[:, :, y1:y2, x1:x2] = images[idx, :, y1:y2, x1:x2]\n","    lam_actual = 1.0 - (x2 - x1) * (y2 - y1) / (W * H)\n","    return images, labels, labels[idx], lam_actual\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# TRAINING (geometric losses + CutMix integrated)\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","def train_epoch(model, loader, optimizer, scheduler, epoch, config, writer):\n","    model.train()\n","    total_loss = 0\n","    total_geo_loss = 0\n","    correct = 0\n","    total = 0\n","\n","    cutmix_alpha = config.get('cutmix_alpha', 1.0)\n","    cutmix_prob = config.get('cutmix_prob', 0.5)\n","    label_smoothing = config.get('label_smoothing', 0.1)\n","    criterion = nn.CrossEntropyLoss(label_smoothing=label_smoothing)\n","\n","    for batch_idx, (images, labels) in enumerate(loader):\n","        images = images.to(device)\n","        labels = labels.to(device)\n","\n","        # CutMix: applied probabilistically per batch\n","        use_cutmix = np.random.rand() < cutmix_prob\n","        if use_cutmix:\n","            images, labels_a, labels_b, lam = cutmix_batch(\n","                images, labels, alpha=cutmix_alpha)\n","            logits = model(images)\n","            ce_loss = lam * criterion(logits, labels_a) + \\\n","                      (1.0 - lam) * criterion(logits, labels_b)\n","            # Accuracy: count correct if matches either label\n","            pred = logits.argmax(1)\n","            correct += (lam * (pred == labels_a).float() +\n","                        (1.0 - lam) * (pred == labels_b).float()).sum().item()\n","        else:\n","            logits = model(images)\n","            ce_loss = criterion(logits, labels)\n","            correct += (logits.argmax(1) == labels).sum().item()\n","\n","        # Geometric regularization — CV target + anchor spread\n","        geo_losses = model.geometric_losses()\n","        geo_loss = geo_losses.get('geo_total', torch.tensor(0.0, device=device))\n","        loss = ce_loss + geo_loss\n","\n","        optimizer.zero_grad()\n","        loss.backward()\n","\n","        # Log gradient norms periodically\n","        if epoch % config['log_grads_every'] == 0 and batch_idx == 0:\n","            log_gradient_norms(model, writer, epoch)\n","\n","        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n","        optimizer.step()\n","        if scheduler is not None:\n","            scheduler.step()\n","\n","        total_loss += ce_loss.item() * images.size(0)\n","        total_geo_loss += geo_loss.item() * images.size(0)\n","        total += images.size(0)\n","\n","    avg_ce = total_loss / total\n","    avg_geo = total_geo_loss / total\n","    return avg_ce, avg_geo, correct / total\n","\n","\n","@torch.no_grad()\n","def evaluate(model, loader):\n","    model.eval()\n","    correct = 0\n","    total = 0\n","    for images, labels in loader:\n","        images = images.to(device)\n","        labels = labels.to(device)\n","        logits = model(images)\n","        correct += (logits.argmax(1) == labels).sum().item()\n","        total += images.size(0)\n","    return correct / total\n","\n","\n","def main():\n","    config = CONFIG.copy()\n","\n","    print(\"=\" * 60)\n","    print(\"  Geometric Transformer — CIFAR-100 (CM-Validated)\")\n","    print(f\"  Input: conv({config['in_channels']}→{config['conv_channels']}) + \"\n","          f\"SVD(rank={config['svd_rank']}) + \"\n","          f\"{config['patch_size']}×{config['patch_size']} patches = \"\n","          f\"{(config['img_size']//config['patch_size'])**2} tokens + CLS\")\n","    print(f\"  Model: d={config['d_model']}, heads={config['n_heads']}, \"\n","          f\"layers={config['n_layers']}, anchors={config['n_anchors']}\")\n","    print(f\"  CM: neighbors={config['cm_neighbors']}, \"\n","          f\"cv_target={config['cv_target']}, \"\n","          f\"cv_weight={config['cv_weight']}, \"\n","          f\"spread_weight={config['spread_weight']}\")\n","    print(f\"  Aug: TrivialAugmentWide + CutMix(α={config['cutmix_alpha']}, \"\n","          f\"p={config['cutmix_prob']}) + \"\n","          f\"RandomErasing(p={config['random_erasing_p']})\")\n","    print(\"=\" * 60)\n","\n","    writer = SummaryWriter(config['log_dir'])\n","    writer.add_text('config', json.dumps(config, indent=2))\n","\n","    print(\"\\nLoading CIFAR-100...\")\n","    train_loader, test_loader = get_dataloaders(config)\n","    print(f\"  Train: {len(train_loader.dataset):,} | Test: {len(test_loader.dataset):,}\")\n","\n","    model = GeoViTClassifier('geo_vit_cifar100', config)\n","    if hasattr(model, 'network_to'):\n","        model.network_to(device=device, strict=False)\n","    else:\n","        model = model.to(device)\n","\n","    n_params = sum(p.numel() for p in model.parameters())\n","    n_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)\n","    print(f\"\\n  Total params:     {n_params:,}\")\n","    print(f\"  Trainable params: {n_trainable:,}\")\n","\n","    model.compile()\n","\n","    for name, module in model.named_children():\n","        n = sum(p.numel() for p in module.parameters())\n","        if n > 0:\n","            print(f\"    {name:<20s}: {n:,}\")\n","\n","    writer.add_scalar('model/total_params', n_params, 0)\n","\n","    # Initial anchor diagnostics\n","    print(f\"\\n  Initial anchor diagnostics:\")\n","    diag = model.anchor_diagnostics()\n","    for layer_name, d in diag.items():\n","        print(f\"    {layer_name}: cv={d['anchor_cv']:.4f}, \"\n","              f\"cm_pos={d['cm_positive_frac']:.3f}, \"\n","              f\"min_dist={d['min_pairwise_dist']:.4f}\")\n","\n","    # Optimizer + scheduler\n","    optimizer = torch.optim.AdamW(\n","        model.parameters(), lr=config['lr'], weight_decay=config['weight_decay'])\n","\n","    total_steps = config['epochs'] * len(train_loader)\n","    warmup_steps = config['warmup_epochs'] * len(train_loader)\n","\n","    def lr_lambda(step):\n","        if step < warmup_steps:\n","            return step / warmup_steps\n","        progress = (step - warmup_steps) / (total_steps - warmup_steps)\n","        return 0.5 * (1 + np.cos(np.pi * progress))\n","\n","    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)\n","\n","    print(f\"\\n{'━'*60}\")\n","    print(f\"  Training for {config['epochs']} epochs\")\n","    print(f\"  Warmup: {config['warmup_epochs']} epochs, \"\n","          f\"LR: {config['lr']}, WD: {config['weight_decay']}\")\n","    print(f\"  Geo reg: cv_w={config['cv_weight']}, spread_w={config['spread_weight']}\")\n","    print(f\"  Aug: TrivialAugmentWide + CutMix(p={config['cutmix_prob']}) + \"\n","          f\"RandomErasing(p={config['random_erasing_p']})\")\n","    print(f\"  TensorBoard: {config['log_dir']}\")\n","    print(f\"  Geo analysis every {config['log_geo_every']} epochs\")\n","    print(f\"{'━'*60}\\n\")\n","\n","    best_acc = 0\n","    save_dir = Path('geo_cifar100'); save_dir.mkdir(exist_ok=True)\n","\n","    for epoch in tqdm(range(config['epochs']), desc=\"Epochs\"):\n","        t0 = time.time()\n","\n","        ce_loss, geo_loss, train_acc = train_epoch(\n","            model, train_loader, optimizer, scheduler, epoch, config, writer)\n","\n","        test_acc = evaluate(model, test_loader)\n","        elapsed = time.time() - t0\n","\n","        lr = optimizer.param_groups[0]['lr']\n","        writer.add_scalar('train/ce_loss', ce_loss, epoch)\n","        writer.add_scalar('train/geo_loss', geo_loss, epoch)\n","        writer.add_scalar('train/total_loss', ce_loss + geo_loss, epoch)\n","        writer.add_scalar('train/accuracy', train_acc, epoch)\n","        writer.add_scalar('test/accuracy', test_acc, epoch)\n","        writer.add_scalar('train/lr', lr, epoch)\n","        writer.add_scalar('train/epoch_time', elapsed, epoch)\n","        writer.add_scalar('gap/train_test', train_acc - test_acc, epoch)\n","\n","        log_weight_norms(model, writer, epoch)\n","\n","        if test_acc > best_acc:\n","            best_acc = test_acc\n","            torch.save({\n","                'state_dict': {k: v.cpu() for k, v in model.state_dict().items()},\n","                'epoch': epoch,\n","                'test_acc': test_acc,\n","                'config': config,\n","            }, save_dir / 'best.pt')\n","\n","        # Full geometric analysis periodically\n","        if epoch % config['log_geo_every'] == 0 or epoch == config['epochs'] - 1:\n","            geo_info = log_geometric_analysis(\n","                model, writer, epoch, test_loader, device, config)\n","\n","            cv_str = ', '.join(f'{cv:.3f}' for cv in geo_info['cv_trajectory'])\n","            cm_str = ', '.join(f'{q:.3f}' for q in geo_info.get('cm_quality_trajectory', []))\n","            res_str = ', '.join(f'{r:.3f}' for r in geo_info.get('res_norms', []))\n","            tqdm.write(\n","                f\"  E{epoch:>3d}  ce={ce_loss:.4f}  geo={geo_loss:.4f}  \"\n","                f\"train={train_acc:.4f}  test={test_acc:.4f}  \"\n","                f\"best={best_acc:.4f}  {elapsed:.1f}s\"\n","                f\"\\n        CV=[{cv_str}]\"\n","                f\"\\n        CM=[{cm_str}]\"\n","                f\"\\n        GR=[{res_str}]\")\n","        elif epoch % 5 == 0:\n","            tqdm.write(\n","                f\"  E{epoch:>3d}  ce={ce_loss:.4f}  geo={geo_loss:.4f}  \"\n","                f\"train={train_acc:.4f}  test={test_acc:.4f}  \"\n","                f\"best={best_acc:.4f}  {elapsed:.1f}s\")\n","\n","    # Final summary\n","    print(f\"\\n{'═'*60}\")\n","    print(f\"  CIFAR-100 RESULTS (CM-Validated)\")\n","    print(f\"{'═'*60}\")\n","    print(f\"  Best test accuracy: {best_acc:.4f} ({best_acc*100:.2f}%)\")\n","    print(f\"  Parameters: {n_params:,}\")\n","    print(f\"  Checkpoint: {save_dir}/best.pt\")\n","    print(f\"  TensorBoard: {config['log_dir']}\")\n","\n","    # Final geometric state + anchor diagnostics\n","    print(f\"\\n  Final geometric state:\")\n","    geo_info = log_geometric_analysis(\n","        model, writer, config['epochs'], test_loader, device, config)\n","\n","    print(f\"\\n  Final anchor diagnostics:\")\n","    diag = model.anchor_diagnostics()\n","    for layer_name, d in diag.items():\n","        print(f\"    {layer_name}: cv={d['anchor_cv']:.4f}, \"\n","              f\"cm_pos={d['cm_positive_frac']:.3f}, \"\n","              f\"cm_mean={d['cm_mean']:.4f}\")\n","\n","    writer.close()\n","    print(f\"\\nDone.\")\n","\n","\n","if __name__ == '__main__':\n","    main()"],"metadata":{"colab":{"base_uri":"https://localhost:8080/","referenced_widgets":["cba1d1d952cc48a99831fc3979d8ed0d","d1d1752b19484f5d82ee9caba836c53e","bfb14422f4e7444fa83296f998bb27f9","70e1b310d0f24624b224ea4585205e98","e204ac6670454441bb2bb313ae4c82c0","e1f6b47538a54273a2e8230ac642d413","d46bfa89a3324ad6b563e2a82ffd152d","5bc38b424a504d84b304ee1e90537191","606bc83980ab41b5b3d3ac6eaa931d09","1dd023b412e64d1293faa213af2a8631","5ff7475e37d44589b3fdfb342f9769c7"]},"id":"PE_JYHY_CZX3","executionInfo":{"status":"ok","timestamp":1774750497546,"user_tz":420,"elapsed":3047551,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"4c0ccf83-ac58-41dd-8d0f-d5c627c92d99"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["Device: cuda\n","  GPU: NVIDIA RTX PRO 6000 Blackwell Server Edition\n","  VRAM: 102.0 GB\n","  Imported from geolip_core (installed)\n","============================================================\n","  Geometric Transformer — CIFAR-100 (CM-Validated)\n","  Input: conv(3→64) + SVD(rank=16) + 4×4 patches = 64 tokens + CLS\n","  Model: d=256, heads=8, layers=8, anchors=512\n","  CM: neighbors=4, cv_target=0.14, cv_weight=0.1, spread_weight=0.01\n","  Aug: TrivialAugmentWide + CutMix(α=1.0, p=0.5) + RandomErasing(p=0.25)\n","============================================================\n","\n","Loading CIFAR-100...\n","  Train: 50,000 | Test: 10,000\n","\n","  Total params:     19,295,276\n","  Trainable params: 19,295,276\n","    components          : 19,295,276\n","\n","  Initial anchor diagnostics:\n","    layer_0: cv=0.0786, cm_pos=1.000, min_dist=0.8744\n","    layer_1: cv=0.0786, cm_pos=1.000, min_dist=0.8738\n","    layer_2: cv=0.0785, cm_pos=1.000, min_dist=0.8677\n","    layer_3: cv=0.0786, cm_pos=1.000, min_dist=0.8716\n","    layer_4: cv=0.0786, cm_pos=1.000, min_dist=0.8746\n","    layer_5: cv=0.0785, cm_pos=1.000, min_dist=0.8729\n","    layer_6: cv=0.0785, cm_pos=1.000, min_dist=0.8747\n","    layer_7: cv=0.0786, cm_pos=1.000, min_dist=0.8727\n","\n","━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\n","  Training for 100 epochs\n","  Warmup: 5 epochs, LR: 0.001, WD: 0.05\n","  Geo reg: cv_w=0.1, spread_w=0.01\n","  Aug: TrivialAugmentWide + CutMix(p=0.5) + RandomErasing(p=0.25)\n","  TensorBoard: runs/geo_cifar100\n","  Geo analysis every 5 epochs\n","━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\n","\n"]},{"output_type":"display_data","data":{"text/plain":["Epochs:   0%|          | 0/100 [00:00<?, ?it/s]"],"application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"cba1d1d952cc48a99831fc3979d8ed0d"}},"metadata":{}},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n","/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n","/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["  E  0  ce=4.5973  geo=0.0007  train=0.0156  test=0.0300  best=0.0300  117.9s\n","        CV=[0.315, 0.319, 0.363, 0.381, 0.386, 0.421, 0.462, 0.533]\n","        CM=[0.881, 0.881, 0.881, 0.881, 0.880, 0.881, 0.881, 0.881]\n","        GR=[7.048, 9.999, 11.276, 13.794, 16.291, 17.840, 18.846, 19.941]\n","  E  5  ce=4.1897  geo=0.0005  train=0.0761  test=0.1042  best=0.1042  29.3s\n","        CV=[0.492, 0.348, 0.349, 0.398, 0.381, 0.399, 0.449, 0.475]\n","        CM=[0.883, 0.874, 0.879, 0.878, 0.878, 0.878, 0.880, 0.877]\n","        GR=[7.052, 10.911, 13.970, 16.664, 19.854, 22.751, 25.318, 25.664]\n","  E 10  ce=3.9923  geo=0.0003  train=0.1176  test=0.1853  best=0.1853  29.3s\n","        CV=[0.472, 0.303, 0.276, 0.305, 0.333, 0.327, 0.544, 0.403]\n","        CM=[0.880, 0.869, 0.877, 0.875, 0.876, 0.875, 0.877, 0.873]\n","        GR=[6.967, 11.161, 14.754, 18.235, 22.352, 25.892, 28.145, 29.686]\n","  E 15  ce=3.8890  geo=0.0003  train=0.1434  test=0.2568  best=0.2568  29.9s\n","        CV=[0.717, 0.291, 0.250, 0.275, 0.299, 0.277, 0.579, 0.361]\n","        CM=[0.878, 0.868, 0.874, 0.872, 0.873, 0.872, 0.874, 0.870]\n","        GR=[6.895, 11.082, 14.676, 18.001, 22.098, 25.520, 27.789, 29.474]\n","  E 20  ce=3.7007  geo=0.0002  train=0.1865  test=0.3028  best=0.3120  29.3s\n","        CV=[0.750, 0.307, 0.239, 0.262, 0.285, 0.273, 0.662, 0.365]\n","        CM=[0.873, 0.864, 0.872, 0.870, 0.868, 0.869, 0.872, 0.867]\n","        GR=[6.795, 10.980, 14.467, 17.952, 21.691, 24.915, 27.035, 28.522]\n","  E 25  ce=3.6969  geo=0.0002  train=0.1961  test=0.3554  best=0.3554  29.3s\n","        CV=[0.764, 0.298, 0.223, 0.256, 0.280, 0.260, 0.646, 0.333]\n","        CM=[0.869, 0.861, 0.870, 0.867, 0.866, 0.867, 0.870, 0.865]\n","        GR=[6.703, 10.867, 14.448, 17.864, 21.492, 24.817, 26.813, 28.449]\n","  E 30  ce=3.4639  geo=0.0002  train=0.2438  test=0.3994  best=0.3994  29.4s\n","        CV=[0.743, 0.287, 0.214, 0.237, 0.251, 0.245, 0.599, 0.336]\n","        CM=[0.866, 0.858, 0.867, 0.864, 0.864, 0.865, 0.868, 0.863]\n","        GR=[6.610, 10.743, 14.318, 17.478, 21.180, 24.504, 26.426, 28.015]\n","  E 35  ce=3.3890  geo=0.0002  train=0.2652  test=0.4408  best=0.4408  29.3s\n","        CV=[0.695, 0.311, 0.210, 0.237, 0.243, 0.236, 0.588, 0.318]\n","        CM=[0.865, 0.857, 0.865, 0.862, 0.862, 0.862, 0.866, 0.861]\n","        GR=[6.537, 10.577, 14.246, 17.315, 20.920, 24.235, 26.305, 27.774]\n","  E 40  ce=3.4295  geo=0.0002  train=0.2637  test=0.4498  best=0.4576  29.3s\n","        CV=[0.725, 0.314, 0.202, 0.227, 0.224, 0.235, 0.625, 0.301]\n","        CM=[0.863, 0.855, 0.863, 0.860, 0.860, 0.860, 0.864, 0.859]\n","        GR=[6.466, 10.393, 13.885, 16.936, 20.531, 23.816, 25.826, 27.316]\n","  E 45  ce=3.2466  geo=0.0002  train=0.3041  test=0.4822  best=0.4882  29.3s\n","        CV=[0.688, 0.295, 0.193, 0.218, 0.224, 0.220, 0.676, 0.298]\n","        CM=[0.862, 0.853, 0.861, 0.859, 0.857, 0.858, 0.862, 0.857]\n","        GR=[6.417, 10.358, 13.793, 16.739, 20.195, 23.405, 25.374, 26.799]\n","  E 50  ce=3.2468  geo=0.0002  train=0.3132  test=0.5142  best=0.5142  29.4s\n","        CV=[0.731, 0.311, 0.193, 0.214, 0.229, 0.223, 0.594, 0.295]\n","        CM=[0.860, 0.851, 0.859, 0.857, 0.856, 0.856, 0.860, 0.856]\n","        GR=[6.364, 10.258, 13.696, 16.632, 20.121, 23.345, 25.393, 26.818]\n","  E 55  ce=3.0758  geo=0.0002  train=0.3527  test=0.5468  best=0.5468  29.3s\n","        CV=[0.709, 0.311, 0.192, 0.214, 0.222, 0.212, 0.623, 0.300]\n","        CM=[0.858, 0.849, 0.858, 0.856, 0.855, 0.855, 0.859, 0.854]\n","        GR=[6.318, 10.216, 13.645, 16.596, 20.189, 23.355, 25.301, 26.637]\n","  E 60  ce=3.0214  geo=0.0002  train=0.3671  test=0.5555  best=0.5555  29.3s\n","        CV=[0.671, 0.317, 0.188, 0.201, 0.209, 0.209, 0.581, 0.285]\n","        CM=[0.857, 0.848, 0.856, 0.855, 0.853, 0.854, 0.857, 0.853]\n","        GR=[6.282, 10.128, 13.507, 16.443, 19.979, 23.034, 24.819, 26.222]\n","  E 65  ce=2.9393  geo=0.0002  train=0.3914  test=0.5809  best=0.5809  29.2s\n","        CV=[0.676, 0.316, 0.188, 0.205, 0.207, 0.212, 0.565, 0.293]\n","        CM=[0.856, 0.847, 0.855, 0.854, 0.852, 0.853, 0.857, 0.853]\n","        GR=[6.250, 10.022, 13.359, 16.273, 19.823, 22.887, 24.732, 26.142]\n","  E 70  ce=2.9535  geo=0.0002  train=0.3963  test=0.5856  best=0.5916  29.3s\n","        CV=[0.701, 0.321, 0.188, 0.206, 0.207, 0.209, 0.560, 0.267]\n","        CM=[0.855, 0.846, 0.855, 0.853, 0.852, 0.852, 0.856, 0.852]\n","        GR=[6.224, 10.036, 13.367, 16.291, 19.847, 22.980, 24.883, 26.290]\n","  E 75  ce=2.6848  geo=0.0002  train=0.4593  test=0.6050  best=0.6050  29.2s\n","        CV=[0.709, 0.319, 0.190, 0.209, 0.202, 0.202, 0.590, 0.282]\n","        CM=[0.855, 0.846, 0.854, 0.853, 0.851, 0.852, 0.856, 0.851]\n","        GR=[6.203, 9.973, 13.263, 16.214, 19.847, 22.978, 24.875, 26.281]\n","  E 80  ce=2.6059  geo=0.0002  train=0.4808  test=0.6135  best=0.6135  30.0s\n","        CV=[0.693, 0.316, 0.189, 0.206, 0.201, 0.206, 0.574, 0.264]\n","        CM=[0.855, 0.846, 0.854, 0.852, 0.851, 0.852, 0.856, 0.851]\n","        GR=[6.191, 9.955, 13.187, 16.131, 19.784, 22.932, 24.876, 26.047]\n","  E 85  ce=2.7445  geo=0.0002  train=0.4526  test=0.6202  best=0.6236  29.2s\n","        CV=[0.704, 0.319, 0.184, 0.203, 0.195, 0.199, 0.580, 0.262]\n","        CM=[0.854, 0.846, 0.853, 0.852, 0.851, 0.852, 0.855, 0.851]\n","        GR=[6.183, 9.950, 13.107, 16.088, 19.764, 22.919, 24.930, 25.980]\n","  E 90  ce=3.0159  geo=0.0002  train=0.3864  test=0.6244  best=0.6245  29.3s\n","        CV=[0.680, 0.326, 0.182, 0.198, 0.204, 0.194, 0.589, 0.279]\n","        CM=[0.854, 0.846, 0.853, 0.853, 0.852, 0.851, 0.855, 0.851]\n","        GR=[6.179, 9.932, 12.974, 16.054, 19.858, 22.976, 24.950, 25.830]\n","  E 95  ce=2.7237  geo=0.0002  train=0.4617  test=0.6264  best=0.6278  29.3s\n","        CV=[0.661, 0.316, 0.183, 0.199, 0.200, 0.202, 0.566, 0.262]\n","        CM=[0.854, 0.847, 0.852, 0.853, 0.852, 0.851, 0.854, 0.851]\n","        GR=[6.177, 9.928, 12.857, 16.027, 19.903, 22.974, 24.872, 25.753]\n","  E 99  ce=2.7514  geo=0.0002  train=0.4538  test=0.6271  best=0.6283  29.3s\n","        CV=[0.642, 0.320, 0.184, 0.197, 0.202, 0.194, 0.586, 0.274]\n","        CM=[0.854, 0.847, 0.852, 0.853, 0.852, 0.851, 0.854, 0.851]\n","        GR=[6.177, 9.930, 12.848, 16.030, 19.912, 22.976, 24.857, 25.741]\n","\n","════════════════════════════════════════════════════════════\n","  CIFAR-100 RESULTS (CM-Validated)\n","════════════════════════════════════════════════════════════\n","  Best test accuracy: 0.6283 (62.83%)\n","  Parameters: 19,295,276\n","  Checkpoint: geo_cifar100/best.pt\n","  TensorBoard: runs/geo_cifar100\n","\n","  Final geometric state:\n","\n","  Final anchor diagnostics:\n","    layer_0: cv=0.1376, cm_pos=0.795, cm_mean=1.3272\n","    layer_1: cv=0.1389, cm_pos=0.746, cm_mean=-1.0980\n","    layer_2: cv=0.1408, cm_pos=0.701, cm_mean=1.6287\n","    layer_3: cv=0.1389, cm_pos=0.799, cm_mean=-0.8710\n","    layer_4: cv=0.1407, cm_pos=0.746, cm_mean=-0.3280\n","    layer_5: cv=0.1403, cm_pos=0.783, cm_mean=0.3428\n","    layer_6: cv=0.1394, cm_pos=0.732, cm_mean=0.6068\n","    layer_7: cv=0.1398, cm_pos=0.699, cm_mean=0.0258\n","\n","Done.\n"]}]},{"cell_type":"code","source":["\"\"\"\n","Geometric Transformer — CIFAR-100 Training with CM-Validated Analysis\n","\n","Changes from previous version:\n","  - CM gate diagnostics per layer: active anchors, gate_mean, cm_positive_frac\n","  - CM quality in geometric residual analysis (replaces blind gate)\n","  - Geometric regularization losses (CV target + anchor spread) in training loop\n","  - Anchor diagnostics via model.anchor_diagnostics()\n","  - CM quality trajectory alongside CV and bridge KL for cooperation analysis\n","\n","TensorBoard logging of every geometric feature element:\n","  - CV (coefficient of variation) per layer — the pentachoron band metric\n","  - CM gate: active anchors, gate mean, cm_positive_frac, quality per position\n","  - Stream agreement/divergence per layer\n","  - Anchor utilization, entropy, spread\n","  - Patchwork activation statistics (from CM-validated triangulation)\n","  - Bridge vs assignment consistency\n","  - Triangulation distance distributions\n","  - SVD spectrum, entropy, novelty\n","  - Quaternion arm norms and composition statistics\n","  - Cayley rotation ‖R-I‖ per layer\n","  - FiLM gamma/beta deviation from identity\n","  - Gate activation statistics\n","  - Gradient norms per component type (including cm_gate)\n","  - Weight norms per component type\n","  - Geometric regularization: CV loss, spread loss per epoch\n","\n","!pip install geolip-core torchvision tqdm tensorboard\n","\"\"\"\n","\n","import torch\n","import torch.nn as nn\n","import torch.nn.functional as F\n","import numpy as np\n","import time, json, math\n","from pathlib import Path\n","from tqdm.auto import tqdm\n","from torch.utils.tensorboard import SummaryWriter\n","\n","device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n","print(f\"Device: {device}\")\n","if device.type == 'cuda':\n","    print(f\"  GPU: {torch.cuda.get_device_name()}\")\n","    print(f\"  VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB\")\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# IMPORT TRANSFORMER\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","# Try geolip_core installed package first, fall back to local file\n","try:\n","    from geolip_core.pipeline.components.geometric_transformer import (\n","        GeometricTransformer, GeometricTransformerLayer,\n","        CayleyOrthogonal, QuaternionCompose, FiLMLayer,\n","        ContentAttention, GeometricAttention, CMValidatedGate,\n","        TorchComponent, BaseTower,\n","        anchor_neighborhood_cm,\n","    )\n","    print(\"  Imported from geolip_core (installed)\")\n","except ImportError:\n","    try:\n","        from geometric_transformer import (\n","            GeometricTransformer, GeometricTransformerLayer,\n","            CayleyOrthogonal, QuaternionCompose, FiLMLayer,\n","            ContentAttention, GeometricAttention, CMValidatedGate,\n","            TorchComponent, BaseTower,\n","            anchor_neighborhood_cm,\n","        )\n","        print(\"  Imported from local geometric_transformer.py\")\n","    except ImportError:\n","        raise ImportError(\n","            \"Cannot find geometric_transformer. Place geometric_transformer.py \"\n","            \"in the working directory or install geolip-core.\")\n","\n","torch.set_float32_matmul_precision('high')\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# CONFIG\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","CONFIG = {\n","    # Model\n","    'd_model': 512,\n","    'n_heads': 8,\n","    'n_layers': 4,\n","    'n_anchors': 128,\n","    'manifold_dim': 256,\n","    'n_comp': 8,\n","    'd_comp': 16,\n","    'context_dim': 64,\n","    'quat_dim': 64,\n","    'dropout': 0.1,\n","    'cm_neighbors': 4,       # CM simplex neighbors\n","\n","    # Input stage\n","    'patch_size': 4,\n","    'img_size': 32,\n","    'in_channels': 3,\n","    'conv_channels': 64,\n","    'svd_rank': 3,\n","\n","    # Training\n","    'epochs': 300,\n","    'batch_size': 1024,\n","    'lr': 1e-3,\n","    'weight_decay': 0.05,\n","    'warmup_epochs': 5,\n","    'label_smoothing': 0.1,\n","    'num_workers': 8,\n","\n","    # Geometric regularization\n","    'cv_target': 0.20,      # pentachoron band center\n","    'cv_weight': 0.1,        # CV loss weight\n","    'spread_weight': 0.01,   # anchor spread loss weight\n","\n","    # Augmentation — tuned for CM gate training\n","    'cutmix_alpha': 1.0,     # CutMix beta distribution α (1.0 = uniform box sizes)\n","    'cutmix_prob': 0.5,      # probability of applying CutMix per batch\n","    'random_erasing_p': 0.25,  # probability of erasing per image\n","\n","    # Data\n","    'num_classes': 100,\n","\n","    # Logging\n","    'log_geo_every': 5,      # full geometric analysis every N epochs\n","    'log_grads_every': 10,   # gradient norms every N epochs\n","    'log_dir': 'runs/geo_cifar100_3',\n","}\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# INPUT STAGE\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","try:\n","    from geolip_core.core.input.svd import SVDObserver\n","    _HAS_SVD = True\n","except ImportError:\n","    _HAS_SVD = False\n","\n","    class SVDObserver(nn.Module):\n","        \"\"\"Fallback SVDObserver.\"\"\"\n","        def __init__(self, in_channels, svd_rank=24):\n","            super().__init__()\n","            self.svd_rank = svd_rank\n","            self.to_svd = nn.Conv2d(in_channels, svd_rank, 1, bias=False)\n","            self.register_buffer('ema_s', torch.ones(svd_rank))\n","            self.register_buffer('ema_vh_flat', torch.eye(svd_rank).reshape(-1))\n","            self.ema_momentum = 0.99\n","\n","        def extract_features(self, S, Vh):\n","            B, k = S.shape\n","            S_safe = S.clamp(min=1e-6)\n","            s_norm = S_safe / (S_safe.sum(dim=-1, keepdim=True) + 1e-8)\n","            vh_diag = Vh.diagonal(dim1=-2, dim2=-1)\n","            vh_offdiag = (Vh.pow(2).sum((-2, -1)) - vh_diag.pow(2).sum(-1)).unsqueeze(-1).clamp(min=0)\n","            s_ent = -(s_norm * torch.log(s_norm.clamp(min=1e-8))).sum(-1, keepdim=True)\n","            out = torch.cat([s_norm, vh_diag, vh_offdiag, s_ent], dim=-1)\n","            return torch.where(torch.isfinite(out), out, torch.zeros_like(out))\n","\n","        def compute_novelty(self, S):\n","            return S - self.ema_s.clone().unsqueeze(0)\n","\n","        def forward(self, x):\n","            B, C, H, W = x.shape\n","            h = self.to_svd(x)\n","            h_flat = h.permute(0, 2, 3, 1).reshape(B, H * W, self.svd_rank)\n","            with torch.amp.autocast('cuda', enabled=False):\n","                with torch.no_grad():\n","                    gram = torch.bmm(h_flat.float().transpose(1, 2), h_flat.float())\n","                    evals, evecs = torch.linalg.eigh(gram)\n","                    evals = evals.flip(-1).clamp(min=1e-12)\n","                    S = evals.sqrt()\n","                    Vh = evecs.flip(-1).transpose(-2, -1)\n","                    S = torch.where(torch.isfinite(S), S, torch.ones_like(S))\n","                    Vh = torch.where(torch.isfinite(Vh), Vh, torch.zeros_like(Vh))\n","            features = self.extract_features(S, Vh)\n","            novelty = self.compute_novelty(S)\n","            return S, Vh, features, novelty\n","\n","        @torch.no_grad()\n","        def update_ema(self, S, Vh):\n","            m = self.ema_momentum\n","            self.ema_s.mul_(m).add_(S.detach().mean(0), alpha=1-m)\n","            self.ema_vh_flat.mul_(m).add_(Vh.detach().mean(0).reshape(-1), alpha=1-m)\n","\n","        @property\n","        def feature_dim(self):\n","            return 2 * self.svd_rank + 2\n","\n","\n","class ConvSVDPatchEmbedding(TorchComponent):\n","    \"\"\"Input stage: conv frontend → SVDObserver → patch tokens.\"\"\"\n","    def __init__(self, name, img_size=32, patch_size=4, in_channels=3,\n","                 conv_channels=64, d_model=256, svd_rank=16):\n","        super().__init__(name)\n","        self.patch_size = patch_size\n","        self.n_patches = (img_size // patch_size) ** 2\n","        self.d_model = d_model\n","        self.svd_rank = svd_rank\n","\n","        self.conv_frontend = nn.Sequential(\n","            nn.Conv2d(in_channels, conv_channels, 3, padding=1, bias=False),\n","            nn.BatchNorm2d(conv_channels), nn.GELU(),\n","            nn.Conv2d(conv_channels, conv_channels, 3, padding=1, bias=False),\n","            nn.BatchNorm2d(conv_channels), nn.GELU(),\n","        )\n","        self.svd_observer = SVDObserver(conv_channels, svd_rank)\n","        self.patch_proj = nn.Conv2d(\n","            conv_channels, d_model, kernel_size=patch_size,\n","            stride=patch_size, bias=False)\n","        self.patch_norm = nn.LayerNorm(d_model)\n","\n","        svd_feat_dim = self.svd_observer.feature_dim\n","        self.svd_to_gamma = nn.Linear(svd_feat_dim, d_model)\n","        self.svd_to_beta = nn.Linear(svd_feat_dim, d_model)\n","        nn.init.zeros_(self.svd_to_gamma.weight); nn.init.ones_(self.svd_to_gamma.bias)\n","        nn.init.zeros_(self.svd_to_beta.weight); nn.init.zeros_(self.svd_to_beta.bias)\n","\n","        self.cls_token = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)\n","        self.pos_embed = nn.Parameter(\n","            torch.randn(1, self.n_patches + 1, d_model) * 0.02)\n","\n","    def forward(self, x):\n","        B = x.shape[0]\n","        feat = self.conv_frontend(x)\n","        S, Vh, svd_features, novelty = self.svd_observer(feat)\n","        tokens = self.patch_proj(feat)\n","        tokens = tokens.flatten(2).transpose(1, 2)\n","        tokens = self.patch_norm(tokens)\n","        gamma = self.svd_to_gamma(svd_features).unsqueeze(1)\n","        beta = self.svd_to_beta(svd_features).unsqueeze(1)\n","        tokens = gamma * tokens + beta\n","        cls = self.cls_token.expand(B, -1, -1)\n","        tokens = torch.cat([cls, tokens], dim=1)\n","        tokens = tokens + self.pos_embed\n","        svd_state = {\n","            'singular_values': S, 'Vh': Vh,\n","            'svd_features': svd_features, 'novelty': novelty,\n","        }\n","        if self.training:\n","            self.svd_observer.update_ema(S, Vh)\n","        return tokens, svd_state\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# CLASSIFIER ( uses GeometricTransformer with CM gates)\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","class GeoViTClassifier(BaseTower):\n","    \"\"\"Geometric Vision Transformer for classification.\n","\n","    Wraps ConvSVDPatchEmbedding + GeometricTransformer + task head.\n","    Exposes geometric_losses() for regularization during training.\n","    \"\"\"\n","    def __init__(self, name, config):\n","        super().__init__(name)\n","        self.config = config\n","\n","        self.attach('patch_embed', ConvSVDPatchEmbedding(\n","            'patch_embed', img_size=config['img_size'],\n","            patch_size=config['patch_size'], in_channels=config['in_channels'],\n","            conv_channels=config['conv_channels'], d_model=config['d_model'],\n","            svd_rank=config['svd_rank'],\n","        ))\n","        self.attach('transformer', GeometricTransformer(\n","            'geo_cifar', d_model=config['d_model'], n_heads=config['n_heads'],\n","            n_layers=config['n_layers'], n_anchors=config['n_anchors'],\n","            manifold_dim=config['manifold_dim'], n_comp=config['n_comp'],\n","            d_comp=config['d_comp'], context_dim=config['context_dim'],\n","            quat_dim=config['quat_dim'], dropout=config['dropout'],\n","            cm_neighbors=config.get('cm_neighbors', 3),\n","        ))\n","        self.attach('head', nn.Sequential(\n","            nn.LayerNorm(config['d_model']),\n","            nn.Linear(config['d_model'], config['d_model']),\n","            nn.GELU(), nn.Dropout(config['dropout']),\n","            nn.Linear(config['d_model'], config['num_classes']),\n","        ))\n","\n","    def forward(self, x, return_geo_state=False):\n","        tokens, svd_state = self['patch_embed'](x)\n","        if return_geo_state:\n","            features, geo_states = self['transformer'](tokens, return_geo_state=True)\n","        else:\n","            features = self['transformer'](tokens)\n","        cls_out = features[:, 0]\n","        logits = self['head'](cls_out)\n","        if return_geo_state:\n","            return logits, geo_states, svd_state\n","        return logits\n","\n","    def geometric_losses(self):\n","        \"\"\"Delegate to transformer's built-in geometric regularization.\"\"\"\n","        return self['transformer'].geometric_losses(\n","            cv_target=self.config.get('cv_target', 0.215),\n","            cv_weight=self.config.get('cv_weight', 0.1),\n","            spread_weight=self.config.get('spread_weight', 0.01),\n","        )\n","\n","    def anchor_diagnostics(self):\n","        \"\"\"Delegate to transformer's anchor diagnostics.\"\"\"\n","        return self['transformer'].anchor_diagnostics()\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# GEOMETRIC ANALYSIS BATTERY ( includes CM diagnostics)\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","@torch.no_grad()\n","def compute_cv(points):\n","    \"\"\"Coefficient of variation on S^(d-1).\n","    CV = std(pairwise_cosine_distances) / mean(pairwise_cosine_distances)\n","    Pentachoron band: CV ∈ [0.20, 0.23].\n","    \"\"\"\n","    points = F.normalize(points.float(), dim=-1)\n","    cos_sim = points @ points.T\n","    n = points.shape[0]\n","    idx = torch.triu_indices(n, n, offset=1, device=points.device)\n","    pairwise_dist = 1.0 - cos_sim[idx[0], idx[1]]\n","    mean_d = pairwise_dist.mean()\n","    std_d = pairwise_dist.std()\n","    cv = (std_d / (mean_d + 1e-8)).item()\n","    return cv, mean_d.item(), std_d.item()\n","\n","\n","@torch.no_grad()\n","def log_geometric_analysis(model, writer, epoch, test_loader, device, config):\n","    \"\"\"Full geometric analysis battery with CM diagnostics.\"\"\"\n","    model.eval()\n","\n","    images, labels = next(iter(test_loader))\n","    images = images[:min(64, images.shape[0])].to(device)\n","    labels = labels[:min(64, labels.shape[0])].to(device)\n","\n","    logits, geo_states, svd_state = model(images, return_geo_state=True)\n","\n","    n_layers = len(geo_states)\n","    pred = logits.argmax(1)\n","    batch_acc = (pred == labels).float().mean().item()\n","    writer.add_scalar('analysis/batch_accuracy', batch_acc, epoch)\n","\n","    # ─── SVD Input Stage ───\n","    S = svd_state['singular_values']\n","    s_norm = S / (S.sum(dim=-1, keepdim=True) + 1e-8)\n","    s_ent = -(s_norm * torch.log(s_norm.clamp(min=1e-8))).sum(-1)\n","    novelty = svd_state['novelty']\n","\n","    writer.add_scalar('svd/entropy_mean', s_ent.mean().item(), epoch)\n","    writer.add_scalar('svd/entropy_std', s_ent.std().item(), epoch)\n","    writer.add_scalar('svd/novelty_norm', novelty.norm(dim=-1).mean().item(), epoch)\n","    writer.add_scalar('svd/top1_ratio', (S[:, 0] / (S.sum(-1) + 1e-8)).mean().item(), epoch)\n","    writer.add_scalar('svd/condition_number',\n","        (S[:, 0] / (S[:, -1].clamp(min=1e-8))).mean().item(), epoch)\n","    for k in range(min(S.shape[1], 5)):\n","        writer.add_scalar(f'svd/S_{k}', S[:, k].mean().item(), epoch)\n","\n","    # SVD FiLM deviation\n","    pe = model['patch_embed']\n","    writer.add_scalar('svd_film/gamma_weight_norm', pe.svd_to_gamma.weight.data.norm().item(), epoch)\n","    writer.add_scalar('svd_film/gamma_bias_dev_from_1',\n","        (pe.svd_to_gamma.bias.data - 1.0).abs().mean().item(), epoch)\n","    writer.add_scalar('svd_film/beta_weight_norm', pe.svd_to_beta.weight.data.norm().item(), epoch)\n","    writer.add_scalar('svd_film/beta_bias_norm', pe.svd_to_beta.bias.data.abs().mean().item(), epoch)\n","\n","    # ─── Anchor Diagnostics (built-in) ───\n","    anchor_diag = model.anchor_diagnostics()\n","    for layer_name, d in anchor_diag.items():\n","        for k, v in d.items():\n","            writer.add_scalar(f'anchor_diag/{layer_name}_{k}', v, epoch)\n","\n","    # ─── Per-Layer Geometric Analysis ───\n","    for i, gs in enumerate(geo_states):\n","        prefix = f'layer_{i}'\n","\n","        # === CV — pentachoron band metric ===\n","        emb = gs['embedding']\n","        # Anchor CV\n","        transformer = model['transformer']\n","        layer = transformer[f'layer_{i}']\n","        anchors = F.normalize(\n","            layer['observer'].association.constellation.anchors, dim=-1)\n","        cv_anchors, mean_d_anchors, std_d_anchors = compute_cv(anchors)\n","        writer.add_scalar(f'{prefix}/cv_anchors', cv_anchors, epoch)\n","        writer.add_scalar(f'{prefix}/anchor_mean_dist', mean_d_anchors, epoch)\n","        writer.add_scalar(f'{prefix}/anchor_std_dist', std_d_anchors, epoch)\n","\n","        # Embedding CV\n","        emb_flat = emb.reshape(-1, emb.shape[-1])\n","        n_sample = min(512, emb_flat.shape[0])\n","        idx = torch.randperm(emb_flat.shape[0], device=device)[:n_sample]\n","        cv_emb, mean_d_emb, std_d_emb = compute_cv(emb_flat[idx])\n","        writer.add_scalar(f'{prefix}/cv_embeddings', cv_emb, epoch)\n","        writer.add_scalar(f'{prefix}/embedding_mean_dist', mean_d_emb, epoch)\n","\n","        # === CM Gate Diagnostics ===\n","        gate_info = gs.get('gate_info', {})\n","        gate_values = gs.get('gate_values')\n","        cm_quality = gs.get('cm_quality')\n","\n","        if gate_info:\n","            writer.add_scalar(f'{prefix}/cm_active_anchors',\n","                gate_info.get('active', 0), epoch)\n","            writer.add_scalar(f'{prefix}/cm_gate_mean',\n","                gate_info.get('gate_mean', 0), epoch)\n","            writer.add_scalar(f'{prefix}/cm_positive_frac',\n","                gate_info.get('cm_positive_frac', 0), epoch)\n","\n","        if gate_values is not None:\n","            gv = gate_values\n","            writer.add_scalar(f'{prefix}/gate_values_min', gv.min().item(), epoch)\n","            writer.add_scalar(f'{prefix}/gate_values_max', gv.max().item(), epoch)\n","            writer.add_scalar(f'{prefix}/gate_values_std', gv.std().item(), epoch)\n","            # Per-anchor gate mean (which anchors are consistently open/closed)\n","            gv_per_anchor = gv.mean(dim=0).mean(dim=0)  # average over B and L\n","            writer.add_scalar(f'{prefix}/gate_anchor_spread',\n","                gv_per_anchor.std().item(), epoch)\n","            # Fraction of positions with >50% anchors open\n","            if gv.dim() == 3:\n","                pos_open_frac = (gv.mean(dim=-1) > 0.5).float().mean().item()\n","            else:\n","                pos_open_frac = (gv > 0.5).float().mean().item()\n","            writer.add_scalar(f'{prefix}/gate_positions_open_frac', pos_open_frac, epoch)\n","\n","        if cm_quality is not None:\n","            writer.add_scalar(f'{prefix}/cm_quality_mean', cm_quality.mean().item(), epoch)\n","            writer.add_scalar(f'{prefix}/cm_quality_std', cm_quality.std().item(), epoch)\n","            writer.add_scalar(f'{prefix}/cm_quality_min', cm_quality.min().item(), epoch)\n","\n","        # === Stream Agreement ===\n","        content = gs['content']\n","        geometric = gs['geometric']\n","        agreement = F.cosine_similarity(\n","            content.reshape(-1, content.shape[-1]),\n","            geometric.reshape(-1, geometric.shape[-1]), dim=-1)\n","        writer.add_scalar(f'{prefix}/stream_agreement_mean', agreement.mean().item(), epoch)\n","        writer.add_scalar(f'{prefix}/stream_agreement_std', agreement.std().item(), epoch)\n","\n","        writer.add_scalar(f'{prefix}/content_norm', content.norm(dim=-1).mean().item(), epoch)\n","        writer.add_scalar(f'{prefix}/geometric_norm', geometric.norm(dim=-1).mean().item(), epoch)\n","\n","        # === Disagreement arm analysis ===\n","        disagree = content - geometric\n","        agree = content * geometric\n","        writer.add_scalar(f'{prefix}/disagree_norm', disagree.norm(dim=-1).mean().item(), epoch)\n","        writer.add_scalar(f'{prefix}/agree_norm', agree.norm(dim=-1).mean().item(), epoch)\n","\n","        # === Anchor Utilization ===\n","        tri = gs['triangulation']\n","        assignment = gs['assignment']\n","        nearest = gs['nearest']\n","        n_anchors = tri.shape[-1]\n","\n","        nearest_flat = nearest.reshape(-1)\n","        counts = torch.bincount(nearest_flat, minlength=n_anchors).float()\n","        total_assignments = counts.sum()\n","\n","        probs = counts / (total_assignments + 1e-8)\n","        anchor_entropy = -(probs * torch.log(probs.clamp(min=1e-8))).sum().item()\n","        max_entropy = math.log(n_anchors)\n","        writer.add_scalar(f'{prefix}/anchor_entropy_normalized',\n","            anchor_entropy / (max_entropy + 1e-8), epoch)\n","        active = (counts > 0).sum().item()\n","        writer.add_scalar(f'{prefix}/anchors_active', active, epoch)\n","        writer.add_scalar(f'{prefix}/anchors_active_frac', active / n_anchors, epoch)\n","        dead = (counts == 0).sum().item()\n","        writer.add_scalar(f'{prefix}/anchors_dead', dead, epoch)\n","\n","        # === Triangulation Statistics ===\n","        writer.add_scalar(f'{prefix}/tri_mean', tri.mean().item(), epoch)\n","        writer.add_scalar(f'{prefix}/tri_std', tri.std().item(), epoch)\n","\n","        # === Soft Assignment Statistics ===\n","        assign_ent = -(assignment * torch.log(assignment.clamp(min=1e-8))).sum(-1)\n","        writer.add_scalar(f'{prefix}/assignment_entropy_mean', assign_ent.mean().item(), epoch)\n","        writer.add_scalar(f'{prefix}/assignment_max_prob',\n","            assignment.max(dim=-1).values.mean().item(), epoch)\n","\n","        # === Patchwork Statistics (now from CM-validated triangulation) ===\n","        pw = gs['patchwork']\n","        writer.add_scalar(f'{prefix}/patchwork_norm', pw.norm(dim=-1).mean().item(), epoch)\n","        writer.add_scalar(f'{prefix}/patchwork_std', pw.std().item(), epoch)\n","        pw_sparsity = (pw.abs() < 0.01).float().mean().item()\n","        writer.add_scalar(f'{prefix}/patchwork_sparsity', pw_sparsity, epoch)\n","\n","        # === Bridge Consistency ===\n","        bridge = gs['bridge']\n","        bridge_soft = F.softmax(bridge, dim=-1)\n","        bridge_assign_kl = F.kl_div(\n","            bridge_soft.log().reshape(-1, n_anchors),\n","            assignment.reshape(-1, n_anchors),\n","            reduction='batchmean', log_target=False)\n","        writer.add_scalar(f'{prefix}/bridge_assignment_kl', bridge_assign_kl.item(), epoch)\n","\n","        # === Quaternion Composition ===\n","        composed = gs['composed']\n","        writer.add_scalar(f'{prefix}/composed_norm', composed.norm(dim=-1).mean().item(), epoch)\n","\n","        # === Geo Context ===\n","        geo_ctx = gs['geo_ctx']\n","        writer.add_scalar(f'{prefix}/geo_ctx_norm', geo_ctx.norm(dim=-1).mean().item(), epoch)\n","\n","        # === Geometric Residual Stream (CM-conditioned) ===\n","        geo_res = gs.get('geo_residual')\n","        if geo_res is not None:\n","            res_norms = geo_res.norm(dim=-1)\n","            writer.add_scalar(f'{prefix}/geo_res_norm', res_norms.mean().item(), epoch)\n","            writer.add_scalar(f'{prefix}/geo_res_std', geo_res.std().item(), epoch)\n","            writer.add_scalar(f'{prefix}/geo_res_sparsity',\n","                (geo_res.abs() < 0.01).float().mean().item(), epoch)\n","            # Cross-position consistency\n","            geo_res_flat = geo_res.reshape(-1, geo_res.shape[-1])\n","            n_s = min(256, geo_res_flat.shape[0])\n","            idx_s = torch.randperm(geo_res_flat.shape[0], device=geo_res.device)[:n_s]\n","            sampled = F.normalize(geo_res_flat[idx_s], dim=-1)\n","            cos_mat = sampled @ sampled.T\n","            triu = torch.triu_indices(n_s, n_s, offset=1, device=geo_res.device)\n","            writer.add_scalar(f'{prefix}/geo_res_consistency',\n","                cos_mat[triu[0], triu[1]].mean().item(), epoch)\n","\n","    # ─── Cayley Rotation Analysis ───\n","    for name, mod in model.named_modules():\n","        if isinstance(mod, CayleyOrthogonal):\n","            R = mod.get_rotation()\n","            I = torch.eye(R.shape[0], device=R.device)\n","            r_dist = (R - I).norm().item()\n","            clean_name = name.replace('.', '_')\n","            writer.add_scalar(f'cayley/{clean_name}_R_minus_I', r_dist, epoch)\n","\n","    # ─── FiLM Layer Analysis ───\n","    film_idx = 0\n","    for name, mod in model.named_modules():\n","        if isinstance(mod, FiLMLayer):\n","            g_b = mod.to_gamma.bias.data\n","            b_b = mod.to_beta.bias.data\n","            writer.add_scalar(f'film/{film_idx}_gamma_dev',\n","                (g_b - 1.0).abs().mean().item(), epoch)\n","            writer.add_scalar(f'film/{film_idx}_beta_dev',\n","                b_b.abs().mean().item(), epoch)\n","            film_idx += 1\n","\n","    # ─── Cross-Layer Trajectories ───\n","    cv_trajectory = []\n","    cm_quality_trajectory = []\n","    res_norms = []\n","    bridge_kls = []\n","\n","    for i, gs in enumerate(geo_states):\n","        # CV\n","        emb = gs['embedding']\n","        emb_flat = emb.reshape(-1, emb.shape[-1])\n","        n_sample = min(512, emb_flat.shape[0])\n","        idx = torch.randperm(emb_flat.shape[0], device=device)[:n_sample]\n","        cv, _, _ = compute_cv(emb_flat[idx])\n","        cv_trajectory.append(cv)\n","\n","        # CM quality\n","        cm_q = gs.get('cm_quality')\n","        if cm_q is not None:\n","            cm_quality_trajectory.append(cm_q.mean().item())\n","\n","        # Geo residual norms\n","        geo_res = gs.get('geo_residual')\n","        if geo_res is not None:\n","            res_norms.append(geo_res.norm(dim=-1).mean().item())\n","\n","        # Bridge KL\n","        n_anchors = gs['assignment'].shape[-1]\n","        bridge_soft = F.softmax(gs['bridge'], dim=-1)\n","        bkl = F.kl_div(\n","            bridge_soft.log().reshape(-1, n_anchors),\n","            gs['assignment'].reshape(-1, n_anchors),\n","            reduction='batchmean', log_target=False).item()\n","        bridge_kls.append(bkl)\n","\n","    # CV trajectory\n","    writer.add_scalar('cv/trajectory_mean', np.mean(cv_trajectory), epoch)\n","    writer.add_scalar('cv/trajectory_std', np.std(cv_trajectory), epoch)\n","    in_band = sum(1 for cv in cv_trajectory if 0.20 <= cv <= 0.23)\n","    writer.add_scalar('cv/layers_in_pentachoron_band', in_band, epoch)\n","    writer.add_scalar('cv/layers_in_band_frac', in_band / len(cv_trajectory), epoch)\n","\n","    # CM quality trajectory\n","    if cm_quality_trajectory:\n","        writer.add_scalar('cm/quality_trajectory_mean',\n","            np.mean(cm_quality_trajectory), epoch)\n","        writer.add_scalar('cm/quality_trajectory_std',\n","            np.std(cm_quality_trajectory), epoch)\n","        writer.add_scalar('cm/quality_min_layer',\n","            np.min(cm_quality_trajectory), epoch)\n","        writer.add_scalar('cm/quality_max_layer',\n","            np.max(cm_quality_trajectory), epoch)\n","\n","    # Geometric residual trajectory\n","    if res_norms:\n","        writer.add_scalar('geo_res/trajectory_start', res_norms[0], epoch)\n","        writer.add_scalar('geo_res/trajectory_end', res_norms[-1], epoch)\n","        writer.add_scalar('geo_res/accumulation_ratio',\n","            res_norms[-1] / (res_norms[0] + 1e-8), epoch)\n","        growth = [res_norms[j+1] - res_norms[j] for j in range(len(res_norms)-1)]\n","        writer.add_scalar('geo_res/growth_mean', np.mean(growth), epoch)\n","        writer.add_scalar('geo_res/growth_std', np.std(growth), epoch)\n","\n","        # Cooperation analysis (includes CM quality)\n","        if len(res_norms) >= 4:\n","            cv_corr = float(np.corrcoef(res_norms, cv_trajectory)[0, 1])\n","            bkl_corr = float(np.corrcoef(res_norms, bridge_kls)[0, 1])\n","            writer.add_scalar('cooperation/geo_res_vs_cv', cv_corr, epoch)\n","            writer.add_scalar('cooperation/geo_res_vs_bridge_kl', bkl_corr, epoch)\n","\n","            if len(cm_quality_trajectory) == len(res_norms):\n","                cm_corr = float(np.corrcoef(\n","                    res_norms, cm_quality_trajectory)[0, 1])\n","                writer.add_scalar('cooperation/geo_res_vs_cm_quality', cm_corr, epoch)\n","                # CM vs CV: do layers with better CM quality also have better CV?\n","                cm_cv_corr = float(np.corrcoef(\n","                    cm_quality_trajectory, cv_trajectory)[0, 1])\n","                writer.add_scalar('cooperation/cm_quality_vs_cv', cm_cv_corr, epoch)\n","\n","    return {\n","        'batch_acc': batch_acc,\n","        'cv_trajectory': cv_trajectory,\n","        'cm_quality_trajectory': cm_quality_trajectory,\n","        'res_norms': res_norms,\n","        'bridge_kls': bridge_kls,\n","    }\n","\n","\n","@torch.no_grad()\n","def log_gradient_norms(model, writer, epoch):\n","    \"\"\"Log gradient norms per component type (includes cm_gate).\"\"\"\n","    type_grads = {}\n","    for name, param in model.named_parameters():\n","        if param.grad is not None:\n","            grad_norm = param.grad.norm().item()\n","            if 'projection' in name and 'proj' in name:\n","                key = 'manifold_proj'\n","            elif 'cm_gate' in name:\n","                key = 'cm_gate'\n","            elif 'observer' in name or 'constellation' in name or 'anchor' in name:\n","                key = 'constellation'\n","            elif 'context' in name:\n","                key = 'geo_context'\n","            elif 'content' in name:\n","                key = 'content_attn'\n","            elif 'geometric' in name and 'film' not in name:\n","                key = 'geo_attn'\n","            elif 'film' in name:\n","                key = 'film'\n","            elif 'rotation' in name or 'cayley' in name or 'A_upper' in name:\n","                key = 'cayley'\n","            elif 'compose' in name or 'quat' in name or 'proj_w' in name:\n","                key = 'quaternion'\n","            elif 'decode' in name:\n","                key = 'decode'\n","            elif 'gate' in name:\n","                key = 'gate'\n","            elif 'conv' in name or 'patch' in name:\n","                key = 'input_stage'\n","            elif 'head' in name:\n","                key = 'head'\n","            elif 'svd' in name:\n","                key = 'svd'\n","            elif 'geo_proj' in name:\n","                key = 'geo_residual_proj'\n","            else:\n","                key = 'other'\n","\n","            if key not in type_grads:\n","                type_grads[key] = []\n","            type_grads[key].append(grad_norm)\n","\n","    for key, norms in type_grads.items():\n","        writer.add_scalar(f'grad_norm/{key}_mean', np.mean(norms), epoch)\n","        writer.add_scalar(f'grad_norm/{key}_max', np.max(norms), epoch)\n","\n","    total = sum(p.grad.norm().item() ** 2\n","                for p in model.parameters() if p.grad is not None) ** 0.5\n","    writer.add_scalar('grad_norm/total', total, epoch)\n","\n","\n","@torch.no_grad()\n","def log_weight_norms(model, writer, epoch):\n","    \"\"\"Log weight norms per component type.\"\"\"\n","    for name, param in model.named_parameters():\n","        if 'A_upper' in name:\n","            clean = name.replace('.', '_')\n","            writer.add_scalar(f'weights/{clean}_norm', param.norm().item(), epoch)\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# DATA\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","def get_dataloaders(config):\n","    import torchvision\n","    import torchvision.transforms as T\n","\n","    # Augmentation pipeline tuned for geometric transformer:\n","    #   TrivialAugmentWide: continuous severity spectrum of geometric + photometric\n","    #     transforms. Exercises CM gate across full quality range — mild distortion\n","    #     keeps CM high, severe distortion creates partially-degenerate simplices.\n","    #   RandomErasing: creates degenerate manifold projections (zero-volume CM simplices).\n","    #     Trains CM gate to close on corrupted regions.\n","    #   CutMix applied at batch level in train_epoch (not here).\n","    train_transform = T.Compose([\n","        T.RandomCrop(32, padding=4),\n","        T.RandomHorizontalFlip(),\n","        T.TrivialAugmentWide(),\n","        T.ToTensor(),\n","        T.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),\n","        T.RandomErasing(p=config.get('random_erasing_p', 0.25),\n","                        scale=(0.02, 0.33), ratio=(0.3, 3.3)),\n","    ])\n","    test_transform = T.Compose([\n","        T.ToTensor(),\n","        T.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),\n","    ])\n","\n","    train_ds = torchvision.datasets.CIFAR100(\n","        root='./data', train=True, download=True, transform=train_transform)\n","    test_ds = torchvision.datasets.CIFAR100(\n","        root='./data', train=False, download=True, transform=test_transform)\n","\n","    train_loader = torch.utils.data.DataLoader(\n","        train_ds, batch_size=config['batch_size'], shuffle=True,\n","        num_workers=config['num_workers'], pin_memory=True, drop_last=True)\n","    test_loader = torch.utils.data.DataLoader(\n","        test_ds, batch_size=config['batch_size'], shuffle=False,\n","        num_workers=config['num_workers'], pin_memory=True)\n","\n","    return train_loader, test_loader\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# CUTMIX — batch-level augmentation for CM gate boundary training\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","def cutmix_batch(images, labels, alpha=1.0):\n","    \"\"\"Apply CutMix to a batch. Returns mixed images + label pairs + lambda.\n","\n","    CutMix replaces a rectangular region of image A with image B.\n","    Positions inside each region have coherent geometry — valid CM simplices.\n","    The boundary between regions has mixed geometric context — the CM gate\n","    should learn to suppress these positions.\n","\n","    Args:\n","        images: (B, C, H, W) batch\n","        labels: (B,) integer labels\n","        alpha: Beta distribution parameter (1.0 = uniform box sizes)\n","\n","    Returns:\n","        images: (B, C, H, W) mixed batch (modified in-place)\n","        labels_a: (B,) labels for region A\n","        labels_b: (B,) labels for region B\n","        lam: float, fraction of image A remaining\n","    \"\"\"\n","    lam = np.random.beta(alpha, alpha)\n","    B = images.size(0)\n","    idx = torch.randperm(B, device=images.device)\n","\n","    H, W = images.shape[2], images.shape[3]\n","    cut_ratio = (1.0 - lam) ** 0.5\n","    cw = int(W * cut_ratio)\n","    ch = int(H * cut_ratio)\n","    cx = np.random.randint(W)\n","    cy = np.random.randint(H)\n","    x1 = max(cx - cw // 2, 0); x2 = min(cx + cw // 2, W)\n","    y1 = max(cy - ch // 2, 0); y2 = min(cy + ch // 2, H)\n","\n","    images[:, :, y1:y2, x1:x2] = images[idx, :, y1:y2, x1:x2]\n","    lam_actual = 1.0 - (x2 - x1) * (y2 - y1) / (W * H)\n","    return images, labels, labels[idx], lam_actual\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# TRAINING (geometric losses + CutMix integrated)\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","def train_epoch(model, loader, optimizer, scheduler, epoch, config, writer):\n","    model.train()\n","    total_loss = 0\n","    total_geo_loss = 0\n","    correct = 0\n","    total = 0\n","\n","    cutmix_alpha = config.get('cutmix_alpha', 1.0)\n","    cutmix_prob = config.get('cutmix_prob', 0.5)\n","    label_smoothing = config.get('label_smoothing', 0.1)\n","    criterion = nn.CrossEntropyLoss(label_smoothing=label_smoothing)\n","\n","    for batch_idx, (images, labels) in enumerate(loader):\n","        images = images.to(device)\n","        labels = labels.to(device)\n","\n","        # CutMix: applied probabilistically per batch\n","        use_cutmix = np.random.rand() < cutmix_prob\n","        if use_cutmix:\n","            images, labels_a, labels_b, lam = cutmix_batch(\n","                images, labels, alpha=cutmix_alpha)\n","            logits = model(images)\n","            ce_loss = lam * criterion(logits, labels_a) + \\\n","                      (1.0 - lam) * criterion(logits, labels_b)\n","            # Accuracy: count correct if matches either label\n","            pred = logits.argmax(1)\n","            correct += (lam * (pred == labels_a).float() +\n","                        (1.0 - lam) * (pred == labels_b).float()).sum().item()\n","        else:\n","            logits = model(images)\n","            ce_loss = criterion(logits, labels)\n","            correct += (logits.argmax(1) == labels).sum().item()\n","\n","        # Geometric regularization — CV target + anchor spread\n","        geo_losses = model.geometric_losses()\n","        geo_loss = geo_losses.get('geo_total', torch.tensor(0.0, device=device))\n","        loss = ce_loss + geo_loss\n","\n","        optimizer.zero_grad()\n","        loss.backward()\n","\n","        # Log gradient norms periodically\n","        if epoch % config['log_grads_every'] == 0 and batch_idx == 0:\n","            log_gradient_norms(model, writer, epoch)\n","\n","        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n","        optimizer.step()\n","        if scheduler is not None:\n","            scheduler.step()\n","\n","        total_loss += ce_loss.item() * images.size(0)\n","        total_geo_loss += geo_loss.item() * images.size(0)\n","        total += images.size(0)\n","\n","    avg_ce = total_loss / total\n","    avg_geo = total_geo_loss / total\n","    return avg_ce, avg_geo, correct / total\n","\n","\n","@torch.no_grad()\n","def evaluate(model, loader):\n","    model.eval()\n","    correct = 0\n","    total = 0\n","    for images, labels in loader:\n","        images = images.to(device)\n","        labels = labels.to(device)\n","        logits = model(images)\n","        correct += (logits.argmax(1) == labels).sum().item()\n","        total += images.size(0)\n","    return correct / total\n","\n","\n","def main():\n","    config = CONFIG.copy()\n","\n","    print(\"=\" * 60)\n","    print(\"  Geometric Transformer — CIFAR-100 (CM-Validated)\")\n","    print(f\"  Input: conv({config['in_channels']}→{config['conv_channels']}) + \"\n","          f\"SVD(rank={config['svd_rank']}) + \"\n","          f\"{config['patch_size']}×{config['patch_size']} patches = \"\n","          f\"{(config['img_size']//config['patch_size'])**2} tokens + CLS\")\n","    print(f\"  Model: d={config['d_model']}, heads={config['n_heads']}, \"\n","          f\"layers={config['n_layers']}, anchors={config['n_anchors']}\")\n","    print(f\"  CM: neighbors={config['cm_neighbors']}, \"\n","          f\"cv_target={config['cv_target']}, \"\n","          f\"cv_weight={config['cv_weight']}, \"\n","          f\"spread_weight={config['spread_weight']}\")\n","    print(f\"  Aug: TrivialAugmentWide + CutMix(α={config['cutmix_alpha']}, \"\n","          f\"p={config['cutmix_prob']}) + \"\n","          f\"RandomErasing(p={config['random_erasing_p']})\")\n","    print(\"=\" * 60)\n","\n","    writer = SummaryWriter(config['log_dir'])\n","    writer.add_text('config', json.dumps(config, indent=2))\n","\n","    print(\"\\nLoading CIFAR-100...\")\n","    train_loader, test_loader = get_dataloaders(config)\n","    print(f\"  Train: {len(train_loader.dataset):,} | Test: {len(test_loader.dataset):,}\")\n","\n","    model = GeoViTClassifier('geo_vit_cifar100', config)\n","    if hasattr(model, 'network_to'):\n","        model.network_to(device=device, strict=False)\n","    else:\n","        model = model.to(device)\n","\n","    n_params = sum(p.numel() for p in model.parameters())\n","    n_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)\n","    print(f\"\\n  Total params:     {n_params:,}\")\n","    print(f\"  Trainable params: {n_trainable:,}\")\n","\n","    model.compile()\n","\n","    for name, module in model.named_children():\n","        n = sum(p.numel() for p in module.parameters())\n","        if n > 0:\n","            print(f\"    {name:<20s}: {n:,}\")\n","\n","    writer.add_scalar('model/total_params', n_params, 0)\n","\n","    # Initial anchor diagnostics\n","    print(f\"\\n  Initial anchor diagnostics:\")\n","    diag = model.anchor_diagnostics()\n","    for layer_name, d in diag.items():\n","        print(f\"    {layer_name}: cv={d['anchor_cv']:.4f}, \"\n","              f\"cm_pos={d['cm_positive_frac']:.3f}, \"\n","              f\"min_dist={d['min_pairwise_dist']:.4f}\")\n","\n","    # Optimizer + scheduler\n","    optimizer = torch.optim.AdamW(\n","        model.parameters(), lr=config['lr'], weight_decay=config['weight_decay'])\n","\n","    total_steps = config['epochs'] * len(train_loader)\n","    warmup_steps = config['warmup_epochs'] * len(train_loader)\n","\n","    def lr_lambda(step):\n","        if step < warmup_steps:\n","            return step / warmup_steps\n","        progress = (step - warmup_steps) / (total_steps - warmup_steps)\n","        return 0.5 * (1 + np.cos(np.pi * progress))\n","\n","    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)\n","\n","    print(f\"\\n{'━'*60}\")\n","    print(f\"  Training for {config['epochs']} epochs\")\n","    print(f\"  Warmup: {config['warmup_epochs']} epochs, \"\n","          f\"LR: {config['lr']}, WD: {config['weight_decay']}\")\n","    print(f\"  Geo reg: cv_w={config['cv_weight']}, spread_w={config['spread_weight']}\")\n","    print(f\"  Aug: TrivialAugmentWide + CutMix(p={config['cutmix_prob']}) + \"\n","          f\"RandomErasing(p={config['random_erasing_p']})\")\n","    print(f\"  TensorBoard: {config['log_dir']}\")\n","    print(f\"  Geo analysis every {config['log_geo_every']} epochs\")\n","    print(f\"{'━'*60}\\n\")\n","\n","    best_acc = 0\n","    save_dir = Path('geo_cifar100'); save_dir.mkdir(exist_ok=True)\n","\n","    for epoch in tqdm(range(config['epochs']), desc=\"Epochs\"):\n","        t0 = time.time()\n","\n","        ce_loss, geo_loss, train_acc = train_epoch(\n","            model, train_loader, optimizer, scheduler, epoch, config, writer)\n","\n","        test_acc = evaluate(model, test_loader)\n","        elapsed = time.time() - t0\n","\n","        lr = optimizer.param_groups[0]['lr']\n","        writer.add_scalar('train/ce_loss', ce_loss, epoch)\n","        writer.add_scalar('train/geo_loss', geo_loss, epoch)\n","        writer.add_scalar('train/total_loss', ce_loss + geo_loss, epoch)\n","        writer.add_scalar('train/accuracy', train_acc, epoch)\n","        writer.add_scalar('test/accuracy', test_acc, epoch)\n","        writer.add_scalar('train/lr', lr, epoch)\n","        writer.add_scalar('train/epoch_time', elapsed, epoch)\n","        writer.add_scalar('gap/train_test', train_acc - test_acc, epoch)\n","\n","        log_weight_norms(model, writer, epoch)\n","\n","        if test_acc > best_acc:\n","            best_acc = test_acc\n","            torch.save({\n","                'state_dict': {k: v.cpu() for k, v in model.state_dict().items()},\n","                'epoch': epoch,\n","                'test_acc': test_acc,\n","                'config': config,\n","            }, save_dir / 'best.pt')\n","\n","        # Full geometric analysis periodically\n","        if epoch % config['log_geo_every'] == 0 or epoch == config['epochs'] - 1:\n","            geo_info = log_geometric_analysis(\n","                model, writer, epoch, test_loader, device, config)\n","\n","            cv_str = ', '.join(f'{cv:.3f}' for cv in geo_info['cv_trajectory'])\n","            cm_str = ', '.join(f'{q:.3f}' for q in geo_info.get('cm_quality_trajectory', []))\n","            res_str = ', '.join(f'{r:.3f}' for r in geo_info.get('res_norms', []))\n","            tqdm.write(\n","                f\"  E{epoch:>3d}  ce={ce_loss:.4f}  geo={geo_loss:.4f}  \"\n","                f\"train={train_acc:.4f}  test={test_acc:.4f}  \"\n","                f\"best={best_acc:.4f}  {elapsed:.1f}s\"\n","                f\"\\n        CV=[{cv_str}]\"\n","                f\"\\n        CM=[{cm_str}]\"\n","                f\"\\n        GR=[{res_str}]\")\n","        elif epoch % 5 == 0:\n","            tqdm.write(\n","                f\"  E{epoch:>3d}  ce={ce_loss:.4f}  geo={geo_loss:.4f}  \"\n","                f\"train={train_acc:.4f}  test={test_acc:.4f}  \"\n","                f\"best={best_acc:.4f}  {elapsed:.1f}s\")\n","\n","    # Final summary\n","    print(f\"\\n{'═'*60}\")\n","    print(f\"  CIFAR-100 RESULTS (CM-Validated)\")\n","    print(f\"{'═'*60}\")\n","    print(f\"  Best test accuracy: {best_acc:.4f} ({best_acc*100:.2f}%)\")\n","    print(f\"  Parameters: {n_params:,}\")\n","    print(f\"  Checkpoint: {save_dir}/best.pt\")\n","    print(f\"  TensorBoard: {config['log_dir']}\")\n","\n","    # Final geometric state + anchor diagnostics\n","    print(f\"\\n  Final geometric state:\")\n","    geo_info = log_geometric_analysis(\n","        model, writer, config['epochs'], test_loader, device, config)\n","\n","    print(f\"\\n  Final anchor diagnostics:\")\n","    diag = model.anchor_diagnostics()\n","    for layer_name, d in diag.items():\n","        print(f\"    {layer_name}: cv={d['anchor_cv']:.4f}, \"\n","              f\"cm_pos={d['cm_positive_frac']:.3f}, \"\n","              f\"cm_mean={d['cm_mean']:.4f}\")\n","\n","    writer.close()\n","    print(f\"\\nDone.\")\n","\n","\n","if __name__ == '__main__':\n","    main()"],"metadata":{"colab":{"base_uri":"https://localhost:8080/","referenced_widgets":["7ba7e8d6232b4b99991f0e44369e88ab","6697a72129524bd0bcd63850b874367c","ff64e3fbea544b56bae47a7a2cd9d54c","3c2531ab76d240ea82c66ba3cdd333d4","bea929296b3648e79b1d201c19f0d172","57192c34c42848cdbfdcb47f8eb6f23f","83d893f9e61f471d8f69fc092b262be6","d73657ec98804890aea6306f5be33bd5","0eeacd094fab458ba80cf4273fc78bfb","8bf14e45224943c38577cb55b46e8933","0fa1fe6d22a744ef81acf99dc3ed7918"]},"id":"ZfPg2N9gkPPM","executionInfo":{"status":"ok","timestamp":1774762872760,"user_tz":420,"elapsed":4828050,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"1bfaf038-cd89-40fa-84cf-9eb191caa45a"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["Device: cuda\n","  GPU: NVIDIA RTX PRO 6000 Blackwell Server Edition\n","  VRAM: 102.0 GB\n","  Imported from geolip_core (installed)\n","============================================================\n","  Geometric Transformer — CIFAR-100 (CM-Validated)\n","  Input: conv(3→64) + SVD(rank=3) + 4×4 patches = 64 tokens + CLS\n","  Model: d=512, heads=8, layers=4, anchors=128\n","  CM: neighbors=4, cv_target=0.2, cv_weight=0.1, spread_weight=0.01\n","  Aug: TrivialAugmentWide + CutMix(α=1.0, p=0.5) + RandomErasing(p=0.25)\n","============================================================\n","\n","Loading CIFAR-100...\n","  Train: 50,000 | Test: 10,000\n","\n","  Total params:     32,970,280\n","  Trainable params: 32,970,280\n","    components          : 32,970,280\n","\n","  Initial anchor diagnostics:\n","    layer_0: cv=0.0239, cm_pos=1.000, min_dist=0.9644\n","    layer_1: cv=0.0243, cm_pos=1.000, min_dist=0.9661\n","    layer_2: cv=0.0243, cm_pos=1.000, min_dist=0.9659\n","    layer_3: cv=0.0245, cm_pos=1.000, min_dist=0.9661\n","\n","━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\n","  Training for 300 epochs\n","  Warmup: 5 epochs, LR: 0.001, WD: 0.05\n","  Geo reg: cv_w=0.1, spread_w=0.01\n","  Aug: TrivialAugmentWide + CutMix(p=0.5) + RandomErasing(p=0.25)\n","  TensorBoard: runs/geo_cifar100_3\n","  Geo analysis every 5 epochs\n","━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\n","\n"]},{"output_type":"display_data","data":{"text/plain":["Epochs:   0%|          | 0/300 [00:00<?, ?it/s]"],"application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"7ba7e8d6232b4b99991f0e44369e88ab"}},"metadata":{}},{"output_type":"stream","name":"stdout","text":["  E  0  ce=4.5294  geo=0.0030  train=0.0281  test=0.0574  best=0.0574  27.4s\n","        CV=[0.302, 0.247, 0.256, 0.273]\n","        CM=[0.881, 0.881, 0.881, 0.881]\n","        GR=[9.970, 14.017, 17.144, 20.017]\n","  E  5  ce=4.1302  geo=0.0006  train=0.0959  test=0.1589  best=0.1589  15.9s\n","        CV=[0.422, 0.412, 0.345, 0.258]\n","        CM=[0.878, 0.879, 0.878, 0.875]\n","        GR=[9.951, 14.896, 19.098, 21.123]\n","  E 10  ce=3.8931  geo=0.0003  train=0.1502  test=0.2748  best=0.2748  15.9s\n","        CV=[0.725, 0.396, 0.271, 0.216]\n","        CM=[0.879, 0.875, 0.876, 0.867]\n","        GR=[9.902, 14.546, 19.015, 21.023]\n","  E 15  ce=3.6642  geo=0.0002  train=0.1958  test=0.3251  best=0.3251  15.9s\n","        CV=[0.733, 0.401, 0.312, 0.221]\n","        CM=[0.878, 0.872, 0.872, 0.854]\n","        GR=[9.825, 15.829, 20.849, 22.090]\n","  E 20  ce=3.6549  geo=0.0002  train=0.2076  test=0.3658  best=0.3733  15.9s\n","        CV=[0.769, 0.328, 0.318, 0.234]\n","        CM=[0.875, 0.871, 0.870, 0.858]\n","        GR=[9.771, 15.852, 20.905, 22.647]\n","  E 25  ce=3.5647  geo=0.0002  train=0.2277  test=0.3805  best=0.3934  15.9s\n","        CV=[0.711, 0.348, 0.421, 0.229]\n","        CM=[0.872, 0.868, 0.864, 0.856]\n","        GR=[9.717, 15.552, 19.346, 21.178]\n","  E 30  ce=3.4316  geo=0.0002  train=0.2562  test=0.4203  best=0.4286  15.9s\n","        CV=[0.638, 0.317, 0.618, 0.271]\n","        CM=[0.870, 0.866, 0.857, 0.852]\n","        GR=[9.594, 15.653, 19.720, 21.399]\n","  E 35  ce=3.3799  geo=0.0002  train=0.2729  test=0.4436  best=0.4499  15.9s\n","        CV=[0.598, 0.310, 0.600, 0.281]\n","        CM=[0.866, 0.860, 0.854, 0.849]\n","        GR=[9.431, 15.008, 19.861, 21.369]\n","  E 40  ce=3.1824  geo=0.0002  train=0.3186  test=0.4722  best=0.4722  15.9s\n","        CV=[0.575, 0.299, 0.687, 0.361]\n","        CM=[0.862, 0.856, 0.845, 0.844]\n","        GR=[9.342, 15.094, 19.425, 21.132]\n","  E 45  ce=3.3307  geo=0.0002  train=0.2916  test=0.4960  best=0.4960  15.9s\n","        CV=[0.557, 0.336, 0.716, 0.463]\n","        CM=[0.858, 0.853, 0.840, 0.838]\n","        GR=[9.219, 14.990, 18.954, 20.568]\n","  E 50  ce=3.3014  geo=0.0002  train=0.2990  test=0.4948  best=0.5077  15.9s\n","        CV=[0.582, 0.367, 0.708, 0.553]\n","        CM=[0.855, 0.849, 0.837, 0.825]\n","        GR=[9.118, 14.246, 19.614, 20.796]\n","  E 55  ce=2.9491  geo=0.0002  train=0.3777  test=0.5449  best=0.5449  15.9s\n","        CV=[0.538, 0.479, 0.728, 0.651]\n","        CM=[0.851, 0.843, 0.827, 0.798]\n","        GR=[9.013, 13.607, 18.208, 19.764]\n","  E 60  ce=3.0646  geo=0.0002  train=0.3604  test=0.5390  best=0.5475  15.9s\n","        CV=[0.554, 0.576, 0.708, 0.620]\n","        CM=[0.845, 0.834, 0.825, 0.774]\n","        GR=[8.928, 13.111, 18.035, 19.635]\n","  E 65  ce=2.9903  geo=0.0002  train=0.3771  test=0.5485  best=0.5500  15.9s\n","        CV=[0.515, 0.610, 0.763, 0.659]\n","        CM=[0.841, 0.826, 0.822, 0.749]\n","        GR=[8.859, 12.565, 17.437, 18.844]\n","  E 70  ce=3.0923  geo=0.0002  train=0.3575  test=0.5486  best=0.5570  15.9s\n","        CV=[0.497, 0.605, 0.755, 0.712]\n","        CM=[0.835, 0.822, 0.817, 0.717]\n","        GR=[8.745, 13.208, 17.593, 18.968]\n","  E 75  ce=2.8893  geo=0.0002  train=0.4065  test=0.5643  best=0.5664  15.9s\n","        CV=[0.509, 0.602, 0.733, 0.656]\n","        CM=[0.831, 0.817, 0.812, 0.712]\n","        GR=[8.611, 12.904, 17.540, 18.776]\n","  E 80  ce=2.7051  geo=0.0002  train=0.4530  test=0.5713  best=0.5854  15.9s\n","        CV=[0.496, 0.681, 0.771, 0.669]\n","        CM=[0.829, 0.809, 0.808, 0.713]\n","        GR=[8.493, 12.319, 17.098, 18.303]\n","  E 85  ce=2.9545  geo=0.0002  train=0.3991  test=0.5862  best=0.5862  15.9s\n","        CV=[0.495, 0.663, 0.702, 0.686]\n","        CM=[0.824, 0.801, 0.808, 0.705]\n","        GR=[8.465, 11.965, 16.710, 18.210]\n","  E 90  ce=2.8081  geo=0.0002  train=0.4362  test=0.5795  best=0.5922  15.9s\n","        CV=[0.487, 0.635, 0.644, 0.648]\n","        CM=[0.819, 0.803, 0.805, 0.703]\n","        GR=[8.346, 12.818, 17.584, 18.952]\n","  E 95  ce=2.5697  geo=0.0002  train=0.4937  test=0.5973  best=0.6033  15.9s\n","        CV=[0.483, 0.640, 0.658, 0.609]\n","        CM=[0.817, 0.792, 0.801, 0.695]\n","        GR=[8.258, 11.401, 15.872, 17.176]\n","  E100  ce=2.7067  geo=0.0002  train=0.4662  test=0.6055  best=0.6076  15.9s\n","        CV=[0.477, 0.625, 0.598, 0.603]\n","        CM=[0.815, 0.790, 0.802, 0.707]\n","        GR=[8.134, 11.059, 16.000, 17.220]\n","  E105  ce=3.0491  geo=0.0002  train=0.3842  test=0.6053  best=0.6204  15.9s\n","        CV=[0.453, 0.626, 0.630, 0.631]\n","        CM=[0.810, 0.777, 0.800, 0.703]\n","        GR=[8.068, 10.362, 15.342, 16.452]\n","  E110  ce=2.7270  geo=0.0002  train=0.4645  test=0.6114  best=0.6204  15.9s\n","        CV=[0.456, 0.632, 0.595, 0.577]\n","        CM=[0.805, 0.774, 0.798, 0.695]\n","        GR=[7.958, 10.222, 14.923, 16.162]\n","  E115  ce=2.6177  geo=0.0002  train=0.4887  test=0.6156  best=0.6237  15.9s\n","        CV=[0.463, 0.616, 0.592, 0.565]\n","        CM=[0.798, 0.771, 0.798, 0.697]\n","        GR=[7.837, 10.057, 15.003, 16.264]\n","  E120  ce=2.6108  geo=0.0002  train=0.4982  test=0.6232  best=0.6268  15.9s\n","        CV=[0.456, 0.620, 0.568, 0.541]\n","        CM=[0.793, 0.767, 0.795, 0.698]\n","        GR=[7.735, 9.763, 14.591, 15.888]\n","  E125  ce=2.5424  geo=0.0002  train=0.5147  test=0.6230  best=0.6291  15.9s\n","        CV=[0.432, 0.632, 0.571, 0.548]\n","        CM=[0.790, 0.767, 0.793, 0.703]\n","        GR=[7.574, 9.682, 14.894, 15.792]\n","  E130  ce=2.5357  geo=0.0002  train=0.5199  test=0.6329  best=0.6373  15.9s\n","        CV=[0.423, 0.597, 0.557, 0.534]\n","        CM=[0.792, 0.756, 0.793, 0.703]\n","        GR=[7.609, 9.011, 14.042, 15.208]\n","  E135  ce=2.5845  geo=0.0002  train=0.5136  test=0.6290  best=0.6373  16.6s\n","        CV=[0.443, 0.579, 0.571, 0.523]\n","        CM=[0.787, 0.751, 0.791, 0.699]\n","        GR=[7.560, 8.984, 13.790, 15.031]\n","  E140  ce=2.4197  geo=0.0002  train=0.5504  test=0.6319  best=0.6383  15.9s\n","        CV=[0.426, 0.593, 0.542, 0.526]\n","        CM=[0.786, 0.751, 0.789, 0.696]\n","        GR=[7.474, 9.391, 14.301, 15.403]\n","  E145  ce=2.5000  geo=0.0002  train=0.5313  test=0.6239  best=0.6383  15.9s\n","        CV=[0.421, 0.558, 0.541, 0.519]\n","        CM=[0.782, 0.745, 0.790, 0.690]\n","        GR=[7.380, 8.630, 13.392, 14.623]\n","  E150  ce=2.3026  geo=0.0002  train=0.5874  test=0.6344  best=0.6444  15.9s\n","        CV=[0.418, 0.557, 0.510, 0.532]\n","        CM=[0.783, 0.741, 0.784, 0.693]\n","        GR=[7.375, 8.594, 13.008, 14.271]\n","  E155  ce=2.3434  geo=0.0002  train=0.5734  test=0.6454  best=0.6487  15.9s\n","        CV=[0.430, 0.539, 0.518, 0.552]\n","        CM=[0.778, 0.740, 0.783, 0.693]\n","        GR=[7.271, 8.651, 13.159, 14.312]\n","  E160  ce=2.2125  geo=0.0002  train=0.6134  test=0.6470  best=0.6487  16.4s\n","        CV=[0.429, 0.526, 0.509, 0.511]\n","        CM=[0.777, 0.739, 0.783, 0.694]\n","        GR=[7.244, 8.680, 13.238, 14.421]\n","  E165  ce=2.1449  geo=0.0002  train=0.6299  test=0.6481  best=0.6526  15.9s\n","        CV=[0.414, 0.505, 0.487, 0.505]\n","        CM=[0.777, 0.737, 0.780, 0.697]\n","        GR=[7.221, 8.553, 12.930, 14.138]\n","  E170  ce=2.2199  geo=0.0002  train=0.6124  test=0.6550  best=0.6550  15.9s\n","        CV=[0.416, 0.521, 0.514, 0.487]\n","        CM=[0.776, 0.736, 0.780, 0.693]\n","        GR=[7.176, 8.430, 12.868, 14.086]\n","  E175  ce=2.4630  geo=0.0002  train=0.5568  test=0.6522  best=0.6581  15.9s\n","        CV=[0.425, 0.506, 0.513, 0.516]\n","        CM=[0.773, 0.736, 0.780, 0.694]\n","        GR=[7.132, 8.750, 13.239, 14.394]\n","  E180  ce=2.3108  geo=0.0002  train=0.5942  test=0.6562  best=0.6591  15.9s\n","        CV=[0.421, 0.520, 0.509, 0.504]\n","        CM=[0.773, 0.735, 0.780, 0.697]\n","        GR=[7.071, 8.788, 13.512, 14.559]\n","  E185  ce=2.4466  geo=0.0002  train=0.5643  test=0.6569  best=0.6591  15.9s\n","        CV=[0.406, 0.505, 0.510, 0.491]\n","        CM=[0.771, 0.734, 0.779, 0.698]\n","        GR=[7.031, 8.665, 13.247, 14.336]\n","  E190  ce=2.2009  geo=0.0002  train=0.6248  test=0.6573  best=0.6641  15.9s\n","        CV=[0.402, 0.504, 0.496, 0.477]\n","        CM=[0.770, 0.735, 0.777, 0.697]\n","        GR=[6.983, 8.707, 13.285, 14.309]\n","  E195  ce=2.3011  geo=0.0002  train=0.5963  test=0.6654  best=0.6654  15.9s\n","        CV=[0.399, 0.510, 0.529, 0.496]\n","        CM=[0.767, 0.730, 0.777, 0.696]\n","        GR=[6.966, 8.302, 12.760, 13.907]\n","  E200  ce=2.1616  geo=0.0002  train=0.6383  test=0.6639  best=0.6701  15.9s\n","        CV=[0.395, 0.503, 0.495, 0.499]\n","        CM=[0.766, 0.733, 0.778, 0.695]\n","        GR=[6.920, 8.509, 13.033, 14.137]\n","  E205  ce=2.3006  geo=0.0002  train=0.6025  test=0.6642  best=0.6701  15.8s\n","        CV=[0.413, 0.499, 0.484, 0.489]\n","        CM=[0.766, 0.734, 0.777, 0.697]\n","        GR=[6.869, 8.414, 12.974, 14.022]\n","  E210  ce=2.1601  geo=0.0002  train=0.6419  test=0.6707  best=0.6707  15.9s\n","        CV=[0.402, 0.478, 0.475, 0.475]\n","        CM=[0.767, 0.731, 0.776, 0.699]\n","        GR=[6.880, 8.239, 12.574, 13.710]\n","  E215  ce=2.4318  geo=0.0002  train=0.5728  test=0.6740  best=0.6740  15.9s\n","        CV=[0.402, 0.490, 0.494, 0.476]\n","        CM=[0.766, 0.732, 0.776, 0.698]\n","        GR=[6.856, 8.328, 12.733, 13.871]\n","  E220  ce=2.2536  geo=0.0002  train=0.6180  test=0.6694  best=0.6740  15.9s\n","        CV=[0.400, 0.493, 0.496, 0.464]\n","        CM=[0.766, 0.731, 0.775, 0.697]\n","        GR=[6.843, 8.182, 12.539, 13.680]\n","  E225  ce=1.9448  geo=0.0002  train=0.7039  test=0.6723  best=0.6743  15.9s\n","        CV=[0.410, 0.503, 0.491, 0.470]\n","        CM=[0.764, 0.732, 0.775, 0.697]\n","        GR=[6.806, 8.205, 12.540, 13.648]\n","  E230  ce=2.2531  geo=0.0002  train=0.6200  test=0.6659  best=0.6743  15.9s\n","        CV=[0.391, 0.488, 0.484, 0.496]\n","        CM=[0.764, 0.732, 0.775, 0.698]\n","        GR=[6.796, 8.189, 12.498, 13.645]\n","  E235  ce=2.0302  geo=0.0002  train=0.6817  test=0.6768  best=0.6768  15.9s\n","        CV=[0.394, 0.478, 0.488, 0.475]\n","        CM=[0.765, 0.731, 0.775, 0.698]\n","        GR=[6.790, 8.268, 12.532, 13.687]\n","  E240  ce=2.0904  geo=0.0002  train=0.6646  test=0.6779  best=0.6791  15.9s\n","        CV=[0.399, 0.499, 0.491, 0.489]\n","        CM=[0.766, 0.732, 0.775, 0.698]\n","        GR=[6.776, 8.351, 12.627, 13.754]\n","  E245  ce=2.0636  geo=0.0002  train=0.6727  test=0.6810  best=0.6810  15.9s\n","        CV=[0.399, 0.471, 0.483, 0.470]\n","        CM=[0.766, 0.733, 0.775, 0.699]\n","        GR=[6.769, 8.416, 12.683, 13.820]\n","  E250  ce=2.1905  geo=0.0002  train=0.6409  test=0.6735  best=0.6810  15.9s\n","        CV=[0.388, 0.478, 0.482, 0.470]\n","        CM=[0.765, 0.733, 0.775, 0.699]\n","        GR=[6.763, 8.403, 12.625, 13.781]\n","  E255  ce=2.1166  geo=0.0002  train=0.6533  test=0.6761  best=0.6810  15.9s\n","        CV=[0.403, 0.496, 0.482, 0.473]\n","        CM=[0.765, 0.733, 0.775, 0.699]\n","        GR=[6.745, 8.426, 12.659, 13.797]\n","  E260  ce=2.4379  geo=0.0002  train=0.5765  test=0.6751  best=0.6810  15.9s\n","        CV=[0.391, 0.481, 0.471, 0.462]\n","        CM=[0.765, 0.734, 0.775, 0.700]\n","        GR=[6.740, 8.436, 12.640, 13.812]\n","  E265  ce=1.9339  geo=0.0002  train=0.7044  test=0.6797  best=0.6810  15.9s\n","        CV=[0.408, 0.484, 0.472, 0.468]\n","        CM=[0.765, 0.734, 0.775, 0.701]\n","        GR=[6.731, 8.445, 12.639, 13.808]\n","  E270  ce=1.9643  geo=0.0002  train=0.6972  test=0.6790  best=0.6810  15.9s\n","        CV=[0.400, 0.478, 0.475, 0.472]\n","        CM=[0.765, 0.734, 0.775, 0.701]\n","        GR=[6.731, 8.440, 12.639, 13.814]\n","  E275  ce=1.8407  geo=0.0002  train=0.7244  test=0.6816  best=0.6816  15.9s\n","        CV=[0.402, 0.483, 0.465, 0.474]\n","        CM=[0.765, 0.734, 0.775, 0.701]\n","        GR=[6.727, 8.453, 12.652, 13.831]\n","  E280  ce=2.0719  geo=0.0002  train=0.6715  test=0.6803  best=0.6825  15.9s\n","        CV=[0.399, 0.476, 0.468, 0.478]\n","        CM=[0.765, 0.734, 0.775, 0.703]\n","        GR=[6.726, 8.450, 12.636, 13.820]\n","  E285  ce=2.0724  geo=0.0002  train=0.6719  test=0.6815  best=0.6825  15.9s\n","        CV=[0.399, 0.478, 0.467, 0.461]\n","        CM=[0.765, 0.734, 0.775, 0.703]\n","        GR=[6.726, 8.428, 12.597, 13.790]\n","  E290  ce=2.0148  geo=0.0002  train=0.6892  test=0.6803  best=0.6825  15.9s\n","        CV=[0.404, 0.498, 0.477, 0.456]\n","        CM=[0.765, 0.734, 0.775, 0.703]\n","        GR=[6.727, 8.433, 12.601, 13.798]\n","  E295  ce=1.9271  geo=0.0002  train=0.7067  test=0.6810  best=0.6825  15.8s\n","        CV=[0.400, 0.482, 0.469, 0.488]\n","        CM=[0.765, 0.734, 0.775, 0.703]\n","        GR=[6.726, 8.433, 12.602, 13.799]\n","  E299  ce=1.9688  geo=0.0002  train=0.7002  test=0.6802  best=0.6825  15.9s\n","        CV=[0.393, 0.492, 0.474, 0.480]\n","        CM=[0.765, 0.734, 0.775, 0.703]\n","        GR=[6.726, 8.437, 12.607, 13.803]\n","\n","════════════════════════════════════════════════════════════\n","  CIFAR-100 RESULTS (CM-Validated)\n","════════════════════════════════════════════════════════════\n","  Best test accuracy: 0.6825 (68.25%)\n","  Parameters: 32,970,280\n","  Checkpoint: geo_cifar100/best.pt\n","  TensorBoard: runs/geo_cifar100_3\n","\n","  Final geometric state:\n","\n","  Final anchor diagnostics:\n","    layer_0: cv=0.1942, cm_pos=0.688, cm_mean=-0.7441\n","    layer_1: cv=0.1957, cm_pos=0.562, cm_mean=-5.3468\n","    layer_2: cv=0.1951, cm_pos=0.680, cm_mean=-1.2126\n","    layer_3: cv=0.1961, cm_pos=0.734, cm_mean=-2.4622\n","\n","Done.\n"]}]},{"cell_type":"markdown","source":["# transformer v3"],"metadata":{"id":"FtkQmzOpAP1p"}},{"cell_type":"code","source":["\"\"\"\n","Geometric Transformer — CIFAR-100 Training with Three-Domain Observer Loss\n","\n","Full observer loss spectrum wired through the CM-gated geometric pipeline:\n","  EXTERNAL:   CE (paired views) + embedding NCE\n","  GEOMETRIC:  patchwork NCE + bridge loss  ← flows through CM gate\n","  INTERNAL:   assign BCE + assign NCE + tri NCE + attraction + CV + spread\n","\n","Changes from previous version:\n","  - Paired augmentation replaces CutMix (two views of each image)\n","  - forward_paired() → compute_loss() replaces single-view + proxy losses\n","  - observer_loss from geolip_core handles the full loss spectrum\n","  - Bridge loss now trained (was computed but never used before)\n","  - Real pentachoron CV replaces pairwise distance proxy\n","  - Plain Adam replaces AdamW (AdamW fights geometric structure)\n","  - NCE memory bank removed (observer_loss handles its own NCE terms)\n","\n","TensorBoard logging preserved from previous version.\n","\n","Requires geolip_core with updated geometric_transformer.py\n","(forward_paired + compute_loss methods).\n","\n","!pip install geolip-core torchvision tqdm tensorboard\n","\"\"\"\n","\n","import torch\n","import torch.nn as nn\n","import torch.nn.functional as F\n","import numpy as np\n","import time, json, math\n","from pathlib import Path\n","from tqdm.auto import tqdm\n","from torch.utils.tensorboard import SummaryWriter\n","\n","device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n","print(f\"Device: {device}\")\n","if device.type == 'cuda':\n","    print(f\"  GPU: {torch.cuda.get_device_name()}\")\n","    print(f\"  VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB\")\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# IMPORT TRANSFORMER\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# IMPORTS\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","from geolip_core.pipeline.components.geometric_transformer import (\n","    GeometricTransformer, GeometricTransformerLayer,\n","    CayleyOrthogonal, QuaternionCompose, FiLMLayer,\n","    ContentAttention, GeometricAttention, CMValidatedGate,\n","    anchor_neighborhood_cm,\n",")\n","from geolip_core.pipeline.observer import TorchComponent\n","\n","torch.set_float32_matmul_precision('high')\n","torch.backends.cudnn.benchmark = True\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# CONFIG\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","CONFIG = {\n","    # Model\n","    'd_model': 384,\n","    'n_heads': 8,\n","    'n_layers': 4,\n","    'n_anchors': 128,\n","    'manifold_dim': 128,\n","    'n_comp': 8,\n","    'd_comp': 32,\n","    'context_dim': 128,\n","    'quat_dim': 64,\n","    'dropout': 0.1,\n","    'cm_neighbors': 3,       # CM simplex neighbors\n","\n","    # Input stage\n","    'patch_size': 4,\n","    'img_size': 32,\n","    'in_channels': 3,\n","    'conv_channels': 64,\n","    'svd_rank': 16,\n","\n","    # Training\n","    'epochs': 100,\n","    'batch_size': 128,           # 2 views concatenated → 1024 effective per forward\n","    'lr': 3e-4,\n","    'weight_decay': 0.005,        # plain Adam — AdamW fights geometric structure\n","    'warmup_epochs': 5,\n","    'label_smoothing': 0.1,\n","    'num_workers': 8,\n","\n","    # Observer loss weights (three_domain_loss canonical recipe)\n","    'cv_target': 0.20,           # pentachoron band center\n","    'w_ce': 0.8,\n","    'w_nce_emb': 0.5,\n","    'w_nce_pw': 1.0,             # patchwork NCE — through gated pipeline\n","    'w_bridge': 1.0,             # bridge loss — through gated pipeline\n","    'w_assign': 0.5,\n","    'w_assign_nce': 0.25,\n","    'w_nce_tri': 0.5,\n","    'w_attract': 0.25,\n","    'w_cv': 0.01,\n","    'w_spread': 0.01,\n","\n","    # Augmentation\n","    'random_erasing_p': 0.25,\n","\n","    # Data\n","    'num_classes': 100,\n","\n","    # Logging\n","    'log_geo_every': 5,      # full geometric analysis every N epochs\n","    'log_grads_every': 10,   # gradient norms every N epochs\n","    'log_dir': 'runs/geo_observer_3domain',\n","}\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# INPUT STAGE\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","from geolip_core.core.input.svd import SVDObserver\n","\n","\n","class ConvSVDPatchEmbedding(TorchComponent):\n","    \"\"\"Input stage: conv frontend → SVDObserver → patch tokens.\"\"\"\n","    def __init__(self, name, img_size=32, patch_size=4, in_channels=3,\n","                 conv_channels=64, d_model=256, svd_rank=16):\n","        super().__init__(name)\n","        self.patch_size = patch_size\n","        self.n_patches = (img_size // patch_size) ** 2\n","        self.d_model = d_model\n","        self.svd_rank = svd_rank\n","\n","        self.conv_frontend = nn.Sequential(\n","            nn.Conv2d(in_channels, conv_channels, 3, padding=1, bias=False),\n","            nn.BatchNorm2d(conv_channels), nn.GELU(),\n","            nn.Conv2d(conv_channels, conv_channels, 3, padding=1, bias=False),\n","            nn.BatchNorm2d(conv_channels), nn.GELU(),\n","        )\n","        self.svd_observer = SVDObserver(conv_channels, svd_rank)\n","        self.patch_proj = nn.Conv2d(\n","            conv_channels, d_model, kernel_size=patch_size,\n","            stride=patch_size, bias=False)\n","        self.patch_norm = nn.LayerNorm(d_model)\n","\n","        svd_feat_dim = self.svd_observer.feature_dim\n","        self.svd_to_gamma = nn.Linear(svd_feat_dim, d_model)\n","        self.svd_to_beta = nn.Linear(svd_feat_dim, d_model)\n","        nn.init.zeros_(self.svd_to_gamma.weight); nn.init.ones_(self.svd_to_gamma.bias)\n","        nn.init.zeros_(self.svd_to_beta.weight); nn.init.zeros_(self.svd_to_beta.bias)\n","\n","        self.cls_token = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)\n","        self.pos_embed = nn.Parameter(\n","            torch.randn(1, self.n_patches + 1, d_model) * 0.02)\n","\n","    @torch.compiler.disable\n","    def _run_svd(self, feat):\n","        \"\"\"SVD decomposition — excluded from torch.compile because\n","        torch.linalg.eigh is not supported in CUDA graph capture.\"\"\"\n","        S, Vh, svd_features, novelty = self.svd_observer(feat)\n","        if self.training:\n","            self.svd_observer.update_ema(S, Vh)\n","        return S, Vh, svd_features, novelty\n","\n","    def forward(self, x):\n","        B = x.shape[0]\n","        feat = self.conv_frontend(x)\n","        S, Vh, svd_features, novelty = self._run_svd(feat)\n","        tokens = self.patch_proj(feat)\n","        tokens = tokens.flatten(2).transpose(1, 2)\n","        tokens = self.patch_norm(tokens)\n","        gamma = self.svd_to_gamma(svd_features).unsqueeze(1)\n","        beta = self.svd_to_beta(svd_features).unsqueeze(1)\n","        tokens = gamma * tokens + beta\n","        cls = self.cls_token.expand(B, -1, -1)\n","        tokens = torch.cat([cls, tokens], dim=1)\n","        tokens = tokens + self.pos_embed\n","        svd_state = {\n","            'singular_values': S, 'Vh': Vh,\n","            'svd_features': svd_features, 'novelty': novelty,\n","        }\n","        return tokens, svd_state\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# CLASSIFIER ( uses GeometricTransformer with CM gates)\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","from geofractal.router.wide_router import WideRouter\n","\n","\n","class ClassificationHead(TorchComponent):\n","    \"\"\"Task head as a proper TorchComponent.\"\"\"\n","    def __init__(self, name, d_model, num_classes, dropout=0.1):\n","        super().__init__(name)\n","        self.net = nn.Sequential(\n","            nn.LayerNorm(d_model),\n","            nn.Linear(d_model, d_model),\n","            nn.GELU(), nn.Dropout(dropout),\n","            nn.Linear(d_model, num_classes),\n","        )\n","\n","    def forward(self, x):\n","        return self.net(x)\n","\n","\n","class GeoViTClassifier(WideRouter):\n","    \"\"\"Geometric Vision Transformer for classification.\n","\n","    WideRouter orchestrator managing:\n","        - patch_embed: ConvSVDPatchEmbedding (input stage)\n","        - transformer: GeometricTransformer (registered tower — 8 layers)\n","        - head: ClassificationHead (task-specific)\n","\n","    forward(x) for inference, forward_paired(v1,v2) + compute_loss() for training.\n","    Call prepare_for_compile() before torch.compile() for optimal graph capture.\n","    \"\"\"\n","    def __init__(self, name, config):\n","        super().__init__(name, strict=False)\n","        self.config = config\n","\n","        self.attach('patch_embed', ConvSVDPatchEmbedding(\n","            'patch_embed', img_size=config['img_size'],\n","            patch_size=config['patch_size'], in_channels=config['in_channels'],\n","            conv_channels=config['conv_channels'], d_model=config['d_model'],\n","            svd_rank=config['svd_rank'],\n","        ))\n","        self.attach('transformer', GeometricTransformer(\n","            'geo_cifar', d_model=config['d_model'], n_heads=config['n_heads'],\n","            n_layers=config['n_layers'], n_anchors=config['n_anchors'],\n","            manifold_dim=config['manifold_dim'], n_comp=config['n_comp'],\n","            d_comp=config['d_comp'], context_dim=config['context_dim'],\n","            quat_dim=config['quat_dim'], dropout=config['dropout'],\n","            cm_neighbors=config.get('cm_neighbors', 3),\n","            nce_bank_size=0,\n","        ))\n","        self.register_tower('transformer')\n","\n","        self.attach('head', ClassificationHead(\n","            'head', config['d_model'], config['num_classes'],\n","            config.get('dropout', 0.1)))\n","\n","    def forward(self, x, x2=None, targets=None, return_geo_state=False):\n","        \"\"\"Unified entry point — compiled by WideRouter.compile().\n","\n","        Training:  model(v1, v2, targets=labels) → (loss, loss_dict)\n","        Inference: model(images)                  → logits\n","        Analysis:  model_raw(images, return_geo_state=True) → (logits, geo_states, svd)\n","        \"\"\"\n","        if x2 is not None:\n","            # ── Training: paired forward through compiled path ──\n","            B = x.shape[0]\n","            v_cat = torch.cat([x, x2], dim=0)\n","            tokens_cat, svd_state = self['patch_embed'](v_cat)\n","            output = self['transformer'].forward_paired(\n","                tokens_cat[:B], tokens_cat[B:])\n","            if targets is not None:\n","                return self._compute_loss(output, targets)\n","            return output\n","\n","        # ── Inference / analysis: single view ──\n","        tokens, svd_state = self['patch_embed'](x)\n","        if return_geo_state:\n","            features, geo_states = self['transformer'](tokens, return_geo_state=True)\n","        else:\n","            features = self['transformer'](tokens)\n","        cls_out = features[:, 0]\n","        logits = self['head'](cls_out)\n","        if return_geo_state:\n","            return logits, geo_states, svd_state\n","        return logits\n","\n","    def _compute_loss(self, output, targets):\n","        \"\"\"Three-domain observer loss — called from forward().\"\"\"\n","        loss_kwargs = {}\n","        for k in ['w_nce_emb', 'w_nce_pw', 'w_bridge', 'w_assign',\n","                   'w_assign_nce', 'w_nce_tri', 'w_attract', 'w_cv', 'w_spread']:\n","            if k in self.config:\n","                loss_kwargs[k] = self.config[k]\n","        if 'cv_target' in self.config:\n","            loss_kwargs['cv_target'] = self.config['cv_target']\n","        return self['transformer'].compute_loss(\n","            output, targets,\n","            w_ce=self.config.get('w_ce', 1.0),\n","            head=self['head'],\n","            **loss_kwargs)\n","\n","    def geometric_losses(self):\n","        \"\"\"Proxy geometric regularization — for backward compat / analysis.\"\"\"\n","        return self['transformer'].geometric_losses(\n","            cv_target=self.config.get('cv_target', 0.22))\n","\n","    def anchor_diagnostics(self):\n","        return self['transformer'].anchor_diagnostics()\n","\n","    def invalidate_caches(self):\n","        \"\"\"Invalidate CM gate caches + tower caches. Call after optimizer.step().\"\"\"\n","        self['transformer'].invalidate_caches()\n","        self.clear_tower_caches()\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# GEOMETRIC ANALYSIS BATTERY ( includes CM diagnostics)\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","@torch.no_grad()\n","def compute_cv(points):\n","    \"\"\"Coefficient of variation on S^(d-1).\n","    CV = std(pairwise_cosine_distances) / mean(pairwise_cosine_distances)\n","    Pentachoron band: CV ∈ [0.20, 0.23].\n","    \"\"\"\n","    points = F.normalize(points.float(), dim=-1)\n","    cos_sim = points @ points.T\n","    n = points.shape[0]\n","    idx = torch.triu_indices(n, n, offset=1, device=points.device)\n","    pairwise_dist = 1.0 - cos_sim[idx[0], idx[1]]\n","    mean_d = pairwise_dist.mean()\n","    std_d = pairwise_dist.std()\n","    cv = (std_d / (mean_d + 1e-8)).item()\n","    return cv, mean_d.item(), std_d.item()\n","\n","\n","@torch.no_grad()\n","def log_geometric_analysis(model, writer, epoch, test_loader, device, config):\n","    \"\"\"Full geometric analysis battery with CM diagnostics.\"\"\"\n","    model.eval()\n","\n","    images, labels = next(iter(test_loader))\n","    images = images[:min(64, images.shape[0])].to(device)\n","    labels = labels[:min(64, labels.shape[0])].to(device)\n","\n","    logits, geo_states, svd_state = model(images, return_geo_state=True)\n","\n","    n_layers = len(geo_states)\n","    pred = logits.argmax(1)\n","    batch_acc = (pred == labels).float().mean().item()\n","    writer.add_scalar('analysis/batch_accuracy', batch_acc, epoch)\n","\n","    # ─── SVD Input Stage ───\n","    S = svd_state['singular_values']\n","    s_norm = S / (S.sum(dim=-1, keepdim=True) + 1e-8)\n","    s_ent = -(s_norm * torch.log(s_norm.clamp(min=1e-8))).sum(-1)\n","    novelty = svd_state['novelty']\n","\n","    writer.add_scalar('svd/entropy_mean', s_ent.mean().item(), epoch)\n","    writer.add_scalar('svd/entropy_std', s_ent.std().item(), epoch)\n","    writer.add_scalar('svd/novelty_norm', novelty.norm(dim=-1).mean().item(), epoch)\n","    writer.add_scalar('svd/top1_ratio', (S[:, 0] / (S.sum(-1) + 1e-8)).mean().item(), epoch)\n","    writer.add_scalar('svd/condition_number',\n","        (S[:, 0] / (S[:, -1].clamp(min=1e-8))).mean().item(), epoch)\n","    for k in range(min(S.shape[1], 5)):\n","        writer.add_scalar(f'svd/S_{k}', S[:, k].mean().item(), epoch)\n","\n","    # SVD FiLM deviation\n","    pe = model['patch_embed']\n","    writer.add_scalar('svd_film/gamma_weight_norm', pe.svd_to_gamma.weight.data.norm().item(), epoch)\n","    writer.add_scalar('svd_film/gamma_bias_dev_from_1',\n","        (pe.svd_to_gamma.bias.data - 1.0).abs().mean().item(), epoch)\n","    writer.add_scalar('svd_film/beta_weight_norm', pe.svd_to_beta.weight.data.norm().item(), epoch)\n","    writer.add_scalar('svd_film/beta_bias_norm', pe.svd_to_beta.bias.data.abs().mean().item(), epoch)\n","\n","    # ─── Anchor Diagnostics (built-in) ───\n","    anchor_diag = model.anchor_diagnostics()\n","    for layer_name, d in anchor_diag.items():\n","        for k, v in d.items():\n","            writer.add_scalar(f'anchor_diag/{layer_name}_{k}', v, epoch)\n","\n","    # ─── Per-Layer Geometric Analysis ───\n","    for i, gs in enumerate(geo_states):\n","        prefix = f'layer_{i}'\n","\n","        # === CV — pentachoron band metric ===\n","        emb = gs['embedding']\n","        # Anchor CV\n","        transformer = model['transformer']\n","        layer = transformer[f'layer_{i}']\n","        anchors = F.normalize(\n","            layer['observer'].association.constellation.anchors, dim=-1)\n","        cv_anchors, mean_d_anchors, std_d_anchors = compute_cv(anchors)\n","        writer.add_scalar(f'{prefix}/cv_anchors', cv_anchors, epoch)\n","        writer.add_scalar(f'{prefix}/anchor_mean_dist', mean_d_anchors, epoch)\n","        writer.add_scalar(f'{prefix}/anchor_std_dist', std_d_anchors, epoch)\n","\n","        # Embedding CV\n","        emb_flat = emb.reshape(-1, emb.shape[-1])\n","        cv_emb, mean_d_emb, std_d_emb = compute_cv(emb_flat)\n","        writer.add_scalar(f'{prefix}/cv_embeddings', cv_emb, epoch)\n","        writer.add_scalar(f'{prefix}/embedding_mean_dist', mean_d_emb, epoch)\n","\n","        # === CM Gate Diagnostics ===\n","        gate_info = gs.get('gate_info', {})\n","        gate_values = gs.get('gate_values')\n","        cm_quality = gs.get('cm_quality')\n","\n","        if gate_info:\n","            writer.add_scalar(f'{prefix}/cm_active_anchors',\n","                gate_info.get('active', 0), epoch)\n","            writer.add_scalar(f'{prefix}/cm_gate_mean',\n","                gate_info.get('gate_mean', 0), epoch)\n","            writer.add_scalar(f'{prefix}/cm_positive_frac',\n","                gate_info.get('cm_positive_frac', 0), epoch)\n","\n","        if gate_values is not None:\n","            gv = gate_values\n","            writer.add_scalar(f'{prefix}/gate_values_min', gv.min().item(), epoch)\n","            writer.add_scalar(f'{prefix}/gate_values_max', gv.max().item(), epoch)\n","            writer.add_scalar(f'{prefix}/gate_values_std', gv.std().item(), epoch)\n","            # Per-anchor gate mean (which anchors are consistently open/closed)\n","            gv_per_anchor = gv.mean(dim=0).mean(dim=0)  # average over B and L\n","            writer.add_scalar(f'{prefix}/gate_anchor_spread',\n","                gv_per_anchor.std().item(), epoch)\n","            # Fraction of positions with >50% anchors open\n","            if gv.dim() == 3:\n","                pos_open_frac = (gv.mean(dim=-1) > 0.5).float().mean().item()\n","            else:\n","                pos_open_frac = (gv > 0.5).float().mean().item()\n","            writer.add_scalar(f'{prefix}/gate_positions_open_frac', pos_open_frac, epoch)\n","\n","        if cm_quality is not None:\n","            writer.add_scalar(f'{prefix}/cm_quality_mean', cm_quality.mean().item(), epoch)\n","            writer.add_scalar(f'{prefix}/cm_quality_std', cm_quality.std().item(), epoch)\n","            writer.add_scalar(f'{prefix}/cm_quality_min', cm_quality.min().item(), epoch)\n","\n","        # === Stream Agreement ===\n","        content = gs['content']\n","        geometric = gs['geometric']\n","        agreement = F.cosine_similarity(\n","            content.reshape(-1, content.shape[-1]),\n","            geometric.reshape(-1, geometric.shape[-1]), dim=-1)\n","        writer.add_scalar(f'{prefix}/stream_agreement_mean', agreement.mean().item(), epoch)\n","        writer.add_scalar(f'{prefix}/stream_agreement_std', agreement.std().item(), epoch)\n","\n","        writer.add_scalar(f'{prefix}/content_norm', content.norm(dim=-1).mean().item(), epoch)\n","        writer.add_scalar(f'{prefix}/geometric_norm', geometric.norm(dim=-1).mean().item(), epoch)\n","\n","        # === Disagreement arm analysis ===\n","        disagree = content - geometric\n","        agree = content * geometric\n","        writer.add_scalar(f'{prefix}/disagree_norm', disagree.norm(dim=-1).mean().item(), epoch)\n","        writer.add_scalar(f'{prefix}/agree_norm', agree.norm(dim=-1).mean().item(), epoch)\n","\n","        # === Anchor Utilization ===\n","        tri = gs['triangulation']\n","        assignment = gs['assignment']\n","        nearest = gs['nearest']\n","        n_anchors = tri.shape[-1]\n","\n","        nearest_flat = nearest.reshape(-1)\n","        counts = torch.bincount(nearest_flat, minlength=n_anchors).float()\n","        total_assignments = counts.sum()\n","\n","        probs = counts / (total_assignments + 1e-8)\n","        anchor_entropy = -(probs * torch.log(probs.clamp(min=1e-8))).sum().item()\n","        max_entropy = math.log(n_anchors)\n","        writer.add_scalar(f'{prefix}/anchor_entropy_normalized',\n","            anchor_entropy / (max_entropy + 1e-8), epoch)\n","        active = (counts > 0).sum().item()\n","        writer.add_scalar(f'{prefix}/anchors_active', active, epoch)\n","        writer.add_scalar(f'{prefix}/anchors_active_frac', active / n_anchors, epoch)\n","        dead = (counts == 0).sum().item()\n","        writer.add_scalar(f'{prefix}/anchors_dead', dead, epoch)\n","\n","        # === Triangulation Statistics ===\n","        writer.add_scalar(f'{prefix}/tri_mean', tri.mean().item(), epoch)\n","        writer.add_scalar(f'{prefix}/tri_std', tri.std().item(), epoch)\n","\n","        # === Soft Assignment Statistics ===\n","        assign_ent = -(assignment * torch.log(assignment.clamp(min=1e-8))).sum(-1)\n","        writer.add_scalar(f'{prefix}/assignment_entropy_mean', assign_ent.mean().item(), epoch)\n","        writer.add_scalar(f'{prefix}/assignment_max_prob',\n","            assignment.max(dim=-1).values.mean().item(), epoch)\n","\n","        # === Patchwork Statistics (now from CM-validated triangulation) ===\n","        pw = gs['patchwork']\n","        writer.add_scalar(f'{prefix}/patchwork_norm', pw.norm(dim=-1).mean().item(), epoch)\n","        writer.add_scalar(f'{prefix}/patchwork_std', pw.std().item(), epoch)\n","        pw_sparsity = (pw.abs() < 0.01).float().mean().item()\n","        writer.add_scalar(f'{prefix}/patchwork_sparsity', pw_sparsity, epoch)\n","\n","        # === Bridge Consistency ===\n","        bridge = gs['bridge']\n","        bridge_soft = F.softmax(bridge, dim=-1)\n","        bridge_assign_kl = F.kl_div(\n","            bridge_soft.log().reshape(-1, n_anchors),\n","            assignment.reshape(-1, n_anchors),\n","            reduction='batchmean', log_target=False)\n","        writer.add_scalar(f'{prefix}/bridge_assignment_kl', bridge_assign_kl.item(), epoch)\n","\n","        # === Quaternion Composition ===\n","        composed = gs['composed']\n","        writer.add_scalar(f'{prefix}/composed_norm', composed.norm(dim=-1).mean().item(), epoch)\n","\n","        # === Geo Context ===\n","        geo_ctx = gs['geo_ctx']\n","        writer.add_scalar(f'{prefix}/geo_ctx_norm', geo_ctx.norm(dim=-1).mean().item(), epoch)\n","\n","        # === Geometric Residual Stream (CM-conditioned) ===\n","        geo_res = gs.get('geo_residual')\n","        if geo_res is not None:\n","            res_norms = geo_res.norm(dim=-1)\n","            writer.add_scalar(f'{prefix}/geo_res_norm', res_norms.mean().item(), epoch)\n","            writer.add_scalar(f'{prefix}/geo_res_std', geo_res.std().item(), epoch)\n","            writer.add_scalar(f'{prefix}/geo_res_sparsity',\n","                (geo_res.abs() < 0.01).float().mean().item(), epoch)\n","            # Cross-position consistency\n","            geo_res_flat = geo_res.reshape(-1, geo_res.shape[-1])\n","            normed = F.normalize(geo_res_flat, dim=-1)\n","            cos_mat = normed @ normed.T\n","            n_pts = normed.shape[0]\n","            triu = torch.triu_indices(n_pts, n_pts, offset=1, device=geo_res.device)\n","            writer.add_scalar(f'{prefix}/geo_res_consistency',\n","                cos_mat[triu[0], triu[1]].mean().item(), epoch)\n","\n","    # ─── Cayley Rotation Analysis ───\n","    for name, mod in model.named_modules():\n","        if isinstance(mod, CayleyOrthogonal):\n","            R = mod.get_rotation()\n","            I = torch.eye(R.shape[0], device=R.device)\n","            r_dist = (R - I).norm().item()\n","            clean_name = name.replace('.', '_')\n","            writer.add_scalar(f'cayley/{clean_name}_R_minus_I', r_dist, epoch)\n","\n","    # ─── FiLM Layer Analysis ───\n","    film_idx = 0\n","    for name, mod in model.named_modules():\n","        if isinstance(mod, FiLMLayer):\n","            g_b = mod.to_gamma.bias.data\n","            b_b = mod.to_beta.bias.data\n","            writer.add_scalar(f'film/{film_idx}_gamma_dev',\n","                (g_b - 1.0).abs().mean().item(), epoch)\n","            writer.add_scalar(f'film/{film_idx}_beta_dev',\n","                b_b.abs().mean().item(), epoch)\n","            film_idx += 1\n","\n","    # ─── Cross-Layer Trajectories ───\n","    cv_trajectory = []\n","    cm_quality_trajectory = []\n","    res_norms = []\n","    bridge_kls = []\n","\n","    for i, gs in enumerate(geo_states):\n","        # CV\n","        emb = gs['embedding']\n","        emb_flat = emb.reshape(-1, emb.shape[-1])\n","        cv, _, _ = compute_cv(emb_flat)\n","        cv_trajectory.append(cv)\n","\n","        # CM quality\n","        cm_q = gs.get('cm_quality')\n","        if cm_q is not None:\n","            cm_quality_trajectory.append(cm_q.mean().item())\n","\n","        # Geo residual norms\n","        geo_res = gs.get('geo_residual')\n","        if geo_res is not None:\n","            res_norms.append(geo_res.norm(dim=-1).mean().item())\n","\n","        # Bridge KL\n","        n_anchors = gs['assignment'].shape[-1]\n","        bridge_soft = F.softmax(gs['bridge'], dim=-1)\n","        bkl = F.kl_div(\n","            bridge_soft.log().reshape(-1, n_anchors),\n","            gs['assignment'].reshape(-1, n_anchors),\n","            reduction='batchmean', log_target=False).item()\n","        bridge_kls.append(bkl)\n","\n","    # CV trajectory\n","    writer.add_scalar('cv/trajectory_mean', np.mean(cv_trajectory), epoch)\n","    writer.add_scalar('cv/trajectory_std', np.std(cv_trajectory), epoch)\n","    in_band = sum(1 for cv in cv_trajectory if 0.20 <= cv <= 0.23)\n","    writer.add_scalar('cv/layers_in_pentachoron_band', in_band, epoch)\n","    writer.add_scalar('cv/layers_in_band_frac', in_band / len(cv_trajectory), epoch)\n","\n","    # CM quality trajectory\n","    if cm_quality_trajectory:\n","        writer.add_scalar('cm/quality_trajectory_mean',\n","            np.mean(cm_quality_trajectory), epoch)\n","        writer.add_scalar('cm/quality_trajectory_std',\n","            np.std(cm_quality_trajectory), epoch)\n","        writer.add_scalar('cm/quality_min_layer',\n","            np.min(cm_quality_trajectory), epoch)\n","        writer.add_scalar('cm/quality_max_layer',\n","            np.max(cm_quality_trajectory), epoch)\n","\n","    # Geometric residual trajectory\n","    if res_norms:\n","        writer.add_scalar('geo_res/trajectory_start', res_norms[0], epoch)\n","        writer.add_scalar('geo_res/trajectory_end', res_norms[-1], epoch)\n","        writer.add_scalar('geo_res/accumulation_ratio',\n","            res_norms[-1] / (res_norms[0] + 1e-8), epoch)\n","        growth = [res_norms[j+1] - res_norms[j] for j in range(len(res_norms)-1)]\n","        writer.add_scalar('geo_res/growth_mean', np.mean(growth), epoch)\n","        writer.add_scalar('geo_res/growth_std', np.std(growth), epoch)\n","\n","        # Cooperation analysis (includes CM quality)\n","        if len(res_norms) >= 4:\n","            cv_corr = float(np.corrcoef(res_norms, cv_trajectory)[0, 1])\n","            bkl_corr = float(np.corrcoef(res_norms, bridge_kls)[0, 1])\n","            writer.add_scalar('cooperation/geo_res_vs_cv', cv_corr, epoch)\n","            writer.add_scalar('cooperation/geo_res_vs_bridge_kl', bkl_corr, epoch)\n","\n","            if len(cm_quality_trajectory) == len(res_norms):\n","                cm_corr = float(np.corrcoef(\n","                    res_norms, cm_quality_trajectory)[0, 1])\n","                writer.add_scalar('cooperation/geo_res_vs_cm_quality', cm_corr, epoch)\n","                # CM vs CV: do layers with better CM quality also have better CV?\n","                cm_cv_corr = float(np.corrcoef(\n","                    cm_quality_trajectory, cv_trajectory)[0, 1])\n","                writer.add_scalar('cooperation/cm_quality_vs_cv', cm_cv_corr, epoch)\n","\n","    return {\n","        'batch_acc': batch_acc,\n","        'cv_trajectory': cv_trajectory,\n","        'cm_quality_trajectory': cm_quality_trajectory,\n","        'res_norms': res_norms,\n","        'bridge_kls': bridge_kls,\n","    }\n","\n","\n","@torch.no_grad()\n","def log_gradient_norms(model, writer, epoch):\n","    \"\"\"Log gradient norms per component type (includes cm_gate).\"\"\"\n","    type_grads = {}\n","    for name, param in model.named_parameters():\n","        if param.grad is not None:\n","            grad_norm = param.grad.norm().item()\n","            if 'projection' in name and 'proj' in name:\n","                key = 'manifold_proj'\n","            elif 'cm_gate' in name:\n","                key = 'cm_gate'\n","            elif 'observer' in name or 'constellation' in name or 'anchor' in name:\n","                key = 'constellation'\n","            elif 'context' in name:\n","                key = 'geo_context'\n","            elif 'content' in name:\n","                key = 'content_attn'\n","            elif 'geometric' in name and 'film' not in name:\n","                key = 'geo_attn'\n","            elif 'film' in name:\n","                key = 'film'\n","            elif 'rotation' in name or 'cayley' in name or 'A_upper' in name:\n","                key = 'cayley'\n","            elif 'compose' in name or 'quat' in name or 'proj_w' in name:\n","                key = 'quaternion'\n","            elif 'decode' in name:\n","                key = 'decode'\n","            elif 'gate' in name:\n","                key = 'gate'\n","            elif 'conv' in name or 'patch' in name:\n","                key = 'input_stage'\n","            elif 'head' in name:\n","                key = 'head'\n","            elif 'svd' in name:\n","                key = 'svd'\n","            elif 'geo_proj' in name:\n","                key = 'geo_residual_proj'\n","            else:\n","                key = 'other'\n","\n","            if key not in type_grads:\n","                type_grads[key] = []\n","            type_grads[key].append(grad_norm)\n","\n","    for key, norms in type_grads.items():\n","        writer.add_scalar(f'grad_norm/{key}_mean', np.mean(norms), epoch)\n","        writer.add_scalar(f'grad_norm/{key}_max', np.max(norms), epoch)\n","\n","    total = sum(p.grad.norm().item() ** 2\n","                for p in model.parameters() if p.grad is not None) ** 0.5\n","    writer.add_scalar('grad_norm/total', total, epoch)\n","\n","\n","@torch.no_grad()\n","def log_weight_norms(model, writer, epoch):\n","    \"\"\"Log weight norms per component type.\"\"\"\n","    for name, param in model.named_parameters():\n","        if 'A_upper' in name:\n","            clean = name.replace('.', '_')\n","            writer.add_scalar(f'weights/{clean}_norm', param.norm().item(), epoch)\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# DATA — paired augmentation for observer loss\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","class PairedTransform:\n","    \"\"\"Returns two independently augmented views of the same image.\"\"\"\n","    def __init__(self, transform):\n","        self.t = transform\n","    def __call__(self, img):\n","        return self.t(img), self.t(img)\n","\n","\n","def get_dataloaders(config):\n","    import torchvision\n","    import torchvision.transforms as T\n","\n","    train_aug = T.Compose([\n","        T.RandomCrop(32, padding=4),\n","        T.RandomHorizontalFlip(),\n","        T.TrivialAugmentWide(),\n","        T.ToTensor(),\n","        T.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),\n","        T.RandomErasing(p=config.get('random_erasing_p', 0.25),\n","                        scale=(0.02, 0.33), ratio=(0.3, 3.3)),\n","    ])\n","    test_transform = T.Compose([\n","        T.ToTensor(),\n","        T.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),\n","    ])\n","\n","    train_ds = torchvision.datasets.CIFAR100(\n","        root='./data', train=True, download=True,\n","        transform=PairedTransform(train_aug))\n","    test_ds = torchvision.datasets.CIFAR100(\n","        root='./data', train=False, download=True,\n","        transform=test_transform)\n","\n","    train_loader = torch.utils.data.DataLoader(\n","        train_ds, batch_size=config['batch_size'], shuffle=True,\n","        num_workers=config['num_workers'], pin_memory=True, drop_last=True)\n","    test_loader = torch.utils.data.DataLoader(\n","        test_ds, batch_size=config['batch_size'] * 2, shuffle=False,\n","        num_workers=config['num_workers'], pin_memory=True)\n","\n","    return train_loader, test_loader\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# TRAINING — three-domain observer loss through gated pipeline\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","def ld_to_scalars(ld):\n","    \"\"\"Convert loss dict values to float scalars for display/logging.\"\"\"\n","    s = {}\n","    for k, v in ld.items():\n","        if isinstance(v, torch.Tensor):\n","            s[k] = v.item() if v.numel() == 1 else v.mean().item()\n","        elif isinstance(v, (float, int)):\n","            s[k] = float(v)\n","    return s\n","\n","\n","def train_epoch(model, loader, optimizer, scheduler, epoch, config, writer):\n","    model.train()\n","    total_loss = 0\n","    correct = 0\n","    total = 0\n","    last_ld = None\n","\n","    for batch_idx, ((v1, v2), labels) in enumerate(loader):\n","        v1, v2, labels = v1.to(device), v2.to(device), labels.to(device)\n","\n","        # Forward + loss through compiled forward()\n","        loss, ld = model(v1, v2, targets=labels)\n","        last_ld = ld\n","\n","        optimizer.zero_grad()\n","        loss.backward()\n","\n","        # Log gradient norms periodically\n","        if epoch % config['log_grads_every'] == 0 and batch_idx == 0:\n","            log_gradient_norms(model, writer, epoch)\n","\n","        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n","        optimizer.step()\n","        if scheduler is not None:\n","            scheduler.step()\n","\n","        # Invalidate CM gate caches — anchors moved\n","        model.invalidate_caches()\n","\n","        total_loss += loss.item() * v1.size(0)\n","        correct += ld.get('acc', 0) * v1.size(0)\n","        total += v1.size(0)\n","\n","    avg_loss = total_loss / total\n","    train_acc = correct / total\n","    return avg_loss, train_acc, last_ld\n","\n","\n","@torch.no_grad()\n","def evaluate(model, loader):\n","    model.eval()\n","    correct = 0\n","    total = 0\n","    for images, labels in loader:\n","        images = images.to(device)\n","        labels = labels.to(device)\n","        logits = model(images)\n","        correct += (logits.argmax(1) == labels).sum().item()\n","        total += images.size(0)\n","    return correct / total\n","\n","\n","def main():\n","    config = CONFIG.copy()\n","\n","    print(\"=\" * 60)\n","    print(\"  Geometric Transformer — CIFAR-100 (Three-Domain Observer Loss)\")\n","    print(f\"  Input: conv({config['in_channels']}→{config['conv_channels']}) + \"\n","          f\"SVD(rank={config['svd_rank']}) + \"\n","          f\"{config['patch_size']}×{config['patch_size']} patches = \"\n","          f\"{(config['img_size']//config['patch_size'])**2} tokens + CLS\")\n","    print(f\"  Model: d={config['d_model']}, heads={config['n_heads']}, \"\n","          f\"layers={config['n_layers']}, anchors={config['n_anchors']}\")\n","    print(f\"  Loss: EXTERNAL(ce={config['w_ce']}, emb_nce={config['w_nce_emb']}) \"\n","          f\"+ GEOMETRIC(pw_nce={config['w_nce_pw']}, bridge={config['w_bridge']}) \"\n","          f\"+ INTERNAL(assign={config['w_assign']}, tri={config['w_nce_tri']}, \"\n","          f\"attract={config['w_attract']})\")\n","    print(f\"  Aug: paired TrivialAugmentWide + RandomErasing(p={config['random_erasing_p']})\")\n","    print(\"=\" * 60)\n","\n","    writer = SummaryWriter(config['log_dir'])\n","    writer.add_text('config', json.dumps(config, indent=2))\n","\n","    print(\"\\nLoading CIFAR-100...\")\n","    train_loader, test_loader = get_dataloaders(config)\n","    print(f\"  Train: {len(train_loader.dataset):,} | Test: {len(test_loader.dataset):,}\")\n","\n","    model = GeoViTClassifier('geo_vit_cifar100', config)\n","    if hasattr(model, 'network_to'):\n","        model.network_to(device=device, strict=False)\n","    else:\n","        model = model.to(device)\n","\n","    # WideRouter compile: discover towers → analyze structure → pre-build\n","    # VMapTowerGroups → torch.compile with CUDA graph optimization\n","    model_raw = model\n","    if device.type == 'cuda':\n","        try:\n","            model_raw.discover_towers()\n","            model = model_raw.compile(mode='default')\n","            print(f\"  WideRouter compiled (default)\")\n","            print(f\"    towers: {model_raw.tower_names}\")\n","            print(f\"    analyzed: {model_raw.objects.get('_analyzed', False)}\")\n","        except Exception as e:\n","            print(f\"  compile failed: {e} — running eager\")\n","            model = model_raw\n","\n","    n_params = sum(p.numel() for p in model.parameters())\n","    n_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)\n","    print(f\"\\n  Total params:     {n_params:,}\")\n","    print(f\"  Trainable params: {n_trainable:,}\")\n","\n","    for name, module in model.named_children():\n","        n = sum(p.numel() for p in module.parameters())\n","        if n > 0:\n","            print(f\"    {name:<20s}: {n:,}\")\n","\n","    writer.add_scalar('model/total_params', n_params, 0)\n","\n","    # Initial anchor diagnostics\n","    print(f\"\\n  Initial anchor diagnostics:\")\n","    diag = model_raw.anchor_diagnostics()\n","    for layer_name, d in diag.items():\n","        print(f\"    {layer_name}: cv={d['anchor_cv']:.4f}, \"\n","              f\"cm_pos={d['cm_positive_frac']:.3f}, \"\n","              f\"min_dist={d['min_pairwise_dist']:.4f}\")\n","\n","    # Optimizer — plain Adam (AdamW fights geometric structure)\n","    optimizer = torch.optim.Adam(\n","        model.parameters(), lr=config['lr'])\n","\n","    total_steps = config['epochs'] * len(train_loader)\n","    warmup_steps = config['warmup_epochs'] * len(train_loader)\n","\n","    def lr_lambda(step):\n","        if step < warmup_steps:\n","            return step / max(1, warmup_steps)\n","        progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)\n","        return 0.5 * (1 + np.cos(np.pi * progress))\n","\n","    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)\n","\n","    print(f\"\\n{'━'*60}\")\n","    print(f\"  Training for {config['epochs']} epochs\")\n","    print(f\"  Warmup: {config['warmup_epochs']} epochs, LR: {config['lr']}\")\n","    print(f\"  Three-domain observer loss (paired views)\")\n","    print(f\"  TensorBoard: {config['log_dir']}\")\n","    print(f\"  Geo analysis every {config['log_geo_every']} epochs\")\n","    print(f\"{'━'*60}\\n\")\n","\n","    best_acc = 0\n","    save_dir = Path('geo_cifar100'); save_dir.mkdir(exist_ok=True)\n","\n","    for epoch in tqdm(range(config['epochs']), desc=\"Epochs\"):\n","        t0 = time.time()\n","\n","        avg_loss, train_acc, last_ld = train_epoch(\n","            model, train_loader, optimizer, scheduler, epoch, config, writer)\n","\n","        test_acc = evaluate(model, test_loader)\n","        elapsed = time.time() - t0\n","\n","        lr = optimizer.param_groups[0]['lr']\n","        writer.add_scalar('train/total_loss', avg_loss, epoch)\n","        writer.add_scalar('train/accuracy', train_acc, epoch)\n","        writer.add_scalar('test/accuracy', test_acc, epoch)\n","        writer.add_scalar('train/lr', lr, epoch)\n","        writer.add_scalar('train/epoch_time', elapsed, epoch)\n","        writer.add_scalar('gap/train_test', train_acc - test_acc, epoch)\n","\n","        # Log individual loss components from last batch\n","        if last_ld is not None:\n","            s = ld_to_scalars(last_ld)\n","            for k, v in s.items():\n","                if isinstance(v, (int, float)) and not math.isnan(v):\n","                    writer.add_scalar(f'loss/{k}', v, epoch)\n","\n","        log_weight_norms(model, writer, epoch)\n","\n","        if test_acc > best_acc:\n","            best_acc = test_acc\n","            torch.save({\n","                'state_dict': {k: v.cpu() for k, v in model.state_dict().items()},\n","                'epoch': epoch,\n","                'test_acc': test_acc,\n","                'config': config,\n","            }, save_dir / 'best.pt')\n","\n","        # Full geometric analysis periodically\n","        if epoch % config['log_geo_every'] == 0 or epoch == config['epochs'] - 1:\n","            geo_info = log_geometric_analysis(\n","                model_raw, writer, epoch, test_loader, device, config)\n","\n","            cv_str = ', '.join(f'{cv:.3f}' for cv in geo_info['cv_trajectory'])\n","            cm_str = ', '.join(f'{q:.3f}' for q in geo_info.get('cm_quality_trajectory', []))\n","\n","            # Three-domain breakdown from last batch\n","            s = ld_to_scalars(last_ld) if last_ld else {}\n","            ext = s.get('loss_external', s.get('loss_task', 0))\n","            geo = s.get('loss_geometric', 0)\n","            intl = s.get('loss_internal', 0)\n","            brg = s.get('bridge_acc', 0)\n","            pw = s.get('nce_pw_acc', 0)\n","\n","            tqdm.write(\n","                f\"  E{epoch:>3d}  L={avg_loss:.3f}\"\n","                f\"  ext={ext:.3f} geo={geo:.3f} int={intl:.3f}\"\n","                f\"  brg={brg*100:.0f}% pw={pw*100:.0f}%\"\n","                f\"  train={train_acc:.4f}  test={test_acc:.4f}\"\n","                f\"  best={best_acc:.4f}  {elapsed:.1f}s\"\n","                f\"\\n        CV=[{cv_str}]\"\n","                f\"\\n        CM=[{cm_str}]\")\n","        elif epoch % 5 == 0:\n","            s = ld_to_scalars(last_ld) if last_ld else {}\n","            tqdm.write(\n","                f\"  E{epoch:>3d}  L={avg_loss:.3f}\"\n","                f\"  train={train_acc:.4f}  test={test_acc:.4f}\"\n","                f\"  best={best_acc:.4f}  {elapsed:.1f}s\")\n","\n","    # Final summary\n","    print(f\"\\n{'═'*60}\")\n","    print(f\"  CIFAR-100 RESULTS (Three-Domain Observer Loss)\")\n","    print(f\"{'═'*60}\")\n","    print(f\"  Best test accuracy: {best_acc:.4f} ({best_acc*100:.2f}%)\")\n","    print(f\"  Parameters: {n_params:,}\")\n","    print(f\"  Checkpoint: {save_dir}/best.pt\")\n","    print(f\"  TensorBoard: {config['log_dir']}\")\n","\n","    # Final geometric state + anchor diagnostics\n","    print(f\"\\n  Final geometric state:\")\n","    geo_info = log_geometric_analysis(\n","        model_raw, writer, config['epochs'], test_loader, device, config)\n","\n","    print(f\"\\n  Final anchor diagnostics:\")\n","    diag = model_raw.anchor_diagnostics()\n","    for layer_name, d in diag.items():\n","        print(f\"    {layer_name}: cv={d['anchor_cv']:.4f}, \"\n","              f\"cm_pos={d['cm_positive_frac']:.3f}, \"\n","              f\"cm_mean={d['cm_mean']:.4f}\")\n","\n","    writer.close()\n","    print(f\"\\nDone.\")\n","\n","\n","if __name__ == '__main__':\n","    main()"],"metadata":{"colab":{"base_uri":"https://localhost:8080/","height":1000,"referenced_widgets":["86d16df51c124161a2728f0003ebed64","b5b3f1a8eec049c78df39ea2e0a878d7","80ab24bf7d1a4ad69b2618334aa90c11","0b8d0ae380024fbb8a0fae99ac4ac102","61cd81b77c3249a98f063dc6f9117df7","d94819a2d40f4687994d35e75011b18c","6efd5fc8aa13488a815d2d74ee9fd2b0","85aa6598a0894f7b8e103c9a79dc897d","288bd07827f54042b5e8e7a342eb586a","bc430ece4e2645a28795a95a2e178010","bb75ffc809d0480e882295dc28ec3f5d"]},"id":"Tgd-MbdcARPa","executionInfo":{"status":"ok","timestamp":1775012731478,"user_tz":420,"elapsed":3430744,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"6c036d18-fb12-4c0f-a602-86e23564aedd"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["Device: cuda\n","  GPU: NVIDIA RTX PRO 6000 Blackwell Server Edition\n","  VRAM: 102.0 GB\n","============================================================\n","  Geometric Transformer — CIFAR-100 (Three-Domain Observer Loss)\n","  Input: conv(3→64) + SVD(rank=16) + 4×4 patches = 64 tokens + CLS\n","  Model: d=384, heads=8, layers=4, anchors=128\n","  Loss: EXTERNAL(ce=0.8, emb_nce=0.5) + GEOMETRIC(pw_nce=1.0, bridge=1.0) + INTERNAL(assign=0.5, tri=0.5, attract=0.25)\n","  Aug: paired TrivialAugmentWide + RandomErasing(p=0.25)\n","============================================================\n","\n","Loading CIFAR-100...\n","  Train: 50,000 | Test: 10,000\n","  WideRouter compiled (default)\n","    towers: ['transformer']\n","    analyzed: True\n","\n","  Total params:     21,360,232\n","  Trainable params: 21,360,232\n","    _orig_mod           : 21,360,232\n","\n","  Initial anchor diagnostics:\n","    layer_0: cv=0.0244, cm_pos=1.000, min_dist=0.9659\n","    layer_1: cv=0.0244, cm_pos=1.000, min_dist=0.9654\n","    layer_2: cv=0.0239, cm_pos=1.000, min_dist=0.9667\n","    layer_3: cv=0.0244, cm_pos=1.000, min_dist=0.9671\n","\n","━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\n","  Training for 100 epochs\n","  Warmup: 5 epochs, LR: 0.0003\n","  Three-domain observer loss (paired views)\n","  TensorBoard: runs/geo_observer_3domain\n","  Geo analysis every 5 epochs\n","━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\n","\n"]},{"output_type":"display_data","data":{"text/plain":["Epochs:   0%|          | 0/100 [00:00<?, ?it/s]"],"application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"86d16df51c124161a2728f0003ebed64"}},"metadata":{}},{"output_type":"stream","name":"stdout","text":["  E  0  L=17.692  ext=4.163 geo=7.518 int=3.647  brg=88% pw=41%  train=0.0385  test=0.0927  best=0.0927  47.3s\n","        CV=[0.342, 0.292, 0.255, 0.265]\n","        CM=[0.882, 0.881, 0.881, 0.884]\n","  E  5  L=10.203  ext=3.329 geo=3.200 int=3.410  brg=94% pw=77%  train=0.1831  test=0.2417  best=0.2417  34.3s\n","        CV=[0.548, 0.572, 0.526, 0.247]\n","        CM=[0.887, 0.887, 0.881, 0.884]\n","  E 10  L=8.918  ext=2.797 geo=2.632 int=3.271  brg=96% pw=87%  train=0.2950  test=0.3701  best=0.3701  33.9s\n","        CV=[0.544, 0.575, 0.552, 0.252]\n","        CM=[0.889, 0.888, 0.877, 0.857]\n","  E 15  L=8.261  ext=2.168 geo=2.617 int=3.258  brg=97% pw=88%  train=0.3783  test=0.4381  best=0.4381  34.3s\n","        CV=[0.555, 0.562, 0.560, 0.248]\n","        CM=[0.890, 0.888, 0.871, 0.793]\n","  E 20  L=7.777  ext=2.127 geo=2.592 int=3.213  brg=96% pw=88%  train=0.4576  test=0.4773  best=0.4800  34.0s\n","        CV=[0.556, 0.534, 0.567, 0.255]\n","        CM=[0.893, 0.888, 0.865, 0.737]\n","  E 25  L=7.409  ext=1.788 geo=2.455 int=3.159  brg=98% pw=91%  train=0.5257  test=0.5054  best=0.5096  33.7s\n","        CV=[0.564, 0.535, 0.550, 0.252]\n","        CM=[0.892, 0.888, 0.863, 0.716]\n","  E 30  L=7.062  ext=1.317 geo=2.198 int=3.047  brg=98% pw=95%  train=0.5966  test=0.5282  best=0.5282  33.8s\n","        CV=[0.576, 0.527, 0.615, 0.254]\n","        CM=[0.894, 0.886, 0.858, 0.711]\n","  E 35  L=6.762  ext=1.292 geo=2.469 int=3.053  brg=95% pw=90%  train=0.6634  test=0.5333  best=0.5349  34.1s\n","        CV=[0.566, 0.536, 0.562, 0.247]\n","        CM=[0.896, 0.886, 0.857, 0.714]\n","  E 40  L=6.553  ext=1.045 geo=2.367 int=3.034  brg=95% pw=92%  train=0.7147  test=0.5481  best=0.5481  33.9s\n","        CV=[0.561, 0.540, 0.559, 0.251]\n","        CM=[0.896, 0.885, 0.854, 0.717]\n","  E 45  L=6.391  ext=0.917 geo=2.587 int=3.049  brg=98% pw=91%  train=0.7528  test=0.5422  best=0.5570  34.2s\n","        CV=[0.567, 0.532, 0.574, 0.261]\n","        CM=[0.897, 0.885, 0.853, 0.725]\n","  E 50  L=6.251  ext=0.842 geo=2.405 int=3.025  brg=98% pw=92%  train=0.7857  test=0.5546  best=0.5570  33.9s\n","        CV=[0.561, 0.541, 0.571, 0.255]\n","        CM=[0.897, 0.884, 0.851, 0.729]\n","  E 55  L=6.126  ext=0.682 geo=2.358 int=3.027  brg=99% pw=95%  train=0.8086  test=0.5525  best=0.5623  34.0s\n","        CV=[0.563, 0.538, 0.565, 0.259]\n","        CM=[0.895, 0.883, 0.853, 0.736]\n","  E 60  L=6.045  ext=0.607 geo=2.243 int=2.931  brg=99% pw=97%  train=0.8319  test=0.5640  best=0.5641  34.3s\n","        CV=[0.565, 0.552, 0.568, 0.258]\n","        CM=[0.895, 0.883, 0.854, 0.742]\n","  E 65  L=5.936  ext=0.549 geo=2.239 int=2.958  brg=98% pw=95%  train=0.8501  test=0.5650  best=0.5702  33.5s\n","        CV=[0.564, 0.551, 0.572, 0.257]\n","        CM=[0.895, 0.883, 0.855, 0.747]\n","  E 70  L=5.876  ext=0.446 geo=2.203 int=2.950  brg=100% pw=98%  train=0.8657  test=0.5698  best=0.5710  34.0s\n","        CV=[0.563, 0.547, 0.585, 0.258]\n","        CM=[0.895, 0.882, 0.855, 0.751]\n","  E 75  L=5.802  ext=0.546 geo=2.428 int=2.909  brg=100% pw=92%  train=0.8800  test=0.5734  best=0.5761  34.4s\n","        CV=[0.560, 0.553, 0.579, 0.260]\n","        CM=[0.895, 0.882, 0.857, 0.758]\n","  E 80  L=5.758  ext=0.482 geo=2.261 int=2.920  brg=98% pw=96%  train=0.8858  test=0.5734  best=0.5804  34.0s\n","        CV=[0.560, 0.555, 0.580, 0.260]\n","        CM=[0.895, 0.882, 0.856, 0.764]\n","  E 85  L=5.701  ext=0.449 geo=2.252 int=2.890  brg=99% pw=95%  train=0.8943  test=0.5790  best=0.5804  33.7s\n","        CV=[0.560, 0.552, 0.576, 0.262]\n","        CM=[0.895, 0.882, 0.856, 0.765]\n","  E 90  L=5.694  ext=0.355 geo=2.237 int=2.927  brg=100% pw=94%  train=0.8974  test=0.5790  best=0.5804  34.2s\n","        CV=[0.560, 0.553, 0.588, 0.262]\n","        CM=[0.895, 0.882, 0.856, 0.766]\n","  E 95  L=5.688  ext=0.325 geo=2.194 int=2.923  brg=99% pw=98%  train=0.9006  test=0.5791  best=0.5808  33.7s\n","        CV=[0.559, 0.554, 0.583, 0.261]\n","        CM=[0.895, 0.882, 0.857, 0.767]\n","  E 99  L=5.698  ext=0.400 geo=2.224 int=2.887  brg=99% pw=94%  train=0.8984  test=0.5803  best=0.5808  34.3s\n","        CV=[0.559, 0.554, 0.583, 0.262]\n","        CM=[0.895, 0.882, 0.857, 0.767]\n","\n","════════════════════════════════════════════════════════════\n","  CIFAR-100 RESULTS (Three-Domain Observer Loss)\n","════════════════════════════════════════════════════════════\n","  Best test accuracy: 0.5808 (58.08%)\n","  Parameters: 21,360,232\n","  Checkpoint: geo_cifar100/best.pt\n","  TensorBoard: runs/geo_observer_3domain\n","\n","  Final geometric state:\n","\n","  Final anchor diagnostics:\n","    layer_0: cv=0.0818, cm_pos=1.000, cm_mean=2.8355\n","    layer_1: cv=0.0640, cm_pos=1.000, cm_mean=2.9221\n","    layer_2: cv=0.0283, cm_pos=1.000, cm_mean=3.2639\n","    layer_3: cv=0.3910, cm_pos=0.172, cm_mean=-4.9637\n","\n","Done.\n"]}]},{"cell_type":"markdown","source":["# flow match transformers"],"metadata":{"id":"W0U7A75M7wVQ"}},{"cell_type":"code","source":["\"\"\"\n","Geometric Transformer — CM-Validated Pipeline\n","==================================================\n","Dual-stream transformer with CM-gated constellation observation,\n","quaternion composition, and per-layer Cayley alignment.\n","\n","CM-validated pipeline changes:\n","    - CM validity gate between association and curation (AnchorGate)\n","    - 5-stream PositionGeometricContext: anchor + structural + history + quality + FLOW\n","    - CM-conditioned geometric residual accumulation (replaces blind learned gate)\n","    - Built-in geometric regularization (CV target + anchor spread)\n","    - Decomposed observer pipeline: association → CM gate → gated curation\n","    - Optional FlowEnsemble: multi-opinion geometric fusion (quat, velocity, orbital, etc.)\n","\n","Pipeline per layer:\n","    1. ManifoldProjection:  h_i → emb_i on S^(d-1) per position\n","    2. ConstellationAssociation: emb_i → raw triangulation, cos, assignment\n","    3. CMValidatedGate: per-anchor CM validity → gate_values (B*L, A)\n","    4. Gated curation: patchwork reads tri * gate_values (validated only)\n","    4.5 FlowEnsemble (optional): multi-opinion geometric predictions\n","    5. PositionGeometricContext: 5 streams → FiLM context (B, L, context_dim)\n","    6. ContentAttention (Stream A): standard MHA\n","    7. GeometricAttention (Stream B): FiLM(Q,K | geo_ctx), V pure\n","    8. CayleyOrthogonal: align B → A basis\n","    9. QuaternionCompose: w=A, i=aligned_B, j=A-B, k=A*B\n","   10. Decode + gated residual\n","   11. CM-conditioned geometric residual write\n","\n","Geometric regularization (call model.geometric_losses() during training):\n","    - CV loss: anchor CV → pentachoron band (0.20-0.23)\n","    - Spread loss: prevent anchor collapse (penalize positive cosine)\n","    These maintain the constellation in the regime where CM validation works.\n","\n","Design principles from Ryan Spearman (ρ=0.309, 76/84 wins):\n","    - FiLM on Q,K ONLY — geometry routes attention, V stays pure\n","    - FiLM on individual arms BEFORE composition, not after\n","    - Quaternion algebra as structural regularizer (non-commutative coupling)\n","    - CayleyOrthogonal guarantees pure rotation (det=1 always)\n","    - Never global average pool — per-position geometric context\n","\n","Author: AbstractPhil + Claude Opus 4.6\n","License: Apache 2.0\n","\"\"\"\n","\n","import math\n","import torch\n","import torch.nn as nn\n","import torch.nn.functional as F\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# IMPORTS — hard dependencies, no fallback stubs\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","from geolip_core.core.associate.constellation import (\n","    ConstellationObserver, ConstellationAssociation, ConstellationCuration,\n","    Constellation, init_anchors_repulsion,\n",")\n","from geolip_core.core.curate.gate import AnchorGate as _GeolipAnchorGate\n","from geolip_core.core.curate.flows import FlowEnsemble\n","import geolip_core.linalg as LA\n","from geolip_core.pipeline.observer import (\n","    TorchComponent, BaseTower, Input, Curation, Distinction,\n",")\n","from geolip_core.core.distinguish.losses import (\n","    observer_loss as _geolip_observer_loss,\n","    ce_loss_paired as _geolip_ce_loss_paired,\n","    spread_loss as _geolip_spread_loss,\n",")\n","\n","# Optional: geofractal WideRouter for compilation\n","try:\n","    from geofractal.router.wide_router import WideRouter\n","    _HAS_WIDE_ROUTER = True\n","except ImportError:\n","    _HAS_WIDE_ROUTER = False\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# CAYLEY-MENGER VALIDITY — geometric quality measurement\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","def pairwise_distances_squared(points):\n","    \"\"\"Batched pairwise squared distances. (B, N, D) → (B, N, N).\"\"\"\n","    gram = torch.bmm(points, points.transpose(1, 2))\n","    diag = gram.diagonal(dim1=-2, dim2=-1)\n","    return diag.unsqueeze(2) + diag.unsqueeze(1) - 2 * gram\n","\n","\n","def cayley_menger_det(points):\n","    \"\"\"Cayley-Menger signed volume² for simplices. (B, K, D) → (B,).\n","\n","    K = number of vertices (k+1 for a k-simplex).\n","    Sign-corrected: positive = valid non-degenerate simplex.\n","    \"\"\"\n","    B, K, D = points.shape\n","    d2 = pairwise_distances_squared(points)\n","    M = torch.zeros(B, K + 1, K + 1, device=points.device, dtype=points.dtype)\n","    M[:, 0, 1:] = 1.0\n","    M[:, 1:, 0] = 1.0\n","    M[:, 1:, 1:] = d2\n","    raw = LA.det(M)\n","    k = K - 1\n","    sign = (-1.0) ** (k + 1)\n","    return sign * raw\n","\n","\n","def anchor_neighborhood_cm(anchors, n_neighbors=3):\n","    \"\"\"Precompute per-anchor CM quality from local neighborhood geometry.\n","\n","    Position-independent. O(A) determinant computations on small matrices.\n","    Each anchor forms a simplex with its k nearest neighbor anchors.\n","    The CM determinant measures local geometric quality — high volume means\n","    the anchor neighborhood is well-conditioned for triangulation.\n","\n","    Args:\n","        anchors: (A, D) normalized anchor positions on S^(d-1)\n","        n_neighbors: neighbors per simplex\n","\n","    Returns:\n","        quality: (A,) signed log-magnitude CM quality per anchor\n","        nn_idx: (A, n_neighbors) neighbor indices\n","    \"\"\"\n","    A, D = anchors.shape\n","    dists = torch.cdist(anchors.unsqueeze(0), anchors.unsqueeze(0)).squeeze(0)\n","    # Mask self-distances without in-place mutation (compile-safe)\n","    self_mask = torch.eye(A, device=anchors.device, dtype=anchors.dtype) * 1e12\n","    dists = dists + self_mask\n","    _, nn_idx = dists.topk(n_neighbors, largest=False)  # (A, n_neighbors)\n","\n","    # Build simplices: [anchor_a, neighbor_1, ..., neighbor_k] — fully vectorized\n","    simplices = torch.cat([\n","        anchors.unsqueeze(1),   # (A, 1, D)\n","        anchors[nn_idx],        # (A, n_neighbors, D)\n","    ], dim=1)                   # (A, K, D)\n","\n","    dets = cayley_menger_det(simplices)  # (A,)\n","    sign = dets.sign()\n","    log_mag = torch.log(dets.abs() + 1e-12)\n","    return sign * log_mag, nn_idx\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# CM VALIDATED GATE — efficient anchor gating for transformer scale\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","class CMValidatedGate(nn.Module):\n","    \"\"\"Anchor gate based on Cayley-Menger validity.\n","\n","    Efficient for transformer scale: anchor CM quality is precomputed O(A²)\n","    and CACHED (only recomputed on invalidate_cache()), then combined with\n","    per-position proximity features through a learned gate.\n","\n","    The gate starts OPEN (bias=+2, sigmoid≈0.88) and learns to CLOSE on\n","    geometrically invalid configurations. Architecture-before-loss: the gate\n","    suppresses degenerate measurements structurally, not through a loss signal.\n","\n","    Gate features per (position, anchor):\n","        - anchor_cm_quality: CM volume of anchor's local neighborhood (cached)\n","        - cos_to_anchor: cosine similarity (position-dependent)\n","\n","    Args:\n","        n_anchors: number of constellation anchors\n","        n_neighbors: neighbors for CM simplex computation\n","    \"\"\"\n","    def __init__(self, n_anchors, n_neighbors=3):\n","        super().__init__()\n","        self.n_anchors = n_anchors\n","        self.n_neighbors = n_neighbors\n","\n","        # Learned gate: [cm_quality, cos_sim] → scalar gate\n","        self.gate_proj = nn.Sequential(\n","            nn.Linear(2, 16),\n","            nn.GELU(),\n","            nn.Linear(16, 1),\n","        )\n","        # Init OPEN — learn to close. sigmoid(2.0) ≈ 0.88\n","        # Small random weight so gradient flows back to gate_proj[0]\n","        nn.init.normal_(self.gate_proj[2].weight, std=0.01)\n","        nn.init.constant_(self.gate_proj[2].bias, 2.0)\n","\n","        # Anchor CM cache — invalidated after optimizer step\n","        self._cached_cm_norm = None\n","\n","    def invalidate_cache(self):\n","        \"\"\"Call after optimizer.step() to recompute anchor CM next forward.\"\"\"\n","        self._cached_cm_norm = None\n","\n","    @torch.compiler.disable\n","    def _get_anchor_cm_norm(self, anchors):\n","        \"\"\"Compute or return cached normalized anchor CM quality.\n","        Excluded from torch.compile — linalg.det not CUDA-graph-safe.\"\"\"\n","        if self._cached_cm_norm is not None:\n","            return self._cached_cm_norm\n","        with torch.no_grad():\n","            anchor_cm, _ = anchor_neighborhood_cm(anchors, self.n_neighbors)\n","            cm_std = anchor_cm.std().clamp(min=1e-8)\n","            self._cached_cm_norm = ((anchor_cm - anchor_cm.mean()) / cm_std).detach()\n","        return self._cached_cm_norm\n","\n","    def _compute_gate(self, anchor_cm_norm, tri):\n","        \"\"\"Compilable gate computation — pure tensor ops, no graph breaks.\"\"\"\n","        N, A = tri.shape\n","        cos_sim = 1.0 - tri  # (N, A)\n","\n","        # Gate features: (N, A, 2)\n","        features = torch.stack([\n","            anchor_cm_norm.unsqueeze(0).expand(N, -1),\n","            cos_sim,\n","        ], dim=-1)\n","\n","        gate_values = torch.sigmoid(self.gate_proj(features).squeeze(-1))\n","\n","        # Diagnostics — pure tensor ops, no graph breaks\n","        gate_info = {\n","            'active': (gate_values.detach() > 0.5).float().sum(-1).mean(),\n","            'gate_mean': gate_values.detach().mean(),\n","            'cm_positive_frac': (anchor_cm_norm > 0).float().mean(),\n","        }\n","\n","        return gate_values, gate_info\n","\n","    def forward(self, embedding, anchors, tri):\n","        \"\"\"Compute per-(position, anchor) gate values.\"\"\"\n","        # Non-compilable: exits torch.compile graph, returns single (A,) tensor\n","        anchor_cm_norm = self._get_anchor_cm_norm(anchors)\n","        # Compilable: re-enters torch.compile graph\n","        return self._compute_gate(anchor_cm_norm, tri)\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# INFONCE MEMORY BANK — contrastive pressure on geometric residual\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","class GeoResidualBank(nn.Module):\n","    \"\"\"Cross-stream contrastive memory bank (CLIP-style).\"\"\"\n","    def __init__(self, proj_dim, bank_size=4096, temperature=0.1):\n","        super().__init__()\n","        self.proj_dim = proj_dim\n","        self.bank_size = bank_size\n","        self.temperature = temperature\n","\n","        self.register_buffer('queue', torch.randn(bank_size, proj_dim))\n","        self.queue = F.normalize(self.queue, dim=-1)\n","        self.register_buffer('queue_ptr', torch.zeros(1, dtype=torch.long))\n","\n","    @torch.no_grad()\n","    def enqueue(self, keys):\n","        B = keys.shape[0]\n","        ptr = int(self.queue_ptr.item())\n","        if ptr + B <= self.bank_size:\n","            self.queue[ptr:ptr + B] = keys\n","        else:\n","            overflow = (ptr + B) - self.bank_size\n","            self.queue[ptr:] = keys[:B - overflow]\n","            self.queue[:overflow] = keys[B - overflow:]\n","        self.queue_ptr[0] = (ptr + B) % self.bank_size\n","\n","    def forward(self, content_proj, geo_proj):\n","        q = F.normalize(content_proj, dim=-1)\n","        k_pos = F.normalize(geo_proj, dim=-1)\n","        k_neg = self.queue.clone().detach()\n","\n","        pos_logits = (q * k_pos).sum(dim=-1, keepdim=True) / self.temperature\n","        neg_logits = q @ k_neg.T / self.temperature\n","\n","        logits = torch.cat([pos_logits, neg_logits], dim=1)\n","        labels = torch.zeros(q.shape[0], dtype=torch.long, device=q.device)\n","\n","        loss = F.cross_entropy(logits, labels)\n","\n","        with torch.no_grad():\n","            acc = (logits.argmax(dim=1) == 0).float().mean()\n","\n","        return loss, acc\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# PROVEN COMPONENTS — from Ryan Spearman (unchanged, tested)\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","class FiLMLayer(TorchComponent):\n","    \"\"\"Feature-wise Linear Modulation. Near-identity-initialized.\n","    gamma ≈ 1 + 0.01·geo_ctx, beta ≈ 0.01·geo_ctx at init.\n","    Gradient flows through to geo_ctx from step 0.\n","    \"\"\"\n","    def __init__(self, name, feature_dim, context_dim):\n","        super().__init__(name)\n","        self.to_gamma = nn.Linear(context_dim, feature_dim)\n","        self.to_beta = nn.Linear(context_dim, feature_dim)\n","        nn.init.normal_(self.to_gamma.weight, std=0.01); nn.init.ones_(self.to_gamma.bias)\n","        nn.init.normal_(self.to_beta.weight, std=0.01); nn.init.zeros_(self.to_beta.bias)\n","\n","    def forward(self, x, ctx):\n","        return self.to_gamma(ctx) * x + self.to_beta(ctx)\n","\n","\n","class CayleyOrthogonal(TorchComponent):\n","    \"\"\"Guaranteed SO(d) rotation via Cayley map. det(Q) = 1 always.\"\"\"\n","    def __init__(self, name, dim):\n","        super().__init__(name)\n","        self.dim = dim\n","        self.A_upper = nn.Parameter(torch.zeros(dim * (dim - 1) // 2) * 0.01)\n","        idx = torch.triu_indices(dim, dim, offset=1)\n","        self.register_buffer('_triu_row', idx[0], persistent=False)\n","        self.register_buffer('_triu_col', idx[1], persistent=False)\n","        self.register_buffer('_eye', torch.eye(dim), persistent=False)\n","\n","    def get_rotation(self):\n","        d = self.dim\n","        A = torch.zeros(d, d, device=self.A_upper.device, dtype=self.A_upper.dtype)\n","        A[self._triu_row, self._triu_col] = self.A_upper\n","        A = A - A.T\n","        return LA.solve(self._eye + A, self._eye - A)\n","\n","    def forward(self, x):\n","        return x @ self.get_rotation().T\n","\n","\n","def quaternion_multiply_batched(q1, q2):\n","    \"\"\"Hamilton product on (B, 4, D) tensors. Fully vectorized.\"\"\"\n","    w1, x1, y1, z1 = q1[:, 0], q1[:, 1], q1[:, 2], q1[:, 3]\n","    w2, x2, y2, z2 = q2[:, 0], q2[:, 1], q2[:, 2], q2[:, 3]\n","    return torch.stack([\n","        w1*w2 - x1*x2 - y1*y2 - z1*z2,\n","        w1*x2 + x1*w2 + y1*z2 - z1*y2,\n","        w1*y2 - x1*z2 + y1*w2 + z1*x2,\n","        w1*z2 + x1*y2 - y1*x2 + z1*w2,\n","    ], dim=1)\n","\n","\n","class QuaternionCompose(TorchComponent):\n","    \"\"\"Four-arm Hamilton product composition. Proven in GeoQuat head.\"\"\"\n","    def __init__(self, name, input_dim, quat_dim=64):\n","        super().__init__(name)\n","        self.quat_dim = quat_dim\n","        self.proj_w = nn.Linear(input_dim, quat_dim)\n","        self.proj_i = nn.Linear(input_dim, quat_dim)\n","        self.proj_j = nn.Linear(input_dim, quat_dim)\n","        self.proj_k = nn.Linear(input_dim, quat_dim)\n","        self.rotation = nn.Parameter(torch.randn(1, 4, quat_dim) * 0.1)\n","\n","    @property\n","    def output_dim(self):\n","        return self.quat_dim * 4\n","\n","    def forward(self, arm_w, arm_i, arm_j, arm_k):\n","        shape = arm_w.shape[:-1]\n","        D = arm_w.shape[-1]\n","        flat = arm_w.dim() > 2\n","        if flat:\n","            arm_w = arm_w.reshape(-1, D); arm_i = arm_i.reshape(-1, D)\n","            arm_j = arm_j.reshape(-1, D); arm_k = arm_k.reshape(-1, D)\n","        q = torch.stack([self.proj_w(arm_w), self.proj_i(arm_i),\n","                         self.proj_j(arm_j), self.proj_k(arm_k)], dim=1)\n","        q = q / (q.norm(dim=1, keepdim=True) + 1e-8)\n","        r = self.rotation.expand(q.shape[0], -1, -1)\n","        r = r / (r.norm(dim=1, keepdim=True) + 1e-8)\n","        composed = quaternion_multiply_batched(r, q)\n","        composed = composed.reshape(q.shape[0], -1)\n","        if flat:\n","            composed = composed.reshape(*shape, -1)\n","        return composed\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# TRANSFORMER-SPECIFIC COMPONENTS\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","class ManifoldProjection(TorchComponent):\n","    \"\"\"Input stage: project transformer hidden states to S^(d-1).\"\"\"\n","    def __init__(self, name, d_model, manifold_dim):\n","        super().__init__(name)\n","        self.proj = nn.Linear(d_model, manifold_dim)\n","        self.norm = nn.LayerNorm(manifold_dim)\n","\n","    def forward(self, hidden_states):\n","        h = self.norm(self.proj(hidden_states))\n","        return F.normalize(h, dim=-1)\n","\n","\n","class PositionGeometricContext(TorchComponent):\n","    \"\"\"Curation stage: 5-stream fusion → FiLM context.\n","\n","    Five streams:\n","        anchor:     cos_to_anchors + assignment + triangulation — WHERE on the manifold\n","        structural: patchwork + embedding — WHAT the local geometry looks like\n","        history:    geo_residual from previous layers — WHAT prior layers observed\n","        quality:    CM gate values per anchor — HOW TRUSTWORTHY is this observation\n","        flow:       FlowEnsemble predictions — WHAT other mathematical lenses see\n","\n","    The flow stream starts at zero (zero-init) and learns to contribute.\n","    Without flows attached, the 5th stream is zeros — equivalent to the\n","    original 4-stream architecture.\n","    \"\"\"\n","    def __init__(self, name, n_anchors, pw_dim, manifold_dim, context_dim):\n","        super().__init__(name)\n","        self.context_dim = context_dim\n","        self.pw_dim = pw_dim\n","        self.manifold_dim = manifold_dim\n","\n","        # WHERE on the manifold\n","        self.anchor_mlp = nn.Sequential(\n","            nn.Linear(n_anchors * 3, context_dim), nn.GELU(), nn.LayerNorm(context_dim))\n","        # WHAT the local geometry looks like\n","        self.struct_mlp = nn.Sequential(\n","            nn.Linear(pw_dim + manifold_dim, context_dim), nn.GELU(), nn.LayerNorm(context_dim))\n","        # WHAT prior layers observed\n","        self.history_mlp = nn.Sequential(\n","            nn.Linear(pw_dim, context_dim), nn.GELU(), nn.LayerNorm(context_dim))\n","        # HOW TRUSTWORTHY — full per-anchor gate profile\n","        self.quality_mlp = nn.Sequential(\n","            nn.Linear(n_anchors, context_dim), nn.GELU(), nn.LayerNorm(context_dim))\n","        # FLOW OPINIONS — anchor-space flow ensemble [N, A] (same shape as gate_values)\n","        # Small init: negligible contribution at start, nonzero gradient path\n","        self.flow_mlp = nn.Sequential(\n","            nn.Linear(n_anchors, context_dim), nn.GELU(), nn.LayerNorm(context_dim))\n","        nn.init.normal_(self.flow_mlp[0].weight, std=0.01)\n","        nn.init.zeros_(self.flow_mlp[0].bias)\n","\n","        # Fuse 5 streams\n","        self.fuse = nn.Sequential(\n","            nn.Linear(context_dim * 5, context_dim), nn.GELU(), nn.LayerNorm(context_dim))\n","\n","    def forward(self, obs_dict, gate_values=None, geo_residual=None, flow_output=None):\n","        \"\"\"\n","        Args:\n","            obs_dict: from decomposed association + gated curation\n","            gate_values: (N, A) CM gate values per anchor, or None\n","            geo_residual: (N, pw_dim) accumulated context, or None for first layer\n","            flow_output: (N, manifold_dim) flow ensemble prediction, or None\n","        Returns:\n","            (N, context_dim) geometric context for FiLM\n","        \"\"\"\n","        anchor_feats = torch.cat([\n","            obs_dict['cos_to_anchors'],\n","            obs_dict['assignment'],\n","            obs_dict['triangulation'],\n","        ], dim=-1)\n","        struct_feats = torch.cat([\n","            obs_dict['patchwork'],\n","            obs_dict['embedding'],\n","        ], dim=-1)\n","\n","        a = self.anchor_mlp(anchor_feats)\n","        s = self.struct_mlp(struct_feats)\n","        h = self.history_mlp(geo_residual) if geo_residual is not None else torch.zeros_like(a)\n","        q = self.quality_mlp(gate_values) if gate_values is not None else torch.zeros_like(a)\n","        f = self.flow_mlp(flow_output) if flow_output is not None else torch.zeros_like(a)\n","\n","        return self.fuse(torch.cat([a, s, h, q, f], dim=-1))\n","\n","\n","class GeometricAttention(TorchComponent):\n","    \"\"\"Attention with FiLM from curated constellation. Stream B.\"\"\"\n","    def __init__(self, name, d_model, n_heads=8, context_dim=128, dropout=0.1):\n","        super().__init__(name)\n","        self.d_model = d_model\n","        self.n_heads = n_heads\n","        self.head_dim = d_model // n_heads\n","        self.scale = self.head_dim ** -0.5\n","\n","        self.w_q = nn.Linear(d_model, d_model)\n","        self.w_k = nn.Linear(d_model, d_model)\n","        self.w_v = nn.Linear(d_model, d_model)\n","        self.w_o = nn.Linear(d_model, d_model)\n","        self.dropout = nn.Dropout(dropout)\n","\n","        self.film_q = FiLMLayer(f'{name}_film_q', d_model, context_dim)\n","        self.film_k = FiLMLayer(f'{name}_film_k', d_model, context_dim)\n","        self.norm = nn.LayerNorm(d_model)\n","\n","        self.ffn1 = nn.Linear(d_model, d_model * 4)\n","        self.film_ffn = FiLMLayer(f'{name}_film_ffn', d_model * 4, context_dim)\n","        self.ffn2 = nn.Linear(d_model * 4, d_model)\n","        self.ffn_drop = nn.Dropout(dropout)\n","        self.ffn_norm = nn.LayerNorm(d_model)\n","\n","    def forward(self, x, geo_ctx, attn_mask=None, key_padding_mask=None):\n","        B, L, D = x.shape\n","        H, HD = self.n_heads, self.head_dim\n","\n","        Q = self.film_q(self.w_q(x), geo_ctx)\n","        K = self.film_k(self.w_k(x), geo_ctx)\n","        V = self.w_v(x)\n","\n","        Q = Q.view(B, L, H, HD).transpose(1, 2)\n","        K = K.view(B, L, H, HD).transpose(1, 2)\n","        V = V.view(B, L, H, HD).transpose(1, 2)\n","\n","        scores = (Q @ K.transpose(-2, -1)) * self.scale\n","        if attn_mask is not None:\n","            scores = scores + attn_mask\n","        if key_padding_mask is not None:\n","            scores = scores.masked_fill(\n","                key_padding_mask.unsqueeze(1).unsqueeze(2), float('-inf'))\n","        attn_out = (self.dropout(F.softmax(scores, dim=-1)) @ V)\n","        attn_out = attn_out.transpose(1, 2).reshape(B, L, D)\n","        x = self.norm(x + self.w_o(attn_out))\n","\n","        h = F.gelu(self.ffn1(x))\n","        h = self.film_ffn(h, geo_ctx)\n","        x = self.ffn_norm(x + self.ffn_drop(self.ffn2(h)))\n","        return x\n","\n","\n","class ContentAttention(TorchComponent):\n","    \"\"\"Standard self-attention. Stream A. No geometric conditioning.\"\"\"\n","    def __init__(self, name, d_model, n_heads=8, dropout=0.1):\n","        super().__init__(name)\n","        self.attn = nn.MultiheadAttention(\n","            d_model, n_heads, dropout=dropout, batch_first=True)\n","        self.norm = nn.LayerNorm(d_model)\n","        self.ffn = nn.Sequential(\n","            nn.Linear(d_model, d_model * 4), nn.GELU(),\n","            nn.Linear(d_model * 4, d_model), nn.Dropout(dropout))\n","        self.ffn_norm = nn.LayerNorm(d_model)\n","\n","    def forward(self, x, attn_mask=None, key_padding_mask=None):\n","        a, _ = self.attn(x, x, x, attn_mask=attn_mask,\n","                         key_padding_mask=key_padding_mask,\n","                         need_weights=False)\n","        x = self.norm(x + a)\n","        x = self.ffn_norm(x + self.ffn(x))\n","        return x\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# LAYER — CM-validated dual-stream with constellation routing + flows\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","class GeometricTransformerLayer(BaseTower):\n","    \"\"\"One layer of the geometric transformer (CM validated + flows).\n","\n","    Pipeline per layer:\n","        1. ManifoldProjection: h → emb on S^(d-1)\n","        2. Association: emb → raw triangulation, cos, assignment\n","        3. CMValidatedGate: per-anchor CM validity → gate_values\n","        4. Gated curation: patchwork reads tri * gate_values\n","        4.5 FlowEnsemble (optional): multi-opinion geometric predictions\n","        5. PositionGeometricContext: 5 streams → FiLM context\n","        6. ContentAttention (Stream A): standard MHA\n","        7. GeometricAttention (Stream B): FiLM(Q,K | geo_ctx)\n","        8. CayleyOrthogonal: align B → A\n","        9. QuaternionCompose: w=A, i=aligned_B, j=A-B, k=A*B\n","       10. Decode + gated residual\n","       11. CM-conditioned geometric residual accumulation\n","\n","    Flows are optional, config-driven, and individually replaceable:\n","        layer['flows'].attach_flow('alignment')\n","        layer['flows'].detach_flow('velocity')\n","    \"\"\"\n","    def __init__(self, name, d_model, n_heads=8, n_anchors=32,\n","                 manifold_dim=256, n_comp=8, d_comp=32,\n","                 context_dim=128, quat_dim=64, dropout=0.1,\n","                 cm_neighbors=3, flow_keys=None, flow_fusion='weighted'):\n","        super().__init__(name)\n","        self.d_model = d_model\n","        self.n_anchors = n_anchors\n","        self.manifold_dim = manifold_dim\n","\n","        # 1. Project to manifold\n","        self.attach('projection', ManifoldProjection(\n","            f'{name}_proj', d_model, manifold_dim))\n","\n","        # 2. Constellation observer (association + curation — called decomposed)\n","        self.attach('observer', ConstellationObserver(\n","            dim=manifold_dim, n_anchors=n_anchors,\n","            n_comp=n_comp, d_comp=d_comp))\n","\n","        # 3. CM validated gate — between association and curation\n","        self.attach('cm_gate', CMValidatedGate(\n","            n_anchors=n_anchors, n_neighbors=cm_neighbors))\n","\n","        # 3.5 Flow ensemble — optional multi-opinion geometric fusion\n","        if flow_keys:\n","            self.attach('flows', FlowEnsemble(\n","                f'{name}_flows', manifold_dim, n_anchors,\n","                flow_keys=flow_keys, fusion=flow_fusion))\n","            # Blend weight: how much flow opinions influence curation\n","            # Starts small → flows fade in as they learn\n","            self.flow_alpha = nn.Parameter(torch.tensor(0.01))\n","\n","        # 4. Fuse observation into FiLM context (5 streams)\n","        pw_dim = self['observer'].curation.patchwork.output_dim\n","        self.attach('context', PositionGeometricContext(\n","            f'{name}_ctx', n_anchors, pw_dim, manifold_dim, context_dim))\n","\n","        # 5. Stream A: content\n","        self.attach('content', ContentAttention(\n","            f'{name}_content', d_model, n_heads, dropout))\n","\n","        # 6. Stream B: geometric\n","        self.attach('geometric', GeometricAttention(\n","            f'{name}_geo', d_model, n_heads, context_dim, dropout))\n","\n","        # 7. Cayley rotation: align B → A\n","        self.attach('rotation', CayleyOrthogonal(f'{name}_cayley', d_model))\n","\n","        # 8. Quaternion composition\n","        self.attach('compose', QuaternionCompose(\n","            f'{name}_quat', d_model, quat_dim))\n","\n","        # 9. Decode + output gate\n","        self.attach('decode', nn.Sequential(\n","            nn.Linear(quat_dim * 4, d_model), nn.GELU(), nn.LayerNorm(d_model)))\n","        self.attach('gate', nn.Sequential(\n","            nn.Linear(d_model * 2, d_model), nn.Sigmoid()))\n","\n","        # 10. Geometric residual projection (no learned gate — CM quality decides)\n","        self._pw_dim = pw_dim\n","        self.attach('geo_proj', nn.Sequential(\n","            nn.Linear(pw_dim, pw_dim), nn.LayerNorm(pw_dim)))\n","\n","    def forward(self, x, geo_residual=None, attn_mask=None, key_padding_mask=None):\n","        \"\"\"\n","        Args:\n","            x: (B, L, D) input hidden states\n","            geo_residual: (B, L, pw_dim) accumulated geometric context,\n","                          or None for first layer\n","\n","        Returns:\n","            x_out: (B, L, D) transformed hidden states\n","            geo_residual_out: (B, L, pw_dim) updated geometric residual\n","            geo_state: dict with full geometric state + CM diagnostics\n","        \"\"\"\n","        B, L, D = x.shape\n","\n","        # ════ 1. Project to manifold ════\n","        emb = self['projection'](x)  # (B, L, manifold_dim)\n","        emb_flat = emb.reshape(B * L, -1)\n","\n","        # ════ 2. Association — raw triangulation ════\n","        a_out = self['observer'].association(emb_flat)\n","\n","        # ════ 3. CM Gate — validate anchor measurements ════\n","        anchors_n = F.normalize(\n","            self['observer'].association.constellation.anchors, dim=-1)\n","        gate_values, gate_info = self['cm_gate'](\n","            emb_flat, anchors_n.detach(), a_out['distances'])\n","\n","        # ════ 4. Gated curation — patchwork reads validated triangulation ════\n","        a_out_gated = dict(a_out)\n","\n","        # ════ 4.5 Flow ensemble — anchor-space geometric opinions ════\n","        flow_opinion = None\n","        if self.has('flows'):\n","            flow_opinion = self['flows'](anchors_n, emb_flat, a_out['distances'])  # [N, A]\n","            # Blend flow opinion into triangulation: raw + alpha*(flow - raw)\n","            # flow_alpha starts at 0.01 → 99% raw, 1% flow opinion\n","            # Gradient: observer_loss → patchwork → distances_weighted → flow_opinion → flows\n","            alpha = self.flow_alpha.sigmoid()\n","            blended_tri = a_out['distances'] + alpha * (flow_opinion - a_out['distances'])\n","            a_out_gated['distances_weighted'] = blended_tri * gate_values\n","        else:\n","            a_out_gated['distances_weighted'] = a_out['distances'] * gate_values\n","\n","        c_out = self['observer'].curation.curate_full(a_out_gated, emb=emb_flat)\n","\n","        # Build observation dict for context\n","        obs = {\n","            'embedding': emb_flat,\n","            'triangulation': a_out['distances'],\n","            'cos_to_anchors': a_out['cos_to_anchors'],\n","            'assignment': a_out['assignment'],\n","            'nearest': a_out['nearest'],\n","            'patchwork': c_out['patchwork'],\n","            'bridge': c_out['bridge'],\n","        }\n","\n","        # ════ 5. Build FiLM context — 5 streams ════\n","        geo_res_flat = geo_residual.reshape(B * L, -1) if geo_residual is not None else None\n","        geo_ctx_flat = self['context'](\n","            obs, gate_values=gate_values, geo_residual=geo_res_flat,\n","            flow_output=flow_opinion)\n","        geo_ctx = geo_ctx_flat.reshape(B, L, -1)\n","\n","        # ════ 6. Stream A: content attention ════\n","        a_out_stream = self['content'](\n","            x, attn_mask=attn_mask, key_padding_mask=key_padding_mask)\n","\n","        # ════ 7. Stream B: geometric attention ════\n","        b_out = self['geometric'](\n","            x, geo_ctx, attn_mask=attn_mask, key_padding_mask=key_padding_mask)\n","\n","        # ════ 8. Cayley rotation: align B → A ════\n","        b_aligned = self['rotation'](b_out)\n","\n","        # ════ 9. Quaternion composition ════\n","        composed = self['compose'](\n","            arm_w=a_out_stream, arm_i=b_aligned,\n","            arm_j=a_out_stream - b_aligned, arm_k=a_out_stream * b_aligned)\n","\n","        # ════ 10. Decode + gated residual ════\n","        decoded = self['decode'](composed)\n","        g = self['gate'](torch.cat([x, decoded], dim=-1))\n","        x_out = g * decoded + (1 - g) * x\n","\n","        # ════ 11. CM-conditioned geometric residual accumulation ════\n","        pw_validated = c_out['patchwork'].reshape(B, L, -1)\n","        cm_quality = gate_values.mean(dim=-1).reshape(B, L, 1)\n","        geo_update = self['geo_proj'](pw_validated)\n","\n","        if geo_residual is None:\n","            geo_residual_out = cm_quality * geo_update\n","        else:\n","            geo_residual_out = geo_residual + cm_quality * geo_update\n","\n","        # ════ Build geo_state dict ════\n","        def _unflatten(t):\n","            if t is None:\n","                return None\n","            if t.dim() == 1:\n","                return t.reshape(B, L)\n","            return t.reshape(B, L, *t.shape[1:])\n","\n","        geo_state = {\n","            'embedding':      emb,\n","            'geo_ctx':        geo_ctx,\n","            'triangulation':  _unflatten(a_out['distances']),\n","            'cos_to_anchors': _unflatten(a_out['cos_to_anchors']),\n","            'assignment':     _unflatten(a_out['assignment']),\n","            'nearest':        _unflatten(a_out['nearest']),\n","            'patchwork':      _unflatten(c_out['patchwork']),\n","            'bridge':         _unflatten(c_out['bridge']),\n","            'gate_values':    _unflatten(gate_values),\n","            'gate_info':      gate_info,\n","            'cm_quality':     cm_quality,\n","            'content':        a_out_stream,\n","            'geometric':      b_out,\n","            'composed':       composed,\n","            'geo_residual':   geo_residual_out,\n","            'flow_opinion':   _unflatten(flow_opinion) if flow_opinion is not None else None,\n","        }\n","\n","        return x_out, geo_residual_out, geo_state\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# FULL MODEL — stack of layers + geometric regularization\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","class GeometricTransformer(BaseTower):\n","    \"\"\"Geometric Transformer — CM-validated dual-stream with optional flows.\n","\n","    Stack of GeometricTransformerLayers with:\n","        - CM-gated observation at every layer\n","        - Optional FlowEnsemble at every layer (config-driven)\n","        - Cross-layer Cayley rotation on hidden states\n","        - Built-in geometric regularization via geometric_losses()\n","    \"\"\"\n","    def __init__(self, name, d_model=512, n_heads=8, n_layers=4,\n","                 n_anchors=32, manifold_dim=256, n_comp=8, d_comp=32,\n","                 context_dim=128, quat_dim=64, dropout=0.1,\n","                 cross_layer_rotation=True, cm_neighbors=3,\n","                 nce_bank_size=4096, nce_temperature=0.1,\n","                 vocab_size=None, max_seq_len=2048,\n","                 flow_keys=None, flow_fusion='weighted'):\n","        super().__init__(name)\n","        self.d_model = d_model\n","        self.n_layers = n_layers\n","        self.n_anchors = n_anchors\n","        self._pw_dim = n_comp * d_comp\n","\n","        if vocab_size is not None:\n","            self.attach('embed', nn.Embedding(vocab_size, d_model))\n","            self.attach('pos_embed', nn.Embedding(max_seq_len, d_model))\n","            self.attach('head', nn.Linear(d_model, vocab_size, bias=False))\n","\n","        for i in range(n_layers):\n","            self.attach(f'layer_{i}', GeometricTransformerLayer(\n","                f'{name}_L{i}', d_model, n_heads, n_anchors,\n","                manifold_dim, n_comp, d_comp, context_dim, quat_dim,\n","                dropout, cm_neighbors,\n","                flow_keys=flow_keys, flow_fusion=flow_fusion))\n","\n","        if cross_layer_rotation and n_layers > 1:\n","            for i in range(n_layers - 1):\n","                self.attach(f'cross_rot_{i}', CayleyOrthogonal(\n","                    f'{name}_xrot_{i}', d_model))\n","\n","        self.attach('final_norm', nn.LayerNorm(d_model))\n","\n","        # Cross-stream contrastive (CLIP-style)\n","        if nce_bank_size > 0:\n","            nce_proj_dim = 128\n","            self.attach('nce_content_proj', nn.Sequential(\n","                nn.Linear(d_model, nce_proj_dim),\n","                nn.GELU(),\n","                nn.Linear(nce_proj_dim, nce_proj_dim),\n","            ))\n","            self.attach('nce_geo_proj', nn.Sequential(\n","                nn.Linear(self._pw_dim, nce_proj_dim),\n","                nn.GELU(),\n","                nn.Linear(nce_proj_dim, nce_proj_dim),\n","            ))\n","            self.attach('nce_bank', GeoResidualBank(\n","                nce_proj_dim, bank_size=nce_bank_size,\n","                temperature=nce_temperature))\n","\n","        self._config = dict(\n","            d_model=d_model, n_heads=n_heads, n_layers=n_layers,\n","            n_anchors=n_anchors, manifold_dim=manifold_dim,\n","            n_comp=n_comp, d_comp=d_comp, context_dim=context_dim,\n","            quat_dim=quat_dim, dropout=dropout,\n","            cross_layer_rotation=cross_layer_rotation,\n","            cm_neighbors=cm_neighbors, vocab_size=vocab_size,\n","            nce_bank_size=nce_bank_size, nce_temperature=nce_temperature,\n","            flow_keys=flow_keys, flow_fusion=flow_fusion,\n","        )\n","\n","    @property\n","    def config(self):\n","        return self._config.copy()\n","\n","    def invalidate_caches(self):\n","        \"\"\"Invalidate all CM gate caches. Call after optimizer.step().\"\"\"\n","        for i in range(self.n_layers):\n","            self[f'layer_{i}']['cm_gate'].invalidate_cache()\n","\n","    def geometric_losses(self, cv_target=0.215, cv_weight=0.1, spread_weight=0.01):\n","        \"\"\"Compute geometric regularization from current anchor geometry.\"\"\"\n","        total_cv = torch.tensor(0.0)\n","        total_spread = torch.tensor(0.0)\n","        n = 0\n","\n","        for i in range(self.n_layers):\n","            layer = self[f'layer_{i}']\n","            anchors = layer['observer'].association.constellation.anchors\n","            anchors_n = F.normalize(anchors, dim=-1)\n","            A = anchors_n.shape[0]\n","\n","            if n == 0:\n","                total_cv = total_cv.to(anchors.device)\n","                total_spread = total_spread.to(anchors.device)\n","\n","            cos = anchors_n @ anchors_n.T\n","            idx = torch.triu_indices(A, A, offset=1, device=cos.device)\n","            pairwise_dist = 1.0 - cos[idx[0], idx[1]]\n","            cv = pairwise_dist.std() / (pairwise_dist.mean() + 1e-8)\n","            total_cv = total_cv + (cv - cv_target).pow(2)\n","\n","            mask = ~torch.eye(A, dtype=torch.bool, device=cos.device)\n","            total_spread = total_spread + F.relu(cos[mask]).mean()\n","\n","            n += 1\n","\n","        losses = {}\n","        if n > 0:\n","            losses['cv'] = cv_weight * total_cv / n\n","            losses['spread'] = spread_weight * total_spread / n\n","            losses['geo_total'] = losses['cv'] + losses['spread']\n","        return losses\n","\n","    def infonce_loss(self, cls_index=0):\n","        \"\"\"Cross-stream contrastive: content queries against decoupled geometry.\"\"\"\n","        if not self.has('nce_bank'):\n","            return {}\n","\n","        hidden = getattr(self, '_last_hidden', None)\n","        geo_residual = getattr(self, '_last_geo_residual', None)\n","        if hidden is None or geo_residual is None:\n","            return {}\n","\n","        content_cls = self['nce_content_proj'](hidden[:, cls_index])\n","        geo_cls = self['nce_geo_proj'](geo_residual[:, cls_index].detach())\n","\n","        loss, acc = self['nce_bank'](content_cls, geo_cls)\n","        return {'nce': loss, 'nce_acc': acc}\n","\n","    @torch.no_grad()\n","    def update_nce_bank(self, cls_index=0):\n","        \"\"\"Enqueue projected geo keys into bank. Call AFTER backward.\"\"\"\n","        if not self.has('nce_bank') or not self.has('nce_geo_proj'):\n","            return\n","\n","        geo_residual = getattr(self, '_last_geo_residual', None)\n","        if geo_residual is None:\n","            return\n","\n","        geo_cls = self['nce_geo_proj'](geo_residual[:, cls_index].detach())\n","        self['nce_bank'].enqueue(F.normalize(geo_cls, dim=-1))\n","\n","    def anchor_diagnostics(self):\n","        \"\"\"Per-layer anchor health diagnostics.\"\"\"\n","        diag = {}\n","        for i in range(self.n_layers):\n","            layer = self[f'layer_{i}']\n","            anchors = layer['observer'].association.constellation.anchors\n","            anchors_n = F.normalize(anchors.detach(), dim=-1)\n","            A = anchors_n.shape[0]\n","\n","            cos = anchors_n @ anchors_n.T\n","            idx = torch.triu_indices(A, A, offset=1, device=cos.device)\n","            pairwise = 1.0 - cos[idx[0], idx[1]]\n","            cv = (pairwise.std() / (pairwise.mean() + 1e-8)).item()\n","\n","            with torch.no_grad():\n","                anchor_cm, _ = anchor_neighborhood_cm(\n","                    anchors_n, layer['cm_gate'].n_neighbors)\n","\n","            diag[f'layer_{i}'] = {\n","                'anchor_cv': cv,\n","                'mean_pairwise_dist': pairwise.mean().item(),\n","                'min_pairwise_dist': pairwise.min().item(),\n","                'cm_positive_frac': (anchor_cm > 0).float().mean().item(),\n","                'cm_mean': anchor_cm.mean().item(),\n","                'cm_std': anchor_cm.std().item(),\n","            }\n","        return diag\n","\n","    def param_report(self):\n","        total = 0\n","        name = getattr(self, '_tower_name', self.__class__.__name__)\n","        print(f\"\\n  {name} — parameter report (CM-validated + flows)\")\n","        print(f\"  {'Component':<35s}  {'Params':>12s}\")\n","        print(f\"  {'─'*35}  {'─'*12}\")\n","        for cname, module in self.named_children():\n","            n = sum(p.numel() for p in module.parameters())\n","            total += n\n","            print(f\"  {cname:<35s}  {n:>12,}\")\n","        print(f\"  {'─'*35}  {'─'*12}\")\n","        print(f\"  {'TOTAL':<35s}  {total:>12,}\")\n","        return total\n","\n","\n","    def forward(self, x, attn_mask=None, key_padding_mask=None,\n","                return_geo_state=False):\n","        if self.has('embed') and x.dtype in (torch.long, torch.int32, torch.int64):\n","            pos = torch.arange(x.shape[1], device=x.device)\n","            x = self['embed'](x) + self['pos_embed'](pos)\n","\n","        geo_states = []\n","        has_xrot = self.has('cross_rot_0')\n","        geo_residual = None\n","\n","        for i in range(self.n_layers):\n","            x, geo_residual, geo_state = self[f'layer_{i}'](\n","                x, geo_residual=geo_residual,\n","                attn_mask=attn_mask, key_padding_mask=key_padding_mask)\n","            if return_geo_state:\n","                geo_states.append(geo_state)\n","            if has_xrot and i < self.n_layers - 1:\n","                x = self[f'cross_rot_{i}'](x)\n","\n","        self._last_geo_residual = geo_residual\n","        self._last_hidden = x\n","\n","        x = self['final_norm'](x)\n","        if self.has('head'):\n","            x = self['head'](x)\n","\n","        return (x, geo_states) if return_geo_state else x\n","\n","    # ── Paired forward + observer loss ──────────────────────────────\n","\n","    def _run_view(self, x, attn_mask=None, key_padding_mask=None):\n","        \"\"\"Run one view through the full pipeline.\n","        Retains ALL layers' geo_states — every layer needs gradient.\n","        \"\"\"\n","        has_xrot = self.has('cross_rot_0')\n","        geo_residual = None\n","\n","        if self.has('embed') and x.dtype in (torch.long, torch.int32, torch.int64):\n","            pos = torch.arange(x.shape[1], device=x.device)\n","            x = self['embed'](x) + self['pos_embed'](pos)\n","\n","        geo_states = []\n","        for i in range(self.n_layers):\n","            x, geo_residual, geo_state = self[f'layer_{i}'](\n","                x, geo_residual=geo_residual,\n","                attn_mask=attn_mask, key_padding_mask=key_padding_mask)\n","            geo_states.append(geo_state)\n","            if has_xrot and i < self.n_layers - 1:\n","                x = self[f'cross_rot_{i}'](x)\n","\n","        x = self['final_norm'](x)\n","        return x, geo_states\n","\n","    def forward_paired(self, x1, x2, cls_index=0,\n","                       attn_mask=None, key_padding_mask=None):\n","        \"\"\"Dual-view forward for observer loss training.\n","\n","        Aggregates geometric observations from ALL layers so every\n","        layer's observer/constellation/flows receive gradient.\n","        Patchwork, bridge, triangulation, etc. are summed across layers —\n","        the observer loss sees the full-stack geometric observation.\n","        \"\"\"\n","        B = x1.shape[0]\n","        x_cat = torch.cat([x1, x2], dim=0)\n","        feat_cat, geo_states = self._run_view(x_cat, attn_mask, key_padding_mask)\n","\n","        c = cls_index\n","        n_layers = len(geo_states)\n","\n","        # Aggregate geometric observations across all layers\n","        def _agg(key):\n","            \"\"\"Sum across layers — every layer gets gradient.\"\"\"\n","            vals = [gs[key] for gs in geo_states if gs.get(key) is not None]\n","            if not vals:\n","                return None\n","            return sum(vals)\n","\n","        emb_agg = _agg('embedding')\n","        pw_agg = _agg('patchwork')\n","        br_agg = _agg('bridge')\n","        asgn_agg = _agg('assignment')\n","        cos_agg = _agg('cos_to_anchors')\n","        tri_agg = _agg('triangulation')\n","        gv_agg = _agg('gate_values')\n","        cm_agg = _agg('cm_quality')\n","\n","        return {\n","            'embedding':      emb_agg[:B, c],\n","            'embedding_aug':  emb_agg[B:, c],\n","            'patchwork1':     pw_agg[:B, c],\n","            'patchwork1_aug': pw_agg[B:, c],\n","            'bridge1':        br_agg[:B, c],\n","            'bridge2':        br_agg[B:, c],\n","            'assign1':        asgn_agg[:B, c],\n","            'assign2':        asgn_agg[B:, c],\n","            'cos1':           cos_agg[:B, c],\n","            'tri1':           tri_agg[:B, c],\n","            'tri2':           tri_agg[B:, c],\n","            'features1':      feat_cat[:B],\n","            'features2':      feat_cat[B:],\n","            'gate_values1':   gv_agg[:B, c],\n","            'gate_values2':   gv_agg[B:, c],\n","            'cm_quality1':    cm_agg[:B],\n","            'cm_quality2':    cm_agg[B:],\n","        }\n","\n","    def compute_loss(self, output, targets, cls_index=0,\n","                     w_ce=1.0, head=None, **loss_kwargs):\n","        final_layer = self[f'layer_{self.n_layers - 1}']\n","        anchors = final_layer['observer'].association.constellation.anchors\n","\n","        obs_loss, ld = _geolip_observer_loss(\n","            output, anchors=anchors, targets=targets,\n","            **loss_kwargs)\n","\n","        if head is not None:\n","            feat1 = output['features1'][:, cls_index]\n","            feat2 = output['features2'][:, cls_index]\n","            logits1 = head(feat1)\n","            logits2 = head(feat2)\n","            l_ce, acc = _geolip_ce_loss_paired(logits1, logits2, targets)\n","            ld['ce'], ld['acc'] = l_ce, acc\n","            ld['logits'] = logits1\n","            loss = w_ce * l_ce + obs_loss\n","            ld['loss_task'] = l_ce.detach()\n","        else:\n","            loss = obs_loss\n","\n","        ld['loss_observer'] = obs_loss.detach()\n","\n","        w_spread = loss_kwargs.get('w_spread', 0.01)\n","        if self.n_layers > 1 and w_spread > 0:\n","            other_spread = torch.tensor(0.0, device=anchors.device)\n","            for i in range(self.n_layers - 1):\n","                layer = self[f'layer_{i}']\n","                layer_anchors = layer['observer'].association.constellation.anchors\n","                other_spread = other_spread + _geolip_spread_loss(layer_anchors)\n","            other_spread = w_spread * other_spread / (self.n_layers - 1)\n","            loss = loss + other_spread\n","            ld['spread_other_layers'] = other_spread.detach()\n","\n","        ld['total'] = loss\n","        return loss, ld\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# FACTORIES\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","def geo_transformer_esm2(name='geo_esm2', n_layers=6, **kw):\n","    \"\"\"Pre-configured for ESM-2 650M (d=1280).\"\"\"\n","    return GeometricTransformer(name, d_model=1280, n_heads=16,\n","        n_layers=n_layers, n_anchors=32, manifold_dim=256,\n","        n_comp=8, d_comp=32, context_dim=128, quat_dim=64, **kw)\n","\n","def geo_transformer_small(name='geo_small', n_layers=4, **kw):\n","    \"\"\"Small config for prototyping.\"\"\"\n","    return GeometricTransformer(name, d_model=256, n_heads=8,\n","        n_layers=n_layers, n_anchors=16, manifold_dim=128,\n","        n_comp=4, d_comp=16, context_dim=64, quat_dim=32, **kw)\n","\n","def geo_transformer_vision(name='geo_vit', n_layers=4, **kw):\n","    \"\"\"For scatter/SVD vision pipeline (patches as tokens).\"\"\"\n","    return GeometricTransformer(name, d_model=384, n_heads=8,\n","        n_layers=n_layers, n_anchors=32, manifold_dim=128,\n","        n_comp=8, d_comp=16, context_dim=64, quat_dim=32, **kw)\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# SELF-TEST\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","if __name__ == '__main__':\n","    print(\"Geometric Transformer — CM Validated — Self-Test\")\n","    print(\"=\" * 60)\n","\n","    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n","\n","    # ── Build small model ──\n","    model = geo_transformer_small('test_cm', n_layers=2)\n","    if hasattr(model, 'network_to'):\n","        model.network_to(device=device, strict=False)\n","    else:\n","        model = model.to(device)\n","    total = model.param_report()\n","\n","    # ── Forward pass ──\n","    B, L, D = 2, 32, 256\n","    x = torch.randn(B, L, D, device=device)\n","    out, geos = model(x, return_geo_state=True)\n","\n","    assert out.shape == (B, L, D), f\"Expected ({B},{L},{D}), got {out.shape}\"\n","    assert len(geos) == 2\n","    print(f\"\\n  Input:  ({B}, {L}, {D})\")\n","    print(f\"  Output: {out.shape}\")\n","    print(f\"  Geo states: {len(geos)} layers\")\n","\n","    # ── Verify CM gate is active ──\n","    for i, gs in enumerate(geos):\n","        gi = gs['gate_info']\n","        cm_q = gs['cm_quality']\n","        gv = gs['gate_values']\n","        print(f\"\\n  Layer {i} CM gate:\")\n","        print(f\"    active anchors:   {gi['active'].item():.1f} / {model.n_anchors}\")\n","        print(f\"    gate mean:        {gi['gate_mean'].item():.4f}\")\n","        print(f\"    cm_positive_frac: {gi['cm_positive_frac'].item():.3f}\")\n","        print(f\"    gate_values:      {gv.shape}  range=[{gv.min():.3f}, {gv.max():.3f}]\")\n","        print(f\"    cm_quality:       {cm_q.shape}  mean={cm_q.mean():.4f}\")\n","\n","    # ── Verify geo_residual continuity ──\n","    gr0 = geos[0]['geo_residual']\n","    gr1 = geos[1]['geo_residual']\n","    print(f\"\\n  Geo residual stream:\")\n","    print(f\"    Layer 0: {gr0.shape}  norm={gr0.norm(dim=-1).mean():.4f}\")\n","    print(f\"    Layer 1: {gr1.shape}  norm={gr1.norm(dim=-1).mean():.4f}\")\n","\n","    # ── Geometric losses ──\n","    geo_losses = model.geometric_losses()\n","    print(f\"\\n  Geometric regularization:\")\n","    for k, v in geo_losses.items():\n","        print(f\"    {k}: {v.item():.6f}\")\n","\n","    # ── Anchor diagnostics ──\n","    diag = model.anchor_diagnostics()\n","    print(f\"\\n  Anchor diagnostics:\")\n","    for layer_name, d in diag.items():\n","        print(f\"    {layer_name}:\")\n","        for k, v in d.items():\n","            print(f\"      {k}: {v:.4f}\")\n","\n","    # ── Verify Cayley rotations ──\n","    print(f\"\\n  Cayley rotations:\")\n","    for name, module in model.named_modules():\n","        if isinstance(module, CayleyOrthogonal):\n","            R = module.get_rotation()\n","            I = torch.eye(R.shape[0], device=R.device)\n","            print(f\"    {name}: ‖RRᵀ-I‖={((R@R.T)-I).norm():.8f}  det={torch.det(R):.4f}\")\n","\n","    # ── Gradient flow through CM gate ──\n","    print(f\"\\n  Gradient flow test:\")\n","    model.zero_grad()\n","    x_grad = torch.randn(B, L, D, device=device, requires_grad=True)\n","    out_grad = model(x_grad)\n","    loss = out_grad.sum()\n","    loss.backward()\n","\n","    # Check gate_proj has gradients\n","    for i in range(model.n_layers):\n","        layer = model[f'layer_{i}']\n","        gate_grads = [p.grad is not None and p.grad.abs().sum() > 0\n","                      for p in layer['cm_gate'].parameters()]\n","        print(f\"    layer_{i} cm_gate grad: {'YES' if all(gate_grads) else 'NO'}\")\n","\n","    # ── Training step simulation ──\n","    print(f\"\\n  Training step simulation:\")\n","    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)\n","    optimizer.zero_grad()\n","\n","    x_train = torch.randn(B, L, D, device=device)\n","    out_train, states = model(x_train, return_geo_state=True)\n","    task_loss = out_train.mean()  # dummy\n","\n","    geo_losses = model.geometric_losses()\n","    total_loss = task_loss + geo_losses.get('geo_total', 0.0)\n","    total_loss.backward()\n","    optimizer.step()\n","    print(f\"    task_loss:  {task_loss.item():.4f}\")\n","    print(f\"    cv_loss:    {geo_losses['cv'].item():.6f}\")\n","    print(f\"    spread_loss:{geo_losses['spread'].item():.6f}\")\n","    print(f\"    total:      {total_loss.item():.4f}\")\n","\n","    # ── Paired forward + observer loss ──\n","    print(f\"\\n  Paired forward + observer loss:\")\n","    model.zero_grad()\n","\n","    x1 = torch.randn(B, L, D, device=device)\n","    x2 = x1 + 0.1 * torch.randn_like(x1)  # view 2 = slight perturbation\n","    targets = torch.randint(0, 10, (B,), device=device)\n","\n","    output = model.forward_paired(x1, x2)\n","    print(f\"    Output keys: {sorted(k for k in output if not k.startswith('geo_'))}\")\n","    for k in ['embedding', 'patchwork1', 'bridge1', 'assign1', 'tri1']:\n","        print(f\"    {k}: {output[k].shape}\")\n","\n","    # Task head for CE\n","    num_classes = 10\n","    head = nn.Linear(D, num_classes).to(device)\n","\n","    loss, ld = model.compute_loss(output, targets, head=head)\n","    print(f\"\\n    Three-domain loss breakdown:\")\n","    for k in ['loss_observer', 'loss_task', 'ce', 'nce_emb', 'nce_pw',\n","               'bridge', 'assign', 'assign_nce', 'nce_tri', 'attract',\n","               'cv', 'spread']:\n","        if k in ld:\n","            v = ld[k]\n","            v = v.item() if isinstance(v, torch.Tensor) else v\n","            print(f\"      {k:16s} = {v:.4f}\")\n","    for k in ['nce_emb_acc', 'nce_pw_acc', 'nce_tri_acc', 'bridge_acc',\n","               'assign_nce_acc', 'acc']:\n","        if k in ld:\n","            v = ld[k]\n","            v = v if isinstance(v, float) else v.item()\n","            print(f\"      {k:16s} = {v*100:.1f}%\")\n","    print(f\"      {'TOTAL':16s} = {loss.item():.4f}\")\n","\n","    # Verify backward through observer loss\n","    loss.backward()\n","    alive_base, dead_base = [], []\n","    for n, p in model.named_parameters():\n","        if p.grad is not None and p.grad.norm() > 0:\n","            alive_base.append(n)\n","        else:\n","            dead_base.append(n)\n","    print(f\"\\n    Gradient flow: {len(alive_base)} params alive, {len(dead_base)} dead\")\n","    if dead_base:\n","        print(f\"\\n    DEAD parameters (base model, paired+observer):\")\n","        for n in dead_base:\n","            print(f\"      {n}\")\n","\n","    # ══════════════════════════════════════════════════════════════\n","    # WIDE ROUTER COMPILATION\n","    # ══════════════════════════════════════════════════════════════\n","\n","    print(f\"\\n{'='*60}\")\n","    print(f\"  WideRouter Compilation\")\n","    print(f\"{'='*60}\")\n","\n","    if _HAS_WIDE_ROUTER:\n","        # Wrap transformer in WideRouter (same pattern as GeoViTClassifier)\n","        router = WideRouter('test_router', strict=False)\n","        router.attach('transformer', model)\n","        router.register_tower('transformer')\n","        router.network_to(device=device, strict=False)\n","\n","        # Discover towers and compile\n","        router.discover_towers()\n","        print(f\"\\n  Towers discovered: {router.tower_names}\")\n","        print(f\"  Analyzed: {router.objects.get('_analyzed', False)}\")\n","\n","        try:\n","            compiled_router = router.compile(mode='default')\n","            print(f\"  WideRouter.compile(mode='default'): OK\")\n","        except Exception as e:\n","            print(f\"  WideRouter.compile: {str(e)[:60]}\")\n","\n","        # Forward through the registered tower directly\n","        with torch.no_grad():\n","            out_via_router = router['transformer'](x)\n","        print(f\"  Forward via router['transformer']: {out_via_router.shape}  OK\")\n","\n","        del router\n","    else:\n","        print(f\"\\n  WideRouter: geofractal not installed\")\n","\n","    print(f\"\\n{'='*60}\")\n","    print(f\"  PASSED — CM-validated pipeline operational\")\n","    print(f\"{'='*60}\")\n","\n","    # ══════════════════════════════════════════════════════════════\n","    # FLOW ENSEMBLE INTEGRATION TESTS\n","    # ══════════════════════════════════════════════════════════════\n","\n","    print(f\"\\n{'='*60}\")\n","    print(f\"  Flow Ensemble Integration\")\n","    print(f\"{'='*60}\")\n","\n","    del model, optimizer\n","    torch.cuda.empty_cache() if device.type == 'cuda' else None\n","\n","    model_f = geo_transformer_small('test_flows', n_layers=2,\n","                                     flow_keys=['quat_lite', 'velocity', 'orbital'])\n","    if hasattr(model_f, 'network_to'):\n","        model_f.network_to(device=device, strict=False)\n","    else:\n","        model_f = model_f.to(device)\n","\n","    total_f = model_f.param_report()\n","    print(f\"\\n  Total params (with flows): {total_f:,}\")\n","\n","    print(f\"\\n  Flow ensemble per layer:\")\n","    for i in range(model_f.n_layers):\n","        layer = model_f[f'layer_{i}']\n","        if layer.has('flows'):\n","            flows = layer['flows']\n","            names = flows.active_flow_names\n","            params = sum(p.numel() for p in flows.parameters())\n","            print(f\"    layer_{i}: {names}  ({params:,} params)\")\n","        else:\n","            print(f\"    layer_{i}: no flows attached\")\n","\n","    x_f = torch.randn(B, L, D, device=device)\n","    out_f, geos_f = model_f(x_f, return_geo_state=True)\n","    assert out_f.shape == (B, L, D)\n","    print(f\"\\n  Forward with flows: {out_f.shape}  OK\")\n","\n","    geo_ctx_0 = geos_f[0]['geo_ctx']\n","    print(f\"  Geo context shape: {geo_ctx_0.shape}  norm={geo_ctx_0.norm(dim=-1).mean():.4f}\")\n","\n","    print(f\"\\n  Flow gradient test (out.sum().backward()):\")\n","    model_f.zero_grad()\n","    x_fg = torch.randn(B, L, D, device=device, requires_grad=True)\n","    out_fg = model_f(x_fg)\n","    out_fg.sum().backward()\n","\n","    alive_simple, dead_simple = [], []\n","    for n, p in model_f.named_parameters():\n","        if p.grad is not None and p.grad.abs().sum() > 0:\n","            alive_simple.append(n)\n","        else:\n","            dead_simple.append(n)\n","    print(f\"    {len(alive_simple)} alive, {len(dead_simple)} dead\")\n","    if dead_simple:\n","        print(f\"\\n    DEAD parameters (out.sum):\")\n","        for n in dead_simple:\n","            print(f\"      {n}\")\n","\n","    print(f\"\\n  Paired forward + observer loss (with flows):\")\n","    model_f.zero_grad()\n","    x1_f = torch.randn(B, L, D, device=device)\n","    x2_f = x1_f + 0.1 * torch.randn_like(x1_f)\n","    targets_f = torch.randint(0, 10, (B,), device=device)\n","\n","    output_f = model_f.forward_paired(x1_f, x2_f)\n","    head_f = nn.Linear(D, num_classes).to(device)\n","    loss_f, ld_f = model_f.compute_loss(output_f, targets_f, head=head_f)\n","    print(f\"    total loss: {loss_f.item():.4f}\")\n","    loss_f.backward()\n","\n","    alive_paired, dead_paired = [], []\n","    for n, p in model_f.named_parameters():\n","        if p.grad is not None and p.grad.abs().sum() > 0:\n","            alive_paired.append(n)\n","        else:\n","            dead_paired.append(n)\n","    print(f\"    {len(alive_paired)} alive, {len(dead_paired)} dead\")\n","    if dead_paired:\n","        print(f\"\\n    DEAD parameters (paired+observer):\")\n","        for n in dead_paired:\n","            print(f\"      {n}\")\n","\n","    print(f\"\\n  Runtime flow management:\")\n","    layer0 = model_f['layer_0']\n","    flows_0 = layer0['flows']\n","    print(f\"    Before:           {flows_0.active_flow_names}\")\n","\n","    flows_0.attach_flow('alignment')\n","    print(f\"    +alignment:       {flows_0.active_flow_names}\")\n","\n","    flows_0.detach_flow('velocity')\n","    print(f\"    -velocity:        {flows_0.active_flow_names}\")\n","\n","    out_swapped = model_f(x_f)\n","    assert out_swapped.shape == (B, L, D)\n","    print(f\"    Forward after swap: {out_swapped.shape}  OK\")\n","\n","    layer1 = model_f['layer_1']\n","    if layer1.has('flows'):\n","        for fn in list(layer1['flows'].active_flow_names):\n","            key = fn.replace('flow_', '')\n","            layer1['flows'].detach_flow(key)\n","        print(f\"    Layer 1 after clear: {layer1['flows'].active_flow_names}\")\n","        out_partial = model_f(x_f)\n","        assert out_partial.shape == (B, L, D)\n","        print(f\"    Forward (L0 flows, L1 empty): {out_partial.shape}  OK\")\n","\n","    print(f\"\\n  Backward compatibility (no flows):\")\n","    model_nf = geo_transformer_small('test_noflows', n_layers=2)\n","    if hasattr(model_nf, 'network_to'):\n","        model_nf.network_to(device=device, strict=False)\n","    else:\n","        model_nf = model_nf.to(device)\n","    out_nf = model_nf(torch.randn(B, L, D, device=device))\n","    assert out_nf.shape == (B, L, D)\n","    print(f\"    Forward (no flows): {out_nf.shape}  OK\")\n","    for i in range(model_nf.n_layers):\n","        assert not model_nf[f'layer_{i}'].has('flows'), f\"layer_{i} should not have flows\"\n","    print(f\"    No flows attached:  OK\")\n","    del model_nf\n","\n","    print(f\"\\n{'='*60}\")\n","    print(f\"  PASSED — CM-validated pipeline operational\")\n","    print(f\"  PASSED — Flow ensemble integration verified\")\n","    print(f\"  PASSED — Flow attach/detach verified\")\n","    print(f\"  PASSED — Backward compatibility verified\")\n","    print(f\"{'='*60}\")"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"aUiHf_oR7xxo","executionInfo":{"status":"ok","timestamp":1775132634560,"user_tz":420,"elapsed":535,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"8a4d8e5f-fe82-4cd9-e239-533a98ec05d6"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["Geometric Transformer — CM Validated — Self-Test\n","============================================================\n","\n","  GeometricTransformer — parameter report (CM-validated + flows)\n","  Component                                  Params\n","  ───────────────────────────────────  ────────────\n","  components                              4,300,066\n","  stages                                          0\n","  ───────────────────────────────────  ────────────\n","  TOTAL                                   4,300,066\n","\n","  Input:  (2, 32, 256)\n","  Output: torch.Size([2, 32, 256])\n","  Geo states: 2 layers\n","\n","  Layer 0 CM gate:\n","    active anchors:   16.0 / 16\n","    gate mean:        0.8814\n","    cm_positive_frac: 0.500\n","    gate_values:      torch.Size([2, 32, 16])  range=[0.881, 0.883]\n","    cm_quality:       torch.Size([2, 32, 1])  mean=0.8814\n","\n","  Layer 1 CM gate:\n","    active anchors:   16.0 / 16\n","    gate mean:        0.8808\n","    cm_positive_frac: 0.562\n","    gate_values:      torch.Size([2, 32, 16])  range=[0.880, 0.882]\n","    cm_quality:       torch.Size([2, 32, 1])  mean=0.8808\n","\n","  Geo residual stream:\n","    Layer 0: torch.Size([2, 32, 64])  norm=7.0513\n","    Layer 1: torch.Size([2, 32, 64])  norm=10.6125\n","\n","  Geometric regularization:\n","    cv: 0.003598\n","    spread: 0.000000\n","    geo_total: 0.003598\n","\n","  Anchor diagnostics:\n","    layer_0:\n","      anchor_cv: 0.0260\n","      mean_pairwise_dist: 1.0664\n","      min_pairwise_dist: 1.0248\n","      cm_positive_frac: 1.0000\n","      cm_mean: 3.6145\n","      cm_std: 0.0238\n","    layer_1:\n","      anchor_cv: 0.0246\n","      mean_pairwise_dist: 1.0662\n","      min_pairwise_dist: 1.0269\n","      cm_positive_frac: 1.0000\n","      cm_mean: 3.6192\n","      cm_std: 0.0170\n","\n","  Cayley rotations:\n","    components.layer_0.components.rotation: ‖RRᵀ-I‖=0.00000000  det=1.0000\n","    components.layer_1.components.rotation: ‖RRᵀ-I‖=0.00000000  det=1.0000\n","    components.cross_rot_0: ‖RRᵀ-I‖=0.00000000  det=1.0000\n","\n","  Gradient flow test:\n","    layer_0 cm_gate grad: YES\n","    layer_1 cm_gate grad: YES\n","\n","  Training step simulation:\n","    task_loss:  0.0000\n","    cv_loss:    0.003598\n","    spread_loss:0.000000\n","    total:      0.0036\n","\n","  Paired forward + observer loss:\n","    Output keys: ['assign1', 'assign2', 'bridge1', 'bridge2', 'cm_quality1', 'cm_quality2', 'cos1', 'embedding', 'embedding_aug', 'features1', 'features2', 'gate_values1', 'gate_values2', 'patchwork1', 'patchwork1_aug', 'tri1', 'tri2']\n","    embedding: torch.Size([2, 128])\n","    patchwork1: torch.Size([2, 64])\n","    bridge1: torch.Size([2, 16])\n","    assign1: torch.Size([2, 16])\n","    tri1: torch.Size([2, 16])\n","\n","    Three-domain loss breakdown:\n","      loss_observer    = 7.2543\n","      loss_task        = 2.3369\n","      ce               = 2.3369\n","      nce_emb          = 0.0000\n","      nce_pw           = 0.6328\n","      bridge           = 5.9117\n","      assign           = 0.2186\n","      assign_nce       = 0.2596\n","      nce_tri          = 0.6607\n","      attract          = 0.8211\n","      cv               = 0.0000\n","      spread           = 0.0000\n","      nce_emb_acc      = 100.0%\n","      nce_pw_acc       = 100.0%\n","      nce_tri_acc      = 100.0%\n","      bridge_acc       = 0.0%\n","      assign_nce_acc   = 100.0%\n","      acc              = 0.0%\n","      TOTAL            = 9.5912\n","\n","    Gradient flow: 225 params alive, 24 dead\n","\n","    DEAD parameters (base model, paired+observer):\n","      components.layer_0.components.context.history_mlp.0.weight\n","      components.layer_0.components.context.history_mlp.0.bias\n","      components.layer_0.components.context.history_mlp.2.weight\n","      components.layer_0.components.context.history_mlp.2.bias\n","      components.layer_0.components.context.flow_mlp.0.weight\n","      components.layer_0.components.context.flow_mlp.0.bias\n","      components.layer_0.components.context.flow_mlp.2.weight\n","      components.layer_0.components.context.flow_mlp.2.bias\n","      components.layer_1.components.context.flow_mlp.0.weight\n","      components.layer_1.components.context.flow_mlp.0.bias\n","      components.layer_1.components.context.flow_mlp.2.weight\n","      components.layer_1.components.context.flow_mlp.2.bias\n","      components.layer_1.components.geo_proj.0.weight\n","      components.layer_1.components.geo_proj.0.bias\n","      components.layer_1.components.geo_proj.1.weight\n","      components.layer_1.components.geo_proj.1.bias\n","      components.nce_content_proj.0.weight\n","      components.nce_content_proj.0.bias\n","      components.nce_content_proj.2.weight\n","      components.nce_content_proj.2.bias\n","      components.nce_geo_proj.0.weight\n","      components.nce_geo_proj.0.bias\n","      components.nce_geo_proj.2.weight\n","      components.nce_geo_proj.2.bias\n","\n","============================================================\n","  WideRouter Compilation\n","============================================================\n","\n","  Towers discovered: ['transformer']\n","  Analyzed: False\n","  WideRouter.compile(mode='default'): OK\n","  Forward via router['transformer']: torch.Size([2, 32, 256])  OK\n","\n","============================================================\n","  PASSED — CM-validated pipeline operational\n","============================================================\n","\n","============================================================\n","  Flow Ensemble Integration\n","============================================================\n","\n","  GeometricTransformer — parameter report (CM-validated + flows)\n","  Component                                  Params\n","  ───────────────────────────────────  ────────────\n","  components                              4,390,764\n","  stages                                          0\n","  ───────────────────────────────────  ────────────\n","  TOTAL                                   4,390,764\n","\n","  Total params (with flows): 4,390,764\n","\n","  Flow ensemble per layer:\n","    layer_0: ['flow_quat_lite', 'flow_velocity', 'flow_orbital']  (45,348 params)\n","    layer_1: ['flow_quat_lite', 'flow_velocity', 'flow_orbital']  (45,348 params)\n","\n","  Forward with flows: torch.Size([2, 32, 256])  OK\n","  Geo context shape: torch.Size([2, 32, 64])  norm=7.9993\n","\n","  Flow gradient test (out.sum().backward()):\n","    259 alive, 20 dead\n","\n","    DEAD parameters (out.sum):\n","      components.layer_0.components.observer.curation.bridge.0.weight\n","      components.layer_0.components.observer.curation.bridge.0.bias\n","      components.layer_0.components.context.history_mlp.0.weight\n","      components.layer_0.components.context.history_mlp.0.bias\n","      components.layer_0.components.context.history_mlp.2.weight\n","      components.layer_0.components.context.history_mlp.2.bias\n","      components.layer_1.components.observer.curation.bridge.0.weight\n","      components.layer_1.components.observer.curation.bridge.0.bias\n","      components.layer_1.components.geo_proj.0.weight\n","      components.layer_1.components.geo_proj.0.bias\n","      components.layer_1.components.geo_proj.1.weight\n","      components.layer_1.components.geo_proj.1.bias\n","      components.nce_content_proj.0.weight\n","      components.nce_content_proj.0.bias\n","      components.nce_content_proj.2.weight\n","      components.nce_content_proj.2.bias\n","      components.nce_geo_proj.0.weight\n","      components.nce_geo_proj.0.bias\n","      components.nce_geo_proj.2.weight\n","      components.nce_geo_proj.2.bias\n","\n","  Paired forward + observer loss (with flows):\n","    total loss: 9.5802\n","    263 alive, 16 dead\n","\n","    DEAD parameters (paired+observer):\n","      components.layer_0.components.context.history_mlp.0.weight\n","      components.layer_0.components.context.history_mlp.0.bias\n","      components.layer_0.components.context.history_mlp.2.weight\n","      components.layer_0.components.context.history_mlp.2.bias\n","      components.layer_1.components.geo_proj.0.weight\n","      components.layer_1.components.geo_proj.0.bias\n","      components.layer_1.components.geo_proj.1.weight\n","      components.layer_1.components.geo_proj.1.bias\n","      components.nce_content_proj.0.weight\n","      components.nce_content_proj.0.bias\n","      components.nce_content_proj.2.weight\n","      components.nce_content_proj.2.bias\n","      components.nce_geo_proj.0.weight\n","      components.nce_geo_proj.0.bias\n","      components.nce_geo_proj.2.weight\n","      components.nce_geo_proj.2.bias\n","\n","  Runtime flow management:\n","    Before:           ['flow_quat_lite', 'flow_velocity', 'flow_orbital']\n","    +alignment:       ['flow_quat_lite', 'flow_velocity', 'flow_orbital', 'flow_alignment']\n","    -velocity:        ['flow_quat_lite', 'flow_orbital', 'flow_alignment']\n","    Forward after swap: torch.Size([2, 32, 256])  OK\n","    Layer 1 after clear: []\n","    Forward (L0 flows, L1 empty): torch.Size([2, 32, 256])  OK\n","\n","  Backward compatibility (no flows):\n","    Forward (no flows): torch.Size([2, 32, 256])  OK\n","    No flows attached:  OK\n","\n","============================================================\n","  PASSED — CM-validated pipeline operational\n","  PASSED — Flow ensemble integration verified\n","  PASSED — Flow attach/detach verified\n","  PASSED — Backward compatibility verified\n","============================================================\n"]}]},{"cell_type":"markdown","source":["# train cifar100 flow match geometric transformer"],"metadata":{"id":"98wsZC2gKNqI"}},{"cell_type":"code","source":["\"\"\"\n","Geometric Transformer — CIFAR-100 Training with Three-Domain Observer Loss\n","\n","Full observer loss spectrum wired through the CM-gated geometric pipeline:\n","  EXTERNAL:   CE (paired views) + embedding NCE\n","  GEOMETRIC:  patchwork NCE + bridge loss  ← flows through CM gate + flow opinions\n","  INTERNAL:   assign BCE + assign NCE + tri NCE + attraction + CV + spread\n","\n","Pipeline:\n","  ConvSVDPatchEmbedding → GeometricTransformer (CM-gated, flows) → ClassificationHead\n","  SVD through geolip_core.linalg (FL kernel, no graph breaks)\n","  WideRouter + CompileRouter stage decomposition\n","  torch.compile inductor fusion on stacked operations\n","\n","!pip install geolip-core torchvision tqdm tensorboard\n","\"\"\"\n","\n","import os, sys\n","import warnings\n","warnings.filterwarnings('ignore')\n","\n","import torch\n","import torch.nn as nn\n","import torch.nn.functional as F\n","import numpy as np\n","import time, json, math\n","from pathlib import Path\n","from tqdm.auto import tqdm\n","from torch.utils.tensorboard import SummaryWriter\n","\n","device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n","print(f\"Device: {device}\")\n","if device.type == 'cuda':\n","    print(f\"  GPU: {torch.cuda.get_device_name()}\")\n","    print(f\"  VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB\")\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# IMPORTS\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","from geolip_core.pipeline.components.geometric_transformer import (\n","    GeometricTransformer, GeometricTransformerLayer,\n","    CayleyOrthogonal, QuaternionCompose, FiLMLayer,\n","    ContentAttention, GeometricAttention, CMValidatedGate,\n","    anchor_neighborhood_cm,\n",")\n","from geolip_core.pipeline.observer import TorchComponent\n","from geofractal.router.wide_router import WideRouter\n","\n","torch.set_float32_matmul_precision('high')\n","torch.backends.cudnn.benchmark = True\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# CONFIG\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","CONFIG = {\n","    # Model\n","    'd_model': 384,\n","    'n_heads': 8,\n","    'n_layers': 8,\n","    'n_anchors': 128,\n","    'manifold_dim': 128,\n","    'n_comp': 8,\n","    'd_comp': 32,\n","    'context_dim': 128,\n","    'quat_dim': 64,\n","    'dropout': 0.1,\n","    'cm_neighbors': 3,\n","\n","    # Flow ensemble\n","    'flow_classes': ['quat_lite', 'velocity', 'orbital'],\n","    'flow_fusion': 'weighted',\n","\n","    # Input stage\n","    'patch_size': 4,\n","    'img_size': 32,\n","    'in_channels': 3,\n","    'conv_channels': 64,\n","    'svd_rank': 12,           # ≤12 for FL eigh (CUDA-graph-safe)\n","\n","    # Training\n","    'epochs': 100,\n","    'batch_size': 128,\n","    'lr': 3e-4,\n","    'weight_decay': 0.005,\n","    'warmup_epochs': 5,\n","    'label_smoothing': 0.1,\n","    'num_workers': 8,\n","\n","    # Observer loss weights — constellation integrity as constraint\n","    'cv_target': 0.20,\n","    'w_ce': 1.0,\n","    'w_nce_emb': 0.5,\n","    'w_nce_pw': 1.0,\n","    'w_bridge': 1.0,\n","    'w_assign': 0.5,\n","    'w_assign_nce': 0.25,\n","    'w_nce_tri': 0.5,\n","    'w_attract': 0.25,\n","    'w_cv': 0.1,         # 10× from 0.01 — prevents constellation collapse\n","    'w_spread': 0.05,    # 5× from 0.01 — anchor repulsion\n","\n","    # Augmentation\n","    'random_erasing_p': 0.25,\n","\n","    # Data\n","    'num_classes': 100,\n","\n","    # Logging\n","    'log_geo_every': 5,\n","    'log_grads_every': 10,\n","    'log_dir': 'runs/geo_observer_flows',\n","\n","    # Precompute\n","    'precompute_freq': 10,  # CM gate + Cayley every N batches\n","}\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# INPUT STAGE — SVD through geolip_core.linalg (no graph breaks)\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","from geolip_core.core.input.svd import SVDObserver\n","\n","\n","class ConvSVDPatchEmbedding(TorchComponent):\n","    \"\"\"Input stage: conv frontend → SVDObserver → patch tokens.\n","    SVD runs through LA.svd(method='gram_eigh') — fully compilable.\"\"\"\n","    def __init__(self, name, img_size=32, patch_size=4, in_channels=3,\n","                 conv_channels=64, d_model=256, svd_rank=16):\n","        super().__init__(name)\n","        self.patch_size = patch_size\n","        self.n_patches = (img_size // patch_size) ** 2\n","        self.d_model = d_model\n","        self.svd_rank = svd_rank\n","\n","        self.conv_frontend = nn.Sequential(\n","            nn.Conv2d(in_channels, conv_channels, 3, padding=1, bias=False),\n","            nn.BatchNorm2d(conv_channels), nn.GELU(),\n","            nn.Conv2d(conv_channels, conv_channels, 3, padding=1, bias=False),\n","            nn.BatchNorm2d(conv_channels), nn.GELU(),\n","        )\n","        self.svd_observer = SVDObserver(conv_channels, svd_rank)\n","        self.patch_proj = nn.Conv2d(\n","            conv_channels, d_model, kernel_size=patch_size,\n","            stride=patch_size, bias=False)\n","        self.patch_norm = nn.LayerNorm(d_model)\n","\n","        svd_feat_dim = self.svd_observer.feature_dim\n","        self.svd_to_gamma = nn.Linear(svd_feat_dim, d_model)\n","        self.svd_to_beta = nn.Linear(svd_feat_dim, d_model)\n","        nn.init.normal_(self.svd_to_gamma.weight, std=0.01)\n","        nn.init.ones_(self.svd_to_gamma.bias)\n","        nn.init.normal_(self.svd_to_beta.weight, std=0.01)\n","        nn.init.zeros_(self.svd_to_beta.bias)\n","\n","        self.cls_token = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)\n","        self.pos_embed = nn.Parameter(\n","            torch.randn(1, self.n_patches + 1, d_model) * 0.02)\n","\n","    def forward(self, x):\n","        B = x.shape[0]\n","        feat = self.conv_frontend(x)\n","        # SVD through geolip_core.linalg — LA.svd(method='gram_eigh')\n","        # Fully compilable, no @torch.compiler.disable needed\n","        S, Vh, svd_features, novelty = self.svd_observer(feat)\n","        # EMA update moved to update_svd_ema() — called outside compiled graph.\n","        # Buffer mutations (self.ema_s, self.ema_vh) break CUDA graph capture.\n","        tokens = self.patch_proj(feat)\n","        tokens = tokens.flatten(2).transpose(1, 2)\n","        tokens = self.patch_norm(tokens)\n","        gamma = self.svd_to_gamma(svd_features).unsqueeze(1)\n","        beta = self.svd_to_beta(svd_features).unsqueeze(1)\n","        tokens = gamma * tokens + beta\n","        cls = self.cls_token.expand(B, -1, -1)\n","        tokens = torch.cat([cls, tokens], dim=1)\n","        tokens = tokens + self.pos_embed\n","        svd_state = {\n","            'singular_values': S, 'Vh': Vh,\n","            'svd_features': svd_features, 'novelty': novelty,\n","        }\n","        return tokens, svd_state\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# CLASSIFIER\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","class ClassificationHead(TorchComponent):\n","    def __init__(self, name, d_model, num_classes, dropout=0.1):\n","        super().__init__(name)\n","        self.net = nn.Sequential(\n","            nn.LayerNorm(d_model),\n","            nn.Linear(d_model, d_model),\n","            nn.GELU(), nn.Dropout(dropout),\n","            nn.Linear(d_model, num_classes),\n","        )\n","\n","    def forward(self, x):\n","        return self.net(x)\n","\n","\n","class GeoViTClassifier(WideRouter):\n","    \"\"\"Geometric Vision Transformer with flow ensemble.\n","\n","    WideRouter orchestrator with CompileRouter stage decomposition:\n","        - patch_embed: ConvSVDPatchEmbedding\n","        - transformer: GeometricTransformer (with flows)\n","        - head: ClassificationHead\n","    \"\"\"\n","    def __init__(self, name, config):\n","        super().__init__(name, strict=False)\n","        self.config = config\n","\n","        self.attach('patch_embed', ConvSVDPatchEmbedding(\n","            'patch_embed', img_size=config['img_size'],\n","            patch_size=config['patch_size'], in_channels=config['in_channels'],\n","            conv_channels=config['conv_channels'], d_model=config['d_model'],\n","            svd_rank=config['svd_rank'],\n","        ))\n","        self.attach('transformer', GeometricTransformer(\n","            'geo_cifar', d_model=config['d_model'], n_heads=config['n_heads'],\n","            n_layers=config['n_layers'], n_anchors=config['n_anchors'],\n","            manifold_dim=config['manifold_dim'], n_comp=config['n_comp'],\n","            d_comp=config['d_comp'], context_dim=config['context_dim'],\n","            quat_dim=config['quat_dim'], dropout=config['dropout'],\n","            cm_neighbors=config.get('cm_neighbors', 3),\n","            nce_bank_size=0,\n","            flow_keys=config.get('flow_classes', None),\n","            flow_fusion=config.get('flow_fusion', 'weighted'),\n","        ))\n","        self.register_tower('transformer')\n","\n","        self.attach('head', ClassificationHead(\n","            'head', config['d_model'], config['num_classes'],\n","            config.get('dropout', 0.1)))\n","\n","    def forward(self, x, x2=None, targets=None, return_geo_state=False):\n","        # NOTE: precompute_cm_gates() must be called BEFORE this forward\n","        # by the training loop. It cannot live here because @torch.compiler.disable\n","        # creates a graph break that splits the CUDA graph into partitions,\n","        # causing backward replay segfaults when partition addresses shift.\n","\n","        if x2 is not None:\n","            B = x.shape[0]\n","            v_cat = torch.cat([x, x2], dim=0)\n","            tokens_cat, svd_state = self['patch_embed'](v_cat)\n","            output = self['transformer'].forward_paired(\n","                tokens_cat[:B], tokens_cat[B:])\n","            if targets is not None:\n","                return self._compute_loss(output, targets)\n","            return output\n","\n","        tokens, svd_state = self['patch_embed'](x)\n","        if return_geo_state:\n","            features, geo_states = self['transformer'](tokens, return_geo_state=True)\n","        else:\n","            features = self['transformer'](tokens)\n","        cls_out = features[:, 0]\n","        logits = self['head'](cls_out)\n","        if return_geo_state:\n","            return logits, geo_states, svd_state\n","        return logits\n","\n","    def _compute_loss(self, output, targets):\n","        \"\"\"Loss computation safe for CUDA graph capture.\n","\n","        Inlines only graph-safe observer losses. Skips entirely:\n","          - cv_loss: @torch.compiler.disable + torch.linalg.det\n","          - spread_loss: boolean indexing (~torch.eye mask)\n","          - knn_accuracy: .item() forces DeviceCopy\n","        These are computed outside via _compute_geo_losses().\n","        \"\"\"\n","        from geolip_core.core.distinguish.losses import (\n","            nce_loss, bridge_loss_paired, assign_bce_loss,\n","            assign_nce_loss, attraction_loss,\n","            ce_loss_paired as _ce_loss,\n","        )\n","        cfg = self.config\n","        final_layer = self['transformer'][f'layer_{self[\"transformer\"].n_layers - 1}']\n","\n","        emb1, emb2 = output['embedding'], output['embedding_aug']\n","\n","        # Embedding NCE\n","        l_nce_emb, nce_emb_acc = nce_loss(emb1, emb2, 0.07, normalize=False)\n","        # Patchwork NCE\n","        l_nce_pw, nce_pw_acc = nce_loss(\n","            output['patchwork1'], output['patchwork1_aug'], 0.1, normalize=True)\n","        # Bridge\n","        l_bridge, bridge_acc = bridge_loss_paired(\n","            output['bridge1'], output['bridge2'],\n","            output['assign1'], output['assign2'])\n","        # Assign BCE\n","        l_assign, assign_ent = assign_bce_loss(output['assign1'], output['cos1'])\n","        # Assign NCE\n","        l_assign_nce, assign_nce_acc = assign_nce_loss(\n","            output['assign1'], output['assign2'], 0.1)\n","        # Tri NCE\n","        l_nce_tri, nce_tri_acc = nce_loss(\n","            output['tri1'], output['tri2'], 0.1, normalize=True)\n","        # Attraction\n","        l_attract, nearest_cos = attraction_loss(output['cos1'])\n","\n","        obs_loss = (\n","            cfg.get('w_nce_emb', 0.5) * l_nce_emb\n","            + cfg.get('w_nce_pw', 1.0) * l_nce_pw\n","            + cfg.get('w_bridge', 1.0) * l_bridge\n","            + cfg.get('w_assign', 0.5) * l_assign\n","            + cfg.get('w_assign_nce', 0.25) * l_assign_nce\n","            + cfg.get('w_nce_tri', 0.5) * l_nce_tri\n","            + cfg.get('w_attract', 0.25) * l_attract\n","        )\n","\n","        # CE loss (graph-safe)\n","        feat1 = output['features1'][:, 0]\n","        feat2 = output['features2'][:, 0]\n","        head = self['head']\n","        logits1 = head(feat1)\n","        logits2 = head(feat2)\n","        l_ce, acc = _ce_loss(logits1, logits2, targets)\n","\n","        loss = cfg.get('w_ce', 1.0) * l_ce + obs_loss\n","\n","        # Single stacked tensor — NO dict crossing compile boundary\n","        # [ce, acc, nce_emb, nce_pw, bridge, assign, assign_nce, nce_tri, attract]\n","        diagnostics = torch.stack([\n","            l_ce.detach(), acc.detach(),\n","            l_nce_emb.detach(), l_nce_pw.detach(),\n","            l_bridge.detach(), l_assign.detach(),\n","            l_assign_nce.detach(), l_nce_tri.detach(),\n","            l_attract.detach(),\n","        ])\n","        return loss, diagnostics\n","\n","    def _compute_geo_losses(self):\n","        \"\"\"CV + spread losses on anchors. Call OUTSIDE compiled graph.\"\"\"\n","        return self['transformer'].geometric_losses(\n","            cv_target=self.config.get('cv_target', 0.20),\n","            cv_weight=self.config.get('w_cv', 0.1),\n","            spread_weight=self.config.get('w_spread', 0.05))\n","\n","    def geometric_losses(self):\n","        return self['transformer'].geometric_losses(\n","            cv_target=self.config.get('cv_target', 0.22))\n","\n","    def anchor_diagnostics(self):\n","        return self['transformer'].anchor_diagnostics()\n","\n","    def precompute_for_step(self):\n","        \"\"\"Precompute cuSOLVER-dependent ops. Call BEFORE compiled forward.\"\"\"\n","        self['transformer'].precompute_cm_gates()\n","\n","    def invalidate_caches(self):\n","        self['transformer'].invalidate_caches()\n","        self.clear_tower_caches()\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# GEOMETRIC ANALYSIS BATTERY\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","@torch.no_grad()\n","def compute_cv(points):\n","    points = F.normalize(points.float(), dim=-1)\n","    cos_sim = points @ points.T\n","    n = points.shape[0]\n","    idx = torch.triu_indices(n, n, offset=1, device=points.device)\n","    pairwise_dist = 1.0 - cos_sim[idx[0], idx[1]]\n","    mean_d = pairwise_dist.mean()\n","    std_d = pairwise_dist.std()\n","    cv = (std_d / (mean_d + 1e-8)).item()\n","    return cv, mean_d.item(), std_d.item()\n","\n","\n","@torch.no_grad()\n","def log_geometric_analysis(model, writer, epoch, test_loader, device, config):\n","    model.eval()\n","\n","    images, labels = next(iter(test_loader))\n","    images = images[:min(64, images.shape[0])].to(device)\n","    labels = labels[:min(64, labels.shape[0])].to(device)\n","\n","    model.precompute_for_step()\n","    logits, geo_states, svd_state = model(images, return_geo_state=True)\n","\n","    n_layers = len(geo_states)\n","    pred = logits.argmax(1)\n","    batch_acc = (pred == labels).float().mean().item()\n","    writer.add_scalar('analysis/batch_accuracy', batch_acc, epoch)\n","\n","    # SVD\n","    S = svd_state['singular_values']\n","    s_norm = S / (S.sum(dim=-1, keepdim=True) + 1e-8)\n","    s_ent = -(s_norm * torch.log(s_norm.clamp(min=1e-8))).sum(-1)\n","    writer.add_scalar('svd/entropy_mean', s_ent.mean().item(), epoch)\n","    writer.add_scalar('svd/top1_ratio', (S[:, 0] / (S.sum(-1) + 1e-8)).mean().item(), epoch)\n","\n","    pe = model['patch_embed']\n","    writer.add_scalar('svd_film/gamma_weight_norm', pe.svd_to_gamma.weight.data.norm().item(), epoch)\n","\n","    # Anchor diagnostics\n","    anchor_diag = model.anchor_diagnostics()\n","    for layer_name, d in anchor_diag.items():\n","        for k, v in d.items():\n","            writer.add_scalar(f'anchor_diag/{layer_name}_{k}', v, epoch)\n","\n","    # Per-layer\n","    transformer = model['transformer']\n","    cv_trajectory = []\n","    cm_quality_trajectory = []\n","    res_norms = []\n","\n","    for i, gs in enumerate(geo_states):\n","        prefix = f'layer_{i}'\n","        layer = transformer[f'layer_{i}']\n","\n","        anchors = F.normalize(\n","            layer['observer'].association.constellation.anchors, dim=-1)\n","        cv_a, _, _ = compute_cv(anchors)\n","        writer.add_scalar(f'{prefix}/cv_anchors', cv_a, epoch)\n","\n","        emb = gs['embedding']\n","        cv_e, _, _ = compute_cv(emb.reshape(-1, emb.shape[-1]))\n","        writer.add_scalar(f'{prefix}/cv_embeddings', cv_e, epoch)\n","        cv_trajectory.append(cv_e)\n","\n","        cm_q = gs.get('cm_quality')\n","        if cm_q is not None:\n","            cm_quality_trajectory.append(cm_q.mean().item())\n","            writer.add_scalar(f'{prefix}/cm_quality_mean', cm_q.mean().item(), epoch)\n","\n","        gate_values = gs.get('gate_values')\n","        if gate_values is not None:\n","            writer.add_scalar(f'{prefix}/gate_values_std', gate_values.std().item(), epoch)\n","\n","        content = gs['content']\n","        geometric = gs['geometric']\n","        agreement = F.cosine_similarity(\n","            content.reshape(-1, content.shape[-1]),\n","            geometric.reshape(-1, geometric.shape[-1]), dim=-1)\n","        writer.add_scalar(f'{prefix}/stream_agreement_mean', agreement.mean().item(), epoch)\n","\n","        tri = gs['triangulation']\n","        nearest = gs['nearest']\n","        n_anchors = tri.shape[-1]\n","        counts = torch.bincount(nearest.reshape(-1), minlength=n_anchors).float()\n","        active = (counts > 0).sum().item()\n","        writer.add_scalar(f'{prefix}/anchors_active', active, epoch)\n","\n","        pw = gs['patchwork']\n","        writer.add_scalar(f'{prefix}/patchwork_norm', pw.norm(dim=-1).mean().item(), epoch)\n","\n","        geo_res = gs.get('geo_residual')\n","        if geo_res is not None:\n","            res_norms.append(geo_res.norm(dim=-1).mean().item())\n","            writer.add_scalar(f'{prefix}/geo_res_norm', res_norms[-1], epoch)\n","\n","        flow_opinion = gs.get('flow_opinion')\n","        if flow_opinion is not None:\n","            writer.add_scalar(f'{prefix}/flow_tri_diff',\n","                (flow_opinion - tri).abs().mean().item(), epoch)\n","\n","        if hasattr(layer, 'flow_alpha'):\n","            writer.add_scalar(f'{prefix}/flow_alpha',\n","                layer.flow_alpha.sigmoid().item(), epoch)\n","\n","    writer.add_scalar('cv/trajectory_mean', np.mean(cv_trajectory), epoch)\n","    in_band = sum(1 for cv in cv_trajectory if 0.20 <= cv <= 0.23)\n","    writer.add_scalar('cv/layers_in_band_frac', in_band / len(cv_trajectory), epoch)\n","\n","    if res_norms:\n","        writer.add_scalar('geo_res/accumulation_ratio',\n","            res_norms[-1] / (res_norms[0] + 1e-8), epoch)\n","\n","    for name, mod in model.named_modules():\n","        if isinstance(mod, CayleyOrthogonal):\n","            R = mod.get_rotation()\n","            I = torch.eye(R.shape[0], device=R.device)\n","            writer.add_scalar(f'cayley/{name.replace(\".\", \"_\")}_R_minus_I',\n","                (R - I).norm().item(), epoch)\n","\n","    film_idx = 0\n","    for name, mod in model.named_modules():\n","        if isinstance(mod, FiLMLayer):\n","            writer.add_scalar(f'film/{film_idx}_gamma_weight_norm',\n","                mod.to_gamma.weight.data.norm().item(), epoch)\n","            film_idx += 1\n","\n","    return {\n","        'batch_acc': batch_acc,\n","        'cv_trajectory': cv_trajectory,\n","        'cm_quality_trajectory': cm_quality_trajectory,\n","        'res_norms': res_norms,\n","    }\n","\n","\n","@torch.no_grad()\n","def log_gradient_norms(model, writer, epoch):\n","    type_grads = {}\n","    for name, param in model.named_parameters():\n","        if param.grad is not None:\n","            grad_norm = param.grad.norm().item()\n","            if 'flows' in name or 'flow_alpha' in name:\n","                key = 'flows'\n","            elif 'cm_gate' in name:\n","                key = 'cm_gate'\n","            elif 'observer' in name or 'constellation' in name:\n","                key = 'constellation'\n","            elif 'context' in name:\n","                key = 'geo_context'\n","            elif 'film' in name:\n","                key = 'film'\n","            elif 'rotation' in name or 'A_upper' in name:\n","                key = 'cayley'\n","            elif 'compose' in name or 'quat' in name:\n","                key = 'quaternion'\n","            elif 'head' in name:\n","                key = 'head'\n","            elif 'conv' in name or 'patch' in name:\n","                key = 'input_stage'\n","            else:\n","                key = 'other'\n","\n","            if key not in type_grads:\n","                type_grads[key] = []\n","            type_grads[key].append(grad_norm)\n","\n","    for key, norms in type_grads.items():\n","        writer.add_scalar(f'grad_norm/{key}_mean', np.mean(norms), epoch)\n","        writer.add_scalar(f'grad_norm/{key}_max', np.max(norms), epoch)\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# DATA\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","class PairedTransform:\n","    def __init__(self, transform):\n","        self.t = transform\n","    def __call__(self, img):\n","        return self.t(img), self.t(img)\n","\n","\n","def get_dataloaders(config):\n","    import torchvision\n","    import torchvision.transforms as T\n","\n","    train_aug = T.Compose([\n","        T.RandomCrop(32, padding=4),\n","        T.RandomHorizontalFlip(),\n","        T.TrivialAugmentWide(),\n","        T.ToTensor(),\n","        T.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),\n","        T.RandomErasing(p=config.get('random_erasing_p', 0.25),\n","                        scale=(0.02, 0.33), ratio=(0.3, 3.3)),\n","    ])\n","    test_transform = T.Compose([\n","        T.ToTensor(),\n","        T.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),\n","    ])\n","\n","    train_ds = torchvision.datasets.CIFAR100(\n","        root='./data', train=True, download=True,\n","        transform=PairedTransform(train_aug))\n","    test_ds = torchvision.datasets.CIFAR100(\n","        root='./data', train=False, download=True,\n","        transform=test_transform)\n","\n","    train_loader = torch.utils.data.DataLoader(\n","        train_ds, batch_size=config['batch_size'], shuffle=True,\n","        num_workers=config['num_workers'], pin_memory=True, drop_last=True)\n","    test_loader = torch.utils.data.DataLoader(\n","        test_ds, batch_size=config['batch_size'] * 2, shuffle=False,\n","        num_workers=config['num_workers'], pin_memory=True)\n","\n","    return train_loader, test_loader\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════════════\n","# TRAINING\n","# ═══════════════════════════════════════════════════════════════════════════════\n","\n","DIAG_KEYS = ['ce', 'acc', 'nce_emb', 'nce_pw', 'bridge', 'assign', 'assign_nce', 'nce_tri', 'attract']\n","\n","def diag_to_dict(diag_tensor):\n","    \"\"\"Unpack diagnostics tensor to dict. Call OUTSIDE compile boundary.\"\"\"\n","    return {k: diag_tensor[i].item() for i, k in enumerate(DIAG_KEYS)}\n","\n","\n","def train_epoch(model, model_raw, loader, optimizer, scheduler, epoch, config, writer, pbar_epoch):\n","    model.train()\n","    total_loss = 0\n","    correct = 0\n","    total = 0\n","    last_diag = None\n","    n_batches = len(loader)\n","    precompute_freq = config.get('precompute_freq', 10)\n","\n","    for batch_idx, ((v1, v2), labels) in enumerate(loader):\n","        v1, v2, labels = v1.to(device), v2.to(device), labels.to(device)\n","\n","        try:\n","            if batch_idx % precompute_freq == 0:\n","                model_raw.invalidate_caches()\n","                model_raw.precompute_for_step()\n","\n","            torch.compiler.cudagraph_mark_step_begin()\n","            loss, diag = model(v1, v2, targets=labels)\n","\n","            optimizer.zero_grad()\n","            loss.backward()\n","\n","            # CV + spread — separate backward through anchor params\n","            geo_losses = model_raw._compute_geo_losses()\n","            if geo_losses:\n","                geo_total = geo_losses.get('geo_total', None)\n","                if geo_total is not None and geo_total.requires_grad:\n","                    geo_total.backward()\n","\n","            last_diag = diag\n","\n","            if epoch % config['log_grads_every'] == 0 and batch_idx == 0:\n","                log_gradient_norms(model, writer, epoch)\n","\n","            # Dynamic grad clip: task loss sets ceiling\n","            clip_val = max(diag[0].item(), 1.0)  # diag[0] = ce\n","            torch.nn.utils.clip_grad_norm_(model.parameters(), clip_val)\n","            optimizer.step()\n","            if scheduler is not None:\n","                scheduler.step()\n","\n","        except Exception as e:\n","            torch.cuda.synchronize()\n","            print(f\"\\n  CRASH at epoch {epoch}, batch {batch_idx}/{n_batches}\")\n","            print(f\"  Error: {str(e)[:300]}\")\n","            import traceback\n","            traceback.print_exc()\n","            raise\n","\n","        total_loss += loss.item() * v1.size(0)\n","        correct += diag[1].item() * v1.size(0)  # diag[1] = acc\n","        total += v1.size(0)\n","\n","        # Update inner progress\n","        if batch_idx % 20 == 0:\n","            pbar_epoch.set_postfix_str(\n","                f\"L={total_loss/total:.2f} acc={correct/total:.3f} \"\n","                f\"[{batch_idx}/{n_batches}]\", refresh=True)\n","\n","    avg_loss = total_loss / total\n","    train_acc = correct / total\n","    return avg_loss, train_acc, last_diag\n","\n","\n","@torch.no_grad()\n","def evaluate(model, model_raw, loader):\n","    model.eval()\n","    correct = 0\n","    total = 0\n","    for images, labels in loader:\n","        images = images.to(device)\n","        labels = labels.to(device)\n","        model_raw.precompute_for_step()\n","        torch.compiler.cudagraph_mark_step_begin()\n","        logits = model(images)\n","        correct += (logits.argmax(1) == labels).sum().item()\n","        total += images.size(0)\n","    return correct / total\n","\n","\n","def main():\n","    config = CONFIG.copy()\n","    flow_str = ', '.join(config.get('flow_classes', [])) or 'none'\n","\n","    print(\"=\" * 60)\n","    print(\"  Geometric Transformer — CIFAR-100\")\n","    print(f\"  Model: d={config['d_model']}, heads={config['n_heads']}, \"\n","          f\"layers={config['n_layers']}, anchors={config['n_anchors']}\")\n","    print(f\"  Flows: [{flow_str}] fusion={config.get('flow_fusion', 'none')}\")\n","    print(f\"  Loss: CE={config['w_ce']}, pw={config['w_nce_pw']}, \"\n","          f\"brg={config['w_bridge']}, cv={config['w_cv']}, spread={config['w_spread']}\")\n","    print(\"=\" * 60)\n","\n","    writer = SummaryWriter(config['log_dir'])\n","    writer.add_text('config', json.dumps(config, indent=2))\n","\n","    print(\"\\nLoading CIFAR-100...\")\n","    train_loader, test_loader = get_dataloaders(config)\n","    print(f\"  Train: {len(train_loader.dataset):,} | Test: {len(test_loader.dataset):,}\")\n","\n","    model = GeoViTClassifier('geo_vit_cifar100', config)\n","    model.network_to(device=device, strict=False)\n","\n","    # ── WideRouter compile ──\n","    model_raw = model\n","    compile_mode = 'default'  # inductor fusion — CUDA graph replay needs Router refactor\n","    if device.type == 'cuda':\n","        try:\n","            towers = model_raw.discover_towers()\n","            print(f\"\\n  WideRouter discovery: {len(towers)} towers\")\n","\n","            from collections import Counter\n","            sig_counts = Counter(model_raw._tower_signature(t) for t in towers)\n","            for sig, count in sig_counts.most_common():\n","                tag = \"BATCH\" if count >= 2 else \"solo\"\n","                print(f\"    [{count}x {tag}] {sig[:55]}\")\n","\n","            # Clear stale inductor cache (old compiled code with graph breaks)\n","            import shutil\n","            cache_dir = '/tmp/torchinductor_root'\n","            if os.path.exists(cache_dir):\n","                shutil.rmtree(cache_dir)\n","                print(f\"  Cleared inductor cache\")\n","\n","            model = model_raw.compile(mode=compile_mode)\n","            print(f\"  WideRouter compiled ({compile_mode})\")\n","            print(f\"  Graph-breaking ops (cv_loss, spread, knn) moved outside compiled forward\")\n","\n","            stats = model_raw.get_wide_stats()\n","            vmap_n = stats.get('vmap_groups', 0)\n","            wide_n = stats.get('wide_primitive_groups', 0)\n","            if vmap_n or wide_n:\n","                print(f\"    vmap groups: {vmap_n}, wide groups: {wide_n}\")\n","\n","            stats = model_raw.get_wide_stats()\n","            vmap_n = stats.get('vmap_groups', 0)\n","            wide_n = stats.get('wide_primitive_groups', 0)\n","            if vmap_n or wide_n:\n","                print(f\"    vmap groups: {vmap_n}, wide groups: {wide_n}\")\n","\n","        except Exception as e:\n","            print(f\"  compile({compile_mode}) failed: {e}\")\n","            print(f\"  Running eager.\")\n","            model = model_raw\n","            compile_mode = 'eager'\n","\n","    n_params = sum(p.numel() for p in model.parameters())\n","    n_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)\n","    print(f\"\\n  Total params:     {n_params:,}\")\n","    print(f\"  Trainable params: {n_trainable:,}\")\n","\n","    for name, module in model.named_children():\n","        n = sum(p.numel() for p in module.parameters())\n","        if n > 0:\n","            print(f\"    {name:<20s}: {n:,}\")\n","\n","    # Flow details\n","    transformer = model_raw['transformer']\n","    for i in range(transformer.n_layers):\n","        layer = transformer[f'layer_{i}']\n","        if layer.has('flows'):\n","            flows = layer['flows']\n","            fp = sum(p.numel() for p in flows.parameters())\n","            print(f\"    layer_{i} flows:     {fp:,}  {flows.active_flow_names}\")\n","\n","    writer.add_scalar('model/total_params', n_params, 0)\n","\n","    print(f\"\\n  Initial anchor diagnostics:\")\n","    diag = model_raw.anchor_diagnostics()\n","    for layer_name, d in diag.items():\n","        print(f\"    {layer_name}: cv={d['anchor_cv']:.4f}, \"\n","              f\"cm_pos={d['cm_positive_frac']:.3f}\")\n","\n","    # Optimizer — plain Adam\n","    optimizer = torch.optim.Adam(model.parameters(), lr=config['lr'])\n","\n","    total_steps = config['epochs'] * len(train_loader)\n","    warmup_steps = config['warmup_epochs'] * len(train_loader)\n","\n","    def lr_lambda(step):\n","        if step < warmup_steps:\n","            return step / max(1, warmup_steps)\n","        progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)\n","        return 0.5 * (1 + np.cos(np.pi * progress))\n","\n","    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)\n","\n","    print(f\"\\n{'━'*60}\")\n","    print(f\"  Training for {config['epochs']} epochs\")\n","    print(f\"  Warmup: {config['warmup_epochs']} epochs, LR: {config['lr']}\")\n","    print(f\"  Geo analysis every {config['log_geo_every']} epochs\")\n","    print(f\"{'━'*60}\\n\")\n","\n","    best_acc = 0\n","    save_dir = Path('geo_cifar100'); save_dir.mkdir(exist_ok=True)\n","    epoch_times = []\n","\n","    pbar = tqdm(range(config['epochs']), desc=\"Training\", unit=\"ep\")\n","\n","    for epoch in pbar:\n","        t0 = time.time()\n","\n","        avg_loss, train_acc, last_diag = train_epoch(\n","            model, model_raw, train_loader, optimizer, scheduler, epoch, config, writer, pbar)\n","\n","        test_acc = evaluate(model, model_raw, test_loader)\n","        elapsed = time.time() - t0\n","        epoch_times.append(elapsed)\n","\n","        lr = optimizer.param_groups[0]['lr']\n","        writer.add_scalar('train/total_loss', avg_loss, epoch)\n","        writer.add_scalar('train/accuracy', train_acc, epoch)\n","        writer.add_scalar('test/accuracy', test_acc, epoch)\n","        writer.add_scalar('train/lr', lr, epoch)\n","        writer.add_scalar('train/epoch_time', elapsed, epoch)\n","\n","        if last_diag is not None:\n","            d = diag_to_dict(last_diag)\n","            for k, v in d.items():\n","                if isinstance(v, (int, float)) and not math.isnan(v):\n","                    writer.add_scalar(f'loss/{k}', v, epoch)\n","\n","        if test_acc > best_acc:\n","            best_acc = test_acc\n","            torch.save({\n","                'state_dict': {k: v.cpu() for k, v in model.state_dict().items()},\n","                'epoch': epoch,\n","                'test_acc': test_acc,\n","                'config': config,\n","            }, save_dir / 'best.pt')\n","\n","        # tqdm postfix — always visible\n","        avg_t = np.mean(epoch_times[-5:])\n","        pbar.set_postfix_str(\n","            f\"L={avg_loss:.2f} tr={train_acc:.3f} te={test_acc:.3f} \"\n","            f\"best={best_acc:.3f} {avg_t:.1f}s/ep\")\n","\n","        # Full geo analysis periodically\n","        if epoch % config['log_geo_every'] == 0 or epoch == config['epochs'] - 1:\n","            geo_info = log_geometric_analysis(\n","                model_raw, writer, epoch, test_loader, device, config)\n","\n","            cv_str = ', '.join(f'{cv:.3f}' for cv in geo_info['cv_trajectory'])\n","            cm_str = ', '.join(f'{q:.3f}' for q in geo_info.get('cm_quality_trajectory', []))\n","            d = diag_to_dict(last_diag) if last_diag is not None else {}\n","\n","            tqdm.write(\n","                f\"  E{epoch:>3d}  L={avg_loss:.3f}\"\n","                f\"  ce={d.get('ce', 0):.3f}  brg={d.get('bridge', 0):.3f}\"\n","                f\"  train={train_acc:.4f}  test={test_acc:.4f}\"\n","                f\"  best={best_acc:.4f}  {elapsed:.1f}s\"\n","                f\"\\n        CV=[{cv_str}]\"\n","                f\"\\n        CM=[{cm_str}]\")\n","\n","    pbar.close()\n","\n","    print(f\"\\n{'═'*60}\")\n","    print(f\"  CIFAR-100 RESULTS\")\n","    print(f\"{'═'*60}\")\n","    print(f\"  Best test accuracy: {best_acc:.4f} ({best_acc*100:.2f}%)\")\n","    print(f\"  Parameters: {n_params:,}\")\n","    print(f\"  Flows: [{flow_str}]\")\n","    print(f\"  Avg epoch time: {np.mean(epoch_times):.1f}s\")\n","    print(f\"  Checkpoint: {save_dir}/best.pt\")\n","\n","    geo_info = log_geometric_analysis(\n","        model_raw, writer, config['epochs'], test_loader, device, config)\n","\n","    diag = model_raw.anchor_diagnostics()\n","    print(f\"\\n  Final anchor diagnostics:\")\n","    for layer_name, d in diag.items():\n","        print(f\"    {layer_name}: cv={d['anchor_cv']:.4f}, \"\n","              f\"cm_pos={d['cm_positive_frac']:.3f}\")\n","\n","    writer.close()\n","    print(f\"\\nDone.\")\n","\n","\n","if __name__ == '__main__':\n","    main()"],"metadata":{"id":"1vclJYU0KTMq","colab":{"base_uri":"https://localhost:8080/","height":529},"executionInfo":{"status":"error","timestamp":1775232032604,"user_tz":420,"elapsed":620,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"9dbfb223-c006-419d-bb73-d9048e66f740"},"execution_count":2,"outputs":[{"output_type":"stream","name":"stdout","text":["Device: cuda\n","  GPU: NVIDIA RTX PRO 6000 Blackwell Server Edition\n","  VRAM: 102.0 GB\n","============================================================\n","  Geometric Transformer — CIFAR-100\n","  Model: d=384, heads=8, layers=8, anchors=128\n","  Flows: [quat_lite, velocity, orbital] fusion=weighted\n","  Loss: CE=1.0, pw=1.0, brg=1.0, cv=0.1, spread=0.05\n","============================================================\n","\n","Loading CIFAR-100...\n"]},{"output_type":"error","ename":"KeyboardInterrupt","evalue":"","traceback":["\u001b[0;31m---------------------------------------------------------------------------\u001b[0m","\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)","\u001b[0;32m/tmp/ipykernel_405529/144620607.py\u001b[0m in \u001b[0;36m<cell line: 0>\u001b[0;34m()\u001b[0m\n\u001b[1;32m    870\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    871\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0m__name__\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m'__main__'\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 872\u001b[0;31m     \u001b[0mmain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m","\u001b[0;32m/tmp/ipykernel_405529/144620607.py\u001b[0m in \u001b[0;36mmain\u001b[0;34m()\u001b[0m\n\u001b[1;32m    688\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    689\u001b[0m     \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"\\nLoading CIFAR-100...\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 690\u001b[0;31m     \u001b[0mtrain_loader\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtest_loader\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mget_dataloaders\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mconfig\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    691\u001b[0m     \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf\"  Train: {len(train_loader.dataset):,} | Test: {len(test_loader.dataset):,}\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    692\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/tmp/ipykernel_405529/144620607.py\u001b[0m in \u001b[0;36mget_dataloaders\u001b[0;34m(config)\u001b[0m\n\u001b[1;32m    562\u001b[0m     ])\n\u001b[1;32m    563\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 564\u001b[0;31m     train_ds = torchvision.datasets.CIFAR100(\n\u001b[0m\u001b[1;32m    565\u001b[0m         \u001b[0mroot\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'./data'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrain\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdownload\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    566\u001b[0m         transform=PairedTransform(train_aug))\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torchvision/datasets/cifar.py\u001b[0m in \u001b[0;36m__init__\u001b[0;34m(self, root, train, transform, target_transform, download)\u001b[0m\n\u001b[1;32m     81\u001b[0m             \u001b[0mfile_path\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mos\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mjoin\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mroot\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbase_folder\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfile_name\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     82\u001b[0m             \u001b[0;32mwith\u001b[0m \u001b[0mopen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfile_path\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"rb\"\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 83\u001b[0;31m                 \u001b[0mentry\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpickle\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mload\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mencoding\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m\"latin1\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     84\u001b[0m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mentry\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"data\"\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     85\u001b[0m                 \u001b[0;32mif\u001b[0m \u001b[0;34m\"labels\"\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mentry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;31mKeyboardInterrupt\u001b[0m: "]}]},{"cell_type":"markdown","source":["# v8"],"metadata":{"id":"zGDUDMszEzQe"}},{"cell_type":"code","source":["\"\"\"\n","Geometric Transformer Redux v8\n","================================\n","Shared constellation across all layers. Momentum-based anchor push.\n","Proper pentachoron CV on embeddings. WideRouter compiled.\n","\n","Architecture:\n","  ConvSVDPatchEmbed → 4× ReduxLayer (shared constellation) → head\n","\n","Each layer:\n","  project → observe (shared anchors) → CM gate → patchwork\n","  → FiLM context → dual attention → gated compose → geo_residual\n","\n","Key differences from v1:\n","  - ONE constellation for all layers (not per-layer)\n","  - Momentum AnchorPush (non-gradient anchor repositioning)\n","  - cv_loss on EMBEDDINGS (pentachoron volumes), not anchor pairwise\n","  - torch.stack diagnostics (no dict crossing compile boundary)\n","  - Dynamic grad clip: max(ce_loss, 1.0)\n","  - Precompute every N batches\n","\n","!pip install geolip-core torchvision tqdm\n","\"\"\"\n","\n","import os, sys, warnings, math, time\n","import numpy as np\n","from pathlib import Path\n","from collections import defaultdict\n","\n","warnings.filterwarnings('ignore')\n","\n","import torch\n","import torch.nn as nn\n","import torch.nn.functional as F\n","from tqdm.auto import tqdm\n","\n","device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n","\n","torch.set_float32_matmul_precision('high')\n","torch.backends.cudnn.benchmark = True\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# IMPORTS from geolip_core — utilities only\n","# ═══════════════════════════════════════════════════════════════════\n","\n","from geolip_core.core.input.svd import SVDObserver\n","from geolip_core.core.associate.constellation import (\n","    ConstellationObserver, ConstellationAssociation, ConstellationCuration,\n",")\n","from geolip_core.core.curate.patchwork import AnchorPush\n","from geolip_core.core.distinguish.losses import (\n","    nce_loss, bridge_loss_paired, assign_bce_loss,\n","    assign_nce_loss, attraction_loss, ce_loss_paired,\n","    cv_loss, spread_loss,\n",")\n","from geolip_core.pipeline.components.geometric_transformer import CMValidatedGate\n","from geolip_core.pipeline.observer import TorchComponent\n","from geofractal.router.wide_router import WideRouter\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# CONFIG\n","# ═══════════════════════════════════════════════════════════════════\n","\n","DEFAULT_CONFIG = {\n","    # Architecture\n","    'd_model': 256,\n","    'n_heads': 8,\n","    'n_layers': 4,\n","    'manifold_dim': 128,\n","    'n_anchors': 128,\n","    'n_comp': 8,\n","    'd_comp': 32,\n","    'context_dim': 128,\n","    'cm_neighbors': 3,\n","    'dropout': 0.1,\n","\n","    # Input\n","    'conv_channels': 64,\n","    'svd_rank': 12,\n","    'patch_size': 4,\n","    'img_size': 32,\n","\n","    # Training\n","    'epochs': 100,\n","    'batch_size': 128,\n","    'lr': 3e-4,\n","    'warmup_epochs': 5,\n","    'num_workers': 4,\n","    'precompute_freq': 10,\n","\n","    # Push\n","    'push_every': 5,      # epochs between anchor pushes\n","    'push_strategy': 'momentum',\n","    'push_accumulate': 2048,  # embeddings to accumulate before push\n","\n","    # Loss\n","    'w_ce': 1.0,\n","    'w_nce_emb': 0.5,\n","    'w_nce_pw': 1.0,\n","    'w_bridge': 1.0,\n","    'w_assign': 0.5,\n","    'w_assign_nce': 0.25,\n","    'w_nce_tri': 0.5,\n","    'w_attract': 0.25,\n","    'w_cv': 0.1,\n","    'w_spread': 0.05,\n","    'cv_target': 0.21,\n","\n","    'num_classes': 100,\n","}\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# INPUT STAGE\n","# ═══════════════════════════════════════════════════════════════════\n","\n","class ConvSVDInput(TorchComponent):\n","    \"\"\"Conv → SVD → patch tokens. FL kernel, compilable.\"\"\"\n","    def __init__(self, name, img_size=32, patch_size=4, in_channels=3,\n","                 conv_channels=64, d_model=256, svd_rank=12):\n","        super().__init__(name)\n","        self.patch_size = patch_size\n","        n_patches = (img_size // patch_size) ** 2\n","\n","        self.conv = nn.Sequential(\n","            nn.Conv2d(in_channels, conv_channels, 3, padding=1, bias=False),\n","            nn.BatchNorm2d(conv_channels), nn.GELU(),\n","            nn.Conv2d(conv_channels, conv_channels, 3, padding=1, bias=False),\n","            nn.BatchNorm2d(conv_channels), nn.GELU(),\n","        )\n","        self.svd = SVDObserver(conv_channels, svd_rank)\n","        self.patch_proj = nn.Conv2d(conv_channels, d_model,\n","                                    kernel_size=patch_size, stride=patch_size, bias=False)\n","        self.patch_norm = nn.LayerNorm(d_model)\n","\n","        feat_dim = self.svd.feature_dim\n","        self.svd_gamma = nn.Linear(feat_dim, d_model)\n","        self.svd_beta = nn.Linear(feat_dim, d_model)\n","        nn.init.normal_(self.svd_gamma.weight, std=0.01)\n","        nn.init.ones_(self.svd_gamma.bias)\n","        nn.init.zeros_(self.svd_beta.bias)\n","\n","        self.cls_token = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)\n","        self.pos_embed = nn.Parameter(torch.randn(1, n_patches + 1, d_model) * 0.02)\n","\n","    def forward(self, x):\n","        B = x.shape[0]\n","        feat = self.conv(x)\n","        S, Vh, svd_feat, novelty = self.svd(feat)\n","        tokens = self.patch_proj(feat).flatten(2).transpose(1, 2)\n","        tokens = self.patch_norm(tokens)\n","        g = self.svd_gamma(svd_feat).unsqueeze(1)\n","        b = self.svd_beta(svd_feat).unsqueeze(1)\n","        tokens = g * tokens + b\n","        tokens = torch.cat([self.cls_token.expand(B, -1, -1), tokens], dim=1)\n","        return tokens + self.pos_embed\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# REDUX LAYER — reads shared constellation\n","# ═══════════════════════════════════════════════════════════════════\n","\n","class ReduxLayer(WideRouter):\n","    \"\"\"Transformer layer observing a shared constellation.\n","\n","    Receives anchors_n and cm_gate externally — does NOT own them.\n","    \"\"\"\n","    def __init__(self, name, d_model=256, n_heads=8, n_anchors=128,\n","                 manifold_dim=128, n_comp=8, d_comp=32,\n","                 context_dim=128, dropout=0.1):\n","        super().__init__(name, strict=False)\n","        self.d_model = d_model\n","        self.manifold_dim = manifold_dim\n","        self.n_anchors = n_anchors\n","        pw_dim = n_comp * d_comp\n","\n","        # Project to manifold (per-layer — each layer sees different aspect)\n","        self.attach('proj', nn.Sequential(\n","            nn.Linear(d_model, manifold_dim), nn.LayerNorm(manifold_dim)))\n","\n","        # Per-layer patchwork (reads shared constellation's triangulation)\n","        from geolip_core.core.curate.patchwork import Patchwork\n","        self.attach('patchwork', Patchwork(\n","            n_anchors=n_anchors, n_comp=n_comp, d_comp=d_comp))\n","        self.attach('bridge_proj', nn.Linear(pw_dim, n_anchors))\n","\n","        # FiLM context: anchor feats + patchwork + history\n","        self.attach('anchor_ctx', nn.Linear(n_anchors * 2, context_dim))  # cos + gate\n","        self.attach('pw_ctx', nn.Linear(pw_dim, context_dim))\n","        self.attach('history_ctx', nn.Linear(pw_dim, context_dim))\n","        self.attach('ctx_fuse', nn.Sequential(\n","            nn.Linear(context_dim * 3, context_dim), nn.GELU()))\n","        self.attach('film_gamma', nn.Linear(context_dim, d_model))\n","        self.attach('film_beta', nn.Linear(context_dim, d_model))\n","\n","        # Dual attention\n","        self.attach('content_attn', nn.MultiheadAttention(\n","            d_model, n_heads, batch_first=True, dropout=dropout))\n","        self.attach('content_norm', nn.LayerNorm(d_model))\n","        self.attach('geo_attn', nn.MultiheadAttention(\n","            d_model, n_heads, batch_first=True, dropout=dropout))\n","        self.attach('geo_norm', nn.LayerNorm(d_model))\n","\n","        # Gated composition\n","        self.attach('gate', nn.Sequential(\n","            nn.Linear(d_model * 2, d_model), nn.Sigmoid()))\n","        self.attach('ffn', nn.Sequential(\n","            nn.Linear(d_model, d_model * 4), nn.GELU(), nn.Dropout(dropout),\n","            nn.Linear(d_model * 4, d_model)))\n","        self.attach('ffn_norm', nn.LayerNorm(d_model))\n","\n","        # Geo residual projection\n","        self._pw_dim = pw_dim\n","        self.attach('geo_proj', nn.Sequential(\n","            nn.Linear(pw_dim, pw_dim), nn.LayerNorm(pw_dim)))\n","\n","    def forward(self, x, anchors_n, gate_values, geo_residual=None):\n","        \"\"\"\n","        Args:\n","            x: (B, L, D) tokens\n","            anchors_n: (A, M) shared normalized anchors\n","            gate_values: (B*L, A) from shared CM gate\n","            geo_residual: (B, L, pw_dim) or None\n","        Returns:\n","            x_out, geo_residual_out, emb_flat, pw_flat, bridge_flat, a_out_dict\n","        \"\"\"\n","        B, L, D = x.shape\n","\n","        # 1. Project to manifold\n","        emb = F.normalize(self['proj'](x), dim=-1)\n","        emb_flat = emb.reshape(B * L, -1)\n","\n","        # 2. Triangulate against shared anchors\n","        cos = emb_flat @ anchors_n.T\n","        distances = 1.0 - cos\n","        assignment = F.softmax(cos / 0.1, dim=-1)\n","\n","        # 3. Gate triangulation\n","        gated_distances = distances * gate_values\n","\n","        # 4. Patchwork curation\n","        a_out = {\n","            'distances': distances,\n","            'distances_weighted': gated_distances,\n","            'cos_to_anchors': cos,\n","            'assignment': assignment,\n","            'nearest': cos.argmax(dim=-1),\n","        }\n","        pw_flat = self['patchwork'](gated_distances)  # (B*L, pw_dim)\n","        bridge_flat = self['bridge_proj'](pw_flat)\n","\n","        # 5. FiLM context\n","        anchor_feats = torch.cat([cos, gate_values], dim=-1)\n","        a_ctx = self['anchor_ctx'](anchor_feats).reshape(B, L, -1)\n","        p_ctx = self['pw_ctx'](pw_flat).reshape(B, L, -1)\n","        h_ctx = self['history_ctx'](geo_residual.reshape(B * L, -1)).reshape(B, L, -1) \\\n","            if geo_residual is not None else torch.zeros_like(a_ctx)\n","        ctx = self['ctx_fuse'](torch.cat([a_ctx, p_ctx, h_ctx], dim=-1))\n","        gamma = self['film_gamma'](ctx)\n","        beta = self['film_beta'](ctx)\n","\n","        # 6. Dual attention\n","        content, _ = self['content_attn'](x, x, x, need_weights=False)\n","        content = self['content_norm'](x + content)\n","        geo_q = x + gamma * x + beta\n","        geo_out, _ = self['geo_attn'](geo_q, x, x, need_weights=False)\n","        geo_out = self['geo_norm'](x + geo_out)\n","\n","        # 7. Gated compose + FFN\n","        g = self['gate'](torch.cat([content, geo_out], dim=-1))\n","        merged = g * geo_out + (1 - g) * content\n","        h = self['ffn'](merged)\n","        x_out = self['ffn_norm'](merged + h)\n","\n","        # 8. Geo residual accumulation\n","        cm_quality = gate_values.mean(dim=-1).reshape(B, L, 1)\n","        geo_update = self['geo_proj'](pw_flat.reshape(B, L, -1))\n","        if geo_residual is None:\n","            geo_residual_out = cm_quality * geo_update\n","        else:\n","            geo_residual_out = geo_residual + cm_quality * geo_update\n","\n","        return x_out, geo_residual_out, emb_flat, pw_flat, bridge_flat, a_out\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# GEOMETRIC TRANSFORMER REDUX — WideRouter\n","# ═══════════════════════════════════════════════════════════════════\n","\n","class GeoTransformerRedux(WideRouter):\n","    \"\"\"Shared constellation, momentum push, proper CV.\n","\n","    All layers observe the SAME constellation. Anchor push moves them\n","    based on accumulated embeddings. CV measured on embedding pentachora.\n","    \"\"\"\n","    def __init__(self, name, config):\n","        super().__init__(name, strict=False)\n","        self.config = config\n","        d = config['d_model']\n","        m = config['manifold_dim']\n","        A = config['n_anchors']\n","        n_layers = config['n_layers']\n","\n","        # Input\n","        self.attach('input', ConvSVDInput(\n","            'input', img_size=config['img_size'], patch_size=config['patch_size'],\n","            conv_channels=config['conv_channels'], d_model=d,\n","            svd_rank=config['svd_rank']))\n","\n","        # SHARED constellation + CM gate (one for all layers)\n","        self.attach('observer', ConstellationObserver(\n","            dim=m, n_anchors=A,\n","            n_comp=config['n_comp'], d_comp=config['d_comp']))\n","        self.attach('cm_gate', CMValidatedGate(A, n_neighbors=config['cm_neighbors']))\n","\n","        # Per-layer transformer blocks (read shared constellation)\n","        for i in range(n_layers):\n","            self.attach(f'layer_{i}', ReduxLayer(\n","                f'L{i}', d, config['n_heads'], A, m,\n","                config['n_comp'], config['d_comp'],\n","                config['context_dim'], config['dropout']))\n","            self.register_tower(f'layer_{i}')\n","\n","        self.attach('final_norm', nn.LayerNorm(d))\n","\n","        # Classification head\n","        pw_dim = config['n_comp'] * config['d_comp']\n","        head_in = d + pw_dim  # cls token + geo_residual pooled\n","        self.attach('head', nn.Sequential(\n","            nn.LayerNorm(head_in), nn.Linear(head_in, d),\n","            nn.GELU(), nn.Dropout(config['dropout']),\n","            nn.Linear(d, config['num_classes'])))\n","\n","        # Manifold projection for shared observation\n","        self.attach('manifold_proj', nn.Sequential(\n","            nn.Linear(d, m), nn.LayerNorm(m)))\n","\n","        self.n_layers = n_layers\n","\n","        # Anchor push (non-gradient momentum updates)\n","        self.push = AnchorPush(\n","            config['push_strategy'], A, m,\n","            decay=0.9, alpha=0.1, beta=0.05, util_floor=0.001)\n","\n","        # Push accumulation buffers\n","        self.register_buffer('_emb_buf', torch.zeros(config['push_accumulate'], m))\n","        self.register_buffer('_lbl_buf', torch.zeros(config['push_accumulate'], dtype=torch.long))\n","        self.register_buffer('_buf_ptr', torch.zeros(1, dtype=torch.long))\n","\n","    @property\n","    def constellation(self):\n","        return self['observer'].association.constellation\n","\n","    @torch.no_grad()\n","    def precompute(self):\n","        \"\"\"CM gate + cache. Call before compiled forward.\"\"\"\n","        anchors_n = F.normalize(self.constellation.anchors, dim=-1)\n","        self['cm_gate'].precompute(anchors_n.detach())\n","\n","    def invalidate(self):\n","        self['cm_gate'].invalidate_cache()\n","\n","    @torch.no_grad()\n","    def accumulate_for_push(self, emb, labels):\n","        \"\"\"Accumulate embeddings for anchor push. Returns True if buffer is full.\"\"\"\n","        B = emb.shape[0]\n","        ptr = int(self._buf_ptr.item())\n","        cap = self._emb_buf.shape[0]\n","        end = min(ptr + B, cap)\n","        n = end - ptr\n","        self._emb_buf[ptr:end] = emb[:n].detach()\n","        self._lbl_buf[ptr:end] = labels[:n].detach()\n","        self._buf_ptr[0] = end\n","        return end >= cap  # True = buffer full, caller should push\n","\n","    @torch.no_grad()\n","    def do_push(self):\n","        \"\"\"Execute momentum anchor push. Returns diagnostics dict.\"\"\"\n","        ptr = int(self._buf_ptr.item())\n","        if ptr < 64:\n","            return {'moved': 0}\n","        dev = self.constellation.anchors.device\n","        emb = self._emb_buf[:ptr].to(dev)\n","        lbl = self._lbl_buf[:ptr].to(dev)\n","        diag = self.push.push(self['observer'], emb, lbl)\n","        self._buf_ptr.zero_()\n","        return diag\n","\n","    def forward(self, x, x2=None, targets=None):\n","        \"\"\"\n","        Single view: returns logits\n","        Paired view with targets: returns (loss, diagnostics_tensor)\n","        \"\"\"\n","        if x2 is not None and targets is not None:\n","            return self._forward_paired(x, x2, targets)\n","\n","        tokens = self['input'](x)\n","        B, L, D = tokens.shape\n","\n","        # Shared observation: project CLS to manifold, observe\n","        anchors_n = F.normalize(self.constellation.anchors, dim=-1)\n","\n","        # Run layers\n","        geo_residual = None\n","        for i in range(self.n_layers):\n","            # Each layer projects tokens to manifold independently\n","            emb_for_gate = F.normalize(self['manifold_proj'](tokens), dim=-1)\n","            flat = emb_for_gate.reshape(B * L, -1)\n","            gate_values, _ = self['cm_gate'](1.0 - flat @ anchors_n.T)\n","\n","            tokens, geo_residual, _, _, _, _ = self[f'layer_{i}'](\n","                tokens, anchors_n, gate_values, geo_residual)\n","\n","        tokens = self['final_norm'](tokens)\n","        cls = tokens[:, 0]\n","        geo_pool = geo_residual[:, 0] if geo_residual is not None else \\\n","            torch.zeros(B, self.config['n_comp'] * self.config['d_comp'], device=x.device)\n","        logits = self['head'](torch.cat([cls, geo_pool], dim=-1))\n","        return logits\n","\n","    def _forward_paired(self, x1, x2, targets):\n","        \"\"\"Paired forward with full observer loss. Returns (loss, diag_tensor).\"\"\"\n","        B = x1.shape[0]\n","        x_cat = torch.cat([x1, x2], dim=0)\n","        tokens = self['input'](x_cat)\n","        BB, L, D = tokens.shape  # BB = 2*B\n","\n","        anchors_n = F.normalize(self.constellation.anchors, dim=-1)\n","\n","        # Run layers — collect LAST layer's observation for loss\n","        geo_residual = None\n","        last_emb = last_pw = last_bridge = last_a_out = None\n","        for i in range(self.n_layers):\n","            emb_for_gate = F.normalize(self['manifold_proj'](tokens), dim=-1)\n","            flat = emb_for_gate.reshape(BB * L, -1)\n","            gate_values, _ = self['cm_gate'](1.0 - flat @ anchors_n.T)\n","\n","            tokens, geo_residual, emb_flat, pw_flat, bridge_flat, a_out = \\\n","                self[f'layer_{i}'](tokens, anchors_n, gate_values, geo_residual)\n","\n","            last_emb = emb_flat\n","            last_pw = pw_flat\n","            last_bridge = bridge_flat\n","            last_a_out = a_out\n","\n","        tokens = self['final_norm'](tokens)\n","\n","        # Split views\n","        cls1, cls2 = tokens[:B, 0], tokens[B:, 0]\n","        geo1 = geo_residual[:B, 0] if geo_residual is not None else torch.zeros_like(cls1[:, :self._pw_dim_])\n","        geo2 = geo_residual[B:, 0] if geo_residual is not None else torch.zeros_like(cls2[:, :self._pw_dim_])\n","\n","        # CLS token embeddings on manifold (for NCE)\n","        emb1 = F.normalize(self['manifold_proj'](cls1), dim=-1)\n","        emb2 = F.normalize(self['manifold_proj'](cls2), dim=-1)\n","\n","        # Last layer observations at CLS position\n","        c = 0  # cls index\n","        N = BB * L\n","        # Reshape to (2B, L, ...) then extract CLS\n","        pw1 = last_pw.reshape(BB, L, -1)[:B, c]\n","        pw2 = last_pw.reshape(BB, L, -1)[B:, c]\n","        brg1 = last_bridge.reshape(BB, L, -1)[:B, c]\n","        brg2 = last_bridge.reshape(BB, L, -1)[B:, c]\n","        assign1 = last_a_out['assignment'].reshape(BB, L, -1)[:B, c]\n","        assign2 = last_a_out['assignment'].reshape(BB, L, -1)[B:, c]\n","        cos1 = last_a_out['cos_to_anchors'].reshape(BB, L, -1)[:B, c]\n","        tri1 = last_a_out['distances'].reshape(BB, L, -1)[:B, c]\n","        tri2 = last_a_out['distances'].reshape(BB, L, -1)[B:, c]\n","\n","        # ═══ LOSSES ═══\n","        cfg = self.config\n","\n","        # CE\n","        logits1 = self['head'](torch.cat([cls1, geo1], dim=-1))\n","        logits2 = self['head'](torch.cat([cls2, geo2], dim=-1))\n","        l_ce, acc = ce_loss_paired(logits1, logits2, targets)\n","\n","        # Embedding NCE\n","        l_nce_emb, _ = nce_loss(emb1, emb2, 0.07, normalize=False)\n","\n","        # Patchwork NCE\n","        l_nce_pw, _ = nce_loss(pw1, pw2, 0.1, normalize=True)\n","\n","        # Bridge\n","        l_bridge, _ = bridge_loss_paired(brg1, brg2, assign1, assign2)\n","\n","        # Assign BCE\n","        l_assign, _ = assign_bce_loss(assign1, cos1)\n","\n","        # Assign NCE\n","        l_assign_nce, _ = assign_nce_loss(assign1, assign2, 0.1)\n","\n","        # Tri NCE\n","        l_nce_tri, _ = nce_loss(tri1, tri2, 0.1, normalize=True)\n","\n","        # Attraction\n","        l_attract, _ = attraction_loss(cos1)\n","\n","        loss = (cfg['w_ce'] * l_ce\n","                + cfg['w_nce_emb'] * l_nce_emb\n","                + cfg['w_nce_pw'] * l_nce_pw\n","                + cfg['w_bridge'] * l_bridge\n","                + cfg['w_assign'] * l_assign\n","                + cfg['w_assign_nce'] * l_assign_nce\n","                + cfg['w_nce_tri'] * l_nce_tri\n","                + cfg['w_attract'] * l_attract)\n","\n","        # Single tensor — no dict crossing compile boundary\n","        diag = torch.stack([\n","            l_ce.detach(), acc.detach(),\n","            l_nce_emb.detach(), l_nce_pw.detach(),\n","            l_bridge.detach(), l_assign.detach(),\n","            l_assign_nce.detach(), l_nce_tri.detach(),\n","            l_attract.detach(),\n","        ])\n","\n","        return loss, diag, cls1  # cls1 has gradient — for CV loss outside compile\n","\n","    @property\n","    def _pw_dim_(self):\n","        return self.config['n_comp'] * self.config['d_comp']\n","\n","    def compute_geo_losses(self, cls_feat=None):\n","        \"\"\"CV (pentachoron on embeddings) + spread (anchor repulsion).\n","        Call OUTSIDE compiled graph.\n","\n","        Args:\n","            cls_feat: (B, d_model) CLS features WITH gradient.\n","                      Projected to manifold here for CV loss.\n","        \"\"\"\n","        losses = {}\n","        cfg = self.config\n","        if cls_feat is not None and cls_feat.shape[0] >= 5:\n","            # Project to manifold OUTSIDE compile — cv_loss uses LA.det\n","            emb = F.normalize(self['manifold_proj'](cls_feat), dim=-1)\n","            l_cv = cv_loss(emb, target=cfg['cv_target'])\n","            losses['cv'] = cfg['w_cv'] * l_cv\n","            losses['cv_raw'] = l_cv.detach()\n","        else:\n","            losses['cv'] = torch.tensor(0.0, device=device)\n","            losses['cv_raw'] = torch.tensor(0.0)\n","\n","        l_spread = spread_loss(self.constellation.anchors)\n","        losses['spread'] = cfg['w_spread'] * l_spread\n","        losses['geo_total'] = losses['cv'] + losses['spread']\n","        return losses\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# DIAGNOSTICS\n","# ═══════════════════════════════════════════════════════════════════\n","\n","DIAG_KEYS = ['ce', 'acc', 'nce_emb', 'nce_pw', 'bridge',\n","             'assign', 'assign_nce', 'nce_tri', 'attract']\n","\n","\n","def diag_to_dict(t):\n","    return {k: t[i].item() for i, k in enumerate(DIAG_KEYS)}\n","\n","\n","@torch.no_grad()\n","def compute_geometric_stats(model):\n","    \"\"\"Full geometric health check.\"\"\"\n","    from geolip_core.core.distinguish.losses import cv_metric\n","    stats = {}\n","    anchors = model.constellation.anchors\n","    anchors_n = F.normalize(anchors, dim=-1)\n","    A = anchors_n.shape[0]\n","    cos = anchors_n @ anchors_n.T\n","    idx = torch.triu_indices(A, A, offset=1, device=cos.device)\n","    pairwise = 1.0 - cos[idx[0], idx[1]]\n","    stats['anchor_cv'] = (pairwise.std() / (pairwise.mean() + 1e-8)).item()\n","    stats['anchor_mean_dist'] = pairwise.mean().item()\n","    stats['anchor_min_dist'] = pairwise.min().item()\n","    stats['cm_pos'] = (model['cm_gate']._cached_cm_norm > 0).float().mean().item()\n","    # Utilization from push\n","    if model.push.strategy == 'momentum' and model.push.accumulator is not None:\n","        stats['push_momentum'] = model.push.accumulator.norm(dim=-1).mean().item()\n","    return stats\n","\n","\n","@torch.no_grad()\n","def compute_gradient_norms(model):\n","    \"\"\"Gradient norms by component type.\"\"\"\n","    norms = defaultdict(list)\n","    for name, p in model.named_parameters():\n","        if p.grad is None:\n","            continue\n","        gn = p.grad.norm().item()\n","        if 'constellation' in name or 'anchor' in name:\n","            norms['constellation'].append(gn)\n","        elif 'cm_gate' in name:\n","            norms['cm_gate'].append(gn)\n","        elif 'content_attn' in name:\n","            norms['content_attn'].append(gn)\n","        elif 'geo_attn' in name:\n","            norms['geo_attn'].append(gn)\n","        elif 'patchwork' in name or 'bridge' in name:\n","            norms['patchwork'].append(gn)\n","        elif 'film' in name or 'ctx' in name:\n","            norms['context'].append(gn)\n","        elif 'head' in name:\n","            norms['head'].append(gn)\n","        elif 'conv' in name or 'input' in name:\n","            norms['input'].append(gn)\n","        else:\n","            norms['other'].append(gn)\n","    return {k: {'mean': np.mean(v), 'max': np.max(v)} for k, v in norms.items()}\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# DRY RUN — noise data, escalating classes\n","# ═══════════════════════════════════════════════════════════════════\n","\n","class PairedTransform:\n","    def __init__(self, t): self.t = t\n","    def __call__(self, img): return self.t(img), self.t(img)\n","\n","\n","def make_noise_data(n_samples, n_classes, img_size=32, channels=3):\n","    \"\"\"Gaussian noise dataset with random labels.\"\"\"\n","    X = torch.randn(n_samples, channels, img_size, img_size)\n","    Y = torch.randint(0, n_classes, (n_samples,))\n","    return torch.utils.data.TensorDataset(X, Y)\n","\n","\n","def make_paired_noise_loader(n_samples, n_classes, batch_size):\n","    \"\"\"Noise data with paired augmentation — shared base, small perturbation.\"\"\"\n","    class PairedNoiseDataset(torch.utils.data.Dataset):\n","        def __init__(self, n, nc):\n","            self.n, self.nc = n, nc\n","            self.labels = torch.randint(0, nc, (n,))\n","            self.base_images = torch.randn(n, 3, 32, 32)\n","        def __len__(self): return self.n\n","        def __getitem__(self, i):\n","            base = self.base_images[i]\n","            v1 = base + 0.1 * torch.randn_like(base)\n","            v2 = base + 0.1 * torch.randn_like(base)\n","            return (v1, v2), self.labels[i]\n","\n","    ds = PairedNoiseDataset(n_samples, n_classes)\n","    return torch.utils.data.DataLoader(\n","        ds, batch_size=batch_size, shuffle=True, num_workers=0, drop_last=True)\n","\n","\n","def dry_run(n_classes, n_epochs, config=None):\n","    \"\"\"Dry run with noise data. Returns training history.\"\"\"\n","    if config is None:\n","        config = DEFAULT_CONFIG.copy()\n","    config['num_classes'] = n_classes\n","\n","    model = GeoTransformerRedux('redux_test', config)\n","    model.network_to(device=device, strict=False)\n","    n_params = sum(p.numel() for p in model.parameters())\n","\n","    print(f\"\\n{'='*60}\")\n","    print(f\"  DRY RUN: {n_classes} classes, {n_epochs} epochs\")\n","    print(f\"  Params: {n_params:,} | d={config['d_model']} layers={config['n_layers']}\")\n","    print(f\"  Shared constellation: {config['n_anchors']} anchors on S^{config['manifold_dim']-1}\")\n","    print(f\"{'='*60}\")\n","\n","    compiled = torch.compile(model, mode='default')\n","    optimizer = torch.optim.Adam(model.parameters(), lr=config['lr'])\n","    loader = make_paired_noise_loader(2048, n_classes, config['batch_size'])\n","\n","    history = defaultdict(list)\n","    precompute_freq = config['precompute_freq']\n","\n","    for epoch in range(n_epochs):\n","        model.train()\n","        epoch_loss = epoch_acc = n = 0\n","        last_diag = None\n","        last_emb = None\n","        t0 = time.time()\n","\n","        for batch_idx, ((v1, v2), labels) in enumerate(loader):\n","            v1, v2, labels = v1.to(device), v2.to(device), labels.to(device)\n","\n","            if batch_idx % precompute_freq == 0:\n","                model.invalidate()\n","                model.precompute()\n","\n","            loss, diag, cls_feat = compiled(v1, v2, targets=labels)\n","\n","            # Geo losses BEFORE backward — combine into single .backward()\n","            geo = model.compute_geo_losses(cls_feat)\n","            total_loss_val = loss + geo['geo_total']\n","\n","            optimizer.zero_grad()\n","            total_loss_val.backward()\n","\n","            # Dynamic clip\n","            clip_val = max(diag[0].item(), 1.0)\n","            torch.nn.utils.clip_grad_norm_(model.parameters(), clip_val)\n","            optimizer.step()\n","\n","            # Accumulate for push (project to manifold, detach)\n","            with torch.no_grad():\n","                emb_for_push = F.normalize(model['manifold_proj'](cls_feat.detach()), dim=-1)\n","            buf_full = model.accumulate_for_push(emb_for_push, labels)\n","            if buf_full:\n","                batch_push = model.do_push()\n","                if batch_push.get('moved', 0) > 0:\n","                    model.invalidate()\n","                    model.precompute()\n","\n","            epoch_loss += loss.item() * v1.size(0)\n","            epoch_acc += diag[1].item() * v1.size(0)\n","            n += v1.size(0)\n","            last_diag = diag\n","            last_emb = emb_for_push\n","\n","        # End-of-epoch push for remaining buffer\n","        push_diag = model.do_push()\n","\n","        elapsed = time.time() - t0\n","        d = diag_to_dict(last_diag)\n","        geo_stats = compute_geometric_stats(model)\n","        grad_norms = compute_gradient_norms(model)\n","        # last_emb is already manifold-projected — compute CV directly\n","        with torch.no_grad():\n","            from geolip_core.core.distinguish.losses import cv_metric\n","            cv_emb_val = cv_metric(last_emb) if last_emb is not None and last_emb.shape[0] >= 5 else 0.0\n","\n","        # Record\n","        history['epoch'].append(epoch)\n","        history['loss'].append(epoch_loss / n)\n","        history['acc'].append(epoch_acc / n)\n","        history['time'].append(elapsed)\n","        for k, v in d.items():\n","            history[f'loss_{k}'].append(v)\n","        for k, v in geo_stats.items():\n","            history[f'geo_{k}'].append(v)\n","        history['cv_emb'].append(cv_emb_val)\n","        if push_diag:\n","            for k, v in push_diag.items():\n","                if isinstance(v, (int, float)):\n","                    history[f'push_{k}'].append(v)\n","\n","        # Print\n","        cfg = config\n","        w_ce = cfg['w_ce'] * d['ce']\n","        w_nce = cfg['w_nce_emb'] * d['nce_emb']\n","        w_pw = cfg['w_nce_pw'] * d['nce_pw']\n","        w_brg = cfg['w_bridge'] * d['bridge']\n","        w_asgn = cfg['w_assign'] * d['assign']\n","        w_anc = cfg['w_assign_nce'] * d['assign_nce']\n","        w_tri = cfg['w_nce_tri'] * d['nce_tri']\n","        w_att = cfg['w_attract'] * d['attract']\n","        w_sum = w_ce + w_nce + w_pw + w_brg + w_asgn + w_anc + w_tri + w_att\n","\n","        push_str = \"\"\n","        if push_diag and push_diag.get('moved', 0) > 0:\n","            push_str = (f\" PUSH: drift={push_diag.get('drift_mean',0):.4f}\"\n","                        f\" active={push_diag.get('n_active',0)}\"\n","                        f\" dead={push_diag.get('dead_count',0)}\")\n","\n","        print(f\"  E{epoch:>3d}  L={epoch_loss/n:.3f} acc={epoch_acc/n:.4f} \"\n","              f\"CV_emb={cv_emb_val:.3f} CV_anc={geo_stats['anchor_cv']:.3f} \"\n","              f\"cm+={geo_stats['cm_pos']:.2f} {elapsed:.1f}s{push_str}\")\n","\n","        if epoch % 5 == 0:\n","            print(f\"    wt: ce={w_ce:.2f} nce={w_nce:.2f} pw={w_pw:.2f} \"\n","                  f\"brg={w_brg:.2f} asgn={w_asgn:.2f} a_nce={w_anc:.2f} \"\n","                  f\"tri={w_tri:.2f} att={w_att:.2f} → {w_sum:.2f}\")\n","            gn = grad_norms\n","            gn_str = ' '.join(f\"{k}={v['mean']:.3f}\" for k, v in sorted(gn.items()))\n","            print(f\"    grad: {gn_str}\")\n","\n","    return history\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# MAIN — escalating dry runs\n","# ═══════════════════════════════════════════════════════════════════\n","\n","if __name__ == '__main__':\n","    print(f\"Device: {device}\")\n","    if device.type == 'cuda':\n","        print(f\"  GPU: {torch.cuda.get_device_name()}\")\n","        print(f\"  VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB\")\n","\n","    # Escalating class count dry runs\n","    for n_classes, n_epochs in [(10, 30), (50, 20), (100, 20)]:\n","        history = dry_run(n_classes, n_epochs)\n","        print(f\"\\n  Summary {n_classes} classes:\"\n","              f\" final_acc={history['acc'][-1]:.4f}\"\n","              f\" CV_emb={history['cv_emb'][-1]:.3f}\"\n","              f\" CV_anc={history['geo_anchor_cv'][-1]:.3f}\\n\")"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"EBTNaq1eE13I","executionInfo":{"status":"ok","timestamp":1775234142360,"user_tz":420,"elapsed":170879,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"4766c38e-94f5-46e0-8818-5fba97c76f32"},"execution_count":7,"outputs":[{"output_type":"stream","name":"stdout","text":["Device: cuda\n","  GPU: NVIDIA RTX PRO 6000 Blackwell Server Edition\n","  VRAM: 102.0 GB\n","\n","============================================================\n","  DRY RUN: 10 classes, 30 epochs\n","  Params: 6,804,747 | d=256 layers=4\n","  Shared constellation: 128 anchors on S^127\n","============================================================\n","  E  0  L=16.388 acc=0.1069 CV_emb=0.206 CV_anc=0.050 cm+=0.54 4.2s\n","    wt: ce=2.31 nce=0.17 pw=4.84 brg=3.17 asgn=0.00 a_nce=1.22 tri=2.42 att=0.10 → 14.25\n","    grad: cm_gate=0.001 constellation=0.005 content_attn=0.089 context=0.001 geo_attn=0.094 head=0.030 input=0.177 other=0.020 patchwork=0.026\n","  E  1  L=12.924 acc=0.1079 CV_emb=0.201 CV_anc=0.060 cm+=0.56 0.8s\n","  E  2  L=11.093 acc=0.1113 CV_emb=0.202 CV_anc=0.076 cm+=0.55 0.8s\n","  E  3  L=11.227 acc=0.1255 CV_emb=0.153 CV_anc=0.099 cm+=0.50 0.8s\n","  E  4  L=14.964 acc=0.1514 CV_emb=0.273 CV_anc=0.129 cm+=0.47 0.8s\n","  E  5  L=14.964 acc=0.1665 CV_emb=0.226 CV_anc=0.168 cm+=0.51 0.8s\n","    wt: ce=2.24 nce=0.03 pw=3.93 brg=4.47 asgn=0.01 a_nce=1.08 tri=2.28 att=0.11 → 14.15\n","    grad: cm_gate=0.002 constellation=0.005 content_attn=0.058 context=0.002 geo_attn=0.070 head=0.022 input=0.265 other=0.007 patchwork=0.006\n","  E  6  L=13.428 acc=0.1831 CV_emb=0.215 CV_anc=0.213 cm+=0.48 0.8s\n","  E  7  L=12.724 acc=0.1855 CV_emb=0.216 CV_anc=0.257 cm+=0.51 0.7s\n","  E  8  L=12.149 acc=0.1958 CV_emb=0.310 CV_anc=0.297 cm+=0.48 0.7s\n","  E  9  L=11.857 acc=0.2090 CV_emb=0.229 CV_anc=0.326 cm+=0.51 0.7s\n","  E 10  L=11.230 acc=0.2324 CV_emb=0.186 CV_anc=0.345 cm+=0.43 0.7s\n","    wt: ce=2.18 nce=0.01 pw=1.71 brg=3.73 asgn=0.01 a_nce=1.07 tri=2.06 att=0.08 → 10.86\n","    grad: cm_gate=0.000 constellation=0.004 content_attn=0.058 context=0.002 geo_attn=0.072 head=0.040 input=0.264 other=0.010 patchwork=0.007\n","  E 11  L=10.841 acc=0.2627 CV_emb=0.237 CV_anc=0.352 cm+=0.37 0.7s\n","  E 12  L=10.553 acc=0.2661 CV_emb=0.200 CV_anc=0.354 cm+=0.45 0.7s\n","  E 13  L=10.200 acc=0.3110 CV_emb=0.167 CV_anc=0.346 cm+=0.43 0.7s\n","  E 14  L=9.932 acc=0.3784 CV_emb=0.176 CV_anc=0.338 cm+=0.55 0.7s\n","  E 15  L=9.655 acc=0.4941 CV_emb=0.172 CV_anc=0.331 cm+=0.50 0.7s\n","    wt: ce=1.51 nce=0.01 pw=0.91 brg=3.78 asgn=0.01 a_nce=1.13 tri=2.02 att=0.07 → 9.44\n","    grad: cm_gate=0.000 constellation=0.002 content_attn=0.043 context=0.003 geo_attn=0.058 head=0.026 input=0.193 other=0.009 patchwork=0.004\n","  E 16  L=9.128 acc=0.6567 CV_emb=0.157 CV_anc=0.330 cm+=0.48 0.7s\n","  E 17  L=8.730 acc=0.7949 CV_emb=0.145 CV_anc=0.335 cm+=0.53 0.7s\n","  E 18  L=8.384 acc=0.8994 CV_emb=0.172 CV_anc=0.344 cm+=0.49 0.7s\n","  E 19  L=8.137 acc=0.9438 CV_emb=0.178 CV_anc=0.354 cm+=0.52 0.7s\n","  E 20  L=7.948 acc=0.9746 CV_emb=0.171 CV_anc=0.361 cm+=0.48 0.7s\n","    wt: ce=0.17 nce=0.02 pw=0.75 brg=3.69 asgn=0.01 a_nce=1.12 tri=2.08 att=0.09 → 7.93\n","    grad: cm_gate=0.001 constellation=0.002 content_attn=0.026 context=0.002 geo_attn=0.034 head=0.007 input=0.125 other=0.005 patchwork=0.003\n","  E 21  L=7.865 acc=0.9849 CV_emb=0.155 CV_anc=0.364 cm+=0.43 0.7s\n","  E 22  L=7.742 acc=0.9888 CV_emb=0.158 CV_anc=0.364 cm+=0.41 0.7s\n","  E 23  L=7.666 acc=0.9893 CV_emb=0.160 CV_anc=0.361 cm+=0.38 0.7s\n","  E 24  L=7.506 acc=0.9937 CV_emb=0.143 CV_anc=0.357 cm+=0.45 0.7s\n","  E 25  L=7.399 acc=0.9927 CV_emb=0.142 CV_anc=0.352 cm+=0.47 0.7s\n","    wt: ce=0.07 nce=0.01 pw=0.64 brg=3.41 asgn=0.01 a_nce=1.10 tri=1.98 att=0.08 → 7.31\n","    grad: cm_gate=0.001 constellation=0.002 content_attn=0.026 context=0.002 geo_attn=0.036 head=0.009 input=0.129 other=0.006 patchwork=0.003\n","  E 26  L=7.462 acc=0.9829 CV_emb=0.147 CV_anc=0.347 cm+=0.46 0.7s\n","  E 27  L=7.322 acc=0.9897 CV_emb=0.142 CV_anc=0.341 cm+=0.53 0.7s\n","  E 28  L=7.214 acc=0.9932 CV_emb=0.137 CV_anc=0.337 cm+=0.51 0.7s\n","  E 29  L=7.199 acc=0.9893 CV_emb=0.145 CV_anc=0.333 cm+=0.53 0.7s\n","\n","  Summary 10 classes: final_acc=0.9893 CV_emb=0.145 CV_anc=0.333\n","\n","\n","============================================================\n","  DRY RUN: 50 classes, 20 epochs\n","  Params: 6,815,027 | d=256 layers=4\n","  Shared constellation: 128 anchors on S^127\n","============================================================\n","  E  0  L=17.739 acc=0.0254 CV_emb=0.298 CV_anc=0.051 cm+=0.50 55.8s\n","    wt: ce=3.97 nce=0.41 pw=4.84 brg=2.58 asgn=0.00 a_nce=1.22 tri=2.42 att=0.10 → 15.54\n","    grad: cm_gate=0.004 constellation=0.005 content_attn=0.242 context=0.001 geo_attn=0.120 head=0.033 input=0.365 other=0.032 patchwork=0.022\n","  E  1  L=16.335 acc=0.0205 CV_emb=0.254 CV_anc=0.064 cm+=0.48 0.9s\n","  E  2  L=15.472 acc=0.0308 CV_emb=0.225 CV_anc=0.084 cm+=0.52 0.9s\n","  E  3  L=16.022 acc=0.0254 CV_emb=0.235 CV_anc=0.107 cm+=0.46 0.9s\n","  E  4  L=15.632 acc=0.0269 CV_emb=0.280 CV_anc=0.140 cm+=0.44 0.9s\n","  E  5  L=15.076 acc=0.0317 CV_emb=0.277 CV_anc=0.176 cm+=0.49 0.8s\n","    wt: ce=3.89 nce=0.09 pw=4.08 brg=2.94 asgn=0.01 a_nce=1.03 tri=2.27 att=0.08 → 14.40\n","    grad: cm_gate=0.002 constellation=0.009 content_attn=0.165 context=0.001 geo_attn=0.090 head=0.038 input=0.434 other=0.022 patchwork=0.012\n","  E  6  L=14.013 acc=0.0400 CV_emb=0.219 CV_anc=0.209 cm+=0.54 0.8s\n","  E  7  L=13.730 acc=0.0386 CV_emb=0.253 CV_anc=0.232 cm+=0.48 0.8s\n","  E  8  L=13.681 acc=0.0464 CV_emb=0.312 CV_anc=0.251 cm+=0.49 0.9s\n","  E  9  L=13.110 acc=0.0488 CV_emb=0.288 CV_anc=0.259 cm+=0.48 0.8s\n","  E 10  L=12.709 acc=0.0522 CV_emb=0.310 CV_anc=0.264 cm+=0.41 0.9s\n","    wt: ce=3.81 nce=0.07 pw=1.61 brg=3.63 asgn=0.01 a_nce=0.98 tri=2.19 att=0.09 → 12.39\n","    grad: cm_gate=0.005 constellation=0.011 content_attn=0.159 context=0.001 geo_attn=0.094 head=0.062 input=0.444 other=0.025 patchwork=0.021\n","  E 11  L=12.256 acc=0.0605 CV_emb=0.224 CV_anc=0.270 cm+=0.45 0.8s\n","  E 12  L=11.810 acc=0.0596 CV_emb=0.268 CV_anc=0.267 cm+=0.41 0.8s\n","  E 13  L=11.563 acc=0.0825 CV_emb=0.221 CV_anc=0.265 cm+=0.41 0.8s\n","  E 14  L=11.063 acc=0.0811 CV_emb=0.252 CV_anc=0.264 cm+=0.37 0.8s\n","  E 15  L=10.786 acc=0.0918 CV_emb=0.232 CV_anc=0.262 cm+=0.39 0.7s\n","    wt: ce=3.71 nce=0.03 pw=0.67 brg=3.14 asgn=0.01 a_nce=0.96 tri=2.21 att=0.09 → 10.81\n","    grad: cm_gate=0.002 constellation=0.007 content_attn=0.131 context=0.002 geo_attn=0.089 head=0.054 input=0.437 other=0.019 patchwork=0.019\n","  E 16  L=10.565 acc=0.1162 CV_emb=0.206 CV_anc=0.260 cm+=0.29 0.7s\n","  E 17  L=10.392 acc=0.1479 CV_emb=0.310 CV_anc=0.258 cm+=0.39 0.7s\n","  E 18  L=10.193 acc=0.1680 CV_emb=0.215 CV_anc=0.258 cm+=0.55 0.7s\n","  E 19  L=9.997 acc=0.2202 CV_emb=0.194 CV_anc=0.257 cm+=0.50 0.8s\n","\n","  Summary 50 classes: final_acc=0.2202 CV_emb=0.194 CV_anc=0.257\n","\n","\n","============================================================\n","  DRY RUN: 100 classes, 20 epochs\n","  Params: 6,827,877 | d=256 layers=4\n","  Shared constellation: 128 anchors on S^127\n","============================================================\n","  E  0  L=18.551 acc=0.0098 CV_emb=0.233 CV_anc=0.052 cm+=0.55 58.6s\n","    wt: ce=4.61 nce=0.24 pw=4.84 brg=2.95 asgn=0.00 a_nce=1.22 tri=2.42 att=0.11 → 16.40\n","    grad: cm_gate=0.003 constellation=0.007 content_attn=0.243 context=0.001 geo_attn=0.165 head=0.066 input=0.419 other=0.045 patchwork=0.039\n","  E  1  L=15.336 acc=0.0107 CV_emb=0.177 CV_anc=0.072 cm+=0.58 0.9s\n","  E  2  L=13.893 acc=0.0132 CV_emb=0.189 CV_anc=0.093 cm+=0.60 0.9s\n","  E  3  L=18.735 acc=0.0151 CV_emb=0.209 CV_anc=0.119 cm+=0.51 0.9s\n","  E  4  L=16.374 acc=0.0205 CV_emb=0.295 CV_anc=0.148 cm+=0.52 1.0s\n","  E  5  L=15.599 acc=0.0200 CV_emb=0.281 CV_anc=0.175 cm+=0.54 0.8s\n","    wt: ce=4.58 nce=0.06 pw=3.58 brg=3.26 asgn=0.01 a_nce=0.98 tri=2.29 att=0.09 → 14.85\n","    grad: cm_gate=0.002 constellation=0.008 content_attn=0.174 context=0.001 geo_attn=0.082 head=0.045 input=0.529 other=0.017 patchwork=0.013\n","  E  6  L=14.591 acc=0.0200 CV_emb=0.198 CV_anc=0.201 cm+=0.55 0.8s\n","  E  7  L=14.371 acc=0.0288 CV_emb=0.246 CV_anc=0.220 cm+=0.59 0.8s\n","  E  8  L=13.838 acc=0.0298 CV_emb=0.213 CV_anc=0.231 cm+=0.61 0.8s\n","  E  9  L=13.288 acc=0.0273 CV_emb=0.228 CV_anc=0.238 cm+=0.59 0.8s\n","  E 10  L=12.685 acc=0.0425 CV_emb=0.209 CV_anc=0.239 cm+=0.48 0.8s\n","    wt: ce=4.48 nce=0.02 pw=1.72 brg=3.20 asgn=0.01 a_nce=0.96 tri=2.14 att=0.08 → 12.60\n","    grad: cm_gate=0.002 constellation=0.007 content_attn=0.162 context=0.001 geo_attn=0.087 head=0.045 input=0.526 other=0.017 patchwork=0.014\n","  E 11  L=12.244 acc=0.0356 CV_emb=0.252 CV_anc=0.238 cm+=0.52 0.7s\n","  E 12  L=11.866 acc=0.0640 CV_emb=0.243 CV_anc=0.236 cm+=0.57 0.7s\n","  E 13  L=11.637 acc=0.0649 CV_emb=0.236 CV_anc=0.233 cm+=0.62 0.8s\n","  E 14  L=11.422 acc=0.0859 CV_emb=0.181 CV_anc=0.232 cm+=0.66 0.7s\n","  E 15  L=11.066 acc=0.0884 CV_emb=0.267 CV_anc=0.232 cm+=0.62 0.7s\n","    wt: ce=4.14 nce=0.03 pw=0.75 brg=2.96 asgn=0.01 a_nce=0.93 tri=2.23 att=0.09 → 11.14\n","    grad: cm_gate=0.004 constellation=0.007 content_attn=0.146 context=0.002 geo_attn=0.093 head=0.055 input=0.487 other=0.019 patchwork=0.018\n","  E 16  L=10.846 acc=0.1143 CV_emb=0.209 CV_anc=0.232 cm+=0.59 0.7s\n","  E 17  L=10.582 acc=0.1523 CV_emb=0.197 CV_anc=0.232 cm+=0.56 0.7s\n","  E 18  L=10.372 acc=0.1963 CV_emb=0.187 CV_anc=0.232 cm+=0.55 0.7s\n","  E 19  L=10.038 acc=0.2612 CV_emb=0.188 CV_anc=0.233 cm+=0.58 0.7s\n","\n","  Summary 100 classes: final_acc=0.2612 CV_emb=0.188 CV_anc=0.233\n","\n"]}]},{"cell_type":"code","source":["\"\"\"\n","Geometric Transformer Redux v8\n","================================\n","Shared constellation across all layers. Momentum-based anchor push.\n","Proper pentachoron CV on embeddings. WideRouter compiled.\n","\n","Architecture:\n","  ConvSVDPatchEmbed → 4× ReduxLayer (shared constellation) → head\n","\n","Each layer:\n","  project → observe (shared anchors) → CM gate → patchwork\n","  → FiLM context → dual attention → gated compose → geo_residual\n","\n","Key differences from v1:\n","  - ONE constellation for all layers (not per-layer)\n","  - Momentum AnchorPush (non-gradient anchor repositioning)\n","  - cv_loss on EMBEDDINGS (pentachoron volumes), not anchor pairwise\n","  - torch.stack diagnostics (no dict crossing compile boundary)\n","  - Dynamic grad clip: max(ce_loss, 1.0)\n","  - Precompute every N batches\n","\n","!pip install geolip-core torchvision tqdm\n","\"\"\"\n","\n","import os, sys, warnings, math, time\n","import numpy as np\n","from pathlib import Path\n","from collections import defaultdict\n","\n","warnings.filterwarnings('ignore')\n","\n","import torch\n","import torch.nn as nn\n","import torch.nn.functional as F\n","from tqdm.auto import tqdm\n","\n","device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n","\n","torch.set_float32_matmul_precision('high')\n","torch.backends.cudnn.benchmark = True\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# IMPORTS from geolip_core — utilities only\n","# ═══════════════════════════════════════════════════════════════════\n","\n","from geolip_core.core.input.svd import SVDObserver\n","from geolip_core.core.associate.constellation import (\n","    ConstellationObserver, ConstellationAssociation, ConstellationCuration,\n",")\n","from geolip_core.core.curate.patchwork import AnchorPush\n","from geolip_core.core.distinguish.losses import (\n","    nce_loss, bridge_loss_paired, assign_bce_loss,\n","    assign_nce_loss, attraction_loss, ce_loss_paired,\n","    cv_loss, spread_loss,\n",")\n","from geolip_core.pipeline.components.geometric_transformer import CMValidatedGate\n","from geolip_core.pipeline.observer import TorchComponent\n","from geofractal.router.wide_router import WideRouter\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# CONFIG\n","# ═══════════════════════════════════════════════════════════════════\n","\n","DEFAULT_CONFIG = {\n","    # Architecture\n","    'd_model': 256,\n","    'n_heads': 8,\n","    'n_layers': 4,\n","    'manifold_dim': 128,\n","    'n_anchors': 256,\n","    'n_comp': 8,\n","    'd_comp': 32,\n","    'context_dim': 128,\n","    'cm_neighbors': 3,\n","    'dropout': 0.1,\n","\n","    # Input\n","    'conv_channels': 128,\n","    'svd_rank': 16,\n","    'patch_size': 16,\n","    'img_size': 32,\n","\n","    # Training\n","    'epochs': 200,\n","    'batch_size': 256,\n","    'lr': 1e-4,\n","    'warmup_epochs': 10,\n","    'num_workers': 8,\n","    'precompute_freq': 10,\n","\n","    # Push\n","    'push_every': 5,\n","    'push_strategy': 'momentum',\n","    'push_accumulate': 8192,  # ~64 batches at B=128 → ~6 pushes/epoch on CIFAR\n","\n","    # Loss\n","    'w_ce': 1.0,\n","    'w_nce_emb': 0.5,\n","    'w_nce_pw': 1.0,\n","    'w_bridge': 1.0,\n","    'w_assign': 0.5,\n","    'w_assign_nce': 0.25,\n","    'w_nce_tri': 0.5,\n","    'w_attract': 0.25,\n","    'w_cv': 0.1,\n","    'w_spread': 0.05,\n","    'cv_target': 0.21,\n","\n","    'num_classes': 100,\n","}\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# INPUT STAGE\n","# ═══════════════════════════════════════════════════════════════════\n","\n","class ConvSVDInput(TorchComponent):\n","    \"\"\"Conv → SVD → patch tokens. FL kernel, compilable.\"\"\"\n","    def __init__(self, name, img_size=32, patch_size=4, in_channels=3,\n","                 conv_channels=64, d_model=256, svd_rank=12):\n","        super().__init__(name)\n","        self.patch_size = patch_size\n","        n_patches = (img_size // patch_size) ** 2\n","\n","        self.conv = nn.Sequential(\n","            nn.Conv2d(in_channels, conv_channels, 3, padding=1, bias=False),\n","            nn.BatchNorm2d(conv_channels), nn.GELU(),\n","            nn.Conv2d(conv_channels, conv_channels, 3, padding=1, bias=False),\n","            nn.BatchNorm2d(conv_channels), nn.GELU(),\n","        )\n","        self.svd = SVDObserver(conv_channels, svd_rank)\n","        self.patch_proj = nn.Conv2d(conv_channels, d_model,\n","                                    kernel_size=patch_size, stride=patch_size, bias=False)\n","        self.patch_norm = nn.LayerNorm(d_model)\n","\n","        feat_dim = self.svd.feature_dim\n","        self.svd_gamma = nn.Linear(feat_dim, d_model)\n","        self.svd_beta = nn.Linear(feat_dim, d_model)\n","        nn.init.normal_(self.svd_gamma.weight, std=0.01)\n","        nn.init.ones_(self.svd_gamma.bias)\n","        nn.init.zeros_(self.svd_beta.bias)\n","\n","        self.cls_token = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)\n","        self.pos_embed = nn.Parameter(torch.randn(1, n_patches + 1, d_model) * 0.02)\n","\n","    def forward(self, x):\n","        B = x.shape[0]\n","        feat = self.conv(x)\n","        S, Vh, svd_feat, novelty = self.svd(feat)\n","        tokens = self.patch_proj(feat).flatten(2).transpose(1, 2)\n","        tokens = self.patch_norm(tokens)\n","        g = self.svd_gamma(svd_feat).unsqueeze(1)\n","        b = self.svd_beta(svd_feat).unsqueeze(1)\n","        tokens = g * tokens + b\n","        tokens = torch.cat([self.cls_token.expand(B, -1, -1), tokens], dim=1)\n","        return tokens + self.pos_embed\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# REDUX LAYER — reads shared constellation\n","# ═══════════════════════════════════════════════════════════════════\n","\n","class ReduxLayer(WideRouter):\n","    \"\"\"Transformer layer observing a shared constellation.\n","\n","    Receives anchors_n and cm_gate externally — does NOT own them.\n","    \"\"\"\n","    def __init__(self, name, d_model=256, n_heads=8, n_anchors=128,\n","                 manifold_dim=128, n_comp=8, d_comp=32,\n","                 context_dim=128, dropout=0.1):\n","        super().__init__(name, strict=False)\n","        self.d_model = d_model\n","        self.manifold_dim = manifold_dim\n","        self.n_anchors = n_anchors\n","        pw_dim = n_comp * d_comp\n","\n","        # Project to manifold (per-layer — each layer sees different aspect)\n","        self.attach('proj', nn.Sequential(\n","            nn.Linear(d_model, manifold_dim), nn.LayerNorm(manifold_dim)))\n","\n","        # Per-layer patchwork (reads shared constellation's triangulation)\n","        from geolip_core.core.curate.patchwork import Patchwork\n","        self.attach('patchwork', Patchwork(\n","            n_anchors=n_anchors, n_comp=n_comp, d_comp=d_comp))\n","        self.attach('bridge_proj', nn.Linear(pw_dim, n_anchors))\n","\n","        # FiLM context: anchor feats + patchwork + history\n","        self.attach('anchor_ctx', nn.Linear(n_anchors * 2, context_dim))  # cos + gate\n","        self.attach('pw_ctx', nn.Linear(pw_dim, context_dim))\n","        self.attach('history_ctx', nn.Linear(pw_dim, context_dim))\n","        self.attach('ctx_fuse', nn.Sequential(\n","            nn.Linear(context_dim * 3, context_dim), nn.GELU()))\n","        self.attach('film_gamma', nn.Linear(context_dim, d_model))\n","        self.attach('film_beta', nn.Linear(context_dim, d_model))\n","\n","        # Dual attention\n","        self.attach('content_attn', nn.MultiheadAttention(\n","            d_model, n_heads, batch_first=True, dropout=dropout))\n","        self.attach('content_norm', nn.LayerNorm(d_model))\n","        self.attach('geo_attn', nn.MultiheadAttention(\n","            d_model, n_heads, batch_first=True, dropout=dropout))\n","        self.attach('geo_norm', nn.LayerNorm(d_model))\n","\n","        # Gated composition\n","        self.attach('gate', nn.Sequential(\n","            nn.Linear(d_model * 2, d_model), nn.Sigmoid()))\n","        self.attach('ffn', nn.Sequential(\n","            nn.Linear(d_model, d_model * 4), nn.GELU(), nn.Dropout(dropout),\n","            nn.Linear(d_model * 4, d_model)))\n","        self.attach('ffn_norm', nn.LayerNorm(d_model))\n","\n","        # Geo residual projection\n","        self._pw_dim = pw_dim\n","        self.attach('geo_proj', nn.Sequential(\n","            nn.Linear(pw_dim, pw_dim), nn.LayerNorm(pw_dim)))\n","\n","    def forward(self, x, anchors_n, gate_values, geo_residual=None):\n","        \"\"\"\n","        Args:\n","            x: (B, L, D) tokens\n","            anchors_n: (A, M) shared normalized anchors\n","            gate_values: (B*L, A) from shared CM gate\n","            geo_residual: (B, L, pw_dim) or None\n","        Returns:\n","            x_out, geo_residual_out, emb_flat, pw_flat, bridge_flat, a_out_dict\n","        \"\"\"\n","        B, L, D = x.shape\n","\n","        # 1. Project to manifold\n","        emb = F.normalize(self['proj'](x), dim=-1, eps=1e-6)\n","        emb_flat = emb.reshape(B * L, -1)\n","\n","        # 2. Triangulate against shared anchors\n","        cos = emb_flat @ anchors_n.T\n","        distances = 1.0 - cos\n","        assignment = F.softmax(cos / 0.1, dim=-1)\n","\n","        # 3. Gate triangulation\n","        gated_distances = distances * gate_values\n","\n","        # 4. Patchwork curation\n","        a_out = {\n","            'distances': distances,\n","            'distances_weighted': gated_distances,\n","            'cos_to_anchors': cos,\n","            'assignment': assignment,\n","            'nearest': cos.argmax(dim=-1),\n","        }\n","        pw_flat = self['patchwork'](gated_distances)  # (B*L, pw_dim)\n","        bridge_flat = self['bridge_proj'](pw_flat)\n","\n","        # 5. FiLM context\n","        anchor_feats = torch.cat([cos, gate_values], dim=-1)\n","        a_ctx = self['anchor_ctx'](anchor_feats).reshape(B, L, -1)\n","        p_ctx = self['pw_ctx'](pw_flat).reshape(B, L, -1)\n","        h_ctx = self['history_ctx'](geo_residual.reshape(B * L, -1)).reshape(B, L, -1) \\\n","            if geo_residual is not None else torch.zeros_like(a_ctx)\n","        ctx = self['ctx_fuse'](torch.cat([a_ctx, p_ctx, h_ctx], dim=-1))\n","        gamma = self['film_gamma'](ctx)\n","        beta = self['film_beta'](ctx)\n","\n","        # 6. Dual attention\n","        content, _ = self['content_attn'](x, x, x, need_weights=False)\n","        content = self['content_norm'](x + content)\n","        geo_q = x + gamma * x + beta\n","        geo_out, _ = self['geo_attn'](geo_q, x, x, need_weights=False)\n","        geo_out = self['geo_norm'](x + geo_out)\n","\n","        # 7. Gated compose + FFN\n","        g = self['gate'](torch.cat([content, geo_out], dim=-1))\n","        merged = g * geo_out + (1 - g) * content\n","        h = self['ffn'](merged)\n","        x_out = self['ffn_norm'](merged + h)\n","\n","        # 8. Geo residual accumulation\n","        cm_quality = gate_values.mean(dim=-1).reshape(B, L, 1)\n","        geo_update = self['geo_proj'](pw_flat.reshape(B, L, -1))\n","        if geo_residual is None:\n","            geo_residual_out = cm_quality * geo_update\n","        else:\n","            geo_residual_out = geo_residual + cm_quality * geo_update\n","\n","        return x_out, geo_residual_out, emb_flat, pw_flat, bridge_flat, a_out\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# GEOMETRIC TRANSFORMER REDUX — WideRouter\n","# ═══════════════════════════════════════════════════════════════════\n","\n","class GeoTransformerRedux(WideRouter):\n","    \"\"\"Shared constellation, momentum push, proper CV.\n","\n","    All layers observe the SAME constellation. Anchor push moves them\n","    based on accumulated embeddings. CV measured on embedding pentachora.\n","    \"\"\"\n","    def __init__(self, name, config):\n","        super().__init__(name, strict=False)\n","        self.config = config\n","        d = config['d_model']\n","        m = config['manifold_dim']\n","        A = config['n_anchors']\n","        n_layers = config['n_layers']\n","\n","        # Input\n","        self.attach('input', ConvSVDInput(\n","            'input', img_size=config['img_size'], patch_size=config['patch_size'],\n","            conv_channels=config['conv_channels'], d_model=d,\n","            svd_rank=config['svd_rank']))\n","\n","        # SHARED constellation + CM gate (one for all layers)\n","        self.attach('observer', ConstellationObserver(\n","            dim=m, n_anchors=A,\n","            n_comp=config['n_comp'], d_comp=config['d_comp']))\n","        self.attach('cm_gate', CMValidatedGate(A, n_neighbors=config['cm_neighbors']))\n","\n","        # Per-layer transformer blocks (read shared constellation)\n","        for i in range(n_layers):\n","            self.attach(f'layer_{i}', ReduxLayer(\n","                f'L{i}', d, config['n_heads'], A, m,\n","                config['n_comp'], config['d_comp'],\n","                config['context_dim'], config['dropout']))\n","            self.register_tower(f'layer_{i}')\n","\n","        self.attach('final_norm', nn.LayerNorm(d))\n","\n","        # Classification head\n","        pw_dim = config['n_comp'] * config['d_comp']\n","        head_in = d + pw_dim  # cls token + geo_residual pooled\n","        self.attach('head', nn.Sequential(\n","            nn.LayerNorm(head_in), nn.Linear(head_in, d),\n","            nn.GELU(), nn.Dropout(config['dropout']),\n","            nn.Linear(d, config['num_classes'])))\n","\n","        # Manifold projection for shared observation\n","        self.attach('manifold_proj', nn.Sequential(\n","            nn.Linear(d, m), nn.LayerNorm(m)))\n","\n","        self.n_layers = n_layers\n","\n","        # Anchor push (non-gradient momentum updates)\n","        self.push = AnchorPush(\n","            config['push_strategy'], A, m,\n","            decay=0.9, alpha=0.05, beta=0.02, util_floor=0.001)\n","\n","        # Push accumulation buffers\n","        self.register_buffer('_emb_buf', torch.zeros(config['push_accumulate'], m))\n","        self.register_buffer('_lbl_buf', torch.zeros(config['push_accumulate'], dtype=torch.long))\n","        self.register_buffer('_buf_ptr', torch.zeros(1, dtype=torch.long))\n","\n","    @property\n","    def constellation(self):\n","        return self['observer'].association.constellation\n","\n","    @torch.no_grad()\n","    def precompute(self):\n","        \"\"\"CM gate + cache. Call before compiled forward.\"\"\"\n","        anchors_n = F.normalize(self.constellation.anchors, dim=-1)\n","        self['cm_gate'].precompute(anchors_n.detach())\n","\n","    def invalidate(self):\n","        self['cm_gate'].invalidate_cache()\n","\n","    @torch.no_grad()\n","    def accumulate_for_push(self, emb, labels):\n","        \"\"\"Accumulate embeddings for anchor push. Returns True if buffer is full.\"\"\"\n","        B = emb.shape[0]\n","        ptr = int(self._buf_ptr.item())\n","        cap = self._emb_buf.shape[0]\n","        end = min(ptr + B, cap)\n","        n = end - ptr\n","        self._emb_buf[ptr:end] = emb[:n].detach()\n","        self._lbl_buf[ptr:end] = labels[:n].detach()\n","        self._buf_ptr[0] = end\n","        return end >= cap  # True = buffer full, caller should push\n","\n","    @torch.no_grad()\n","    def do_push(self):\n","        \"\"\"Execute momentum anchor push. Returns diagnostics dict.\"\"\"\n","        ptr = int(self._buf_ptr.item())\n","        if ptr < 64:\n","            return {'moved': 0}\n","        dev = self.constellation.anchors.device\n","        emb = self._emb_buf[:ptr].to(dev)\n","        lbl = self._lbl_buf[:ptr].to(dev)\n","        diag = self.push.push(self['observer'], emb, lbl)\n","        self._buf_ptr.zero_()\n","        return diag\n","\n","    def forward(self, x, x2=None, targets=None):\n","        \"\"\"\n","        Single view: returns logits\n","        Paired view with targets: returns (loss, diagnostics_tensor)\n","        \"\"\"\n","        if x2 is not None and targets is not None:\n","            return self._forward_paired(x, x2, targets)\n","\n","        tokens = self['input'](x)\n","        B, L, D = tokens.shape\n","\n","        # Shared observation: project CLS to manifold, observe\n","        anchors_n = F.normalize(self.constellation.anchors, dim=-1)\n","\n","        # Run layers\n","        geo_residual = None\n","        for i in range(self.n_layers):\n","            # Each layer projects tokens to manifold independently\n","            emb_for_gate = F.normalize(self['manifold_proj'](tokens), dim=-1, eps=1e-6)\n","            flat = emb_for_gate.reshape(B * L, -1)\n","            gate_values, _ = self['cm_gate'](1.0 - flat @ anchors_n.T)\n","\n","            tokens, geo_residual, _, _, _, _ = self[f'layer_{i}'](\n","                tokens, anchors_n, gate_values, geo_residual)\n","\n","        tokens = self['final_norm'](tokens)\n","        cls = tokens[:, 0]\n","        geo_pool = geo_residual[:, 0] if geo_residual is not None else \\\n","            torch.zeros(B, self.config['n_comp'] * self.config['d_comp'], device=x.device)\n","        logits = self['head'](torch.cat([cls, geo_pool], dim=-1))\n","        return logits\n","\n","    def _forward_paired(self, x1, x2, targets):\n","        \"\"\"Paired forward with full observer loss. Returns (loss, diag_tensor).\"\"\"\n","        B = x1.shape[0]\n","        x_cat = torch.cat([x1, x2], dim=0)\n","        tokens = self['input'](x_cat)\n","        BB, L, D = tokens.shape  # BB = 2*B\n","\n","        anchors_n = F.normalize(self.constellation.anchors, dim=-1)\n","\n","        # Run layers — collect LAST layer's observation for loss\n","        geo_residual = None\n","        last_emb = last_pw = last_bridge = last_a_out = None\n","        for i in range(self.n_layers):\n","            emb_for_gate = F.normalize(self['manifold_proj'](tokens), dim=-1, eps=1e-6)\n","            flat = emb_for_gate.reshape(BB * L, -1)\n","            gate_values, _ = self['cm_gate'](1.0 - flat @ anchors_n.T)\n","\n","            tokens, geo_residual, emb_flat, pw_flat, bridge_flat, a_out = \\\n","                self[f'layer_{i}'](tokens, anchors_n, gate_values, geo_residual)\n","\n","            last_emb = emb_flat\n","            last_pw = pw_flat\n","            last_bridge = bridge_flat\n","            last_a_out = a_out\n","\n","        tokens = self['final_norm'](tokens)\n","\n","        # Split views\n","        cls1, cls2 = tokens[:B, 0], tokens[B:, 0]\n","        geo1 = geo_residual[:B, 0] if geo_residual is not None else torch.zeros_like(cls1[:, :self._pw_dim_])\n","        geo2 = geo_residual[B:, 0] if geo_residual is not None else torch.zeros_like(cls2[:, :self._pw_dim_])\n","\n","        # CLS token embeddings on manifold (for NCE)\n","        emb1 = F.normalize(self['manifold_proj'](cls1), dim=-1, eps=1e-6)\n","        emb2 = F.normalize(self['manifold_proj'](cls2), dim=-1, eps=1e-6)\n","\n","        # Last layer observations at CLS position\n","        c = 0  # cls index\n","        N = BB * L\n","        # Reshape to (2B, L, ...) then extract CLS\n","        pw1 = last_pw.reshape(BB, L, -1)[:B, c]\n","        pw2 = last_pw.reshape(BB, L, -1)[B:, c]\n","        brg1 = last_bridge.reshape(BB, L, -1)[:B, c]\n","        brg2 = last_bridge.reshape(BB, L, -1)[B:, c]\n","        assign1 = last_a_out['assignment'].reshape(BB, L, -1)[:B, c]\n","        assign2 = last_a_out['assignment'].reshape(BB, L, -1)[B:, c]\n","        cos1 = last_a_out['cos_to_anchors'].reshape(BB, L, -1)[:B, c]\n","        tri1 = last_a_out['distances'].reshape(BB, L, -1)[:B, c]\n","        tri2 = last_a_out['distances'].reshape(BB, L, -1)[B:, c]\n","\n","        # ═══ LOSSES ═══\n","        cfg = self.config\n","\n","        # CE\n","        logits1 = self['head'](torch.cat([cls1, geo1], dim=-1))\n","        logits2 = self['head'](torch.cat([cls2, geo2], dim=-1))\n","        l_ce, acc = ce_loss_paired(logits1, logits2, targets)\n","\n","        # Embedding NCE\n","        l_nce_emb, _ = nce_loss(emb1, emb2, 0.1, normalize=False)\n","\n","        # Patchwork NCE\n","        l_nce_pw, _ = nce_loss(pw1, pw2, 0.1, normalize=True)\n","\n","        # Bridge\n","        l_bridge, _ = bridge_loss_paired(brg1, brg2, assign1, assign2)\n","\n","        # Assign BCE\n","        l_assign, _ = assign_bce_loss(assign1, cos1)\n","\n","        # Assign NCE\n","        l_assign_nce, _ = assign_nce_loss(assign1, assign2, 0.1)\n","\n","        # Tri NCE\n","        l_nce_tri, _ = nce_loss(tri1, tri2, 0.1, normalize=True)\n","\n","        # Attraction\n","        l_attract, _ = attraction_loss(cos1)\n","\n","        loss = (cfg['w_ce'] * l_ce\n","                + cfg['w_nce_emb'] * l_nce_emb\n","                + cfg['w_nce_pw'] * l_nce_pw\n","                + cfg['w_bridge'] * l_bridge\n","                + cfg['w_assign'] * l_assign\n","                + cfg['w_assign_nce'] * l_assign_nce\n","                + cfg['w_nce_tri'] * l_nce_tri\n","                + cfg['w_attract'] * l_attract)\n","\n","        # Single tensor — no dict crossing compile boundary\n","        diag = torch.stack([\n","            l_ce.detach(), acc.detach(),\n","            l_nce_emb.detach(), l_nce_pw.detach(),\n","            l_bridge.detach(), l_assign.detach(),\n","            l_assign_nce.detach(), l_nce_tri.detach(),\n","            l_attract.detach(),\n","        ])\n","\n","        return loss, diag, cls1  # cls1 has gradient — for CV loss outside compile\n","\n","    @property\n","    def _pw_dim_(self):\n","        return self.config['n_comp'] * self.config['d_comp']\n","\n","    def compute_geo_losses(self, cls_feat=None):\n","        \"\"\"CV (pentachoron on embeddings) + spread (anchor repulsion).\n","        Call OUTSIDE compiled graph.\n","\n","        Args:\n","            cls_feat: (B, d_model) CLS features WITH gradient.\n","                      Projected to manifold here for CV loss.\n","        \"\"\"\n","        losses = {}\n","        cfg = self.config\n","        if cls_feat is not None and cls_feat.shape[0] >= 5:\n","            # Project to manifold OUTSIDE compile — cv_loss uses LA.det\n","            emb = F.normalize(self['manifold_proj'](cls_feat), dim=-1)\n","            l_cv = cv_loss(emb, target=cfg['cv_target'])\n","            losses['cv'] = cfg['w_cv'] * l_cv\n","            losses['cv_raw'] = l_cv.detach()\n","        else:\n","            losses['cv'] = torch.tensor(0.0, device=device)\n","            losses['cv_raw'] = torch.tensor(0.0)\n","\n","        l_spread = spread_loss(self.constellation.anchors)\n","        losses['spread'] = cfg['w_spread'] * l_spread\n","        losses['geo_total'] = losses['cv'] + losses['spread']\n","        return losses\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# DIAGNOSTICS\n","# ═══════════════════════════════════════════════════════════════════\n","\n","DIAG_KEYS = ['ce', 'acc', 'nce_emb', 'nce_pw', 'bridge',\n","             'assign', 'assign_nce', 'nce_tri', 'attract']\n","\n","\n","def diag_to_dict(t):\n","    return {k: t[i].item() for i, k in enumerate(DIAG_KEYS)}\n","\n","\n","@torch.no_grad()\n","def compute_geometric_stats(model):\n","    \"\"\"Full geometric health check.\"\"\"\n","    from geolip_core.core.distinguish.losses import cv_metric\n","    stats = {}\n","    anchors = model.constellation.anchors\n","    anchors_n = F.normalize(anchors, dim=-1)\n","    A = anchors_n.shape[0]\n","    cos = anchors_n @ anchors_n.T\n","    idx = torch.triu_indices(A, A, offset=1, device=cos.device)\n","    pairwise = 1.0 - cos[idx[0], idx[1]]\n","    stats['anchor_cv'] = (pairwise.std() / (pairwise.mean() + 1e-8)).item()\n","    stats['anchor_mean_dist'] = pairwise.mean().item()\n","    stats['anchor_min_dist'] = pairwise.min().item()\n","    stats['cm_pos'] = (model['cm_gate']._cached_cm_norm > 0).float().mean().item()\n","    # Utilization from push\n","    if model.push.strategy == 'momentum' and model.push.accumulator is not None:\n","        stats['push_momentum'] = model.push.accumulator.norm(dim=-1).mean().item()\n","    return stats\n","\n","\n","@torch.no_grad()\n","def compute_gradient_norms(model):\n","    \"\"\"Gradient norms by component type.\"\"\"\n","    norms = defaultdict(list)\n","    for name, p in model.named_parameters():\n","        if p.grad is None:\n","            continue\n","        gn = p.grad.norm().item()\n","        if 'constellation' in name or 'anchor' in name:\n","            norms['constellation'].append(gn)\n","        elif 'cm_gate' in name:\n","            norms['cm_gate'].append(gn)\n","        elif 'content_attn' in name:\n","            norms['content_attn'].append(gn)\n","        elif 'geo_attn' in name:\n","            norms['geo_attn'].append(gn)\n","        elif 'patchwork' in name or 'bridge' in name:\n","            norms['patchwork'].append(gn)\n","        elif 'film' in name or 'ctx' in name:\n","            norms['context'].append(gn)\n","        elif 'head' in name:\n","            norms['head'].append(gn)\n","        elif 'conv' in name or 'input' in name:\n","            norms['input'].append(gn)\n","        else:\n","            norms['other'].append(gn)\n","    return {k: {'mean': np.mean(v), 'max': np.max(v)} for k, v in norms.items()}\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# DRY RUN — noise data, escalating classes\n","# ═══════════════════════════════════════════════════════════════════\n","\n","class PairedTransform:\n","    def __init__(self, t): self.t = t\n","    def __call__(self, img): return self.t(img), self.t(img)\n","\n","\n","def make_noise_data(n_samples, n_classes, img_size=32, channels=3):\n","    \"\"\"Gaussian noise dataset with random labels.\"\"\"\n","    X = torch.randn(n_samples, channels, img_size, img_size)\n","    Y = torch.randint(0, n_classes, (n_samples,))\n","    return torch.utils.data.TensorDataset(X, Y)\n","\n","\n","def make_paired_noise_loader(n_samples, n_classes, batch_size):\n","    \"\"\"Noise data with paired augmentation — shared base, small perturbation.\"\"\"\n","    class PairedNoiseDataset(torch.utils.data.Dataset):\n","        def __init__(self, n, nc):\n","            self.n, self.nc = n, nc\n","            self.labels = torch.randint(0, nc, (n,))\n","            self.base_images = torch.randn(n, 3, 32, 32)\n","        def __len__(self): return self.n\n","        def __getitem__(self, i):\n","            base = self.base_images[i]\n","            v1 = base + 0.1 * torch.randn_like(base)\n","            v2 = base + 0.1 * torch.randn_like(base)\n","            return (v1, v2), self.labels[i]\n","\n","    ds = PairedNoiseDataset(n_samples, n_classes)\n","    return torch.utils.data.DataLoader(\n","        ds, batch_size=batch_size, shuffle=True, num_workers=0, drop_last=True)\n","\n","\n","def dry_run(n_classes, n_epochs, config=None):\n","    \"\"\"Dry run with noise data. Returns training history.\"\"\"\n","    if config is None:\n","        config = DEFAULT_CONFIG.copy()\n","    config['num_classes'] = n_classes\n","\n","    model = GeoTransformerRedux('redux_test', config)\n","    model.network_to(device=device, strict=False)\n","    n_params = sum(p.numel() for p in model.parameters())\n","\n","    print(f\"\\n{'='*60}\")\n","    print(f\"  DRY RUN: {n_classes} classes, {n_epochs} epochs\")\n","    print(f\"  Params: {n_params:,} | d={config['d_model']} layers={config['n_layers']}\")\n","    print(f\"  Shared constellation: {config['n_anchors']} anchors on S^{config['manifold_dim']-1}\")\n","    print(f\"{'='*60}\")\n","\n","    compiled = torch.compile(model, mode='default')\n","    optimizer = torch.optim.Adam(model.parameters(), lr=config['lr'])\n","    loader = make_paired_noise_loader(2048, n_classes, config['batch_size'])\n","\n","    history = defaultdict(list)\n","    precompute_freq = config['precompute_freq']\n","\n","    for epoch in range(n_epochs):\n","        model.train()\n","        epoch_loss = epoch_acc = n = 0\n","        last_diag = None\n","        last_emb = None\n","        t0 = time.time()\n","\n","        for batch_idx, ((v1, v2), labels) in enumerate(loader):\n","            v1, v2, labels = v1.to(device), v2.to(device), labels.to(device)\n","\n","            if batch_idx % precompute_freq == 0:\n","                model.invalidate()\n","                model.precompute()\n","\n","            loss, diag, cls_feat = compiled(v1, v2, targets=labels)\n","\n","            # Geo losses BEFORE backward — combine into single .backward()\n","            geo = model.compute_geo_losses(cls_feat)\n","            total_loss_val = loss + geo['geo_total']\n","\n","            optimizer.zero_grad()\n","            total_loss_val.backward()\n","\n","            # Dynamic clip\n","            clip_val = max(diag[0].item(), 1.0)\n","            torch.nn.utils.clip_grad_norm_(model.parameters(), clip_val)\n","            optimizer.step()\n","\n","            # Accumulate for push (project to manifold, detach)\n","            with torch.no_grad():\n","                emb_for_push = F.normalize(model['manifold_proj'](cls_feat.detach()), dim=-1)\n","            buf_full = model.accumulate_for_push(emb_for_push, labels)\n","            if buf_full:\n","                batch_push = model.do_push()\n","                if batch_push.get('moved', 0) > 0:\n","                    model.invalidate()\n","                    model.precompute()\n","\n","            epoch_loss += loss.item() * v1.size(0)\n","            epoch_acc += diag[1].item() * v1.size(0)\n","            n += v1.size(0)\n","            last_diag = diag\n","            last_emb = emb_for_push\n","\n","        # End-of-epoch push for remaining buffer\n","        push_diag = model.do_push()\n","\n","        elapsed = time.time() - t0\n","        d = diag_to_dict(last_diag)\n","        geo_stats = compute_geometric_stats(model)\n","        grad_norms = compute_gradient_norms(model)\n","        # last_emb is already manifold-projected — compute CV directly\n","        with torch.no_grad():\n","            from geolip_core.core.distinguish.losses import cv_metric\n","            cv_emb_val = cv_metric(last_emb) if last_emb is not None and last_emb.shape[0] >= 5 else 0.0\n","\n","        # Record\n","        history['epoch'].append(epoch)\n","        history['loss'].append(epoch_loss / n)\n","        history['acc'].append(epoch_acc / n)\n","        history['time'].append(elapsed)\n","        for k, v in d.items():\n","            history[f'loss_{k}'].append(v)\n","        for k, v in geo_stats.items():\n","            history[f'geo_{k}'].append(v)\n","        history['cv_emb'].append(cv_emb_val)\n","        if push_diag:\n","            for k, v in push_diag.items():\n","                if isinstance(v, (int, float)):\n","                    history[f'push_{k}'].append(v)\n","\n","        # Print\n","        cfg = config\n","        w_ce = cfg['w_ce'] * d['ce']\n","        w_nce = cfg['w_nce_emb'] * d['nce_emb']\n","        w_pw = cfg['w_nce_pw'] * d['nce_pw']\n","        w_brg = cfg['w_bridge'] * d['bridge']\n","        w_asgn = cfg['w_assign'] * d['assign']\n","        w_anc = cfg['w_assign_nce'] * d['assign_nce']\n","        w_tri = cfg['w_nce_tri'] * d['nce_tri']\n","        w_att = cfg['w_attract'] * d['attract']\n","        w_sum = w_ce + w_nce + w_pw + w_brg + w_asgn + w_anc + w_tri + w_att\n","\n","        push_str = \"\"\n","        if push_diag and push_diag.get('moved', 0) > 0:\n","            push_str = (f\" PUSH: drift={push_diag.get('drift_mean',0):.4f}\"\n","                        f\" active={push_diag.get('n_active',0)}\"\n","                        f\" dead={push_diag.get('dead_count',0)}\")\n","\n","        print(f\"  E{epoch:>3d}  L={epoch_loss/n:.3f} acc={epoch_acc/n:.4f} \"\n","              f\"CV_emb={cv_emb_val:.3f} CV_anc={geo_stats['anchor_cv']:.3f} \"\n","              f\"cm+={geo_stats['cm_pos']:.2f} {elapsed:.1f}s{push_str}\")\n","\n","        if epoch % 5 == 0:\n","            print(f\"    wt: ce={w_ce:.2f} nce={w_nce:.2f} pw={w_pw:.2f} \"\n","                  f\"brg={w_brg:.2f} asgn={w_asgn:.2f} a_nce={w_anc:.2f} \"\n","                  f\"tri={w_tri:.2f} att={w_att:.2f} → {w_sum:.2f}\")\n","            gn = grad_norms\n","            gn_str = ' '.join(f\"{k}={v['mean']:.3f}\" for k, v in sorted(gn.items()))\n","            print(f\"    grad: {gn_str}\")\n","\n","    return history\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# MAIN — escalating dry runs\n","# ═══════════════════════════════════════════════════════════════════\n","\n","def get_cifar100_loaders(config):\n","    import torchvision\n","    import torchvision.transforms as T\n","\n","    train_aug = T.Compose([\n","        T.RandomCrop(32, padding=4), T.RandomHorizontalFlip(),\n","        T.TrivialAugmentWide(), T.ToTensor(),\n","        T.Normalize((0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761)),\n","        T.RandomErasing(p=0.25),\n","    ])\n","    test_t = T.Compose([\n","        T.ToTensor(),\n","        T.Normalize((0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761)),\n","    ])\n","\n","    train_ds = torchvision.datasets.CIFAR100(\n","        './data', train=True, download=True, transform=PairedTransform(train_aug))\n","    test_ds = torchvision.datasets.CIFAR100(\n","        './data', train=False, download=True, transform=test_t)\n","\n","    train_loader = torch.utils.data.DataLoader(\n","        train_ds, batch_size=config['batch_size'], shuffle=True,\n","        num_workers=config['num_workers'], pin_memory=True, drop_last=True)\n","    test_loader = torch.utils.data.DataLoader(\n","        test_ds, batch_size=config['batch_size'] * 2, shuffle=False,\n","        num_workers=config['num_workers'], pin_memory=True)\n","    return train_loader, test_loader\n","\n","\n","@torch.no_grad()\n","def evaluate(model, model_raw, loader):\n","    model.eval()\n","    model_raw.invalidate()\n","    model_raw.precompute()\n","    correct = total = 0\n","    for images, labels in loader:\n","        images, labels = images.to(device), labels.to(device)\n","        logits = model(images)\n","        correct += (logits.argmax(1) == labels).sum().item()\n","        total += images.size(0)\n","    return correct / total\n","\n","\n","def train_cifar(config=None):\n","    \"\"\"Full CIFAR-100 training with geometric tracking.\"\"\"\n","    if config is None:\n","        config = DEFAULT_CONFIG.copy()\n","    #config['num_classes'] = 100\n","    #config['epochs'] = 200\n","    #config['batch_size'] = 128\n","\n","    print(f\"\\n{'='*60}\")\n","    print(f\"  GeoTransformer Redux — CIFAR-100\")\n","    print(f\"  d={config['d_model']}, layers={config['n_layers']}, \"\n","          f\"anchors={config['n_anchors']}, manifold={config['manifold_dim']}\")\n","    print(f\"  Shared constellation, momentum push, pentachoron CV\")\n","    print(f\"{'='*60}\")\n","\n","    train_loader, test_loader = get_cifar100_loaders(config)\n","    print(f\"  Train: {len(train_loader.dataset):,} | Test: {len(test_loader.dataset):,}\")\n","\n","    model = GeoTransformerRedux('redux_cifar', config)\n","    model.network_to(device=device, strict=False)\n","    n_params = sum(p.numel() for p in model.parameters())\n","    print(f\"  Params: {n_params:,}\")\n","\n","    model_raw = model\n","    compiled = torch.compile(model, mode='default')\n","    print(f\"  Compiled (default)\")\n","\n","    optimizer = torch.optim.Adam(model.parameters(), lr=config['lr'])\n","    total_steps = config['epochs'] * len(train_loader)\n","    warmup_steps = config['warmup_epochs'] * len(train_loader)\n","\n","    def lr_lambda(step):\n","        if step < warmup_steps:\n","            return step / max(1, warmup_steps)\n","        progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)\n","        return 0.5 * (1 + math.cos(math.pi * progress))\n","\n","    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)\n","    precompute_freq = config['precompute_freq']\n","\n","    print(f\"  Training {config['epochs']} epochs\")\n","    print(f\"{'━'*60}\\n\")\n","\n","    best_acc = 0\n","    save_dir = Path('geo_redux'); save_dir.mkdir(exist_ok=True)\n","    pbar = tqdm(range(config['epochs']), desc=\"Training\", unit=\"ep\")\n","\n","    for epoch in pbar:\n","        model.train()\n","        total_loss = correct = total = 0\n","        last_diag = None\n","        last_emb = None\n","        push_diag = {}\n","        t0 = time.time()\n","\n","        for batch_idx, ((v1, v2), labels) in enumerate(train_loader):\n","            v1, v2, labels = v1.to(device), v2.to(device), labels.to(device)\n","\n","            if batch_idx % precompute_freq == 0:\n","                model_raw.invalidate()\n","                model_raw.precompute()\n","\n","            loss, diag, cls_feat = compiled(v1, v2, targets=labels)\n","\n","            # Detailed NaN diagnostics\n","            if torch.isnan(loss) or torch.isnan(diag).any():\n","                d = {k: diag[i].item() for i, k in enumerate(DIAG_KEYS)}\n","                nan_losses = [k for k, v in d.items() if math.isnan(v)]\n","                # Check parameters for NaN\n","                nan_params = []\n","                for pname, p in model_raw.named_parameters():\n","                    if torch.isnan(p).any():\n","                        nan_params.append(pname)\n","                # Check anchors\n","                anc = model_raw.constellation.anchors\n","                anc_nan = torch.isnan(anc).any().item()\n","                anc_norm = anc.norm(dim=-1)\n","                # Check CM gate buffer\n","                cm_nan = torch.isnan(model_raw['cm_gate']._cached_cm_norm).any().item()\n","                tqdm.write(\n","                    f\"\\n  ═══ NaN at E{epoch}B{batch_idx} ═══\"\n","                    f\"\\n    loss={loss.item():.6f}\"\n","                    f\"\\n    diag: {' '.join(f'{k}={v:.4f}' for k,v in d.items())}\"\n","                    f\"\\n    NaN losses: {nan_losses}\"\n","                    f\"\\n    NaN params ({len(nan_params)}): {nan_params[:5]}\"\n","                    f\"\\n    anchors: nan={anc_nan} norm_min={anc_norm.min():.4f} norm_max={anc_norm.max():.4f}\"\n","                    f\"\\n    CM buffer nan={cm_nan}\"\n","                    f\"\\n    last push drift={push_diag.get('drift_mean', 'n/a') if push_diag else 'none'}\"\n","                )\n","                return best_acc\n","\n","            # Spread + CV on embeddings\n","            geo = model_raw.compute_geo_losses(cls_feat)\n","            geo_total = geo['geo_total']\n","\n","            # Check geo for NaN\n","            if torch.isnan(geo_total):\n","                tqdm.write(\n","                    f\"\\n  ═══ NaN in geo_losses at E{epoch}B{batch_idx} ═══\"\n","                    f\"\\n    cv={geo.get('cv', 'n/a')}, spread={geo.get('spread', 'n/a')}\"\n","                    f\"\\n    Skipping geo this batch\"\n","                )\n","                total_loss_val = loss\n","            else:\n","                total_loss_val = loss + geo_total\n","\n","            optimizer.zero_grad()\n","            total_loss_val.backward()\n","\n","            # Post-backward NaN check\n","            grad_nan = False\n","            for pname, p in model_raw.named_parameters():\n","                if p.grad is not None and torch.isnan(p.grad).any():\n","                    grad_nan = True\n","                    tqdm.write(f\"  NaN grad in {pname} at E{epoch}B{batch_idx}\")\n","                    break\n","            if grad_nan:\n","                optimizer.zero_grad()\n","                model_raw.invalidate()\n","                model_raw.precompute()\n","                continue\n","\n","            #clip_val = min(max(diag[0].item(), 1.0), 2.0)\n","            torch.nn.utils.clip_grad_norm_(model.parameters(), 1)\n","            optimizer.step()\n","            scheduler.step()\n","\n","            # Push accumulation\n","            with torch.no_grad():\n","                emb_for_push = F.normalize(\n","                    model_raw['manifold_proj'](cls_feat.detach()), dim=-1, eps=1e-6)\n","            buf_full = model_raw.accumulate_for_push(emb_for_push, labels)\n","            if buf_full:\n","                push_diag = model_raw.do_push()\n","                if push_diag.get('moved', 0) > 0:\n","                    model_raw.invalidate()\n","                    model_raw.precompute()\n","\n","            total_loss += loss.item() * v1.size(0)\n","            correct += diag[1].item() * v1.size(0)\n","            total += v1.size(0)\n","            last_diag = diag\n","            last_emb = emb_for_push\n","\n","        # End-of-epoch push for remaining buffer\n","        push_diag = model_raw.do_push()\n","\n","        train_loss = total_loss / total\n","        train_acc = correct / total\n","        elapsed = time.time() - t0\n","\n","        test_acc = evaluate(compiled, model_raw, test_loader)\n","        if test_acc > best_acc:\n","            best_acc = test_acc\n","            torch.save({\n","                'state_dict': {k: v.cpu() for k, v in model.state_dict().items()},\n","                'epoch': epoch, 'test_acc': test_acc, 'config': config,\n","            }, save_dir / 'best.pt')\n","\n","        pbar.set_postfix_str(\n","            f\"L={train_loss:.2f} tr={train_acc:.3f} te={test_acc:.3f} \"\n","            f\"best={best_acc:.3f} {elapsed:.1f}s\")\n","\n","        if epoch % 5 == 0 or epoch == config['epochs'] - 1:\n","            d = diag_to_dict(last_diag)\n","            geo_stats = compute_geometric_stats(model_raw)\n","            with torch.no_grad():\n","                from geolip_core.core.distinguish.losses import cv_metric\n","                cv_emb = cv_metric(last_emb) if last_emb is not None else 0.0\n","\n","            cfg = config\n","            w_ce = cfg['w_ce'] * d['ce']\n","            w_nce = cfg['w_nce_emb'] * d['nce_emb']\n","            w_pw = cfg['w_nce_pw'] * d['nce_pw']\n","            w_brg = cfg['w_bridge'] * d['bridge']\n","\n","            push_str = \"\"\n","            if push_diag and push_diag.get('moved', 0) > 0:\n","                push_str = (f\" PUSH: drift={push_diag.get('drift_mean',0):.4f}\"\n","                            f\" active={push_diag.get('n_active',0)}\")\n","\n","            tqdm.write(\n","                f\"  E{epoch:>3d}  L={train_loss:.3f} \"\n","                f\"train={train_acc:.4f} test={test_acc:.4f} best={best_acc:.4f} \"\n","                f\"CV_e={cv_emb:.3f} CV_a={geo_stats['anchor_cv']:.3f} \"\n","                f\"cm+={geo_stats['cm_pos']:.2f} {elapsed:.1f}s{push_str}\"\n","                f\"\\n    wt: ce={w_ce:.2f} nce={w_nce:.2f} pw={w_pw:.2f} \"\n","                f\"brg={w_brg:.2f} \"\n","                f\"sum={w_ce+w_nce+w_pw+w_brg+cfg['w_assign']*d['assign']+cfg['w_assign_nce']*d['assign_nce']+cfg['w_nce_tri']*d['nce_tri']+cfg['w_attract']*d['attract']:.2f}\")\n","\n","    pbar.close()\n","    print(f\"\\n{'═'*60}\")\n","    print(f\"  CIFAR-100: best={best_acc:.4f} ({best_acc*100:.2f}%) | {n_params:,} params\")\n","    print(f\"{'═'*60}\")\n","    return best_acc\n","\n","\n","if __name__ == '__main__':\n","    print(f\"Device: {device}\")\n","    if device.type == 'cuda':\n","        print(f\"  GPU: {torch.cuda.get_device_name()}\")\n","        print(f\"  VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB\")\n","\n","    # Detect notebook vs CLI\n","    try:\n","        import argparse\n","        parser = argparse.ArgumentParser()\n","        parser.add_argument('--mode', default='cifar', choices=['cifar', 'dry'])\n","        args = parser.parse_args()\n","        mode = args.mode\n","    except SystemExit:\n","        mode = 'cifar'  # default in notebooks\n","\n","    if mode == 'dry':\n","        for n_classes, n_epochs in [(10, 30), (50, 20), (100, 20)]:\n","            history = dry_run(n_classes, n_epochs)\n","            print(f\"\\n  Summary {n_classes} classes:\"\n","                  f\" final_acc={history['acc'][-1]:.4f}\"\n","                  f\" CV_emb={history['cv_emb'][-1]:.3f}\"\n","                  f\" CV_anc={history['geo_anchor_cv'][-1]:.3f}\\n\")\n","    else:\n","        train_cifar()"],"metadata":{"colab":{"base_uri":"https://localhost:8080/","height":662,"referenced_widgets":["8638148368d34ba78e9451e5d38ea54f","3fe3ced2159249f2abdd190983a48dd6","1e30a8e7935241f9a58dede6e148656b","a0153bdb2cde41d485969562898cd3d2","5d7a4a42979c4622ad3c62594e48c786","f629365478a84dca84b8f64287eba2dd","54d2555ec2f94980bff8fd90a6d705fb","d67d42e594774904b7cd4367478da24c","02e36a82c7224d1ab00b1f54f5c78900","614125e4c7174d0983d0106216228730","9ea9af4b1a3c4255800e4260fde64577"]},"id":"MjY3DkkiSwNq","executionInfo":{"status":"error","timestamp":1775263597434,"user_tz":420,"elapsed":13213,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"8916f69a-0c78-42e0-c352-81f34e6a2d62"},"execution_count":2,"outputs":[{"output_type":"stream","name":"stdout","text":["Device: cuda\n","  GPU: NVIDIA RTX PRO 6000 Blackwell Server Edition\n","  VRAM: 102.0 GB\n","\n","============================================================\n","  GeoTransformer Redux — CIFAR-100\n","  d=256, layers=4, anchors=256, manifold=128\n","  Shared constellation, momentum push, pentachoron CV\n","============================================================\n"]},{"output_type":"stream","name":"stderr","text":["usage: colab_kernel_launcher.py [-h] [--mode {cifar,dry}]\n","colab_kernel_launcher.py: error: unrecognized arguments: -f /root/.local/share/jupyter/runtime/kernel-75cd7e9d-8e97-4606-91b3-7c80cc793bcc.json\n"]},{"output_type":"stream","name":"stdout","text":["  Train: 50,000 | Test: 10,000\n","  Params: 15,409,829\n","  Compiled (default)\n","  Training 200 epochs\n","━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\n","\n"]},{"output_type":"display_data","data":{"text/plain":["Training:   0%|          | 0/200 [00:00<?, ?ep/s]"],"application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"8638148368d34ba78e9451e5d38ea54f"}},"metadata":{}},{"output_type":"stream","name":"stdout","text":["  NaN grad in components.input.pos_embed at E0B21\n"]},{"output_type":"error","ename":"KeyboardInterrupt","evalue":"","traceback":["\u001b[0;31m---------------------------------------------------------------------------\u001b[0m","\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)","\u001b[0;32m/tmp/ipykernel_10985/2567062161.py\u001b[0m in \u001b[0;36m<cell line: 0>\u001b[0;34m()\u001b[0m\n\u001b[1;32m   1041\u001b[0m                   f\" CV_anc={history['geo_anchor_cv'][-1]:.3f}\\n\")\n\u001b[1;32m   1042\u001b[0m     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1043\u001b[0;31m         \u001b[0mtrain_cifar\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m","\u001b[0;32m/tmp/ipykernel_10985/2567062161.py\u001b[0m in \u001b[0;36mtrain_cifar\u001b[0;34m(config)\u001b[0m\n\u001b[1;32m    932\u001b[0m             \u001b[0mgrad_nan\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mFalse\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    933\u001b[0m             \u001b[0;32mfor\u001b[0m \u001b[0mpname\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mp\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mmodel_raw\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnamed_parameters\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 934\u001b[0;31m                 \u001b[0;32mif\u001b[0m \u001b[0mp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgrad\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0misnan\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgrad\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0many\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    935\u001b[0m                     \u001b[0mgrad_nan\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mTrue\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    936\u001b[0m                     \u001b[0mtqdm\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mwrite\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf\"  NaN grad in {pname} at E{epoch}B{batch_idx}\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;31mKeyboardInterrupt\u001b[0m: "]}]},{"cell_type":"markdown","source":["# geo cifar100"],"metadata":{"id":"dZ4m4MNuwYU6"}},{"cell_type":"code","source":["\"\"\"\n","GeoLIP Fast — Multi-scale SVD + Constellation MLP for CIFAR-100\n","\n","No transformers, no routers, no towers. Pure nn.Module.\n","Uses geolip_core utilities: SVDObserver, constellation losses, CM functions.\n","\n","Multi-scale input:\n","  32×32×3  → conv → SVD → spectral features  (global composition)\n","  16×16×4  → conv → SVD → spectral features  (mid-level structure)\n","  8×8×16   → conv → SVD → spectral features  (local texture)\n","\n","Pipeline:\n","  Multi-scale conv → SVD per scale → concat features → MLP → manifold\n","  → triangulate against constellation → patchwork → predict\n","\n","!pip install geolip-core torchvision tqdm\n","\"\"\"\n","\n","import os, warnings\n","warnings.filterwarnings('ignore')\n","\n","import torch\n","import torch.nn as nn\n","import torch.nn.functional as F\n","import numpy as np\n","import time, math\n","from pathlib import Path\n","from tqdm.auto import tqdm\n","\n","device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n","print(f\"Device: {device}\")\n","if device.type == 'cuda':\n","    print(f\"  GPU: {torch.cuda.get_device_name()}\")\n","    print(f\"  VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB\")\n","\n","torch.set_float32_matmul_precision('high')\n","torch.backends.cudnn.benchmark = True\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# IMPORTS from geolip_core — utilities only, no pipeline classes\n","# ═══════════════════════════════════════════════════════════════════\n","\n","from geolip_core.core.input.svd import SVDObserver\n","from geolip_core.core.distinguish.losses import (\n","    nce_loss, bridge_loss_paired, assign_bce_loss,\n","    assign_nce_loss, attraction_loss, ce_loss_paired,\n",")\n","from geolip_core.pipeline.components.geometric_transformer import (\n","    anchor_neighborhood_cm, CMValidatedGate,\n",")\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# CONFIG\n","# ═══════════════════════════════════════════════════════════════════\n","\n","CONFIG = {\n","    # Architecture\n","    'd_model': 256,\n","    'manifold_dim': 128,\n","    'n_anchors': 128,\n","    'svd_rank': 12,\n","    'conv_channels': 64,\n","    'mlp_hidden': 512,\n","    'cm_neighbors': 3,\n","\n","    # Training\n","    'epochs': 200,\n","    'batch_size': 256,\n","    'lr': 3e-4,\n","    'warmup_epochs': 5,\n","    'num_workers': 8,\n","\n","    # Loss weights\n","    'w_ce': 1.0,\n","    'w_nce_emb': 0.5,\n","    'w_nce_pw': 0.75,\n","    'w_bridge': 0.75,\n","    'w_assign': 0.5,\n","    'w_assign_nce': 0.25,\n","    'w_nce_tri': 0.5,\n","    'w_attract': 0.25,\n","    'w_cv': 0.05,\n","    'w_spread': 0.05,\n","    'cv_target': 0.21,\n","\n","    'num_classes': 100,\n","}\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# MULTI-SCALE SVD ENCODER\n","# ═══════════════════════════════════════════════════════════════════\n","\n","class MultiScaleSVDEncoder(nn.Module):\n","    \"\"\"Multi-scale conv → SVD spectral features.\n","\n","    Scale 0: 32×32 whole image → global spectral signature\n","    Scale 1: 16×16 patches (2×2 grid) → mid-level structure\n","    Scale 2: 8×8 patches (4×4 grid) → local texture\n","\n","    Each scale: conv block → SVDObserver → spectral features\n","    Output: concatenated spectral features from all scales\n","    \"\"\"\n","    def __init__(self, in_channels=3, conv_channels=64, svd_rank=12):\n","        super().__init__()\n","        self.svd_rank = svd_rank\n","\n","        # Shared conv stem\n","        self.stem = nn.Sequential(\n","            nn.Conv2d(in_channels, conv_channels, 3, padding=1, bias=False),\n","            nn.BatchNorm2d(conv_channels), nn.GELU(),\n","            nn.Conv2d(conv_channels, conv_channels, 3, padding=1, bias=False),\n","            nn.BatchNorm2d(conv_channels), nn.GELU(),\n","        )\n","\n","        # Per-scale refinement + SVD\n","        # Scale 0: 32×32 full — 1 token\n","        self.scale0_conv = nn.Sequential(\n","            nn.Conv2d(conv_channels, conv_channels, 3, padding=1, bias=False),\n","            nn.BatchNorm2d(conv_channels), nn.GELU())\n","        self.scale0_svd = SVDObserver(conv_channels, svd_rank)\n","\n","        # Scale 1: 16×16 patches — 4 tokens (2×2 unfold)\n","        self.scale1_conv = nn.Sequential(\n","            nn.Conv2d(conv_channels, conv_channels, 3, padding=1, bias=False),\n","            nn.BatchNorm2d(conv_channels), nn.GELU())\n","        self.scale1_svd = SVDObserver(conv_channels, svd_rank)\n","\n","        # Scale 2: 8×8 patches — 16 tokens (4×4 unfold)\n","        self.scale2_conv = nn.Sequential(\n","            nn.Conv2d(conv_channels, conv_channels, 3, padding=1, bias=False),\n","            nn.BatchNorm2d(conv_channels), nn.GELU())\n","        self.scale2_svd = SVDObserver(conv_channels, svd_rank)\n","\n","        # Feature dim per SVD: 2*rank + 2\n","        self.feat_per_svd = 2 * svd_rank + 2\n","        # Total: 1×feat + 4×feat + 16×feat = 21×feat\n","        self.n_tokens = 21\n","        self.output_dim = self.n_tokens * self.feat_per_svd\n","\n","    def _patch_svd(self, feat, patch_h, patch_w, conv, svd):\n","        \"\"\"Unfold into patches, run conv+SVD on each.\"\"\"\n","        B, C, H, W = feat.shape\n","        ph, pw = H // patch_h, W // patch_w\n","        # Unfold into (B*n_patches, C, ph, pw)\n","        patches = feat.unfold(2, ph, ph).unfold(3, pw, pw)\n","        patches = patches.contiguous().reshape(B * patch_h * patch_w, C, ph, pw)\n","        patches = conv(patches)\n","        S, Vh, features, novelty = svd(patches)\n","        # (B*n_patches, feat_dim) → (B, n_patches, feat_dim)\n","        features = features.reshape(B, patch_h * patch_w, -1)\n","        S = S.reshape(B, patch_h * patch_w, -1)\n","        return features, S\n","\n","    def forward(self, x):\n","        \"\"\"(B, 3, 32, 32) → (B, output_dim), all_S\"\"\"\n","        B = x.shape[0]\n","        feat = self.stem(x)  # (B, C, 32, 32)\n","\n","        # Scale 0: full 32×32\n","        s0_feat = self.scale0_conv(feat)\n","        S0, Vh0, svd_feat0, nov0 = self.scale0_svd(s0_feat)\n","        # (B, feat_per_svd) → (B, 1, feat_per_svd)\n","        svd_feat0 = svd_feat0.unsqueeze(1)\n","\n","        # Scale 1: 2×2 grid of 16×16 patches → (B, 4, feat_per_svd)\n","        svd_feat1, S1 = self._patch_svd(feat, 2, 2, self.scale1_conv, self.scale1_svd)\n","\n","        # Scale 2: 4×4 grid of 8×8 patches → (B, 16, feat_per_svd)\n","        svd_feat2, S2 = self._patch_svd(feat, 4, 4, self.scale2_conv, self.scale2_svd)\n","\n","        # Concat all scales: (B, 21, feat_per_svd)\n","        all_tokens = torch.cat([svd_feat0, svd_feat1, svd_feat2], dim=1)\n","        all_flat = all_tokens.reshape(B, -1)  # (B, 21 * feat_per_svd)\n","\n","        return all_flat, S0\n","\n","    @torch.no_grad()\n","    def update_ema(self, x):\n","        \"\"\"Run forward again to update SVD EMA. Call after backward.\"\"\"\n","        feat = self.stem(x)\n","        s0 = self.scale0_conv(feat)\n","        S0, Vh0, _, _ = self.scale0_svd(s0)\n","        self.scale0_svd.update_ema(S0, Vh0)\n","\n","        # Scale 1\n","        B, C, H, W = feat.shape\n","        p1 = feat.unfold(2, 16, 16).unfold(3, 16, 16)\n","        p1 = p1.contiguous().reshape(B * 4, C, 16, 16)\n","        p1 = self.scale1_conv(p1)\n","        S1, Vh1, _, _ = self.scale1_svd(p1)\n","        self.scale1_svd.update_ema(S1, Vh1)\n","\n","        # Scale 2\n","        p2 = feat.unfold(2, 8, 8).unfold(3, 8, 8)\n","        p2 = p2.contiguous().reshape(B * 16, C, 8, 8)\n","        p2 = self.scale2_conv(p2)\n","        S2, Vh2, _, _ = self.scale2_svd(p2)\n","        self.scale2_svd.update_ema(S2, Vh2)\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# CONSTELLATION + PATCHWORK (inline, no router)\n","# ═══════════════════════════════════════════════════════════════════\n","\n","class Constellation(nn.Module):\n","    \"\"\"Learnable anchor set on S^(d-1) with CM-validated gate.\"\"\"\n","    def __init__(self, n_anchors, manifold_dim, cm_neighbors=3):\n","        super().__init__()\n","        self.n_anchors = n_anchors\n","        self.manifold_dim = manifold_dim\n","        self.anchors = nn.Parameter(torch.randn(n_anchors, manifold_dim) * 0.1)\n","        self.cm_gate = CMValidatedGate(n_anchors, n_neighbors=cm_neighbors)\n","\n","    @torch.no_grad()\n","    def precompute(self):\n","        anchors_n = F.normalize(self.anchors, dim=-1)\n","        self.cm_gate.precompute(anchors_n.detach())\n","\n","    def invalidate(self):\n","        self.cm_gate.invalidate_cache()\n","\n","    def forward(self, emb):\n","        \"\"\"(B, manifold_dim) → distances, cos, assignment, gate_values\"\"\"\n","        anchors_n = F.normalize(self.anchors, dim=-1)\n","        emb_n = F.normalize(emb, dim=-1)\n","        cos = emb_n @ anchors_n.T           # (B, A)\n","        distances = 1.0 - cos\n","        assignment = F.softmax(cos / 0.1, dim=-1)\n","        gate_values, gate_info = self.cm_gate(distances)\n","        gated_distances = distances * gate_values\n","        return {\n","            'cos': cos, 'distances': distances,\n","            'gated_distances': gated_distances,\n","            'assignment': assignment, 'gate_values': gate_values,\n","            'nearest': cos.argmax(dim=-1),\n","        }\n","\n","\n","class SimplePatchwork(nn.Module):\n","    \"\"\"Lightweight patchwork: gated distances → compact features.\"\"\"\n","    def __init__(self, n_anchors, d_out):\n","        super().__init__()\n","        self.proj = nn.Sequential(\n","            nn.Linear(n_anchors, d_out), nn.LayerNorm(d_out), nn.GELU(),\n","            nn.Linear(d_out, d_out))\n","\n","    def forward(self, gated_distances):\n","        return self.proj(gated_distances)\n","\n","\n","class BridgePredictor(nn.Module):\n","    \"\"\"Predict assignment from patchwork features.\"\"\"\n","    def __init__(self, pw_dim, n_anchors):\n","        super().__init__()\n","        self.proj = nn.Linear(pw_dim, n_anchors)\n","\n","    def forward(self, pw_features):\n","        return self.proj(pw_features)\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# FULL MODEL\n","# ═══════════════════════════════════════════════════════════════════\n","\n","class GeoFastClassifier(nn.Module):\n","    \"\"\"Multi-scale SVD → MLP → constellation → predict.\n","\n","    No transformers, no attention, no routers.\n","    \"\"\"\n","    def __init__(self, config):\n","        super().__init__()\n","        self.config = config\n","        d = config['d_model']\n","        m = config['manifold_dim']\n","        A = config['n_anchors']\n","        pw_dim = A  # patchwork output = n_anchors\n","\n","        # Multi-scale SVD encoder\n","        self.encoder = MultiScaleSVDEncoder(\n","            conv_channels=config['conv_channels'],\n","            svd_rank=config['svd_rank'])\n","\n","        # MLP: SVD features → d_model\n","        self.backbone = nn.Sequential(\n","            nn.Linear(self.encoder.output_dim, d),\n","            nn.LayerNorm(d), nn.GELU(), nn.Dropout(0.1),\n","            nn.Linear(d, d),\n","            nn.LayerNorm(d), nn.GELU(), nn.Dropout(0.1),\n","        )\n","\n","        # Project to manifold\n","        self.to_manifold = nn.Sequential(\n","            nn.Linear(d, m), nn.LayerNorm(m))\n","\n","        # Constellation + patchwork + bridge\n","        self.constellation = Constellation(A, m, config['cm_neighbors'])\n","        self.patchwork = SimplePatchwork(A, pw_dim)\n","        self.bridge = BridgePredictor(pw_dim, A)\n","\n","        # Classification head from concatenated features\n","        head_in = d + pw_dim  # backbone features + patchwork\n","        self.head = nn.Sequential(\n","            nn.Linear(head_in, d), nn.GELU(), nn.Dropout(0.1),\n","            nn.Linear(d, config['num_classes']))\n","\n","    def precompute(self):\n","        self.constellation.precompute()\n","\n","    def invalidate(self):\n","        self.constellation.invalidate()\n","\n","    def _encode(self, x):\n","        \"\"\"Image → (backbone_feat, emb, constellation_out, pw, bridge)\"\"\"\n","        svd_flat, S = self.encoder(x)\n","        feat = self.backbone(svd_flat)\n","        emb = self.to_manifold(feat)\n","        c_out = self.constellation(emb)\n","        pw = self.patchwork(c_out['gated_distances'])\n","        brg = self.bridge(pw)\n","        return feat, emb, c_out, pw, brg\n","\n","    def forward(self, x, x2=None, targets=None):\n","        if x2 is not None and targets is not None:\n","            return self._forward_paired(x, x2, targets)\n","\n","        feat, emb, c_out, pw, brg = self._encode(x)\n","        combined = torch.cat([feat, pw], dim=-1)\n","        logits = self.head(combined)\n","        return logits\n","\n","    def _forward_paired(self, x1, x2, targets):\n","        \"\"\"Paired forward for observer loss training.\"\"\"\n","        feat1, emb1, c1, pw1, brg1 = self._encode(x1)\n","        feat2, emb2, c2, pw2, brg2 = self._encode(x2)\n","\n","        # CE\n","        combined1 = torch.cat([feat1, pw1], dim=-1)\n","        combined2 = torch.cat([feat2, pw2], dim=-1)\n","        logits1 = self.head(combined1)\n","        logits2 = self.head(combined2)\n","        l_ce, acc = ce_loss_paired(logits1, logits2, targets)\n","\n","        # Embedding NCE\n","        l_nce_emb, _ = nce_loss(emb1, emb2, 0.07, normalize=True)\n","\n","        # Patchwork NCE\n","        l_nce_pw, _ = nce_loss(pw1, pw2, 0.1, normalize=True)\n","\n","        # Bridge\n","        l_bridge, _ = bridge_loss_paired(brg1, brg2, c1['assignment'], c2['assignment'])\n","\n","        # Assign BCE\n","        l_assign, _ = assign_bce_loss(c1['assignment'], c1['cos'])\n","\n","        # Assign NCE\n","        l_assign_nce, _ = assign_nce_loss(c1['assignment'], c2['assignment'], 0.1)\n","\n","        # Tri NCE\n","        l_nce_tri, _ = nce_loss(c1['distances'], c2['distances'], 0.1, normalize=True)\n","\n","        # Attraction\n","        l_attract, _ = attraction_loss(c1['cos'])\n","\n","        cfg = self.config\n","        loss = (cfg['w_ce'] * l_ce\n","                + cfg['w_nce_emb'] * l_nce_emb\n","                + cfg['w_nce_pw'] * l_nce_pw\n","                + cfg['w_bridge'] * l_bridge\n","                + cfg['w_assign'] * l_assign\n","                + cfg['w_assign_nce'] * l_assign_nce\n","                + cfg['w_nce_tri'] * l_nce_tri\n","                + cfg['w_attract'] * l_attract)\n","\n","        # Single tensor crossing compile boundary — no dict overhead\n","        # [ce, acc, nce_emb, nce_pw, bridge, assign, assign_nce, nce_tri, attract]\n","        diagnostics = torch.stack([\n","            l_ce.detach(), acc.detach(),\n","            l_nce_emb.detach(), l_nce_pw.detach(),\n","            l_bridge.detach(), l_assign.detach(),\n","            l_assign_nce.detach(), l_nce_tri.detach(),\n","            l_attract.detach(),\n","        ])\n","\n","        return loss, diagnostics\n","\n","    def compute_geo_losses(self):\n","        \"\"\"CV + spread on anchors. Call outside compiled graph.\"\"\"\n","        anchors = self.constellation.anchors\n","        anchors_n = F.normalize(anchors, dim=-1)\n","        A = anchors_n.shape[0]\n","        cos = anchors_n @ anchors_n.T\n","        idx = torch.triu_indices(A, A, offset=1, device=cos.device)\n","        pairwise = 1.0 - cos[idx[0], idx[1]]\n","        cv = pairwise.std() / (pairwise.mean() + 1e-8)\n","        l_cv = (cv - self.config['cv_target']).pow(2)\n","        mask = ~torch.eye(A, dtype=torch.bool, device=cos.device)\n","        l_spread = F.relu(cos[mask]).mean()\n","        return {\n","            'cv': self.config['w_cv'] * l_cv,\n","            'spread': self.config['w_spread'] * l_spread,\n","            'cv_raw': cv.detach(),\n","            'geo_total': self.config['w_cv'] * l_cv + self.config['w_spread'] * l_spread,\n","        }\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# DATA\n","# ═══════════════════════════════════════════════════════════════════\n","\n","class PairedTransform:\n","    def __init__(self, t): self.t = t\n","    def __call__(self, img): return self.t(img), self.t(img)\n","\n","\n","def get_dataloaders(config):\n","    import torchvision\n","    import torchvision.transforms as T\n","\n","    train_aug = T.Compose([\n","        T.RandomCrop(32, padding=4), T.RandomHorizontalFlip(),\n","        T.TrivialAugmentWide(), T.ToTensor(),\n","        T.Normalize((0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761)),\n","        T.RandomErasing(p=0.25),\n","    ])\n","    test_t = T.Compose([\n","        T.ToTensor(),\n","        T.Normalize((0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761)),\n","    ])\n","\n","    train_ds = torchvision.datasets.CIFAR100(\n","        './data', train=True, download=True, transform=PairedTransform(train_aug))\n","    test_ds = torchvision.datasets.CIFAR100(\n","        './data', train=False, download=True, transform=test_t)\n","\n","    train_loader = torch.utils.data.DataLoader(\n","        train_ds, batch_size=config['batch_size'], shuffle=True,\n","        num_workers=config['num_workers'], pin_memory=True, drop_last=True)\n","    test_loader = torch.utils.data.DataLoader(\n","        test_ds, batch_size=config['batch_size'] * 2, shuffle=False,\n","        num_workers=config['num_workers'], pin_memory=True)\n","    return train_loader, test_loader\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# TRAINING\n","# ═══════════════════════════════════════════════════════════════════\n","\n","@torch.no_grad()\n","def evaluate(model, loader):\n","    model.eval()\n","    model.invalidate()\n","    model.precompute()\n","    correct = total = 0\n","    for images, labels in loader:\n","        images, labels = images.to(device), labels.to(device)\n","        logits = model(images)\n","        correct += (logits.argmax(1) == labels).sum().item()\n","        total += images.size(0)\n","    return correct / total\n","\n","\n","def main():\n","    config = CONFIG.copy()\n","    print(\"=\" * 60)\n","    print(\"  GeoLIP Fast — Multi-scale SVD + Constellation MLP\")\n","    print(f\"  d={config['d_model']}, manifold={config['manifold_dim']}, \"\n","          f\"anchors={config['n_anchors']}, svd_rank={config['svd_rank']}\")\n","    print(f\"  21 tokens: 1×32² + 4×16² + 16×8²\")\n","    print(\"=\" * 60)\n","\n","    train_loader, test_loader = get_dataloaders(config)\n","    print(f\"  Train: {len(train_loader.dataset):,} | Test: {len(test_loader.dataset):,}\")\n","\n","    model = GeoFastClassifier(config).to(device)\n","    n_params = sum(p.numel() for p in model.parameters())\n","    print(f\"  Params: {n_params:,}\")\n","    print(f\"  Encoder output: {model.encoder.output_dim} \"\n","          f\"({model.encoder.n_tokens} tokens × {model.encoder.feat_per_svd} SVD features)\")\n","\n","    # Compile\n","    compiled = torch.compile(model, mode='default')\n","    print(f\"  Compiled (default)\")\n","\n","    optimizer = torch.optim.Adam(model.parameters(), lr=config['lr'])\n","    total_steps = config['epochs'] * len(train_loader)\n","    warmup_steps = config['warmup_epochs'] * len(train_loader)\n","\n","    def lr_lambda(step):\n","        if step < warmup_steps:\n","            return step / max(1, warmup_steps)\n","        progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)\n","        return 0.5 * (1 + math.cos(math.pi * progress))\n","\n","    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)\n","\n","    precompute_freq = 10  # precompute CM every N batches\n","\n","    print(f\"\\n  Training {config['epochs']} epochs, precompute every {precompute_freq} batches\")\n","    print(f\"{'━' * 60}\\n\")\n","\n","    best_acc = 0\n","    save_dir = Path('geo_fast'); save_dir.mkdir(exist_ok=True)\n","    pbar = tqdm(range(config['epochs']), desc=\"Training\", unit=\"ep\")\n","\n","    # Diagnostics index: [ce, acc, nce_emb, nce_pw, bridge, assign, assign_nce, nce_tri, attract]\n","    DIAG_KEYS = ['ce', 'acc', 'nce_emb', 'nce_pw', 'bridge', 'assign', 'assign_nce', 'nce_tri', 'attract']\n","\n","    for epoch in pbar:\n","        model.train()\n","        total_loss = correct = total = 0\n","        last_diag = None\n","        t0 = time.time()\n","\n","        for batch_idx, ((v1, v2), labels) in enumerate(train_loader):\n","            v1, v2, labels = v1.to(device), v2.to(device), labels.to(device)\n","\n","            if batch_idx % precompute_freq == 0:\n","                model.invalidate()\n","                model.precompute()\n","\n","            loss, diag = compiled(v1, v2, targets=labels)\n","\n","            optimizer.zero_grad()\n","            loss.backward()\n","\n","            # Geo losses (CV + spread)\n","            geo = model.compute_geo_losses()\n","            geo_total = geo['geo_total']\n","            if geo_total.requires_grad:\n","                geo_total.backward()\n","\n","            # Dynamic clip: task loss scale sets the ceiling\n","            clip_val = max(diag[0].item(), 1.0)  # diag[0] = ce\n","            torch.nn.utils.clip_grad_norm_(model.parameters(), clip_val)\n","            optimizer.step()\n","            scheduler.step()\n","\n","            total_loss += loss.item() * v1.size(0)\n","            correct += diag[1].item() * v1.size(0)  # diag[1] = acc\n","            total += v1.size(0)\n","            last_diag = diag\n","\n","        train_loss = total_loss / total\n","        train_acc = correct / total\n","        elapsed = time.time() - t0\n","\n","        test_acc = evaluate(model, test_loader)\n","        if test_acc > best_acc:\n","            best_acc = test_acc\n","            torch.save({\n","                'state_dict': model.state_dict(),\n","                'epoch': epoch, 'test_acc': test_acc, 'config': config,\n","            }, save_dir / 'best.pt')\n","\n","        pbar.set_postfix_str(\n","            f\"L={train_loss:.2f} tr={train_acc:.3f} te={test_acc:.3f} \"\n","            f\"best={best_acc:.3f} {elapsed:.1f}s\")\n","\n","        if (epoch % 10 == 0 or epoch == config['epochs'] - 1) and last_diag is not None:\n","            geo = model.compute_geo_losses()\n","            d = {k: last_diag[i].item() for i, k in enumerate(DIAG_KEYS)}\n","            cfg = config\n","            # Weighted contributions\n","            w_ce = cfg['w_ce'] * d['ce']\n","            w_nce_e = cfg['w_nce_emb'] * d['nce_emb']\n","            w_nce_p = cfg['w_nce_pw'] * d['nce_pw']\n","            w_brg = cfg['w_bridge'] * d['bridge']\n","            w_asgn = cfg['w_assign'] * d['assign']\n","            w_anc = cfg['w_assign_nce'] * d['assign_nce']\n","            w_tri = cfg['w_nce_tri'] * d['nce_tri']\n","            w_att = cfg['w_attract'] * d['attract']\n","            tqdm.write(\n","                f\"  E{epoch:>3d}  L={train_loss:.3f} train={train_acc:.4f} \"\n","                f\"test={test_acc:.4f} best={best_acc:.4f} \"\n","                f\"CV={geo['cv_raw'].item():.3f} {elapsed:.1f}s\"\n","                f\"\\n    raw:  ce={d['ce']:.3f} nce_emb={d['nce_emb']:.3f} \"\n","                f\"nce_pw={d['nce_pw']:.3f} brg={d['bridge']:.3f} \"\n","                f\"assign={d['assign']:.3f} assign_nce={d['assign_nce']:.3f} \"\n","                f\"nce_tri={d['nce_tri']:.3f} attract={d['attract']:.3f}\"\n","                f\"\\n    wt'd: ce={w_ce:.2f} nce_e={w_nce_e:.2f} nce_p={w_nce_p:.2f} \"\n","                f\"brg={w_brg:.2f} asgn={w_asgn:.2f} a_nce={w_anc:.2f} \"\n","                f\"tri={w_tri:.2f} att={w_att:.2f} \"\n","                f\"sum={w_ce+w_nce_e+w_nce_p+w_brg+w_asgn+w_anc+w_tri+w_att:.2f}\")\n","\n","    pbar.close()\n","    print(f\"\\n{'═' * 60}\")\n","    print(f\"  RESULTS: {best_acc:.4f} ({best_acc*100:.2f}%) | {n_params:,} params\")\n","    print(f\"{'═' * 60}\")\n","\n","\n","if __name__ == '__main__':\n","    main()"],"metadata":{"colab":{"base_uri":"https://localhost:8080/","height":1000,"referenced_widgets":["74ac9e73987f44b6bb162f0f9ba9c846","f41e6f87ec564b1c972a55f84ba7238a","fd296886a4ae41818fc8f64787521637","abc21c48f62b435a984e971dc69c8cb2","5201c3b55b9d42138b84be227d368237","4f7527e7353449319bd4370b6cd3184b","6718ec41897c4c01ae79ec4a6a5ecd4d","5ea7f24370ac4900809d0b9e52bf8017","f473b773987d4b848b1760bb55fd45ec","f21f834c60a44e148fc08227a306dc09","f9408474f42248e79577eb763e72029a"]},"id":"cUkIfLCswZe9","executionInfo":{"status":"error","timestamp":1775228772117,"user_tz":420,"elapsed":643984,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"3dba5a8e-48fb-4a48-a1bf-9a316c512b4a"},"execution_count":6,"outputs":[{"output_type":"stream","name":"stdout","text":["Device: cuda\n","  GPU: NVIDIA RTX PRO 6000 Blackwell Server Edition\n","  VRAM: 102.0 GB\n","============================================================\n","  GeoLIP Fast — Multi-scale SVD + Constellation MLP\n","  d=256, manifold=128, anchors=128, svd_rank=12\n","  21 tokens: 1×32² + 4×16² + 16×8²\n","============================================================\n","  Train: 50,000 | Test: 10,000\n","  Params: 582,629\n","  Encoder output: 546 (21 tokens × 26 SVD features)\n","  Compiled (default)\n","\n","  Training 200 epochs, precompute every 10 batches\n","━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\n","\n"]},{"output_type":"display_data","data":{"text/plain":["Training:   0%|          | 0/200 [00:00<?, ?ep/s]"],"application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"74ac9e73987f44b6bb162f0f9ba9c846"}},"metadata":{}},{"output_type":"stream","name":"stdout","text":["  E  0  L=19.501 train=0.0106 test=0.0137 best=0.0137 CV=0.088 192.3s\n","    raw:  ce=4.615 nce_emb=5.649 nce_pw=5.545 brg=3.449 assign=0.014 assign_nce=5.568 nce_tri=5.545 attract=0.516\n","    wt'd: ce=4.61 nce_e=2.82 nce_p=4.16 brg=2.59 asgn=0.01 a_nce=1.39 tri=2.77 att=0.13 sum=18.49\n","  E 10  L=15.599 train=0.0303 test=0.0308 best=0.0308 CV=0.255 3.3s\n","    raw:  ce=4.472 nce_emb=5.402 nce_pw=5.417 brg=0.302 assign=0.001 assign_nce=5.552 nce_tri=5.508 attract=0.050\n","    wt'd: ce=4.47 nce_e=2.70 nce_p=4.06 brg=0.23 asgn=0.00 a_nce=1.39 tri=2.75 att=0.01 sum=15.62\n","  E 20  L=16.564 train=0.0372 test=0.0334 best=0.0355 CV=0.407 3.3s\n","    raw:  ce=4.401 nce_emb=5.339 nce_pw=5.368 brg=1.842 assign=0.011 assign_nce=5.408 nce_tri=5.451 attract=0.052\n","    wt'd: ce=4.40 nce_e=2.67 nce_p=4.03 brg=1.38 asgn=0.01 a_nce=1.35 tri=2.73 att=0.01 sum=16.57\n","  E 30  L=16.832 train=0.0408 test=0.0392 best=0.0414 CV=0.457 3.3s\n","    raw:  ce=4.399 nce_emb=5.386 nce_pw=5.392 brg=2.317 assign=0.012 assign_nce=5.457 nce_tri=5.437 attract=0.057\n","    wt'd: ce=4.40 nce_e=2.69 nce_p=4.04 brg=1.74 asgn=0.01 a_nce=1.36 tri=2.72 att=0.01 sum=16.98\n","  E 40  L=16.872 train=0.0439 test=0.0419 best=0.0456 CV=0.476 3.3s\n","    raw:  ce=4.344 nce_emb=5.272 nce_pw=5.285 brg=2.527 assign=0.012 assign_nce=5.384 nce_tri=5.381 attract=0.063\n","    wt'd: ce=4.34 nce_e=2.64 nce_p=3.96 brg=1.89 asgn=0.01 a_nce=1.35 tri=2.69 att=0.02 sum=16.90\n","  E 50  L=16.685 train=0.0478 test=0.0467 best=0.0510 CV=0.489 3.3s\n","    raw:  ce=4.294 nce_emb=5.184 nce_pw=5.197 brg=2.389 assign=0.012 assign_nce=5.282 nce_tri=5.346 attract=0.062\n","    wt'd: ce=4.29 nce_e=2.59 nce_p=3.90 brg=1.79 asgn=0.01 a_nce=1.32 tri=2.67 att=0.02 sum=16.59\n","  E 60  L=16.656 train=0.0501 test=0.0484 best=0.0510 CV=0.496 3.3s\n","    raw:  ce=4.344 nce_emb=5.175 nce_pw=5.186 brg=2.387 assign=0.012 assign_nce=5.317 nce_tri=5.356 attract=0.066\n","    wt'd: ce=4.34 nce_e=2.59 nce_p=3.89 brg=1.79 asgn=0.01 a_nce=1.33 tri=2.68 att=0.02 sum=16.64\n","  E 70  L=16.634 train=0.0509 test=0.0498 best=0.0558 CV=0.497 3.3s\n","    raw:  ce=4.215 nce_emb=5.136 nce_pw=5.166 brg=2.427 assign=0.012 assign_nce=5.272 nce_tri=5.323 attract=0.068\n","    wt'd: ce=4.21 nce_e=2.57 nce_p=3.87 brg=1.82 asgn=0.01 a_nce=1.32 tri=2.66 att=0.02 sum=16.48\n","  E 80  L=16.604 train=0.0525 test=0.0548 best=0.0563 CV=0.490 3.3s\n","    raw:  ce=4.318 nce_emb=5.308 nce_pw=5.291 brg=2.470 assign=0.012 assign_nce=5.378 nce_tri=5.388 attract=0.072\n","    wt'd: ce=4.32 nce_e=2.65 nce_p=3.97 brg=1.85 asgn=0.01 a_nce=1.34 tri=2.69 att=0.02 sum=16.85\n"]},{"output_type":"error","ename":"KeyboardInterrupt","evalue":"","traceback":["\u001b[0;31m---------------------------------------------------------------------------\u001b[0m","\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)","\u001b[0;32m/tmp/ipykernel_371939/2959789835.py\u001b[0m in \u001b[0;36m<cell line: 0>\u001b[0;34m()\u001b[0m\n\u001b[1;32m    592\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    593\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0m__name__\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m'__main__'\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 594\u001b[0;31m     \u001b[0mmain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m","\u001b[0;32m/tmp/ipykernel_371939/2959789835.py\u001b[0m in \u001b[0;36mmain\u001b[0;34m()\u001b[0m\n\u001b[1;32m    547\u001b[0m         \u001b[0melapsed\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtime\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtime\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m-\u001b[0m \u001b[0mt0\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    548\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 549\u001b[0;31m         \u001b[0mtest_acc\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mevaluate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtest_loader\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    550\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mtest_acc\u001b[0m \u001b[0;34m>\u001b[0m \u001b[0mbest_acc\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    551\u001b[0m             \u001b[0mbest_acc\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtest_acc\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/utils/_contextlib.py\u001b[0m in \u001b[0;36mdecorate_context\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    122\u001b[0m         \u001b[0;31m# pyrefly: ignore [bad-context-manager]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    123\u001b[0m         \u001b[0;32mwith\u001b[0m \u001b[0mctx_factory\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 124\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mfunc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    125\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    126\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mdecorate_context\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/tmp/ipykernel_371939/2959789835.py\u001b[0m in \u001b[0;36mevaluate\u001b[0;34m(model, loader)\u001b[0m\n\u001b[1;32m    456\u001b[0m     \u001b[0;32mfor\u001b[0m \u001b[0mimages\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlabels\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mloader\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    457\u001b[0m         \u001b[0mimages\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlabels\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mimages\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlabels\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 458\u001b[0;31m         \u001b[0mlogits\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mimages\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    459\u001b[0m         \u001b[0mcorrect\u001b[0m \u001b[0;34m+=\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mlogits\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0margmax\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0mlabels\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msum\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mitem\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    460\u001b[0m         \u001b[0mtotal\u001b[0m \u001b[0;34m+=\u001b[0m \u001b[0mimages\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msize\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1774\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_compiled_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[misc]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1775\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1776\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1777\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1778\u001b[0m     \u001b[0;31m# torchrec tests the code consistency with the following code\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1785\u001b[0m                 \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_pre_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_hooks\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1786\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1787\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1788\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1789\u001b[0m         \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/tmp/ipykernel_371939/2959789835.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, x, x2, targets)\u001b[0m\n\u001b[1;32m    326\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_forward_paired\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx2\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtargets\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    327\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 328\u001b[0;31m         \u001b[0mfeat\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0memb\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mc_out\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpw\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbrg\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_encode\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    329\u001b[0m         \u001b[0mcombined\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mfeat\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpw\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdim\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    330\u001b[0m         \u001b[0mlogits\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mhead\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcombined\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/tmp/ipykernel_371939/2959789835.py\u001b[0m in \u001b[0;36m_encode\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m    314\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_encode\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    315\u001b[0m         \u001b[0;34m\"\"\"Image → (backbone_feat, emb, constellation_out, pw, bridge)\"\"\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 316\u001b[0;31m         \u001b[0msvd_flat\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mS\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mencoder\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    317\u001b[0m         \u001b[0mfeat\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackbone\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msvd_flat\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    318\u001b[0m         \u001b[0memb\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto_manifold\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfeat\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1774\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_compiled_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[misc]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1775\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1776\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1777\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1778\u001b[0m     \u001b[0;31m# torchrec tests the code consistency with the following code\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1785\u001b[0m                 \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_pre_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_hooks\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1786\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1787\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1788\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1789\u001b[0m         \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/tmp/ipykernel_371939/2959789835.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m    161\u001b[0m         \u001b[0;31m# Scale 0: full 32×32\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    162\u001b[0m         \u001b[0ms0_feat\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mscale0_conv\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfeat\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 163\u001b[0;31m         \u001b[0mS0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mVh0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msvd_feat0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnov0\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mscale0_svd\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0ms0_feat\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    164\u001b[0m         \u001b[0;31m# (B, feat_per_svd) → (B, 1, feat_per_svd)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    165\u001b[0m         \u001b[0msvd_feat0\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0msvd_feat0\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0munsqueeze\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1774\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_compiled_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[misc]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1775\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1776\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1777\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1778\u001b[0m     \u001b[0;31m# torchrec tests the code consistency with the following code\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1785\u001b[0m                 \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_pre_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_hooks\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1786\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1787\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1788\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1789\u001b[0m         \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/geolip_core/core/input/svd.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m    103\u001b[0m         \u001b[0mh\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto_svd\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    104\u001b[0m         \u001b[0mh_flat\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mh\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpermute\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m2\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m3\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mreshape\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mB\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mH\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0mW\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msvd_rank\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 105\u001b[0;31m         \u001b[0mS\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mVh\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_decompose\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mh_flat\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    106\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mS\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mVh\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mextract_features\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mS\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mVh\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcompute_novelty\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mS\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    107\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/geolip_core/core/input/svd.py\u001b[0m in \u001b[0;36m_decompose\u001b[0;34m(self, matrix)\u001b[0m\n\u001b[1;32m     60\u001b[0m             \u001b[0mA\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmatrix\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfloat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     61\u001b[0m             \u001b[0mG\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mA\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtranspose\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m@\u001b[0m \u001b[0mA\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 62\u001b[0;31m             \u001b[0meigenvalues\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0meigenvectors\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mLA\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0meigh\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mG\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmethod\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'fl'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     63\u001b[0m             \u001b[0;31m# Detach — FL deflation loop mutates cl in-place (version conflict)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     64\u001b[0m             \u001b[0meigenvalues\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0meigenvalues\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdetach\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/geolip_core/linalg/eigh.py\u001b[0m in \u001b[0;36meigh\u001b[0;34m(A, method)\u001b[0m\n\u001b[1;32m    181\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    182\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mmethod\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m'fl'\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mmethod\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m'auto'\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mbackend\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0muse_fl_eigh\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mn\u001b[0m \u001b[0;34m<=\u001b[0m \u001b[0m_FL_MAX_N\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mA\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mis_cuda\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 183\u001b[0;31m         \u001b[0mvals\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mvecs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mFLEigh\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mA\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    184\u001b[0m     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    185\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mmethod\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m'auto'\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mA\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mis_cuda\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1774\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_compiled_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[misc]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1775\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1776\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1777\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1778\u001b[0m     \u001b[0;31m# torchrec tests the code consistency with the following code\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1785\u001b[0m                 \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_pre_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_hooks\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1786\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1787\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1788\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1789\u001b[0m         \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/geolip_core/linalg/eigh.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, A)\u001b[0m\n\u001b[1;32m     91\u001b[0m                     \u001b[0md2\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0md2\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0mz\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mdp\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     92\u001b[0m                     \u001b[0mdp\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdp\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0mz\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mpv\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 93\u001b[0;31m                     \u001b[0mpv\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpv\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0mz\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mcl\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mj\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     94\u001b[0m                 \u001b[0mok\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpv\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mabs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m>\u001b[0m \u001b[0;36m1e-30\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     95\u001b[0m                 \u001b[0mps\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mwhere\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mok\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpv\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mones_like\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpv\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;31mKeyboardInterrupt\u001b[0m: "]}]},{"cell_type":"markdown","source":["# smaller prototype"],"metadata":{"id":"U68DMM5murm_"}},{"cell_type":"code","source":["\"\"\"\n","Router→Tower→Component prototype WITH wide ensemble + FL eigh + SVD.\n","\n","Proper hierarchy for WideRouter batching:\n","  MiniClassifier (WideRouter)\n","    └── patch_embed (SVDPatchEmbed — component)\n","    └── transformer (MiniTransformerRouter — WideRouter)\n","         ├── layer_0 (MiniGeoLayer — WideRouter)  ← Router with parallel flows\n","         │    ├── flow_0 (FlowTower — BaseTower)  ← [3x BATCH] vmappable\n","         │    ├── flow_1 (FlowTower — BaseTower)\n","         │    └── flow_2 (FlowTower — BaseTower)\n","         ├── layer_1 (MiniGeoLayer — WideRouter)\n","         │    ├── flow_0..2\n","         ├── layer_2, layer_3 ...\n","         └── final_norm\n","\n","Layers are sequential. Flows within each layer are PARALLEL — the ensemble\n","that WideRouter discovers and batches via vmap.\n","\"\"\"\n","\n","import os\n","os.environ['CUDA_LAUNCH_BLOCKING'] = '1'\n","import sys\n","_stderr_log = open('/tmp/proto_stderr.log', 'w')\n","sys.stderr = _stderr_log\n","import warnings\n","warnings.filterwarnings('ignore')\n","import faulthandler\n","_crash_log = open('/tmp/proto_crash.log', 'w')\n","faulthandler.enable(file=_crash_log, all_threads=True)\n","\n","import torch\n","import torch.nn as nn\n","import torch.nn.functional as F\n","import shutil\n","\n","device = torch.device('cuda')\n","torch.set_float32_matmul_precision('high')\n","\n","from geofractal.router.wide_router import WideRouter\n","from geofractal.router.base_tower import BaseTower\n","from geofractal.router.components.torch_component import TorchComponent\n","\n","import geolip_core.linalg as LA\n","from geolip_core.linalg.svd import gram_fl_eigh_svd\n","from geolip_core.linalg.eigh import FLEigh\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# ACTUAL CMValidatedGate + support functions (from geometric_transformer.py)\n","# Raw nn.Module, NOT TorchComponent — exactly as in real model\n","# ═══════════════════════════════════════════════════════════════════\n","\n","def pairwise_distances_squared(points):\n","    gram = torch.bmm(points, points.transpose(1, 2))\n","    diag = gram.diagonal(dim1=-2, dim2=-1)\n","    return diag.unsqueeze(2) + diag.unsqueeze(1) - 2 * gram\n","\n","def cayley_menger_det(points):\n","    B, K, D = points.shape\n","    d2 = pairwise_distances_squared(points)\n","    M = torch.zeros(B, K + 1, K + 1, device=points.device, dtype=points.dtype)\n","    M[:, 0, 1:] = 1.0; M[:, 1:, 0] = 1.0; M[:, 1:, 1:] = d2\n","    raw = LA.det(M)\n","    k = K - 1\n","    sign = (-1.0) ** (k + 1)\n","    return sign * raw\n","\n","def anchor_neighborhood_cm(anchors, n_neighbors=3):\n","    A, D = anchors.shape\n","    dists = torch.cdist(anchors.unsqueeze(0), anchors.unsqueeze(0)).squeeze(0)\n","    self_mask = torch.eye(A, device=anchors.device, dtype=anchors.dtype) * 1e12\n","    dists = dists + self_mask\n","    _, nn_idx = dists.topk(n_neighbors, largest=False)\n","    simplices = torch.cat([anchors.unsqueeze(1), anchors[nn_idx]], dim=1)\n","    dets = cayley_menger_det(simplices)\n","    sign = dets.sign()\n","    log_mag = torch.log(dets.abs() + 1e-12)\n","    return sign * log_mag, nn_idx\n","\n","class CMValidatedGate(nn.Module):\n","    \"\"\"ACTUAL CMValidatedGate — raw nn.Module, NOT TorchComponent.\"\"\"\n","    def __init__(self, n_anchors, n_neighbors=3):\n","        super().__init__()\n","        self.n_anchors = n_anchors\n","        self.n_neighbors = n_neighbors\n","        self.gate_proj = nn.Sequential(\n","            nn.Linear(2, 16), nn.GELU(), nn.Linear(16, 1))\n","        nn.init.normal_(self.gate_proj[2].weight, std=0.01)\n","        nn.init.constant_(self.gate_proj[2].bias, 2.0)\n","        self.register_buffer('_cached_cm_norm', torch.zeros(n_anchors), persistent=False)\n","        self._cache_warm = False\n","\n","    def invalidate_cache(self):\n","        self._cache_warm = False\n","\n","    def precompute(self, anchors):\n","        if self._cache_warm:\n","            return\n","        with torch.no_grad():\n","            anchor_cm, _ = anchor_neighborhood_cm(anchors, self.n_neighbors)\n","            cm_std = anchor_cm.std().clamp(min=1e-8)\n","            new_val = ((anchor_cm - anchor_cm.mean()) / cm_std).detach()\n","            self._cached_cm_norm.copy_(new_val)\n","            self._cache_warm = True\n","\n","    def forward(self, tri):\n","        N, A = tri.shape\n","        features = torch.stack([\n","            self._cached_cm_norm.unsqueeze(0).expand(N, -1),\n","            1.0 - tri,\n","        ], dim=-1)\n","        gate_values = torch.sigmoid(self.gate_proj(features).squeeze(-1))\n","        gate_info = {\n","            'active': (gate_values.detach() > 0.5).float().sum(-1).mean(),\n","            'gate_mean': gate_values.detach().mean(),\n","            'cm_positive_frac': (self._cached_cm_norm > 0).float().mean(),\n","        }\n","        return gate_values, gate_info\n","\n","class GeoResidualBank(nn.Module):\n","    \"\"\"ACTUAL GeoResidualBank — raw nn.Module, NOT TorchComponent.\"\"\"\n","    def __init__(self, proj_dim, bank_size=4096, temperature=0.1):\n","        super().__init__()\n","        self.proj_dim = proj_dim\n","        self.bank_size = bank_size\n","        self.temperature = temperature\n","        self.register_buffer('queue', torch.randn(bank_size, proj_dim))\n","        self.queue = F.normalize(self.queue, dim=-1)\n","        self.register_buffer('queue_ptr', torch.zeros(1, dtype=torch.long))\n","\n","    @torch.no_grad()\n","    def enqueue(self, keys):\n","        B = keys.shape[0]\n","        ptr = int(self.queue_ptr.item())\n","        if ptr + B <= self.bank_size:\n","            self.queue[ptr:ptr + B] = keys\n","        else:\n","            overflow = (ptr + B) - self.bank_size\n","            self.queue[ptr:] = keys[:B - overflow]\n","            self.queue[:overflow] = keys[B - overflow:]\n","        self.queue_ptr[0] = (ptr + B) % self.bank_size\n","\n","    def forward(self, content_proj, geo_proj):\n","        q = F.normalize(content_proj, dim=-1)\n","        k_pos = F.normalize(geo_proj, dim=-1)\n","        k_neg = self.queue.clone().detach()\n","        pos_logits = (q * k_pos).sum(dim=-1, keepdim=True) / self.temperature\n","        neg_logits = q @ k_neg.T / self.temperature\n","        logits = torch.cat([pos_logits, neg_logits], dim=1)\n","        labels = torch.zeros(q.shape[0], dtype=torch.long, device=q.device)\n","        loss = F.cross_entropy(logits, labels)\n","        with torch.no_grad():\n","            acc = (logits.argmax(dim=1) == 0).float().mean()\n","        return loss, acc\n","\n","class DeepObserver(nn.Module):\n","    \"\"\"Observer with deep attribute chain — raw nn.Module.\"\"\"\n","    def __init__(self, n_anchors, manifold_dim):\n","        super().__init__()\n","        self.association = nn.Module()\n","        self.association.constellation = nn.Module()\n","        self.association.constellation.anchors = nn.Parameter(\n","            torch.randn(n_anchors, manifold_dim) * 0.1)\n","\n","    def forward(self, emb):\n","        anchors_n = F.normalize(self.association.constellation.anchors, dim=-1)\n","        cos = emb @ anchors_n.T\n","        distances = 1.0 - cos\n","        assignment = F.softmax(cos / 0.1, dim=-1)\n","        return {\n","            'distances': distances, 'cos_to_anchors': cos,\n","            'assignment': assignment, 'nearest': cos.argmax(dim=-1),\n","        }\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# COMPONENTS — stateless leaves\n","# ═══════════════════════════════════════════════════════════════════\n","\n","class SVDPatchEmbed(TorchComponent):\n","    \"\"\"Conv → FL SVD → patch tokens. EMA from Router.\"\"\"\n","    def __init__(self, name, d_model=256, svd_rank=8,\n","                 img_size=32, patch_size=4, in_ch=3, conv_ch=32):\n","        super().__init__(name)\n","        self.svd_rank = svd_rank\n","        n_patches = (img_size // patch_size) ** 2\n","        self.conv = nn.Sequential(\n","            nn.Conv2d(in_ch, conv_ch, 3, padding=1, bias=False),\n","            nn.BatchNorm2d(conv_ch), nn.GELU())\n","        self.to_svd = nn.Conv2d(conv_ch, svd_rank, 1, bias=False)\n","        self.patch_proj = nn.Conv2d(conv_ch, d_model, patch_size,\n","                                     stride=patch_size, bias=False)\n","        self.norm = nn.LayerNorm(d_model)\n","        feat_dim = 2 * svd_rank + 2\n","        self.svd_gamma = nn.Linear(feat_dim, d_model)\n","        self.svd_beta = nn.Linear(feat_dim, d_model)\n","        nn.init.ones_(self.svd_gamma.bias); nn.init.zeros_(self.svd_beta.bias)\n","        self.cls_token = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)\n","        self.pos_embed = nn.Parameter(torch.randn(1, n_patches + 1, d_model) * 0.02)\n","\n","    def _svd_features(self, S, Vh):\n","        S_s = S.clamp(min=1e-6)\n","        s_n = S_s / (S_s.sum(-1, keepdim=True) + 1e-8)\n","        vh_d = Vh.diagonal(dim1=-2, dim2=-1)\n","        vh_o = (Vh.pow(2).sum((-2,-1)) - vh_d.pow(2).sum(-1)).unsqueeze(-1).clamp(min=0)\n","        s_e = -(s_n * torch.log(s_n.clamp(min=1e-8))).sum(-1, keepdim=True)\n","        o = torch.cat([s_n, vh_d, vh_o, s_e], dim=-1)\n","        return torch.where(torch.isfinite(o), o, torch.zeros_like(o))\n","\n","    def forward(self, x, ema_s):\n","        B = x.shape[0]\n","        feat = self.conv(x)\n","        h = self.to_svd(feat).permute(0,2,3,1).reshape(B, -1, self.svd_rank)\n","        with torch.amp.autocast('cuda', enabled=False):\n","            U, S, Vh = gram_fl_eigh_svd(h.float())\n","        S = S.clamp(min=1e-6)\n","        sf = self._svd_features(S, Vh)\n","        g = self.svd_gamma(sf).unsqueeze(1)\n","        b = self.svd_beta(sf).unsqueeze(1)\n","        tok = self.norm(self.patch_proj(feat).flatten(2).transpose(1,2))\n","        tok = g * tok + b\n","        cls = self.cls_token.expand(B, -1, -1)\n","        tok = torch.cat([cls, tok], dim=1) + self.pos_embed\n","        return tok, S.detach(), Vh.detach()\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# TOWER — Flow (parallel, vmappable)\n","# ═══════════════════════════════════════════════════════════════════\n","\n","class FlowTower(BaseTower):\n","    \"\"\"Single geometric flow — produces an opinion from anchors + embeddings.\n","\n","    Uses FLEigh for spectral analysis. Fully compilable.\n","    Structurally identical across flows → WideRouter can vmap batch them.\n","    \"\"\"\n","    def __init__(self, name, n_anchors=32, manifold_dim=64, d_out=64):\n","        super().__init__(name, strict=False)\n","        self.manifold_dim = manifold_dim\n","        self.n_neighbors = 5\n","\n","        # Spectral analysis via FL eigh\n","        self.proj = nn.Sequential(\n","            nn.Linear(n_anchors, d_out * 2), nn.GELU(),\n","            nn.Linear(d_out * 2, d_out), nn.LayerNorm(d_out))\n","\n","    def forward(self, distances, anchors_n):\n","        \"\"\"(N, A) distances + (A, M) anchors → (N, d_out) opinion.\n","\n","        Uses FLEigh on anchor neighborhood Gram matrices.\n","        \"\"\"\n","        N, A = distances.shape\n","\n","        # Build local neighborhoods\n","        _, topk_idx = distances.topk(self.n_neighbors, dim=-1, largest=False)\n","        neighbors = anchors_n[topk_idx.reshape(-1)].reshape(\n","            N, self.n_neighbors, -1)\n","        gram = torch.bmm(neighbors, neighbors.transpose(1, 2))\n","\n","        # Spectral decomposition — FL eigh (compilable)\n","        eigenvalues, eigenvectors = FLEigh()(gram)\n","\n","        # Spectral-weighted distances\n","        spectral_weight = eigenvalues[:, -1:] / (eigenvalues.sum(-1, keepdim=True) + 1e-8)\n","        weighted = distances * spectral_weight\n","\n","        return self.proj(weighted)\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# LAYER (Router with parallel flow ensemble)\n","# ═══════════════════════════════════════════════════════════════════\n","\n","class MiniGeoLayer(WideRouter):\n","    \"\"\"Geometric layer with geo_residual accumulation stream.\n","\n","    geo_residual flows across layers:\n","      Layer 0: geo_residual=None → creates from patchwork\n","      Layer N: receives accumulated residual, adds own patchwork contribution\n","    History MLP processes prior layers' observations for FiLM context.\n","    \"\"\"\n","    def __init__(self, name, d_model=256, n_heads=4, n_anchors=32,\n","                 manifold_dim=64, n_flows=3, flow_dim=64):\n","        super().__init__(name, strict=False)\n","        self.d_model = d_model\n","        self.n_anchors = n_anchors\n","        self.manifold_dim = manifold_dim\n","        self.n_flows = n_flows\n","        pw_dim = n_anchors  # simplified patchwork dim\n","\n","        # Manifold projection\n","        self.attach('proj', nn.Sequential(\n","            nn.Linear(d_model, manifold_dim), nn.LayerNorm(manifold_dim)))\n","\n","        # RAW nn.Module components (NOT TorchComponent)\n","        self.attach('observer', DeepObserver(n_anchors, manifold_dim))\n","        self.attach('cm_gate', CMValidatedGate(n_anchors, n_neighbors=3))\n","\n","        # Parallel flow ensemble\n","        for i in range(n_flows):\n","            flow = FlowTower(f'flow_{i}', n_anchors, manifold_dim, flow_dim)\n","            self.attach(f'flow_{i}', flow)\n","            self.register_tower(f'flow_{i}')\n","        self.attach('flow_fuse', nn.Linear(n_flows * flow_dim, d_model))\n","\n","        # Cayley rotation\n","        n_upper = manifold_dim * (manifold_dim - 1) // 2\n","        self.A_upper = nn.Parameter(torch.zeros(n_upper) * 0.01)\n","        self.register_buffer('_cached_rotation', torch.eye(manifold_dim), persistent=False)\n","        self._cache_warm = False\n","\n","        # Dual attention streams\n","        self.attach('content_attn', nn.MultiheadAttention(\n","            d_model, n_heads, batch_first=True))\n","        self.attach('content_norm', nn.LayerNorm(d_model))\n","        self.attach('geo_attn', nn.MultiheadAttention(\n","            d_model, n_heads, batch_first=True))\n","        self.attach('geo_norm', nn.LayerNorm(d_model))\n","\n","        # Gated residual composition\n","        self.attach('gate', nn.Sequential(nn.Linear(d_model * 2, d_model), nn.Sigmoid()))\n","        self.attach('ffn', nn.Sequential(\n","            nn.Linear(d_model, d_model * 4), nn.GELU(),\n","            nn.Linear(d_model * 4, d_model)))\n","        self.attach('ffn_norm', nn.LayerNorm(d_model))\n","\n","        # FiLM context — 3 streams: anchor, flow, history\n","        self.attach('anchor_mlp', nn.Linear(n_anchors * 3, d_model))  # cos+assign+tri\n","        self.attach('flow_mlp', nn.Linear(d_model, d_model))\n","        self.attach('history_mlp', nn.Linear(pw_dim, d_model))\n","        self.attach('film_fuse', nn.Sequential(\n","            nn.Linear(d_model * 3, d_model), nn.GELU()))\n","        self.attach('film_gamma', nn.Linear(d_model, d_model))\n","        self.attach('film_beta', nn.Linear(d_model, d_model))\n","\n","        # Geo residual projection — CM-gated patchwork → residual\n","        self.attach('geo_proj', nn.Sequential(\n","            nn.Linear(pw_dim, pw_dim), nn.LayerNorm(pw_dim)))\n","\n","    def forward(self, x, geo_residual=None):\n","        \"\"\"Returns (x_out, geo_residual_out, geo_state).\n","\n","        geo_residual: (B, L, pw_dim) from previous layer, or None for first.\n","        \"\"\"\n","        B, L, D = x.shape\n","\n","        # ═ 1. Project + rotate (buffer read) ═\n","        emb = F.normalize(self['proj'](x), dim=-1)\n","        emb_rot = emb @ self._cached_rotation.T\n","\n","        # ═ 2. Observer + CM gate ═\n","        flat = emb_rot.reshape(-1, self.manifold_dim)\n","        a_out = self['observer'](flat)\n","        distances = a_out['distances']\n","        cos = a_out['cos_to_anchors']\n","        assignment = a_out['assignment']\n","        gate_values, gate_info = self['cm_gate'](distances)\n","        gated_distances = distances * gate_values\n","\n","        # ═ 3. Flow ensemble ═\n","        anchors_n = F.normalize(\n","            self['observer'].association.constellation.anchors, dim=-1)\n","        flow_opinions = []\n","        for i in range(self.n_flows):\n","            opinion = self[f'flow_{i}'](gated_distances, anchors_n)\n","            flow_opinions.append(opinion)\n","        combined = torch.cat(flow_opinions, dim=-1)\n","        fused = self['flow_fuse'](combined).reshape(B, L, D)\n","\n","        # ═ 4. FiLM context — 3 streams ═\n","        anchor_feats = torch.cat([\n","            cos.reshape(B, L, -1),\n","            assignment.reshape(B, L, -1),\n","            distances.reshape(B, L, -1),\n","        ], dim=-1)\n","        a_ctx = self['anchor_mlp'](anchor_feats)\n","        f_ctx = self['flow_mlp'](fused)\n","        h_ctx = self['history_mlp'](geo_residual) if geo_residual is not None \\\n","            else torch.zeros(B, L, D, device=x.device)\n","        ctx = self['film_fuse'](torch.cat([a_ctx, f_ctx, h_ctx], dim=-1))\n","        gamma = self['film_gamma'](ctx)\n","        beta = self['film_beta'](ctx)\n","\n","        # ═ 5. Dual attention + gated composition ═\n","        content, _ = self['content_attn'](x, x, x, need_weights=False)\n","        content = self['content_norm'](x + content)\n","        geo_out, _ = self['geo_attn'](x + gamma * x + beta, x, x, need_weights=False)\n","        geo_out = self['geo_norm'](x + geo_out)\n","        g = self['gate'](torch.cat([content, geo_out], dim=-1))\n","        merged = g * geo_out + (1 - g) * content\n","\n","        # ═ 6. FFN ═\n","        h = self['ffn'](merged)\n","        x_out = self['ffn_norm'](merged + h)\n","\n","        # ═ 7. Geo residual accumulation ═\n","        pw_features = gated_distances.reshape(B, L, -1)  # (B, L, n_anchors)\n","        pw_proj = self['geo_proj'](pw_features)\n","        geo_residual_out = pw_proj if geo_residual is None else geo_residual + pw_proj\n","\n","        # geo_state\n","        geo_state = {\n","            'embedding': emb_rot.reshape(B, L, -1),\n","            'triangulation': distances.reshape(B, L, -1),\n","            'cos_to_anchors': cos.reshape(B, L, -1),\n","            'assignment': assignment.reshape(B, L, -1),\n","            'nearest': cos.reshape(B, L, -1).argmax(-1),\n","            'patchwork': fused,\n","            'content': content,\n","            'geometric': geo_out,\n","            'gate_values': gate_values.reshape(B, L, -1),\n","            'geo_residual': geo_residual_out,\n","            'flow_opinion': combined.reshape(B, L, -1),\n","        }\n","        return x_out, geo_residual_out, geo_state\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# TRANSFORMER ROUTER — sequential layers, manages state\n","# ═══════════════════════════════════════════════════════════════════\n","\n","class MiniTransformerRouter(WideRouter):\n","    \"\"\"Orchestrates layers sequentially. Manages rotations + SVD EMA.\"\"\"\n","    def __init__(self, name, d_model=256, n_heads=4, n_layers=4,\n","                 n_anchors=32, manifold_dim=64, svd_rank=8, n_flows=3):\n","        super().__init__(name, strict=False)\n","        self.n_layers = n_layers\n","        self.manifold_dim = manifold_dim\n","\n","        for i in range(n_layers):\n","            layer = MiniGeoLayer(\n","                f'layer_{i}', d_model, n_heads, n_anchors,\n","                manifold_dim, n_flows=n_flows)\n","            self.attach(f'layer_{i}', layer)\n","            self.register_tower(f'layer_{i}')\n","\n","        self.attach('final_norm', nn.LayerNorm(d_model))\n","\n","        # NO rotation buffers on Router — they live on the layers now\n","        # (matches real model where CayleyOrthogonal has _cached_rotation)\n","\n","        # SVD EMA\n","        self.register_buffer('_ema_s', torch.ones(svd_rank), persistent=False)\n","        self._ema_momentum = 0.99\n","\n","    @torch.no_grad()\n","    def precompute(self):\n","        \"\"\"Write rotations + gate quality to LAYER buffers.\"\"\"\n","        for i in range(self.n_layers):\n","            layer = self[f'layer_{i}']\n","            d = layer.manifold_dim\n","            A = torch.zeros(d, d, device=layer.A_upper.device)\n","            idx = torch.triu_indices(d, d, offset=1)\n","            A[idx[0], idx[1]] = layer.A_upper\n","            A = A - A.T\n","            I = torch.eye(d, device=A.device)\n","            R = torch.linalg.solve(I + A, I - A)\n","            layer._cached_rotation.copy_(R)\n","            layer._cache_warm = True\n","\n","            # CM gate precompute — writes to raw nn.Module buffer\n","            anchors_n = F.normalize(\n","                layer['observer'].association.constellation.anchors, dim=-1)\n","            layer['cm_gate'].precompute(anchors_n)\n","\n","    def invalidate(self):\n","        \"\"\"Invalidate ALL caches on layers.\"\"\"\n","        for i in range(self.n_layers):\n","            layer = self[f'layer_{i}']\n","            layer._cache_warm = False\n","            layer['cm_gate'].invalidate_cache()\n","\n","    @torch.no_grad()\n","    def update_svd_ema(self, S):\n","        m = self._ema_momentum\n","        self._ema_s.mul_(m).add_(S.mean(0), alpha=1-m)\n","\n","    def forward(self, tokens):\n","        \"\"\"Accumulates geo_states + geo_residual across layers.\"\"\"\n","        geo_states = []\n","        geo_residual = None\n","        for i in range(self.n_layers):\n","            tokens, geo_residual, geo_state = self[f'layer_{i}'](\n","                tokens, geo_residual=geo_residual)\n","            geo_states.append(geo_state)\n","        return self['final_norm'](tokens), geo_states\n","\n","    def forward_paired(self, x1, x2):\n","        \"\"\"Paired forward — matches real model's forward_paired.\"\"\"\n","        B = x1.shape[0]\n","        x_cat = torch.cat([x1, x2], dim=0)\n","        feat_cat, geo_states = self.forward(x_cat)\n","\n","        # Build output dict from FINAL layer's geo_state (matches real model)\n","        gs = geo_states[-1]\n","        c = 0  # cls_index\n","        return {\n","            'embedding':      gs['embedding'][:B, c],\n","            'embedding_aug':  gs['embedding'][B:, c],\n","            'patchwork1':     gs['patchwork'][:B, c],\n","            'patchwork1_aug': gs['patchwork'][B:, c],\n","            'assign1':        gs['assignment'][:B, c],\n","            'assign2':        gs['assignment'][B:, c],\n","            'cos1':           gs['cos_to_anchors'][:B, c],\n","            'tri1':           gs['triangulation'][:B, c],\n","            'tri2':           gs['triangulation'][B:, c],\n","            'features1':      feat_cat[:B],\n","            'features2':      feat_cat[B:],\n","        }\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# TOP-LEVEL\n","# ═══════════════════════════════════════════════════════════════════\n","\n","class MiniClassifier(WideRouter):\n","    def __init__(self, d_model=256, n_heads=4, n_layers=4,\n","                 n_anchors=32, manifold_dim=64, svd_rank=8, n_flows=3):\n","        super().__init__('mini_geo', strict=False)\n","        self.d_model = d_model\n","        self.n_anchors = n_anchors\n","        self.manifold_dim = manifold_dim\n","        self.n_layers = n_layers\n","\n","        self.attach('patch_embed', SVDPatchEmbed('pe', d_model=d_model, svd_rank=svd_rank))\n","        self.attach('transformer', MiniTransformerRouter(\n","            'xfmr', d_model, n_heads, n_layers, n_anchors, manifold_dim,\n","            svd_rank, n_flows))\n","        self.attach('head', nn.Sequential(nn.LayerNorm(d_model), nn.Linear(d_model, 100)))\n","        self.register_tower('transformer')\n","\n","    def forward(self, x1, x2=None, targets=None):\n","        ema_s = self['transformer']._ema_s\n","\n","        if x2 is not None:\n","            B = x1.shape[0]\n","            x_cat = torch.cat([x1, x2], dim=0)\n","            tokens, S, Vh = self['patch_embed'](x_cat, ema_s)\n","\n","            # forward_paired — accumulates geo_states, builds output dict\n","            output = self['transformer'].forward_paired(tokens[:B], tokens[B:])\n","\n","            if targets is not None:\n","                loss, ld = self._observer_loss(output, targets)\n","                return loss, ld, S.detach()\n","            return output, S.detach()\n","\n","        tokens, S, Vh = self['patch_embed'](x1, ema_s)\n","        feat, _ = self['transformer'](tokens)\n","        return self['head'](feat[:, 0]), S.detach()\n","\n","    def _observer_loss(self, output, targets):\n","        \"\"\"Loss from output dict — matches real model's compute_loss pattern.\n","\n","        Uses DEEP NESTED ATTRIBUTE ACCESS to get anchors:\n","          self['transformer']['layer_N'].anchors\n","        Uses IMPORTED LOSS FUNCTIONS from geolip_core.\n","        \"\"\"\n","        from geolip_core.core.distinguish.losses import (\n","            nce_loss, bridge_loss_paired, assign_bce_loss,\n","            assign_nce_loss, attraction_loss,\n","            ce_loss_paired,\n","        )\n","\n","        # CE via head — uses features from output dict\n","        feat1 = output['features1'][:, 0]\n","        feat2 = output['features2'][:, 0]\n","        l_ce, acc = ce_loss_paired(\n","            self['head'](feat1), self['head'](feat2), targets)\n","\n","        # Embedding NCE — from output dict\n","        l_nce_emb, nce_acc = nce_loss(\n","            output['embedding'], output['embedding_aug'], 0.07, normalize=False)\n","\n","        # Patchwork NCE\n","        l_nce_pw, _ = nce_loss(\n","            output['patchwork1'], output['patchwork1_aug'], 0.1, normalize=True)\n","\n","        # Bridge — DEEP NESTED ACCESS matching real model:\n","        # self['transformer']['layer_N']['observer'].association.constellation.anchors\n","        last_idx = self.n_layers - 1\n","        final_anchors = F.normalize(\n","            self['transformer'][f'layer_{last_idx}']['observer'].association.constellation.anchors,\n","            dim=-1)\n","\n","        # Recompute assignment from anchors (like real model)\n","        cos1 = output['embedding'] @ final_anchors.T\n","        assign1 = F.softmax(cos1 / 0.1, dim=-1)\n","        cos2 = output['embedding_aug'] @ final_anchors.T\n","        assign2 = F.softmax(cos2 / 0.1, dim=-1)\n","        bridge1 = output['patchwork1'][:, :self.n_anchors]  # (B, n_anchors)\n","        bridge2 = output['patchwork1_aug'][:, :self.n_anchors]\n","\n","        l_bridge, bridge_acc = bridge_loss_paired(\n","            bridge1, bridge2, assign1, assign2)\n","\n","        # Assign BCE\n","        l_assign, assign_ent = assign_bce_loss(assign1, cos1)\n","\n","        # Assign NCE\n","        l_assign_nce, _ = assign_nce_loss(assign1, assign2, 0.1)\n","\n","        # Tri NCE\n","        l_nce_tri, _ = nce_loss(output['tri1'], output['tri2'], 0.1, normalize=True)\n","\n","        # Attraction\n","        l_attract, nearest_cos = attraction_loss(output['cos1'])\n","\n","        loss = (1.0 * l_ce + 0.5 * l_nce_emb + 1.0 * l_nce_pw\n","                + 1.0 * l_bridge + 0.5 * l_assign + 0.25 * l_assign_nce\n","                + 0.5 * l_nce_tri + 0.25 * l_attract)\n","\n","        ld = {\n","            'ce': l_ce, 'acc': acc,\n","            'nce_emb': l_nce_emb, 'nce_pw': l_nce_pw,\n","            'bridge': l_bridge, 'assign': l_assign,\n","            'nce_tri': l_nce_tri, 'attract': l_attract,\n","            'total': loss,\n","        }\n","        return loss, ld\n","\n","    def precompute(self):\n","        self['transformer'].precompute()\n","\n","    def invalidate(self):\n","        self['transformer'].invalidate()\n","\n","    def update_ema(self, S):\n","        self['transformer'].update_svd_ema(S)\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# RUN\n","# ═══════════════════════════════════════════════════════════════════\n","\n","B = 64\n","\n","print(\"=\" * 60)\n","print(\"  Router→Tower prototype: wide ensemble + FL eigh + SVD\")\n","print(\"=\" * 60)\n","\n","model = MiniClassifier(d_model=384, n_layers=8, n_anchors=128,\n","                       manifold_dim=128, svd_rank=16, n_flows=3)\n","model.network_to(device=device, strict=False)\n","n_params = sum(p.numel() for p in model.parameters())\n","print(f\"  Params: {n_params:,}\")\n","\n","# Discovery — should find flow ensembles\n","print(f\"\\n  Tower discovery:\")\n","for level_name in ['mini_geo (top)', 'transformer', 'layer_0']:\n","    if level_name == 'mini_geo (top)':\n","        router = model\n","    elif level_name == 'transformer':\n","        router = model['transformer']\n","    else:\n","        router = model['transformer']['layer_0']\n","\n","    towers = router.discover_towers()\n","    print(f\"    {level_name}: {len(towers)} towers\")\n","    from collections import Counter\n","    sigs = Counter(router._tower_signature(t) for t in towers)\n","    for sig, count in sigs.most_common():\n","        tag = \"BATCH\" if count >= 2 else \"solo\"\n","        print(f\"      [{count}x {tag}] {sig[:50]}\")\n","\n","# Clear cache\n","cache_dir = '/tmp/torchinductor_root'\n","if os.path.exists(cache_dir):\n","    shutil.rmtree(cache_dir)\n","\n","compiled = model.compile(mode='reduce-overhead')\n","print(f\"\\n  Compiled (reduce-overhead)\")\n","\n","optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)\n","x1 = torch.randn(B, 3, 32, 32, device=device)\n","x2 = torch.randn(B, 3, 32, 32, device=device)\n","targets = torch.randint(0, 100, (B,), device=device)\n","\n","print(f\"\\n  Training steps (paired views + full observer loss):\")\n","for step in range(5):\n","    _crash_log.write(f\"S{step} precompute\\n\"); _crash_log.flush()\n","    model.precompute()\n","\n","    _crash_log.write(f\"S{step} forward\\n\"); _crash_log.flush()\n","    torch.compiler.cudagraph_mark_step_begin()\n","    loss, ld, S = compiled(x1, x2, targets=targets)\n","\n","    _crash_log.write(f\"S{step} backward\\n\"); _crash_log.flush()\n","    optimizer.zero_grad()\n","    loss.backward()\n","\n","    _crash_log.write(f\"S{step} ema+step\\n\"); _crash_log.flush()\n","    model.update_ema(S)\n","    optimizer.step()\n","    model.invalidate()  # clear caches after param update\n","\n","    ce = ld['ce'].item()\n","    nce = ld['nce_emb'].item()\n","    brg = ld['bridge'].item()\n","    print(f\"    Step {step}: loss={loss.item():.3f} ce={ce:.3f} nce={nce:.3f} bridge={brg:.3f}\")\n","\n","print(f\"\\n  ALL 5 STEPS PASSED\")\n","print(f\"  InfoNCE (embedding + patchwork + tri): ✓\")\n","print(f\"  Bridge loss (soft CE): ✓\")\n","print(f\"  Assignment BCE + scatter: ✓\")\n","print(f\"  Attraction loss: ✓\")\n","print(f\"  Paired CE: ✓\")\n","print(f\"  FL SVD + FLEigh spectral: ✓\")\n","print(f\"  reduce-overhead backward replay: ✓\")\n","print(f\"  Memory: {torch.cuda.max_memory_allocated()/1e9:.2f} GB\")\n","print(\"=\" * 60)"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"1xOdUv01utMV","executionInfo":{"status":"ok","timestamp":1775183138263,"user_tz":420,"elapsed":285294,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"33d7b15e-a7df-45f9-a9ce-762363dcf859"},"execution_count":6,"outputs":[{"output_type":"stream","name":"stdout","text":["============================================================\n","  Router→Tower prototype: wide ensemble + FL eigh + SVD\n","============================================================\n","  Params: 32,172,812\n","\n","  Tower discovery:\n","    mini_geo (top): 1 towers\n","      [1x solo] MiniTransformerRouter_0_9_31882504_ModuleDict\n","    transformer: 8 towers\n","      [8x BATCH] MiniGeoLayer_0_21_3985217_ModuleDict\n","    layer_0: 3 towers\n","      [3x BATCH] FlowTower_0_0_24896_ModuleDict_ModuleList_Sequenti\n","\n","  Compiled (reduce-overhead)\n","\n","  Training steps (paired views + full observer loss):\n","    Step 0: loss=19.328 ce=4.742 nce=4.213 bridge=4.989\n","    Step 1: loss=18.693 ce=4.135 nce=4.187 bridge=5.019\n","    Step 2: loss=18.547 ce=4.006 nce=4.175 bridge=5.053\n","    Step 3: loss=18.319 ce=3.946 nce=4.159 bridge=4.922\n","    Step 4: loss=18.071 ce=3.907 nce=4.160 bridge=4.723\n","\n","  ALL 5 STEPS PASSED\n","  InfoNCE (embedding + patchwork + tri): ✓\n","  Bridge loss (soft CE): ✓\n","  Assignment BCE + scatter: ✓\n","  Attraction loss: ✓\n","  Paired CE: ✓\n","  FL SVD + FLEigh spectral: ✓\n","  reduce-overhead backward replay: ✓\n","  Memory: 7.04 GB\n","============================================================\n"]}]},{"cell_type":"markdown","source":["# graph analysis"],"metadata":{"id":"X5K51SXfMUbq"}},{"cell_type":"code","source":["\"\"\"\n","Inductor crash diagnostic — captures generated code around segfault line.\n","\n","Runs B0 (triggers compilation), saves the generated inductor file,\n","reads the crash lines, then attempts B1.\n","\"\"\"\n","import os, glob, shutil\n","os.environ['CUDA_LAUNCH_BLOCKING'] = '1'\n","\n","import faulthandler\n","_crash_log = open('/tmp/geo_crash.log', 'w')\n","faulthandler.enable(file=_crash_log, all_threads=True)\n","\n","import torch, torch.nn as nn, torch.nn.functional as F\n","device = torch.device('cuda')\n","torch.set_float32_matmul_precision('high')\n","\n","# Clear stale inductor cache\n","cache_dir = '/tmp/torchinductor_root'\n","if os.path.exists(cache_dir):\n","    shutil.rmtree(cache_dir)\n","    print(\"Cleared inductor cache\")\n","\n","from geolip_core.pipeline.components.geometric_transformer import GeometricTransformer\n","from geolip_core.pipeline.observer import TorchComponent\n","from geolip_core.core.input.svd import SVDObserver\n","from geofractal.router.wide_router import WideRouter\n","\n","# ── Minimal reproduction of GeoViTClassifier ──\n","class ConvSVDPatchEmbedding(TorchComponent):\n","    def __init__(self, name, d_model=384, svd_rank=12):\n","        super().__init__(name)\n","        self.conv_frontend = nn.Sequential(\n","            nn.Conv2d(3, 64, 3, padding=1, bias=False),\n","            nn.BatchNorm2d(64), nn.GELU(),\n","            nn.Conv2d(64, 64, 3, padding=1, bias=False),\n","            nn.BatchNorm2d(64), nn.GELU(),\n","        )\n","        self.svd_observer = SVDObserver(64, svd_rank)\n","        self.patch_proj = nn.Conv2d(64, d_model, kernel_size=4, stride=4, bias=False)\n","        self.patch_norm = nn.LayerNorm(d_model)\n","        svd_feat_dim = self.svd_observer.feature_dim\n","        self.svd_to_gamma = nn.Linear(svd_feat_dim, d_model)\n","        self.svd_to_beta = nn.Linear(svd_feat_dim, d_model)\n","        nn.init.normal_(self.svd_to_gamma.weight, std=0.01)\n","        nn.init.ones_(self.svd_to_gamma.bias)\n","        nn.init.normal_(self.svd_to_beta.weight, std=0.01)\n","        nn.init.zeros_(self.svd_to_beta.bias)\n","        self.cls_token = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)\n","        self.pos_embed = nn.Parameter(torch.randn(1, 65, d_model) * 0.02)\n","\n","    def forward(self, x):\n","        B = x.shape[0]\n","        feat = self.conv_frontend(x)\n","        S, Vh, svd_features, novelty = self.svd_observer(feat)\n","        tokens = self.patch_proj(feat).flatten(2).transpose(1, 2)\n","        tokens = self.patch_norm(tokens)\n","        gamma = self.svd_to_gamma(svd_features).unsqueeze(1)\n","        beta = self.svd_to_beta(svd_features).unsqueeze(1)\n","        tokens = gamma * tokens + beta\n","        cls = self.cls_token.expand(B, -1, -1)\n","        tokens = torch.cat([cls, tokens], dim=1)\n","        return tokens + self.pos_embed, {}\n","\n","class TestClassifier(WideRouter):\n","    def __init__(self):\n","        super().__init__('test', strict=False)\n","        self.attach('patch_embed', ConvSVDPatchEmbedding('pe'))\n","        self.attach('transformer', GeometricTransformer(\n","            'geo', d_model=384, n_heads=8, n_layers=8, n_anchors=128,\n","            manifold_dim=128, n_comp=8, d_comp=32, context_dim=128,\n","            quat_dim=64, dropout=0.0, cm_neighbors=3, nce_bank_size=0,\n","            flow_keys=['quat_lite', 'velocity', 'orbital'], flow_fusion='weighted',\n","        ))\n","        self.attach('head', nn.Sequential(\n","            nn.LayerNorm(384), nn.Linear(384, 384), nn.GELU(), nn.Linear(384, 100)))\n","        self.register_tower('transformer')\n","\n","    def forward(self, x, x2=None, targets=None):\n","        if x2 is not None:\n","            B = x.shape[0]\n","            v_cat = torch.cat([x, x2], dim=0)\n","            tokens_cat, _ = self['patch_embed'](v_cat)\n","            output = self['transformer'].forward_paired(tokens_cat[:B], tokens_cat[B:])\n","            if targets is not None:\n","                return self['transformer'].compute_loss(\n","                    output, targets, w_ce=1.0, head=self['head'])\n","            return output\n","        tokens, _ = self['patch_embed'](x)\n","        feat = self['transformer'](tokens)\n","        return self['head'](feat[:, 0])\n","\n","    def precompute(self):\n","        self['transformer'].precompute_cm_gates()\n","\n","    def invalidate(self):\n","        self['transformer'].invalidate_caches()\n","\n","model = TestClassifier()\n","model.network_to(device=device, strict=False)\n","optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)\n","\n","B = 128\n","print(f\"Params: {sum(p.numel() for p in model.parameters()):,}\")\n","\n","# Compile\n","compiled = model.compile(mode='reduce-overhead')\n","print(\"Compiled (reduce-overhead)\")\n","\n","# ── B0: First training step (triggers compilation) ──\n","print(\"\\n=== B0: First training step ===\")\n","x1 = torch.randn(B, 3, 32, 32, device=device)\n","x2 = torch.randn(B, 3, 32, 32, device=device)\n","targets = torch.randint(0, 100, (B,), device=device)\n","\n","model.precompute()\n","torch.compiler.cudagraph_mark_step_begin()\n","loss, ld = compiled(x1, x2, targets=targets)\n","print(f\"  Forward OK — loss={loss.item():.4f}\")\n","\n","optimizer.zero_grad()\n","loss.backward()\n","print(f\"  Backward OK\")\n","\n","optimizer.step()\n","model.invalidate()\n","print(f\"  Step + invalidate OK\")\n","torch.cuda.synchronize()\n","\n","# ── Capture the generated inductor code ──\n","print(\"\\n=== Capturing generated inductor code ===\")\n","gen_files = sorted(glob.glob('/tmp/torchinductor_root/**/*.py', recursive=True),\n","                   key=os.path.getsize, reverse=True)\n","\n","crash_file = None\n","for f in gen_files[:5]:\n","    size = os.path.getsize(f)\n","    lines = sum(1 for _ in open(f))\n","    print(f\"  {os.path.basename(f)}: {lines} lines, {size/1024:.0f} KB\")\n","    if lines > 20000:\n","        crash_file = f\n","\n","if crash_file:\n","    print(f\"\\n  Reading crash file: {os.path.basename(crash_file)}\")\n","    with open(crash_file) as cf:\n","        all_lines = cf.readlines()\n","    total = len(all_lines)\n","    print(f\"  Total lines: {total}\")\n","\n","    # Save the full file\n","    out_path = '/mnt/user-data/outputs/inductor_generated.py'\n","    shutil.copy(crash_file, out_path)\n","    print(f\"  Full file saved to inductor_generated.py\")\n","\n","    # Read around the crash lines\n","    for target_line in [25559, 42025]:\n","        if target_line <= total:\n","            start = max(0, target_line - 30)\n","            end = min(total, target_line + 10)\n","            print(f\"\\n  === Lines {start+1}-{end} (around line {target_line}) ===\")\n","            for i in range(start, end):\n","                marker = \" >>>\" if i == target_line - 1 else \"    \"\n","                print(f\"  {marker} {i+1:>6}: {all_lines[i].rstrip()}\")\n","\n","    # Find all 'partition_' function definitions\n","    print(f\"\\n  === Partition functions ===\")\n","    for i, line in enumerate(all_lines):\n","        if 'def partition_' in line or 'def call(' in line:\n","            print(f\"    {i+1:>6}: {line.rstrip()}\")\n","\n","    # Find what ops are near the crash\n","    print(f\"\\n  === .copy_ and buffer ops near crash ===\")\n","    for target in [25559]:\n","        start = max(0, target - 100)\n","        end = min(total, target + 10)\n","        for i in range(start, end):\n","            line = all_lines[i]\n","            if any(kw in line for kw in ['copy_', 'register_buffer', '_cached',\n","                                          '_cache_warm', 'solve', 'det', 'eigh',\n","                                          'mutation', 'DeviceCopy', 'offset']):\n","                print(f\"    {i+1:>6}: {line.rstrip()}\")\n","else:\n","    print(\"  No large generated file found\")\n","\n","# ── B1: Second training step (triggers crash) ──\n","print(\"\\n=== B1: Second training step (expect crash here) ===\")\n","x1 = torch.randn(B, 3, 32, 32, device=device)\n","x2 = torch.randn(B, 3, 32, 32, device=device)\n","targets = torch.randint(0, 100, (B,), device=device)\n","\n","model.precompute()\n","torch.compiler.cudagraph_mark_step_begin()\n","\n","_crash_log.write(\"B1 forward\\n\"); _crash_log.flush()\n","loss, ld = compiled(x1, x2, targets=targets)\n","print(f\"  Forward OK — loss={loss.item():.4f}\")\n","\n","_crash_log.write(\"B1 backward\\n\"); _crash_log.flush()\n","optimizer.zero_grad()\n","loss.backward()\n","print(f\"  Backward OK\")\n","\n","_crash_log.write(\"B1 step\\n\"); _crash_log.flush()\n","optimizer.step()\n","model.invalidate()\n","print(f\"  Step + invalidate OK\")\n","\n","print(\"\\n=== B1 survived! Trying B2... ===\")\n","x1 = torch.randn(B, 3, 32, 32, device=device)\n","x2 = torch.randn(B, 3, 32, 32, device=device)\n","targets = torch.randint(0, 100, (B,), device=device)\n","\n","model.precompute()\n","torch.compiler.cudagraph_mark_step_begin()\n","loss, ld = compiled(x1, x2, targets=targets)\n","optimizer.zero_grad()\n","loss.backward()\n","optimizer.step()\n","model.invalidate()\n","print(f\"  B2 OK — loss={loss.item():.4f}\")\n","\n","print(\"\\nAll steps passed.\")"],"metadata":{"colab":{"base_uri":"https://localhost:8080/","height":1000},"id":"mgWPLzHfCrxa","executionInfo":{"status":"error","timestamp":1775172883685,"user_tz":420,"elapsed":411609,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"30776152-bc08-428c-c6fc-3590b98deafc"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["Cleared inductor cache\n","Params: 42,867,988\n","Compiled (reduce-overhead)\n","\n","=== B0: First training step ===\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n","/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:7627: UserWarning: \n","Online softmax is disabled on the fly since Inductor decides to\n","split the reduction. Cut an issue to PyTorch if this is an\n","important use case and you want to speed it up with online\n","softmax.\n","\n","  warnings.warn(\n"]},{"output_type":"stream","name":"stdout","text":["  Forward OK — loss=21.2552\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:7627: UserWarning: \n","Online softmax is disabled on the fly since Inductor decides to\n","split the reduction. Cut an issue to PyTorch if this is an\n","important use case and you want to speed it up with online\n","softmax.\n","\n","  warnings.warn(\n"]},{"output_type":"stream","name":"stdout","text":["  Backward OK\n","  Step + invalidate OK\n","\n","=== Capturing generated inductor code ===\n","  c7itseqjueaeuxnhckyh4l3j6ltbojztkwxszd7qsruiulignqcg.py: 59891 lines, 5118 KB\n","  cyiy7pduxhxebl2tzlt754pdelyn6ab5cpj5ylo5rzadxndibsjd.py: 44263 lines, 3945 KB\n","  cr7uxvfrxg2h5bxmbo26rs2k6chluxxfvwlhh7iutte4ffncdmgq.py: 1350 lines, 100 KB\n","  c7sjofxkxjestfz7atkeh242dyueemztfukcepe24bc7dxzmmjj6.py: 1128 lines, 87 KB\n","  c6zazyknob4x5mhymze2wvetwcrwwse3b233rdm4alacbj5hwado.py: 375 lines, 25 KB\n","\n","  Reading crash file: cyiy7pduxhxebl2tzlt754pdelyn6ab5cpj5ylo5rzadxndibsjd.py\n","  Total lines: 44263\n"]},{"output_type":"error","ename":"FileNotFoundError","evalue":"[Errno 2] No such file or directory: '/mnt/user-data/outputs/inductor_generated.py'","traceback":["\u001b[0;31m---------------------------------------------------------------------------\u001b[0m","\u001b[0;31mFileNotFoundError\u001b[0m                         Traceback (most recent call last)","\u001b[0;32m/tmp/ipykernel_76804/911878686.py\u001b[0m in \u001b[0;36m<cell line: 0>\u001b[0;34m()\u001b[0m\n\u001b[1;32m    150\u001b[0m     \u001b[0;31m# Save the full file\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    151\u001b[0m     \u001b[0mout_path\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m'/mnt/user-data/outputs/inductor_generated.py'\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 152\u001b[0;31m     \u001b[0mshutil\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcopy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcrash_file\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mout_path\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    153\u001b[0m     \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf\"  Full file saved to inductor_generated.py\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    154\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/lib/python3.12/shutil.py\u001b[0m in \u001b[0;36mcopy\u001b[0;34m(src, dst, follow_symlinks)\u001b[0m\n\u001b[1;32m    433\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mos\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0misdir\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdst\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    434\u001b[0m         \u001b[0mdst\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mos\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mjoin\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdst\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mos\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbasename\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msrc\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 435\u001b[0;31m     \u001b[0mcopyfile\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msrc\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdst\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfollow_symlinks\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mfollow_symlinks\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    436\u001b[0m     \u001b[0mcopymode\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msrc\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdst\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfollow_symlinks\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mfollow_symlinks\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    437\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mdst\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/lib/python3.12/shutil.py\u001b[0m in \u001b[0;36mcopyfile\u001b[0;34m(src, dst, follow_symlinks)\u001b[0m\n\u001b[1;32m    260\u001b[0m         \u001b[0;32mwith\u001b[0m \u001b[0mopen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msrc\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'rb'\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mfsrc\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    261\u001b[0m             \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 262\u001b[0;31m                 \u001b[0;32mwith\u001b[0m \u001b[0mopen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdst\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'wb'\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mfdst\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    263\u001b[0m                     \u001b[0;31m# macOS\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    264\u001b[0m                     \u001b[0;32mif\u001b[0m \u001b[0m_HAS_FCOPYFILE\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;31mFileNotFoundError\u001b[0m: [Errno 2] No such file or directory: '/mnt/user-data/outputs/inductor_generated.py'"]}]},{"cell_type":"code","source":["import glob, os\n","\n","gen_files = sorted(glob.glob('/tmp/torchinductor_root/**/*.py', recursive=True),\n","                   key=os.path.getsize, reverse=True)\n","\n","for f in gen_files[:3]:\n","    lines = open(f).readlines()\n","    n = len(lines)\n","    print(f\"\\n{'='*70}\")\n","    print(f\"FILE: {os.path.basename(f)} — {n} lines\")\n","\n","    # Find partitions\n","    for i, line in enumerate(lines):\n","        if 'def partition_' in line or 'def call(' in line:\n","            print(f\"  {i+1:>6}: {line.rstrip()}\")\n","\n","    # Print first 80 lines of partition_0\n","    for i, line in enumerate(lines):\n","        if 'def partition_0(' in line:\n","            end = min(n, i + 80)\n","            print(f\"\\n  --- partition_0 (lines {i+1}-{end}) ---\")\n","            for j in range(i, end):\n","                print(f\"  {j+1:>6}: {lines[j].rstrip()}\")\n","            break"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"HF4QWNrYkFvm","executionInfo":{"status":"ok","timestamp":1775173133225,"user_tz":420,"elapsed":150,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"d9acd6cc-0411-4bbe-f7ab-96a83cc9e95c"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["\n","======================================================================\n","FILE: c7itseqjueaeuxnhckyh4l3j6ltbojztkwxszd7qsruiulignqcg.py — 59891 lines\n","   47561: def partition_0(args):\n","   58566:     def call(self, args):\n","\n","  --- partition_0 (lines 47561-47640) ---\n","   47561: def partition_0(args):\n","   47562:     primals_2, primals_1, primals_3, primals_5, primals_6, primals_7, primals_8, primals_9, primals_11, primals_12, primals_13, primals_14, primals_15, primals_17, primals_21, primals_20, primals_22, primals_24, primals_18, primals_19, primals_23, primals_25, primals_27, primals_26, primals_30, primals_28, primals_29, primals_31, primals_32, primals_33, primals_34, primals_37, primals_36, primals_39, primals_38, primals_41, primals_40, primals_43, primals_42, primals_44, primals_46, primals_45, primals_48, primals_47, primals_35, primals_49, primals_50, primals_52, primals_51, primals_54, primals_53, primals_55, primals_56, primals_58, primals_57, primals_60, primals_59, primals_61, primals_62, primals_64, primals_63, primals_66, primals_65, primals_67, primals_68, primals_70, primals_69, primals_72, primals_71, primals_73, primals_74, primals_76, primals_75, primals_78, primals_77, primals_79, primals_80, primals_82, primals_81, primals_84, primals_83, primals_85, primals_86, primals_88, primals_87, primals_90, primals_89, primals_91, primals_92, primals_94, primals_93, primals_96, primals_95, primals_97, primals_98, primals_102, primals_101, primals_103, primals_104, primals_106, primals_105, primals_107, primals_108, primals_110, primals_109, primals_111, primals_112, primals_114, primals_113, primals_115, primals_116, primals_118, primals_117, primals_119, primals_120, primals_122, primals_121, primals_124, primals_123, primals_134, primals_133, primals_136, primals_135, primals_137, primals_138, primals_140, primals_139, primals_142, primals_141, primals_143, primals_144, primals_145, primals_146, primals_148, primals_147, primals_125, primals_126, primals_149, primals_150, primals_128, primals_127, primals_129, primals_152, primals_151, primals_154, primals_153, primals_155, primals_156, primals_157, primals_158, primals_159, primals_160, primals_161, primals_130, primals_131, primals_132, primals_163, primals_162, primals_165, primals_164, primals_167, primals_166, primals_169, primals_168, primals_170, primals_172, primals_171, primals_173, primals_174, primals_176, primals_175, primals_178, primals_177, primals_181, primals_183, primals_182, primals_186, primals_184, primals_185, primals_187, primals_188, primals_189, primals_190, primals_193, primals_192, primals_195, primals_194, primals_197, primals_196, primals_199, primals_198, primals_200, primals_202, primals_201, primals_204, primals_203, primals_191, primals_205, primals_206, primals_208, primals_207, primals_210, primals_209, primals_211, primals_212, primals_214, primals_213, primals_216, primals_215, primals_217, primals_218, primals_220, primals_219, primals_222, primals_221, primals_223, primals_224, primals_226, primals_225, primals_228, primals_227, primals_229, primals_230, primals_232, primals_231, primals_234, primals_233, primals_235, primals_236, primals_238, primals_237, primals_240, primals_239, primals_241, primals_242, primals_244, primals_243, primals_246, primals_245, primals_247, primals_248, primals_250, primals_249, primals_252, primals_251, primals_253, primals_254, primals_338, primals_337, primals_179, primals_180, primals_339, primals_340, primals_258, primals_257, primals_259, primals_260, primals_262, primals_261, primals_263, primals_264, primals_266, primals_265, primals_267, primals_268, primals_270, primals_269, primals_271, primals_272, primals_274, primals_273, primals_275, primals_276, primals_278, primals_277, primals_279, primals_280, primals_282, primals_281, primals_283, primals_294, primals_293, primals_296, primals_295, primals_297, primals_298, primals_300, primals_299, primals_302, primals_301, primals_303, primals_304, primals_305, primals_306, primals_307, primals_284, primals_308, primals_285, primals_286, primals_309, primals_310, primals_288, primals_287, primals_289, primals_312, primals_311, primals_314, primals_313, primals_315, primals_316, primals_317, primals_318, primals_319, primals_320, primals_321, primals_290, primals_291, primals_292, primals_323, primals_322, primals_325, primals_324, primals_327, primals_326, primals_329, primals_328, primals_330, primals_332, primals_331, primals_333, primals_334, primals_336, primals_335, primals_341, primals_343, primals_342, primals_346, primals_344, primals_345, primals_347, primals_348, primals_349, primals_350, primals_353, primals_352, primals_355, primals_354, primals_357, primals_356, primals_359, primals_358, primals_360, primals_362, primals_361, primals_364, primals_363, primals_351, primals_365, primals_366, primals_368, primals_367, primals_370, primals_369, primals_371, primals_372, primals_374, primals_373, primals_376, primals_375, primals_377, primals_378, primals_380, primals_379, primals_382, primals_381, primals_383, primals_384, primals_386, primals_385, primals_388, primals_387, primals_389, primals_390, primals_392, primals_391, primals_394, primals_393, primals_395, primals_396, primals_398, primals_397, primals_400, primals_399, primals_401, primals_402, primals_404, primals_403, primals_406, primals_405, primals_407, primals_408, primals_410, primals_409, primals_412, primals_411, primals_413, primals_414, primals_418, primals_417, primals_419, primals_420, primals_422, primals_421, primals_423, primals_424, primals_426, primals_425, primals_427, primals_428, primals_430, primals_429, primals_431, primals_432, primals_434, primals_433, primals_435, primals_436, primals_438, primals_437, primals_439, primals_440, primals_442, primals_441, primals_443, primals_454, primals_453, primals_456, primals_455, primals_457, primals_458, primals_460, primals_459, primals_462, primals_461, primals_463, primals_464, primals_465, primals_466, primals_467, primals_444, primals_468, primals_445, primals_446, primals_469, primals_470, primals_448, primals_447, primals_449, primals_472, primals_471, primals_474, primals_473, primals_475, primals_476, primals_477, primals_478, primals_479, primals_480, primals_481, primals_450, primals_451, primals_452, primals_483, primals_482, primals_485, primals_484, primals_487, primals_486, primals_489, primals_488, primals_490, primals_492, primals_491, primals_493, primals_494, primals_496, primals_495, primals_498, primals_497, primals_501, primals_503, primals_502, primals_506, primals_504, primals_505, primals_507, primals_508, primals_509, primals_510, primals_513, primals_512, primals_515, primals_514, primals_517, primals_516, primals_519, primals_518, primals_520, primals_522, primals_521, primals_524, primals_523, primals_511, primals_525, primals_526, primals_528, primals_527, primals_530, primals_529, primals_531, primals_532, primals_534, primals_533, primals_536, primals_535, primals_537, primals_538, primals_540, primals_539, primals_542, primals_541, primals_543, primals_544, primals_546, primals_545, primals_548, primals_547, primals_549, primals_550, primals_552, primals_551, primals_554, primals_553, primals_555, primals_556, primals_558, primals_557, primals_560, primals_559, primals_561, primals_562, primals_564, primals_563, primals_566, primals_565, primals_567, primals_568, primals_570, primals_569, primals_572, primals_571, primals_573, primals_574, primals_658, primals_657, primals_499, primals_500, primals_659, primals_660, primals_578, primals_577, primals_579, primals_580, primals_582, primals_581, primals_583, primals_584, primals_586, primals_585, primals_587, primals_588, primals_590, primals_589, primals_591, primals_592, primals_594, primals_593, primals_595, primals_596, primals_598, primals_597, primals_599, primals_600, primals_602, primals_601, primals_603, primals_614, primals_613, primals_616, primals_615, primals_617, primals_618, primals_620, primals_619, primals_622, primals_621, primals_623, primals_624, primals_625, primals_626, primals_627, primals_604, primals_628, primals_605, primals_606, primals_629, primals_630, primals_608, primals_607, primals_609, primals_632, primals_631, primals_634, primals_633, primals_635, primals_636, primals_637, primals_638, primals_639, primals_640, primals_641, primals_610, primals_611, primals_612, primals_643, primals_642, primals_645, primals_644, primals_647, primals_646, primals_649, primals_648, primals_650, primals_652, primals_651, primals_653, primals_654, primals_656, primals_655, primals_661, primals_663, primals_662, primals_666, primals_664, primals_665, primals_667, primals_668, primals_669, primals_670, primals_673, primals_672, primals_675, primals_674, primals_677, primals_676, primals_679, primals_678, primals_680, primals_682, primals_681, primals_684, primals_683, primals_671, primals_685, primals_686, primals_688, primals_687, primals_690, primals_689, primals_691, primals_692, primals_694, primals_693, primals_696, primals_695, primals_697, primals_698, primals_700, primals_699, primals_702, primals_701, primals_703, primals_704, primals_706, primals_705, primals_708, primals_707, primals_709, primals_710, primals_712, primals_711, primals_714, primals_713, primals_715, primals_716, primals_718, primals_717, primals_720, primals_719, primals_721, primals_722, primals_724, primals_723, primals_726, primals_725, primals_727, primals_728, primals_730, primals_729, primals_732, primals_731, primals_733, primals_734, primals_738, primals_737, primals_739, primals_740, primals_742, primals_741, primals_743, primals_744, primals_746, primals_745, primals_747, primals_748, primals_750, primals_749, primals_751, primals_752, primals_754, primals_753, primals_755, primals_756, primals_758, primals_757, primals_759, primals_760, primals_762, primals_761, primals_763, primals_774, primals_773, primals_776, primals_775, primals_777, primals_778, primals_780, primals_779, primals_782, primals_781, primals_783, primals_784, primals_785, primals_786, primals_787, primals_764, primals_788, primals_765, primals_766, primals_789, primals_790, primals_768, primals_767, primals_769, primals_792, primals_791, primals_794, primals_793, primals_795, primals_796, primals_797, primals_798, primals_799, primals_800, primals_801, primals_770, primals_771, primals_772, primals_803, primals_802, primals_805, primals_804, primals_807, primals_806, primals_809, primals_808, primals_810, primals_812, primals_811, primals_813, primals_814, primals_816, primals_815, primals_818, primals_817, primals_821, primals_823, primals_822, primals_826, primals_824, primals_825, primals_827, primals_828, primals_829, primals_830, primals_833, primals_832, primals_835, primals_834, primals_837, primals_836, primals_839, primals_838, primals_840, primals_842, primals_841, primals_844, primals_843, primals_831, primals_845, primals_846, primals_848, primals_847, primals_850, primals_849, primals_851, primals_852, primals_854, primals_853, primals_856, primals_855, primals_857, primals_858, primals_860, primals_859, primals_862, primals_861, primals_863, primals_864, primals_866, primals_865, primals_868, primals_867, primals_869, primals_870, primals_872, primals_871, primals_874, primals_873, primals_875, primals_876, primals_878, primals_877, primals_880, primals_879, primals_881, primals_882, primals_884, primals_883, primals_886, primals_885, primals_887, primals_888, primals_890, primals_889, primals_892, primals_891, primals_893, primals_894, primals_978, primals_977, primals_819, primals_820, primals_979, primals_980, primals_898, primals_897, primals_899, primals_900, primals_902, primals_901, primals_903, primals_904, primals_906, primals_905, primals_907, primals_908, primals_910, primals_909, primals_911, primals_912, primals_914, primals_913, primals_915, primals_916, primals_918, primals_917, primals_919, primals_920, primals_922, primals_921, primals_923, primals_934, primals_933, primals_936, primals_935, primals_937, primals_938, primals_940, primals_939, primals_942, primals_941, primals_943, primals_944, primals_945, primals_946, primals_947, primals_924, primals_948, primals_925, primals_926, primals_949, primals_950, primals_928, primals_927, primals_929, primals_952, primals_951, primals_954, primals_953, primals_955, primals_956, primals_957, primals_958, primals_959, primals_960, primals_961, primals_930, primals_931, primals_932, primals_963, primals_962, primals_965, primals_964, primals_967, primals_966, primals_969, primals_968, primals_970, primals_972, primals_971, primals_973, primals_974, primals_976, primals_975, primals_981, primals_983, primals_982, primals_986, primals_984, primals_985, primals_987, primals_988, primals_989, primals_990, primals_993, primals_992, primals_995, primals_994, primals_997, primals_996, primals_999, primals_998, primals_1000, primals_1002, primals_1001, primals_1004, primals_1003, primals_991, primals_1005, primals_1006, primals_1008, primals_1007, primals_1010, primals_1009, primals_1011, primals_1012, primals_1014, primals_1013, primals_1016, primals_1015, primals_1017, primals_1018, primals_1020, primals_1019, primals_1022, primals_1021, primals_1023, primals_1024, primals_1026, primals_1025, primals_1028, primals_1027, primals_1029, primals_1030, primals_1032, primals_1031, primals_1034, primals_1033, primals_1035, primals_1036, primals_1038, primals_1037, primals_1040, primals_1039, primals_1041, primals_1042, primals_1044, primals_1043, primals_1046, primals_1045, primals_1047, primals_1048, primals_1050, primals_1049, primals_1052, primals_1051, primals_1053, primals_1054, primals_1058, primals_1057, primals_1059, primals_1060, primals_1062, primals_1061, primals_1063, primals_1064, primals_1066, primals_1065, primals_1067, primals_1068, primals_1070, primals_1069, primals_1071, primals_1072, primals_1074, primals_1073, primals_1075, primals_1076, primals_1078, primals_1077, primals_1079, primals_1080, primals_1082, primals_1081, primals_1083, primals_1094, primals_1093, primals_1096, primals_1095, primals_1097, primals_1098, primals_1100, primals_1099, primals_1102, primals_1101, primals_1103, primals_1104, primals_1105, primals_1106, primals_1107, primals_1084, primals_1108, primals_1085, primals_1086, primals_1109, primals_1110, primals_1088, primals_1087, primals_1089, primals_1112, primals_1111, primals_1114, primals_1113, primals_1115, primals_1116, primals_1117, primals_1118, primals_1119, primals_1120, primals_1121, primals_1090, primals_1091, primals_1092, primals_1123, primals_1122, primals_1125, primals_1124, primals_1127, primals_1126, primals_1129, primals_1128, primals_1130, primals_1132, primals_1131, primals_1133, primals_1134, primals_1136, primals_1135, primals_1138, primals_1137, primals_1139, primals_1140, primals_1141, primals_1143, primals_1142, primals_1146, primals_1144, primals_1145, primals_1147, primals_1148, primals_1149, primals_1150, primals_1153, primals_1152, primals_1155, primals_1154, primals_1157, primals_1156, primals_1159, primals_1158, primals_1160, primals_1162, primals_1161, primals_1164, primals_1163, primals_1151, primals_1165, primals_1166, primals_1168, primals_1167, primals_1170, primals_1169, primals_1171, primals_1172, primals_1174, primals_1173, primals_1176, primals_1175, primals_1177, primals_1178, primals_1180, primals_1179, primals_1182, primals_1181, primals_1183, primals_1184, primals_1186, primals_1185, primals_1188, primals_1187, primals_1189, primals_1190, primals_1192, primals_1191, primals_1194, primals_1193, primals_1195, primals_1196, primals_1198, primals_1197, primals_1200, primals_1199, primals_1201, primals_1202, primals_1204, primals_1203, primals_1206, primals_1205, primals_1207, primals_1208, primals_1210, primals_1209, primals_1212, primals_1211, primals_1213, primals_1214, primals_1216, primals_1215, primals_1218, primals_1217, primals_1219, primals_1220, primals_1222, primals_1221, primals_1223, primals_1224, primals_1226, primals_1225, primals_1227, primals_1228, primals_1230, primals_1229, primals_1231, primals_1232, primals_1234, primals_1233, primals_1235, primals_1236, primals_1238, primals_1237, primals_1239, primals_1240, primals_1242, primals_1241, primals_1243, primals_1254, primals_1253, primals_1256, primals_1255, primals_1257, primals_1258, primals_1260, primals_1259, primals_1262, primals_1261, primals_1263, primals_1264, primals_1265, primals_1266, primals_1267, primals_1244, primals_1268, primals_1245, primals_1246, primals_1269, primals_1270, primals_1248, primals_1247, primals_1249, primals_1272, primals_1271, primals_1274, primals_1273, primals_1275, primals_1276, primals_1277, primals_1278, primals_1279, primals_1280, primals_1281, primals_1250, primals_1251, primals_1252, primals_1283, primals_1282, primals_1285, primals_1284, primals_1287, primals_1286, primals_1289, primals_1288, primals_1290, primals_1292, primals_1291, primals_1293, primals_1294, primals_1296, primals_1295, primals_1301, primals_1302, primals_4, primals_10 = args\n","   47563:     args.clear()\n","   47564:     assert_size_stride(primals_2, (128, 3, 32, 32), (3072, 1024, 32, 1))\n","   47565:     assert_size_stride(primals_1, (128, 3, 32, 32), (3072, 1024, 32, 1))\n","   47566:     assert_size_stride(primals_3, (64, 3, 3, 3), (27, 9, 3, 1))\n","   47567:     assert_size_stride(primals_5, (64, ), (1, ))\n","   47568:     assert_size_stride(primals_6, (64, ), (1, ))\n","   47569:     assert_size_stride(primals_7, (64, ), (1, ))\n","   47570:     assert_size_stride(primals_8, (64, ), (1, ))\n","   47571:     assert_size_stride(primals_9, (64, 64, 3, 3), (576, 9, 3, 1))\n","   47572:     assert_size_stride(primals_11, (64, ), (1, ))\n","   47573:     assert_size_stride(primals_12, (64, ), (1, ))\n","   47574:     assert_size_stride(primals_13, (64, ), (1, ))\n","   47575:     assert_size_stride(primals_14, (64, ), (1, ))\n","   47576:     assert_size_stride(primals_15, (12, 64, 1, 1), (64, 1, 1, 1))\n","   47577:     assert_size_stride(primals_17, (384, 64, 4, 4), (1024, 16, 4, 1))\n","   47578:     assert_size_stride(primals_21, (384, ), (1, ))\n","   47579:     assert_size_stride(primals_20, (384, 26), (26, 1))\n","   47580:     assert_size_stride(primals_22, (384, 26), (26, 1))\n","   47581:     assert_size_stride(primals_24, (1, 1, 384), (384, 384, 1))\n","   47582:     assert_size_stride(primals_18, (384, ), (1, ))\n","   47583:     assert_size_stride(primals_19, (384, ), (1, ))\n","   47584:     assert_size_stride(primals_23, (384, ), (1, ))\n","   47585:     assert_size_stride(primals_25, (1, 65, 384), (24960, 384, 1))\n","   47586:     assert_size_stride(primals_27, (128, ), (1, ))\n","   47587:     assert_size_stride(primals_26, (128, 384), (384, 1))\n","   47588:     assert_size_stride(primals_30, (128, 128), (128, 1))\n","   47589:     assert_size_stride(primals_28, (128, ), (1, ))\n","   47590:     assert_size_stride(primals_29, (128, ), (1, ))\n","   47591:     assert_size_stride(primals_31, (128, ), (1, ))\n","   47592:     assert_size_stride(primals_32, (16, 2), (2, 1))\n","   47593:     assert_size_stride(primals_33, (16, ), (1, ))\n","   47594:     assert_size_stride(primals_34, (1, 16), (16, 1))\n","   47595:     assert_size_stride(primals_37, (64, ), (1, ))\n","   47596:     assert_size_stride(primals_36, (64, 128), (128, 1))\n","   47597:     assert_size_stride(primals_39, (4, ), (1, ))\n","   47598:     assert_size_stride(primals_38, (4, 64), (64, 1))\n","   47599:     assert_size_stride(primals_41, (128, ), (1, ))\n","   47600:     assert_size_stride(primals_40, (128, 256), (256, 1))\n","   47601:     assert_size_stride(primals_43, (128, ), (1, ))\n","   47602:     assert_size_stride(primals_42, (128, 128), (128, 1))\n","   47603:     assert_size_stride(primals_44, (), ())\n","   47604:     assert_size_stride(primals_46, (12, ), (1, ))\n","   47605:     assert_size_stride(primals_45, (12, 128), (128, 1))\n","   47606:     assert_size_stride(primals_48, (128, ), (1, ))\n","   47607:     assert_size_stride(primals_47, (128, 12), (12, 1))\n","   47608:     assert_size_stride(primals_35, (1, ), (1, ))\n","   47609:     assert_size_stride(primals_49, (3, ), (1, ))\n","   47610:     assert_size_stride(primals_50, (), ())\n","   47611:     assert_size_stride(primals_52, (64, ), (1, ))\n","   47612:     assert_size_stride(primals_51, (64, 16), (16, 1))\n","   47613:     assert_size_stride(primals_54, (32, ), (1, ))\n","   47614:     assert_size_stride(primals_53, (32, 64), (64, 1))\n","   47615:     assert_size_stride(primals_55, (32, ), (1, ))\n","   47616:     assert_size_stride(primals_56, (32, ), (1, ))\n","   47617:     assert_size_stride(primals_58, (64, ), (1, ))\n","   47618:     assert_size_stride(primals_57, (64, 16), (16, 1))\n","   47619:     assert_size_stride(primals_60, (32, ), (1, ))\n","   47620:     assert_size_stride(primals_59, (32, 64), (64, 1))\n","   47621:     assert_size_stride(primals_61, (32, ), (1, ))\n","   47622:     assert_size_stride(primals_62, (32, ), (1, ))\n","   47623:     assert_size_stride(primals_64, (64, ), (1, ))\n","   47624:     assert_size_stride(primals_63, (64, 16), (16, 1))\n","   47625:     assert_size_stride(primals_66, (32, ), (1, ))\n","   47626:     assert_size_stride(primals_65, (32, 64), (64, 1))\n","   47627:     assert_size_stride(primals_67, (32, ), (1, ))\n","   47628:     assert_size_stride(primals_68, (32, ), (1, ))\n","   47629:     assert_size_stride(primals_70, (64, ), (1, ))\n","   47630:     assert_size_stride(primals_69, (64, 16), (16, 1))\n","   47631:     assert_size_stride(primals_72, (32, ), (1, ))\n","   47632:     assert_size_stride(primals_71, (32, 64), (64, 1))\n","   47633:     assert_size_stride(primals_73, (32, ), (1, ))\n","   47634:     assert_size_stride(primals_74, (32, ), (1, ))\n","   47635:     assert_size_stride(primals_76, (64, ), (1, ))\n","   47636:     assert_size_stride(primals_75, (64, 16), (16, 1))\n","   47637:     assert_size_stride(primals_78, (32, ), (1, ))\n","   47638:     assert_size_stride(primals_77, (32, 64), (64, 1))\n","   47639:     assert_size_stride(primals_79, (32, ), (1, ))\n","   47640:     assert_size_stride(primals_80, (32, ), (1, ))\n","\n","======================================================================\n","FILE: cyiy7pduxhxebl2tzlt754pdelyn6ab5cpj5ylo5rzadxndibsjd.py — 44263 lines\n","   22418: def partition_0(args):\n","   41747:     def call(self, args):\n","\n","  --- partition_0 (lines 22418-22497) ---\n","   22418: def partition_0(args):\n","   22419:     tangents_7, primals_1301, mul_16366, div_1355, mm_55, addmm_430, getitem_421, rsqrt_168, primals_1293, primals_1294, addmm_431, primals_1295, view_709, primals_1291, cat_98, primals_1290, cat_97, primals_1288, view_702, primals_1286, view_701, primals_1284, mm_62, primals_1282, add_13221, view_697, primals_1281, primals_1279, mul_16337, div_1366, primals_1277, view_695, addmm_422, addmm_423, primals_1271, primals_1269, mul_16331, div_1367, primals_1275, view_665, primals_1273, view_689, primals_1267, view_687, div_1352, permute_657, permute_658, permute_659, primals_1265, primals_1263, addmm_417, addmm_418, primals_1261, primals_1259, primals_1257, addmm_414, addmm_415, primals_1255, primals_1253, primals_1251, mul_16326, div_1368, primals_1249, view_661, addmm_412, primals_1247, primals_1245, mul_16321, div_1369, view_659, primals_1243, view_657, view_654, view_655, view_656, getitem_408, getitem_409, getitem_410, getitem_411, view_648, primals_1242, primals_1239, addmm_410, getitem_407, rsqrt_163, primals_1237, primals_1223, addmm_406, getitem_399, rsqrt_159, primals_1221, primals_1235, addmm_409, getitem_405, rsqrt_162, primals_1233, primals_1231, addmm_408, getitem_403, rsqrt_161, primals_1229, tangents_3, tangents_4, primals_1215, primals_1213, addmm_403, getitem_395, rsqrt_157, primals_1211, addmm_402, primals_1209, primals_1207, addmm_401, getitem_393, rsqrt_156, primals_1205, addmm_400, primals_1203, primals_1201, addmm_399, getitem_391, rsqrt_155, primals_1199, addmm_398, primals_1197, primals_1195, addmm_397, getitem_389, rsqrt_154, primals_1193, addmm_396, primals_1191, primals_1189, addmm_395, getitem_387, rsqrt_153, primals_1187, addmm_394, primals_1185, primals_1183, addmm_393, getitem_385, rsqrt_152, primals_1181, addmm_392, primals_1179, primals_1177, addmm_391, getitem_383, rsqrt_151, primals_1175, addmm_390, primals_1173, primals_1171, addmm_389, getitem_381, rsqrt_150, primals_1169, addmm_388, primals_1167, full_default_875, sigmoid_28, primals_1166, tangents_8, tangents_9, add_13179, primals_1165, div_1204, div_1212, pow_269, where_3326, primals_1160, addmm_385, sum_261, div_1211, pow_267, primals_1158, addmm_384, primals_1156, div_1210, pow_265, view_637, slice_scatter_7, index_56, index_57, index_58, index_59, addmm_383, index_63, index_61, primals_1154, addmm_382, primals_1152, primals_1219, addmm_405, getitem_397, rsqrt_158, primals_1217, primals_1150, addmm_default, primals_1148, tangents_5, div_1206, tangents_6, addmm_387, tangents_1, tangents_2, addmm_379, getitem_373, rsqrt_149, primals_1144, primals_1145, pow_255, primals_1142, cat_96, sum_281, primals_1227, addmm_407, getitem_401, rsqrt_160, primals_1225, view_646, cat_95, cat_94, cat_92, pow_284, slice_786, pow_283, slice_785, pow_282, slice_784, pow_281, slice_783, pow_280, slice_782, pow_279, slice_781, pow_278, slice_780, pow_277, slice_779, sub_2533, cat_91, sub_2275, primals_1163, view_643, getitem_379, sub_2531, mm_default_281, mm_59, mm_default_280, squeeze_dim_554, squeeze_dim_537, squeeze_dim_534, squeeze_dim_533, squeeze_dim_530, squeeze_dim_529, squeeze_dim_526, squeeze_dim_524, squeeze_43, pow_275, unsqueeze_1128, unsqueeze_1084, select_scatter_943, full_default_441, full_443, add_13117, sub_2526, select_scatter_944, add_13141, sub_2527, add_13142, add_13144, add_13146, add_13148, add_13150, add_13152, add_13154, add_13156, add_13158, gt_1122, where_3304, div_1389, add_13162, add_13160, gt_1121, where_3302, select_scatter_1033, where_3300, gt_1120, full_default_439, convert_element_type_79, full_default_443, add_11793, add_11792, add_11791, add_11790, add_11789, add_11788, add_11787, add_11786, add_11785, add_11784, add_11783, add_11794, pow_271, addmm_386, primals_1161, mul_14553, cat_90, mul_14544, sub_2278, mean_38, pow_261, pow_257, primals_1146, view_640, view_638, view_633, primals_1141, view_622, primals_1139, addmm_378, getitem_371, rsqrt_148, primals_1140, cat_80, mm_47, addmm_376, getitem_369, rsqrt_147, primals_1133, primals_1134, addmm_377, primals_1135, view_619, primals_1131, cat_86, primals_1130, cat_85, primals_1128, view_612, primals_1126, view_611, primals_1124, mm_54, primals_1122, add_11741, view_607, primals_1121, primals_1119, mul_14508, div_1460, primals_1117, view_605, addmm_368, addmm_369, primals_1111, primals_1109, mul_14502, div_1461, primals_1115, view_575, primals_1113, view_599, primals_1107, view_597, div_1200, permute_953, permute_954, permute_955, primals_1105, primals_1103, addmm_363, addmm_364, primals_1101, primals_1099, primals_1097, addmm_360, addmm_361, primals_1095, primals_1093, primals_1091, mul_14497, div_1462, primals_1089, view_571, addmm_358, primals_1087, primals_1085, mul_14492, div_1463, view_569, primals_1083, view_567, view_564, view_565, view_566, getitem_356, getitem_357, getitem_358, getitem_359, view_558, primals_1082, primals_1079, addmm_356, getitem_355, rsqrt_142, primals_1077, primals_1063, addmm_352, getitem_347, rsqrt_138, primals_1061, primals_1137, primals_1029, addmm_341, getitem_335, rsqrt_132, primals_1027, addmm_340, primals_1025, primals_1023, addmm_339, getitem_333, rsqrt_131, primals_1021, addmm_338, primals_1019, primals_1017, addmm_337, getitem_331, rsqrt_130, primals_1015, addmm_336, primals_1013, primals_1011, addmm_335, getitem_329, rsqrt_129, primals_1009, addmm_334, primals_1007, primals_1053, addmm_349, getitem_343, rsqrt_136, primals_1051, addmm_348, primals_1049, primals_1047, addmm_347, getitem_341, rsqrt_135, primals_1045, addmm_346, primals_1043, primals_1041, addmm_345, getitem_339, rsqrt_134, primals_1039, addmm_344, primals_1037, primals_1035, addmm_343, getitem_337, rsqrt_133, primals_1033, addmm_342, primals_1031, sigmoid_24, primals_1006, primals_1075, addmm_355, getitem_353, rsqrt_141, primals_1073, primals_1005, div_1052, div_1060, pow_234, where_3361, primals_1000, addmm_331, sum_228, div_1059, pow_232, primals_998, addmm_330, primals_996, div_1058, pow_230, slice_scatter_6, index_48, index_49, index_50, index_51, addmm_329, index_55, index_53, primals_994, addmm_328, primals_992, primals_1071, addmm_354, getitem_351, rsqrt_140, primals_1069, primals_1059, addmm_351, getitem_345, rsqrt_137, primals_1057, addmm_333, add_11699, primals_990, addmm_default_1, primals_988, mm_48, amax_default_3, sum_221, div_1051, pow_220, addmm_325, getitem_321, rsqrt_128, primals_984, primals_985, primals_982, cat_84, sum_248, primals_1067, addmm_353, getitem_349, rsqrt_139, primals_1065, view_556, cat_83, cat_82, pow_249, slice_697, pow_248, slice_696, pow_247, slice_695, pow_246, slice_694, pow_245, slice_693, pow_244, slice_692, pow_243, slice_691, pow_242, slice_690, sub_2244, cat_79, primals_1003, view_553, getitem_327, sub_2242, mm_default_299, mm_51, mm_default_298, squeeze_dim_590, squeeze_dim_469, squeeze_dim_466, squeeze_dim_465, squeeze_dim_462, squeeze_dim_461, squeeze_dim_456, squeeze_38, pow_240, unsqueeze_1002, unsqueeze_958, select_scatter_828, add_11637, sub_2237, select_scatter_829, add_11661, sub_2238, add_11662, add_11664, add_11666, add_11668, add_11670, add_11672, add_11674, add_11676, add_11678, gt_997, where_2937, div_1483, add_11682, add_11680, gt_996, where_2935, select_scatter_918, where_2933, gt_995, convert_element_type_69, add_10313, add_10312, add_10311, add_10310, add_10309, add_10308, add_10307, add_10306, add_10305, add_10304, add_10303, add_10314, pow_236, addmm_332, primals_1001, mul_12724, cat_78, mul_12715, sub_1989, mean_33, pow_226, pow_222, primals_986, view_550, view_548, view_543, primals_981, view_532, primals_979, addmm_324, getitem_319, rsqrt_127, primals_980, cat_68, mm_39, addmm_322, getitem_317, rsqrt_126, primals_973, primals_974, addmm_323, primals_975, view_529, primals_971, cat_74, primals_970, cat_73, primals_968, view_522, primals_966, view_521, primals_964, mm_46, primals_962, add_10261, view_517, primals_961, primals_959, mul_12679, div_1554, primals_957, view_515, addmm_314, addmm_315, primals_951, primals_949, mul_12673, div_1555, primals_955, view_485, primals_953, view_509, primals_947, view_507, div_1048, permute_1245, permute_1246, permute_1247, primals_945, primals_943, addmm_309, addmm_310, primals_941, primals_939, primals_937, addmm_306, addmm_307, primals_935, primals_933, primals_931, mul_12668, div_1556, primals_929, view_481, addmm_304, primals_927, primals_925, mul_12663, div_1557, view_479, primals_923, view_477, view_474, view_475, view_476, getitem_304, getitem_305, getitem_306, getitem_307, view_468, primals_922, primals_919, addmm_302, getitem_303, rsqrt_121, primals_917, primals_903, addmm_298, getitem_295, rsqrt_117, primals_901, primals_915, addmm_301, getitem_301, rsqrt_120, primals_913, primals_977, primals_893, addmm_295, getitem_291, rsqrt_115, primals_891, addmm_294, primals_889, primals_887, addmm_293, getitem_289, rsqrt_114, primals_885, addmm_292, primals_883, primals_881, addmm_291, getitem_287, rsqrt_113, primals_879, addmm_290, primals_877, primals_875, addmm_289, getitem_285, rsqrt_112, primals_873, addmm_288, primals_871, primals_869, addmm_287, getitem_283, rsqrt_111, primals_867, addmm_286, primals_865, primals_863, addmm_285, getitem_281, rsqrt_110, primals_861, addmm_284, primals_859, primals_857, addmm_283, getitem_279, rsqrt_109, primals_855, addmm_282, primals_853, primals_851, addmm_281, getitem_277, rsqrt_108, primals_849, addmm_280, primals_847, sigmoid_20, primals_846, primals_845, div_900, div_908, pow_199, where_3396, primals_840, addmm_277, sum_195, div_907, pow_197, primals_838, addmm_276, primals_836, div_906, pow_195, slice_scatter_5, index_40, index_41, index_42, index_43, addmm_275, index_47, index_45, primals_834, addmm_274, primals_832, primals_899, addmm_297, getitem_293, rsqrt_116, primals_897, primals_911, addmm_300, getitem_299, rsqrt_119, primals_909, addmm_279, add_10219, primals_830, addmm_default_2, primals_828, mm_40, amax_default_5, sum_188, div_899, pow_185, addmm_271, getitem_269, rsqrt_107, primals_824, primals_825, primals_822, cat_72, sum_215, primals_907, addmm_299, getitem_297, rsqrt_118, primals_905, view_466, cat_71, cat_70, pow_214, slice_608, pow_213, slice_607, pow_212, slice_606, pow_211, slice_605, pow_210, slice_604, pow_209, slice_603, pow_208, slice_602, pow_207, slice_601, sub_1955, cat_67, primals_843, view_463, getitem_275, sub_1953, mm_default_317, mm_43, mm_default_316, squeeze_dim_626, squeeze_dim_401, squeeze_dim_398, squeeze_dim_397, squeeze_dim_394, squeeze_dim_393, squeeze_dim_388, squeeze_33, pow_205, unsqueeze_876, unsqueeze_832, select_scatter_713, add_10157, sub_1948, select_scatter_714, add_10181, sub_1949, add_10182, add_10184, add_10186, add_10188, add_10190, add_10192, add_10194, add_10196, add_10198, gt_872, where_2570, div_1577, add_10202, add_10200, gt_871, where_2568, select_scatter_803, where_2566, gt_870, convert_element_type_59, add_8833, add_8832, add_8831, add_8830, add_8829, add_8828, add_8827, add_8826, add_8825, add_8824, add_8823, add_8834, pow_201, addmm_278, primals_841, mul_10895, cat_66, mul_10886, sub_1700, mean_28, pow_191, pow_187, primals_826, view_460, view_458, view_453, primals_821, view_442, primals_819, addmm_270, getitem_267, rsqrt_106, primals_820, cat_56, mm_31, addmm_268, getitem_265, rsqrt_105, primals_813, primals_814, addmm_269, primals_815, view_439, primals_811, cat_62, primals_810, cat_61, primals_808, view_432, primals_806, view_431, primals_804, mm_38, primals_802, add_8781, view_427, primals_801, primals_799, mul_10850, div_1648, primals_797, view_425, addmm_260, addmm_261, primals_791, primals_789, mul_10844, div_1649, primals_795, view_395, primals_793, view_419, primals_787, view_417, div_896, permute_1537, permute_1538, permute_1539, primals_785, primals_783, addmm_255, addmm_256, primals_781, primals_779, primals_777, addmm_252, addmm_253, primals_775, primals_773, primals_771, mul_10839, div_1650, primals_769, view_391, addmm_250, primals_767, primals_765, mul_10834, div_1651, view_389, primals_763, view_387, view_384, view_385, view_386, getitem_252, getitem_253, getitem_254, getitem_255, view_378, primals_762, primals_759, addmm_248, getitem_251, rsqrt_100, primals_757, primals_743, addmm_244, getitem_243, rsqrt_96, primals_741, primals_755, addmm_247, getitem_249, rsqrt_99, primals_753, primals_817, primals_733, addmm_241, getitem_239, rsqrt_94, primals_731, addmm_240, primals_729, primals_727, addmm_239, getitem_237, rsqrt_93, primals_725, addmm_238, primals_723, primals_721, addmm_237, getitem_235, rsqrt_92, primals_719, addmm_236, primals_717, primals_715, addmm_235, getitem_233, rsqrt_91, primals_713, addmm_234, primals_711, primals_709, addmm_233, getitem_231, rsqrt_90, primals_707, addmm_232, primals_705, primals_703, addmm_231, getitem_229, rsqrt_89, primals_701, addmm_230, primals_699, primals_697, addmm_229, getitem_227, rsqrt_88, primals_695, addmm_228, primals_693, primals_691, addmm_227, getitem_225, rsqrt_87, primals_689, addmm_226, primals_687, sigmoid_16, primals_686, primals_685, div_748, div_756, pow_164, where_3431, primals_680, addmm_223, sum_162, div_755, pow_162, primals_678, addmm_222, primals_676, div_754, pow_160, slice_scatter_4, index_32, index_33, index_34, index_35, addmm_221, index_39, index_37, primals_674, addmm_220, primals_672, primals_739, addmm_243, getitem_241, rsqrt_95, primals_737, primals_751, addmm_246, getitem_247, rsqrt_98, primals_749, addmm_225, add_8739, primals_670, addmm_default_3, primals_668, mm_32, amax_default_7, sum_155, div_747, pow_150, addmm_217, getitem_217, rsqrt_86, primals_664, primals_665, primals_662, cat_60, sum_182, primals_747, addmm_245, getitem_245, rsqrt_97, primals_745, view_376, cat_59, cat_58, pow_179, slice_519, pow_178, slice_518, pow_177, slice_517, pow_176, slice_516, pow_175, slice_515, pow_174, slice_514, pow_173, slice_513, pow_172, slice_512, sub_1666, cat_55, primals_683, view_373, getitem_223, sub_1664, mm_default_335, mm_35, mm_default_334, squeeze_dim_662, squeeze_dim_333, squeeze_dim_330, squeeze_dim_329, squeeze_dim_326, squeeze_dim_325, squeeze_dim_320, squeeze_28, pow_170, unsqueeze_750, unsqueeze_706, select_scatter_598, add_8677, sub_1659, select_scatter_599, add_8701, sub_1660, add_8702, add_8704, add_8706, add_8708, add_8710, add_8712, add_8714, add_8716, add_8718, gt_747, where_2203, div_1671, add_8722, add_8720, gt_746, where_2201, select_scatter_688, where_2199, gt_745, convert_element_type_49, add_7353, add_7352, add_7351, add_7350, add_7349, add_7348, add_7347, add_7346, add_7345, add_7344, add_7343, add_7354, pow_166, addmm_224, primals_681, mul_9066, cat_54, mul_9057, sub_1411, mean_23, pow_156, pow_152, primals_666, view_370, view_368, view_363, primals_661, view_352, addmm_216, getitem_215, rsqrt_85, primals_659, primals_660, cat_44, mm_23, addmm_214, getitem_213, rsqrt_84, primals_653, primals_654, addmm_215, primals_655, view_349, primals_651, cat_50, primals_650, cat_49, primals_648, view_342, primals_646, view_341, primals_644, mm_30, primals_642, add_7301, view_337, primals_641, primals_639, mul_9021, div_1742, primals_637, view_335, addmm_206, addmm_207, primals_631, primals_629, mul_9015, div_1743, primals_635, view_305, primals_633, view_329, primals_627, view_327, div_744, permute_1829, permute_1830, permute_1831, primals_625, primals_623, addmm_201, addmm_202, primals_621, primals_619, primals_617, addmm_198, addmm_199, primals_615, primals_613, primals_611, mul_9010, div_1744, primals_609, view_301, addmm_196, primals_607, primals_605, mul_9005, div_1745, view_299, primals_603, view_297, view_294, view_295, view_296, getitem_200, getitem_201, getitem_202, getitem_203, view_288, primals_602, primals_599, addmm_194, getitem_199, rsqrt_79, primals_597, primals_583, addmm_190, getitem_191, rsqrt_75, primals_581, primals_595, addmm_193, getitem_197, rsqrt_78, primals_593, primals_657, primals_573, addmm_187, getitem_187, rsqrt_73, primals_571, addmm_186, primals_569, primals_567, addmm_185, getitem_185, rsqrt_72, primals_565, addmm_184, primals_563, primals_561, addmm_183, getitem_183, rsqrt_71, primals_559, addmm_182, primals_557, primals_555, addmm_181, getitem_181, rsqrt_70, primals_553, addmm_180, primals_551, primals_549, addmm_179, getitem_179, rsqrt_69, primals_547, addmm_178, primals_545, primals_543, addmm_177, getitem_177, rsqrt_68, primals_541, addmm_176, primals_539, primals_537, addmm_175, getitem_175, rsqrt_67, primals_535, addmm_174, primals_533, primals_531, addmm_173, getitem_173, rsqrt_66, primals_529, addmm_172, primals_527, sigmoid_12, primals_526, primals_525, div_596, div_604, pow_129, where_3466, primals_520, addmm_169, sum_129, div_603, pow_127, primals_518, addmm_168, primals_516, div_602, pow_125, slice_scatter_3, index_24, index_25, index_26, index_27, addmm_167, index_31, index_29, primals_514, addmm_166, primals_512, primals_579, addmm_189, getitem_189, rsqrt_74, primals_577, primals_591, addmm_192, getitem_195, rsqrt_77, primals_589, addmm_171, add_7259, primals_510, addmm_default_4, primals_508, mm_24, amax_default_9, sum_122, div_595, pow_115, addmm_163, getitem_165, rsqrt_65, primals_504, primals_505, primals_502, cat_48, sum_149, primals_587, addmm_191, getitem_193, rsqrt_76, primals_585, view_286, cat_47, cat_46, primals_501, mm_15, addmm_160, getitem_161, rsqrt_63, primals_493, primals_494, addmm_161, primals_495, primals_491, primals_490, cat_37, primals_488, primals_486, primals_484, add_5821, primals_481, primals_479, mul_7192, div_1836, primals_477, primals_475, addmm_152, addmm_153, primals_473, primals_471, primals_469, mul_7186, div_1837, primals_467, permute_2121, div_592, permute_2122, primals_463, addmm_147, addmm_148, primals_461, permute_2123, primals_457, addmm_144, addmm_145, primals_455, primals_439, addmm_140, getitem_147, rsqrt_58, primals_437, primals_427, addmm_137, getitem_141, rsqrt_55, primals_425, view_262, primals_499, addmm_162, getitem_163, rsqrt_64, primals_500, view_172, primals_339, addmm_108, getitem_111, rsqrt_43, primals_340, pow_144, slice_430, pow_143, slice_429, pow_142, slice_428, pow_141, slice_427, pow_140, slice_426, pow_139, slice_425, pow_138, slice_424, pow_137, slice_423, sub_1377, cat_43, primals_523, view_283, getitem_171, sub_1375, mm_default_353, mm_27, mm_default_352, squeeze_dim_698, squeeze_dim_265, squeeze_dim_262, squeeze_dim_261, squeeze_dim_258, squeeze_dim_257, squeeze_dim_252, squeeze_23, pow_135, unsqueeze_624, unsqueeze_580, select_scatter_483, add_7197, sub_1370, select_scatter_484, add_7221, sub_1371, add_7222, add_7224, add_7226, add_7228, add_7230, add_7232, add_7234, add_7236, add_7238, gt_622, where_1836, div_1765, add_7242, add_7240, gt_621, where_1834, select_scatter_573, where_1832, gt_620, convert_element_type_39, add_5873, add_5872, add_5871, add_5870, add_5869, add_5868, add_5867, add_5866, add_5865, add_5864, add_5863, add_5874, pow_131, addmm_170, primals_521, mul_7237, cat_42, mul_7228, sub_1122, mean_18, pow_121, pow_117, primals_506, view_280, view_278, view_273, cat_32, view_259, cat_38, view_252, view_251, mm_22, primals_482, view_247, view_245, view_215, view_239, view_237, primals_465, primals_459, primals_453, primals_451, mul_7181, div_1838, primals_449, view_211, addmm_142, primals_447, primals_445, mul_7176, div_1839, view_209, primals_443, view_207, view_204, view_205, view_206, getitem_148, getitem_149, getitem_150, getitem_151, view_198, primals_442, primals_423, addmm_136, getitem_139, rsqrt_54, primals_421, primals_435, addmm_139, getitem_145, rsqrt_57, primals_433, primals_497, primals_413, addmm_133, getitem_135, rsqrt_52, primals_411, addmm_132, primals_409, primals_407, addmm_131, getitem_133, rsqrt_51, primals_405, addmm_130, primals_403, primals_401, addmm_129, getitem_131, rsqrt_50, primals_399, addmm_128, primals_397, primals_395, addmm_127, getitem_129, rsqrt_49, primals_393, addmm_126, primals_391, primals_389, addmm_125, getitem_127, rsqrt_48, primals_387, addmm_124, primals_385, primals_383, addmm_123, getitem_125, rsqrt_47, primals_381, addmm_122, primals_379, primals_377, addmm_121, getitem_123, rsqrt_46, primals_375, addmm_120, primals_373, primals_371, addmm_119, getitem_121, rsqrt_45, primals_369, addmm_118, primals_367, sigmoid_8, primals_366, primals_365, div_444, div_452, pow_94, where_3501, primals_360, addmm_115, sum_96, div_451, pow_92, primals_358, addmm_114, primals_356, div_450, pow_90, slice_scatter_2, index_16, index_17, index_18, index_19, addmm_113, index_23, index_21, primals_354, addmm_112, primals_352, primals_419, addmm_135, getitem_137, rsqrt_53, primals_417, primals_431, addmm_138, getitem_143, rsqrt_56, primals_429, addmm_117, add_5779, primals_350, addmm_default_5, primals_348, mm_16, amax_default_11, sum_89, div_443, pow_80, addmm_109, getitem_113, rsqrt_44, primals_344, primals_345, primals_342, cat_36, sum_116, view_196, cat_35, cat_34, pow_109, slice_341, pow_108, slice_340, pow_107, slice_339, pow_106, slice_338, pow_105, slice_337, pow_104, slice_336, pow_103, slice_335, pow_102, slice_334, sub_1088, cat_31, primals_363, view_193, getitem_119, sub_1086, mm_default_371, mm_19, mm_default_370, squeeze_dim_734, squeeze_dim_197, squeeze_dim_194, squeeze_dim_193, squeeze_dim_190, squeeze_dim_189, squeeze_dim_184, squeeze_18, pow_100, unsqueeze_498, unsqueeze_454, select_scatter_368, add_5717, sub_1081, select_scatter_369, add_5741, sub_1082, add_5742, add_5744, add_5746, add_5748, add_5750, add_5752, add_5754, add_5756, add_5758, gt_497, where_1469, div_1859, add_5762, add_5760, gt_496, where_1467, select_scatter_458, where_1465, gt_495, convert_element_type_29, add_4393, add_4392, add_4391, add_4390, add_4389, add_4388, add_4387, add_4386, add_4385, add_4384, add_4383, add_4394, pow_96, addmm_116, primals_361, mul_5408, cat_30, mul_5399, sub_833, mean_13, pow_86, pow_82, primals_346, view_190, view_188, view_183, primals_341, cat_20, mm_7, addmm_106, getitem_109, rsqrt_42, primals_333, primals_334, addmm_107, primals_335, view_169, primals_331, cat_26, primals_330, cat_25, primals_328, view_162, primals_326, view_161, primals_324, mm_14, primals_322, add_4341, view_157, primals_321, primals_319, mul_5363, div_1930, primals_317, view_155, addmm_98, addmm_99, primals_311, primals_309, mul_5357, div_1931, primals_315, view_125, primals_313, view_149, primals_307, view_147, div_440, permute_2413, permute_2414, permute_2415, primals_305, primals_303, addmm_93, addmm_94, primals_301, primals_299, primals_297, addmm_90, addmm_91, primals_295, primals_293, primals_291, mul_5352, div_1932, primals_289, view_121, addmm_88, primals_287, primals_285, mul_5347, div_1933, view_119, primals_283, view_117, view_114, view_115, view_116, getitem_96, getitem_97, getitem_98, getitem_99, view_108, primals_282, primals_279, addmm_86, getitem_95, rsqrt_37, primals_277, primals_263, addmm_82, getitem_87, rsqrt_33, primals_261, primals_275, addmm_85, getitem_93, rsqrt_36, primals_273, primals_337, primals_253, addmm_79, getitem_83, rsqrt_31, primals_251, addmm_78, primals_249, primals_247, addmm_77, getitem_81, rsqrt_30, primals_245, addmm_76, primals_243, primals_241, addmm_75, getitem_79, rsqrt_29, primals_239, addmm_74, primals_237, primals_235, addmm_73, getitem_77, rsqrt_28, primals_233, addmm_72, primals_231, primals_229, addmm_71, getitem_75, rsqrt_27, primals_227, addmm_70, primals_225, primals_223, addmm_69, getitem_73, rsqrt_26, primals_221, addmm_68, primals_219, primals_217, addmm_67, getitem_71, rsqrt_25, primals_215, addmm_66, primals_213, primals_211, addmm_65, getitem_69, rsqrt_24, primals_209, addmm_64, primals_207, sigmoid_4, primals_206, primals_205, div_292, div_300, pow_59, where_3536, primals_200, addmm_61, sum_63, div_299, pow_57, primals_198, addmm_60, primals_196, div_298, pow_55, slice_scatter_1, index_8, index_9, index_10, index_11, addmm_59, index_15, index_13, primals_194, addmm_58, primals_192, primals_259, addmm_81, getitem_85, rsqrt_32, primals_257, primals_271, addmm_84, getitem_91, rsqrt_35, primals_269, addmm_63, add_4299, primals_190, addmm_default_6, primals_188, mm_8, amax_default_13, sum_56, div_291, pow_45, addmm_55, getitem_61, rsqrt_23, primals_184, primals_185, primals_182, cat_24, sum_83, primals_267, addmm_83, getitem_89, rsqrt_34, primals_265, view_106, cat_23, cat_22, pow_74, slice_252, pow_73, slice_251, pow_72, slice_250, pow_71, slice_249, pow_70, slice_248, pow_69, slice_247, pow_68, slice_246, pow_67, slice_245, sub_799, cat_19, primals_203, view_103, getitem_67, sub_797, mm_default_389, mm_11, mm_default_388, squeeze_dim_770, squeeze_dim_129, squeeze_dim_126, squeeze_dim_125, squeeze_dim_122, squeeze_dim_121, squeeze_dim_116, squeeze_13, pow_65, unsqueeze_372, unsqueeze_328, select_scatter_253, add_4237, sub_792, select_scatter_254, add_4261, sub_793, add_4262, add_4264, add_4266, add_4268, add_4270, add_4272, add_4274, add_4276, add_4278, gt_372, where_1102, div_1953, add_4282, add_4280, gt_371, where_1100, select_scatter_343, where_1098, gt_370, convert_element_type_19, add_2913, add_2912, add_2911, add_2910, add_2909, add_2908, add_2907, add_2906, add_2905, add_2904, add_2903, add_2914, pow_61, addmm_62, primals_201, mul_3579, cat_18, mul_3570, sub_544, mean_8, pow_51, pow_47, primals_186, view_100, view_98, view_93, primals_181, view_82, primals_179, addmm_54, getitem_59, rsqrt_22, primals_180, cat_8, cat_3, addmm_52, getitem_57, rsqrt_21, primals_173, primals_174, addmm_53, primals_175, view_79, primals_171, cat_14, primals_170, cat_13, primals_168, view_72, primals_166, view_71, primals_164, mm_6, primals_162, add_2862, view_67, primals_161, primals_159, mul_3534, div_2024, primals_157, view_65, addmm_44, addmm_45, primals_151, primals_149, addmm_43, getitem_53, rsqrt_19, primals_155, view_35, primals_153, view_59, primals_147, view_57, div_288, permute_2705, permute_2706, permute_2707, primals_145, primals_143, addmm_39, addmm_40, primals_141, primals_139, primals_137, addmm_36, addmm_37, primals_135, primals_133, primals_131, mul_3523, div_2026, primals_129, view_31, addmm_34, primals_127, primals_125, addmm_33, getitem_49, rsqrt_17, view_29, primals_123, view_27, view_24, view_25, view_26, getitem_44, getitem_45, getitem_46, getitem_47, view_18, primals_122, primals_119, addmm_32, getitem_43, rsqrt_16, primals_117, cat_12, primals_115, addmm_31, getitem_41, rsqrt_15, primals_113, sum_50, primals_111, addmm_30, getitem_39, rsqrt_14, primals_109, sigmoid, primals_107, addmm_29, getitem_37, rsqrt_13, primals_105, cat_11, primals_103, addmm_28, getitem_35, rsqrt_12, primals_101, cat_10, primals_177, primals_97, addmm_26, getitem_33, rsqrt_11, primals_95, pow_39, addmm_25, primals_93, slice_163, primals_91, addmm_24, getitem_31, rsqrt_10, primals_89, pow_38, addmm_23, primals_87, slice_162, primals_85, addmm_22, getitem_29, rsqrt_9, primals_83, pow_37, addmm_21, primals_81, slice_161, primals_79, addmm_20, getitem_27, rsqrt_8, primals_77, pow_36, addmm_19, primals_75, slice_160, primals_73, addmm_18, getitem_25, rsqrt_7, primals_71, addmm_17, primals_69, primals_67, addmm_16, getitem_23, rsqrt_6, primals_65, addmm_15, primals_63, primals_61, addmm_14, getitem_21, rsqrt_5, primals_59, addmm_13, primals_57, primals_55, addmm_12, getitem_19, rsqrt_4, primals_53, addmm_11, primals_51, primals_50, pow_35, slice_159, pow_34, slice_158, pow_33, slice_157, pow_32, slice_156, sub_511, primals_49, cat_7, mm, addmm_10, primals_47, view_14, getitem_17, sub_509, mm_default_407, mm_3, mm_default_406, squeeze_dim_806, squeeze_dim_61, squeeze_dim_58, squeeze_dim_57, squeeze_dim_54, squeeze_dim_53, squeeze_dim_48, squeeze_8, pow_30, unsqueeze_246, unsqueeze_202, select_scatter_138, add_2761, sub_504, select_scatter_139, add_2785, sub_505, add_2786, add_2788, add_2790, add_2792, add_2794, add_2796, add_2798, add_2800, add_2802, gt_247, div_2046, add_2804, add_2806, gt_246, select_scatter_228, where_731, gt_245, convert_element_type_9, add_1437, add_1436, add_1435, add_1434, add_1433, add_1432, add_1431, add_1430, add_1429, add_1428, add_1427, add_1438, pow_26, addmm_9, primals_45, div_140, div_148, pow_24, where_3571, primals_44, addmm_8, sum_30, div_147, pow_22, primals_42, mul_1755, addmm_7, primals_40, cat_6, add_2823, div_146, pow_20, addmm_6, slice_scatter, cat_5, div_139, primals_38, addmm_5, primals_36, primals_34, addmm_default_7, primals_32, amax_default_15, sum_23, pow_10, addmm_2, getitem_11, rsqrt_3, primals_28, primals_29, mul_1746, sub_256, mean_3, pow_16, pow_12, primals_30, view_11, view_9, primals_26, convolution_3, getitem_9, rsqrt_2, primals_18, primals_19, where_369, addmm, mul_19, primals_17, convolution_1, getitem_3, rsqrt_1, primals_13, primals_14, mul_9, primals_9, convolution, getitem_1, rsqrt, primals_7, primals_8, cat, primals_3 = args\n","   22420:     args.clear()\n","   22421:     assert_size_stride(tangents_7, (256, 65, 384), (24960, 384, 1))\n","   22422:     assert_size_stride(primals_1301, (384, ), (1, ))\n","   22423:     assert_size_stride(mul_16366, (256, 65, 384), (24960, 384, 1))\n","   22424:     assert_size_stride(div_1355, (256, 65, 1), (65, 1, 1))\n","   22425:     assert_size_stride(mm_55, (16640, 384), (384, 1))\n","   22426:     assert_size_stride(addmm_430, (16640, 384), (384, 1))\n","   22427:     assert_size_stride(getitem_421, (256, 65, 1), (65, 1, 1))\n","   22428:     assert_size_stride(rsqrt_168, (256, 65, 1), (65, 1, 1))\n","   22429:     assert_size_stride(primals_1293, (384, ), (1, ))\n","   22430:     assert_size_stride(primals_1294, (384, ), (1, ))\n","   22431:     assert_size_stride(addmm_431, (16640, 384), (384, 1))\n","   22432:     assert_size_stride(primals_1295, (384, 768), (768, 1))\n","   22433:     assert_size_stride(view_709, (16640, 768), (768, 1))\n","   22434:     assert_size_stride(primals_1291, (384, 256), (256, 1))\n","   22435:     assert_size_stride(cat_98, (16640, 256), (256, 1))\n","   22436:     assert_size_stride(primals_1290, (1, 4, 64), (256, 64, 1))\n","   22437:     assert_size_stride(cat_97, (16640, 256), (256, 1))\n","   22438:     assert_size_stride(primals_1288, (64, 384), (384, 1))\n","   22439:     assert_size_stride(view_702, (16640, 384), (384, 1))\n","   22440:     assert_size_stride(primals_1286, (64, 384), (384, 1))\n","   22441:     assert_size_stride(view_701, (16640, 384), (384, 1))\n","   22442:     assert_size_stride(primals_1284, (64, 384), (384, 1))\n","   22443:     assert_size_stride(mm_62, (16640, 384), (384, 1))\n","   22444:     assert_size_stride(primals_1282, (64, 384), (384, 1))\n","   22445:     assert_size_stride(add_13221, (256, 65, 384), (24960, 384, 1))\n","   22446:     assert_size_stride(view_697, (16640, 384), (384, 1))\n","   22447:     assert_size_stride(primals_1281, (384, 384), (384, 1))\n","   22448:     assert_size_stride(primals_1279, (384, ), (1, ))\n","   22449:     assert_size_stride(mul_16337, (256, 65, 384), (24960, 384, 1))\n","   22450:     assert_size_stride(div_1366, (256, 65, 1), (65, 1, 1))\n","   22451:     assert_size_stride(primals_1277, (384, 1536), (1536, 1))\n","   22452:     assert_size_stride(view_695, (16640, 1536), (1536, 1))\n","   22453:     assert_size_stride(addmm_422, (16640, 1536), (1536, 1))\n","   22454:     assert_size_stride(addmm_423, (16640, 1536), (1536, 1))\n","   22455:     assert_size_stride(primals_1271, (1536, 384), (384, 1))\n","   22456:     assert_size_stride(primals_1269, (384, ), (1, ))\n","   22457:     assert_size_stride(mul_16331, (256, 65, 384), (24960, 384, 1))\n","   22458:     assert_size_stride(div_1367, (256, 65, 1), (65, 1, 1))\n","   22459:     assert_size_stride(primals_1275, (1536, 128), (128, 1))\n","   22460:     assert_size_stride(view_665, (16640, 128), (128, 1))\n","   22461:     assert_size_stride(primals_1273, (1536, 128), (128, 1))\n","   22462:     assert_size_stride(view_689, (16640, 384), (384, 1))\n","   22463:     assert_size_stride(primals_1267, (384, 384), (384, 1))\n","   22464:     assert_size_stride(view_687, (16640, 384), (384, 1))\n","   22465:     assert_size_stride(div_1352, (256, 8, 65, 65), (34048, 4256, 65, 1))\n","   22466:     assert_size_stride(permute_657, (2048, 48, 65), (3136, 1, 48))\n","   22467:     assert_size_stride(permute_658, (2048, 48, 65), (3136, 1, 48))\n","   22468:     assert_size_stride(permute_659, (2048, 65, 48), (3136, 1, 65))\n","   22469:     assert_size_stride(primals_1265, (384, 384), (384, 1))\n","   22470:     assert_size_stride(primals_1263, (384, 128), (128, 1))\n","   22471:     assert_size_stride(addmm_417, (16640, 384), (384, 1))\n","   22472:     assert_size_stride(addmm_418, (16640, 384), (384, 1))\n","   22473:     assert_size_stride(primals_1261, (384, 128), (128, 1))\n","   22474:     assert_size_stride(primals_1259, (384, 384), (384, 1))\n","   22475:     assert_size_stride(primals_1257, (384, 128), (128, 1))\n","   22476:     assert_size_stride(addmm_414, (16640, 384), (384, 1))\n","   22477:     assert_size_stride(addmm_415, (16640, 384), (384, 1))\n","   22478:     assert_size_stride(primals_1255, (384, 128), (128, 1))\n","   22479:     assert_size_stride(primals_1253, (384, 384), (384, 1))\n","   22480:     assert_size_stride(primals_1251, (384, ), (1, ))\n","   22481:     assert_size_stride(mul_16326, (256, 65, 384), (24960, 384, 1))\n","   22482:     assert_size_stride(div_1368, (256, 65, 1), (65, 1, 1))\n","   22483:     assert_size_stride(primals_1249, (384, 1536), (1536, 1))\n","   22484:     assert_size_stride(view_661, (16640, 1536), (1536, 1))\n","   22485:     assert_size_stride(addmm_412, (16640, 1536), (1536, 1))\n","   22486:     assert_size_stride(primals_1247, (1536, 384), (384, 1))\n","   22487:     assert_size_stride(primals_1245, (384, ), (1, ))\n","   22488:     assert_size_stride(mul_16321, (256, 65, 384), (24960, 384, 1))\n","   22489:     assert_size_stride(div_1369, (256, 65, 1), (65, 1, 1))\n","   22490:     assert_size_stride(view_659, (16640, 384), (384, 1))\n","   22491:     assert_size_stride(primals_1243, (384, 384), (384, 1))\n","   22492:     assert_size_stride(view_657, (16640, 384), (384, 1))\n","   22493:     assert_size_stride(view_654, (256, 8, 65, 48), (384, 48, 98304, 1))\n","   22494:     assert_size_stride(view_655, (256, 8, 65, 48), (384, 48, 98304, 1))\n","   22495:     assert_size_stride(view_656, (256, 8, 65, 48), (384, 48, 98304, 1))\n","   22496:     assert_size_stride(getitem_408, (256, 8, 65, 48), (24960, 48, 384, 1))\n","   22497:     assert_size_stride(getitem_409, (256, 8, 96), (768, 96, 1))\n","\n","======================================================================\n","FILE: cr7uxvfrxg2h5bxmbo26rs2k6chluxxfvwlhh7iutte4ffncdmgq.py — 1350 lines\n","    1137: def partition_0(args):\n","    1316:     def call(self, args):\n","\n","  --- partition_0 (lines 1137-1216) ---\n","    1137: def partition_0(args):\n","    1138:     primals_7, primals_8, primals_5, primals_9, primals_6, primals_1, primals_2, primals_10, primals_11, primals_3, primals_4 = args\n","    1139:     args.clear()\n","    1140:     assert_size_stride(primals_7, (128, 128), (8320, 1))\n","    1141:     assert_size_stride(primals_8, (128, 128), (8320, 1))\n","    1142:     assert_size_stride(primals_5, (128, 128), (8320, 1))\n","    1143:     assert_size_stride(primals_9, (128, 128), (24960, 1))\n","    1144:     assert_size_stride(primals_6, (128, 128), (8320, 1))\n","    1145:     assert_size_stride(primals_1, (128, 128), (8320, 1))\n","    1146:     assert_size_stride(primals_2, (128, 128), (8320, 1))\n","    1147:     assert_size_stride(primals_10, (128, 128), (8320, 1))\n","    1148:     assert_size_stride(primals_11, (128, 128), (8320, 1))\n","    1149:     assert_size_stride(primals_3, (128, 256), (16640, 1))\n","    1150:     assert_size_stride(primals_4, (128, 256), (16640, 1))\n","    1151:     with torch.cuda._DeviceGuard(0):\n","    1152:         torch.cuda.set_device(0)\n","    1153:         buf31 = empty_strided_cuda((128, 1), (1, 128), torch.float32)\n","    1154:         buf32 = reinterpret_tensor(buf31, (128, 1), (1, 1), 0); del buf31  # reuse\n","    1155:         buf30 = empty_strided_cuda((128, 1), (1, 1), torch.float32)\n","    1156:         buf33 = empty_strided_cuda((128, ), (1, ), torch.float32)\n","    1157:         # Topologically Sorted Source Nodes: [log_softmax_1, mul_1, sum_2], Original ATen: [aten._log_softmax, prims.prepare_softmax_online, aten.sub, aten.mul, aten.sum]\n","    1158:         stream0 = get_raw_stream(0)\n","    1159:         triton_per_fused__log_softmax_mul_prepare_softmax_online_sub_sum_0.run(buf32, primals_7, primals_8, buf30, buf33, 128, 128, stream=stream0)\n","    1160:         buf23 = empty_strided_cuda((128, 1), (1, 128), torch.float32)\n","    1161:         buf24 = reinterpret_tensor(buf23, (128, 1), (1, 1), 0); del buf23  # reuse\n","    1162:         buf61 = empty_strided_cuda((128, ), (1, ), torch.float32)\n","    1163:         buf62 = empty_strided_cuda((128, ), (1, ), torch.int64)\n","    1164:         buf36 = empty_strided_cuda((128, 128), (128, 1), torch.bool)\n","    1165:         buf22 = empty_strided_cuda((128, 1), (1, 1), torch.float32)\n","    1166:         buf25 = empty_strided_cuda((128, ), (1, ), torch.float32)\n","    1167:         buf27 = empty_strided_cuda((128, ), (1, ), torch.int64)\n","    1168:         buf28 = empty_strided_cuda((128, ), (1, ), torch.int64)\n","    1169:         buf37 = empty_strided_cuda((128, ), (1, ), torch.float32)\n","    1170:         buf39 = empty_strided_cuda((128, ), (1, ), torch.float32)\n","    1171:         # Topologically Sorted Source Nodes: [labels, log_softmax, mul, sum_1, argmax_2, argmax_3, nearest, unsqueeze, clamp, loss_4, clamp_1, log, mul_2, sum_3, max_1], Original ATen: [aten.arange, aten._log_softmax, prims.prepare_softmax_online, aten.sub, aten.mul, aten.sum, aten.argmax, aten.unsqueeze, aten.view, aten.expand, aten.eq, aten.scalar_tensor, aten.where, aten.clamp, aten.binary_cross_entropy, aten.log, aten.max]\n","    1172:         stream0 = get_raw_stream(0)\n","    1173:         triton_per_fused__log_softmax_arange_argmax_binary_cross_entropy_clamp_eq_expand_log_max_mul_prepare_softmax_online_scalar_tensor_sub_sum_unsqueeze_view_where_1.run(buf24, primals_5, primals_9, primals_6, buf61, buf62, buf36, buf22, buf25, buf27, buf28, buf37, buf39, 128, 128, stream=stream0)\n","    1174:         del primals_9\n","    1175:         buf29 = empty_strided_cuda((), (), torch.float32)\n","    1176:         buf70 = buf29; del buf29  # reuse\n","    1177:         # Topologically Sorted Source Nodes: [eq_2, float_3, acc_2], Original ATen: [aten.eq, aten._to_copy, aten.mean]\n","    1178:         stream0 = get_raw_stream(0)\n","    1179:         triton_per_fused__to_copy_eq_mean_2.run(buf70, buf27, buf28, 1, 128, stream=stream0)\n","    1180:         del buf27\n","    1181:         buf26 = empty_strided_cuda((), (), torch.float32)\n","    1182:         buf69 = buf26; del buf26  # reuse\n","    1183:         # Topologically Sorted Source Nodes: [mean_2, loss_2, mean_4, loss_3, add, l_bridge], Original ATen: [aten.mean, aten.neg, aten.add, aten.div]\n","    1184:         stream0 = get_raw_stream(0)\n","    1185:         triton_per_fused_add_div_mean_neg_3.run(buf69, buf25, buf33, 1, 128, stream=stream0)\n","    1186:         buf38 = empty_strided_cuda((), (), torch.float32)\n","    1187:         buf71 = buf38; del buf38  # reuse\n","    1188:         # Topologically Sorted Source Nodes: [clamp, loss_4], Original ATen: [aten.scalar_tensor, aten.where, aten.clamp, aten.binary_cross_entropy]\n","    1189:         stream0 = get_raw_stream(0)\n","    1190:         triton_per_fused_binary_cross_entropy_clamp_scalar_tensor_where_4.run(buf71, buf37, 1, 128, stream=stream0)\n","    1191:         buf40 = empty_strided_cuda((), (), torch.float32)\n","    1192:         buf72 = buf40; del buf40  # reuse\n","    1193:         # Topologically Sorted Source Nodes: [mean_6, entropy], Original ATen: [aten.mean, aten.neg]\n","    1194:         stream0 = get_raw_stream(0)\n","    1195:         triton_per_fused_mean_neg_5.run(buf72, buf39, 1, 128, stream=stream0)\n","    1196:         buf63 = empty_strided_cuda((), (), torch.float32)\n","    1197:         buf64 = empty_strided_cuda((), (), torch.float32)\n","    1198:         buf77 = buf63; del buf63  # reuse\n","    1199:         buf78 = buf64; del buf64  # reuse\n","    1200:         # Topologically Sorted Source Nodes: [sub, loss_7, nearest_cos_1], Original ATen: [aten.rsub, aten.mean]\n","    1201:         stream0 = get_raw_stream(0)\n","    1202:         triton_per_fused_mean_rsub_6.run(buf77, buf78, buf61, 1, 128, stream=stream0)\n","    1203:         buf0 = empty_strided_cuda((128, 128), (128, 1), torch.float32)\n","    1204:         # Topologically Sorted Source Nodes: [getattr_1, matmul], Original ATen: [aten.permute, aten.mm]\n","    1205:         extern_kernels.mm(primals_1, reinterpret_tensor(primals_2, (128, 128), (1, 8320), 0), out=buf0)\n","    1206:         buf1 = reinterpret_tensor(buf61, (128, 1), (1, 1), 0); del buf61  # reuse\n","    1207:         buf2 = reinterpret_tensor(buf39, (128, 1), (1, 128), 0); del buf39  # reuse\n","    1208:         buf7 = buf28; del buf28  # reuse\n","    1209:         buf3 = reinterpret_tensor(buf2, (128, 1), (1, 1), 0); del buf2  # reuse\n","    1210:         # Topologically Sorted Source Nodes: [sim, loss, argmax], Original ATen: [aten.div, aten.mul, aten.amax, aten.sub, aten._log_softmax, aten.argmax]\n","    1211:         stream0 = get_raw_stream(0)\n","    1212:         triton_per_fused__log_softmax_amax_argmax_div_mul_sub_7.run(buf3, buf0, buf1, buf7, 128, 128, stream=stream0)\n","    1213:         buf8 = empty_strided_cuda((), (), torch.float32)\n","    1214:         buf66 = buf8; del buf8  # reuse\n","    1215:         # Topologically Sorted Source Nodes: [labels, eq, float_1, acc], Original ATen: [aten.arange, aten.eq, aten._to_copy, aten.mean]\n","    1216:         stream0 = get_raw_stream(0)\n"]}]},{"cell_type":"markdown","source":["# eigh compiled"],"metadata":{"id":"o4bkrLRFLPxO"}},{"cell_type":"code","source":["\"\"\"\n","CompiledEigh: torch.compile(fullgraph=True) drop-in for torch.linalg.eigh.\n","\n","Eliminates all graph breaks, device-host syncs, and dynamic allocation.\n","Output contract matches torch.linalg.eigh:\n","  eigenvalues:  [*, n] real, ascending\n","  eigenvectors: [*, n, n] orthonormal columns\n","\n","Author: AbstractPhil / GeoLIP project\n","\"\"\"\n","\n","import math\n","import torch\n","import torch.nn as nn\n","from torch import Tensor\n","from typing import Tuple, Optional\n","\n","\n","# =============================================================================\n","# Constants\n","# =============================================================================\n","\n","DEFAULT_MAX_NEWTON: int = 8\n","DEFAULT_MAX_JACOBI_SWEEPS: int = 10  # 10 sweeps saturates for n <= 16\n","JACOBI_THRESHOLD: int = 16\n","\n","\n","# =============================================================================\n","# Atom: 2x2 Symmetric Eigenproblem\n","# =============================================================================\n","\n","def eigh_2x2(a: Tensor, b: Tensor, c: Tensor, eps: float = 1e-30\n","             ) -> Tuple[Tensor, Tensor, Tensor, Tensor]:\n","    \"\"\"\n","    Closed-form eigendecomposition of batched 2x2 symmetric matrices.\n","    Returns: (lambda1, lambda2, cos_theta, sin_theta), lambda1 <= lambda2.\n","    \"\"\"\n","    trace = a + c\n","    diff = a - c\n","    two_b = 2.0 * b\n","    hyp = torch.sqrt(diff * diff + two_b * two_b + eps)\n","\n","    lambda1 = 0.5 * (trace - hyp)\n","    lambda2 = 0.5 * (trace + hyp)\n","\n","    vx = two_b\n","    vy = lambda2 - a\n","    norm_v = torch.sqrt(vx * vx + vy * vy + eps)\n","    cos_theta = vy / norm_v\n","    sin_theta = vx / norm_v\n","\n","    return lambda1, lambda2, cos_theta, sin_theta\n","\n","\n","# =============================================================================\n","# Utility: Newton-Schulz Orthogonalization (all-bmm, GPU-native)\n","# =============================================================================\n","\n","def orthogonalize_ns(V: Tensor, n_iter: int = 2) -> Tensor:\n","    \"\"\"\n","    Re-orthogonalize columns of V via Newton-Schulz iteration.\n","\n","    Computes V @ (V^T V)^{-1/2} using the coupled iteration:\n","        X_0 = I,  Y_0 = V^T V\n","        X_{k+1} = 0.5 * X_k @ (3I - Y_k)\n","        Y_{k+1} = 0.5 * (3I - Y_k) @ Y_k\n","\n","    X converges to (V^T V)^{-1/2}, Y converges to I.\n","    Cubically convergent when V^T V ≈ I.\n","\n","    Convergence from ||V^T V - I|| = ε:\n","        1 iteration: error → O(ε²)  ≈ 1e-6 from 1e-3\n","        2 iterations: error → O(ε⁴) ≈ 1e-12 from 1e-3\n","\n","    All ops are bmm — fully compiled, no sequential column processing.\n","\n","    V: [B, n, n] (square, columns are approximate eigenvectors)\n","    Returns: [B, n, n] with orthonormal columns\n","    \"\"\"\n","    B, n, m = V.shape\n","    I_n = torch.eye(m, device=V.device, dtype=V.dtype).unsqueeze(0).expand(B, -1, -1)\n","\n","    # Z = V^T V ≈ I\n","    Y = torch.bmm(V.transpose(-2, -1), V)\n","    X = I_n.clone()\n","\n","    for _ in range(n_iter):\n","        T = 3.0 * I_n - Y           # (3I - Y_k)\n","        X = 0.5 * torch.bmm(X, T)   # X_{k+1}\n","        Y = 0.5 * torch.bmm(T, Y)   # Y_{k+1}\n","\n","    return torch.bmm(V, X)\n","\n","\n","# =============================================================================\n","# Phase 1: Householder Tridiagonalization\n","# =============================================================================\n","\n","class HouseholderTridiagonalizer(nn.Module):\n","    \"\"\"\n","    Reduces batched symmetric A to tridiagonal T = Q^T A Q via\n","    Householder reflections. Fixed loop bounds, compilable.\n","    \"\"\"\n","\n","    def __init__(self, max_n: int, eps: float = 1e-30):\n","        super().__init__()\n","        self.max_n = max_n\n","        self.eps = eps\n","\n","    def forward(self, A: Tensor, d: Tensor, e: Tensor,\n","                reflectors: Tensor) -> None:\n","        B, n, _ = A.shape\n","        eps = self.eps\n","\n","        for k in range(n - 2):\n","            tail_len = n - k - 1\n","            x = A[:, k + 1:, k].clone()\n","\n","            sigma = torch.sqrt((x * x).sum(dim=-1, keepdim=True) + eps)\n","            sign_x0 = torch.where(x[:, 0:1] >= 0,\n","                                  torch.ones_like(sigma),\n","                                  -torch.ones_like(sigma))\n","            alpha = -sign_x0 * sigma\n","\n","            v = x.clone()\n","            v[:, 0:1] = v[:, 0:1] - alpha\n","            v_norm = torch.sqrt((v * v).sum(dim=-1, keepdim=True) + eps)\n","            v = v / v_norm\n","\n","            reflectors[k, :, :tail_len] = v\n","            if tail_len < n:\n","                reflectors[k, :, tail_len:] = 0.0\n","\n","            sub_A = A[:, k + 1:, k + 1:]\n","            v_col = v.unsqueeze(-1)\n","\n","            p = torch.bmm(sub_A, v_col).squeeze(-1)\n","            vtp = (v * p).sum(dim=-1, keepdim=True)\n","            q = p - vtp * v\n","\n","            q_col = q.unsqueeze(-1)\n","            q_row = q.unsqueeze(-2)\n","            v_row = v.unsqueeze(-2)\n","\n","            A[:, k + 1:, k + 1:] -= 2.0 * (v_col @ q_row + q_col @ v_row)\n","            A[:, k, k + 1] = alpha.squeeze(-1)\n","            A[:, k + 1, k] = alpha.squeeze(-1)\n","\n","        for i in range(n):\n","            d[:, i] = A[:, i, i]\n","        for i in range(n - 1):\n","            e[:, i] = A[:, i, i + 1]\n","\n","\n","# =============================================================================\n","# Phase 2a: Secular Equation Newton Solver (Fixed Budget)\n","# =============================================================================\n","\n","class SecularNewtonSolver(nn.Module):\n","\n","    def __init__(self, max_newton: int = DEFAULT_MAX_NEWTON,\n","                 eps: float = 1e-30, tol: float = 1e-7):\n","        super().__init__()\n","        self.max_newton = max_newton\n","        self.eps = eps\n","        self.tol = tol\n","\n","    def forward(self, delta: Tensor, z_sq: Tensor,\n","                rho: Tensor, mask: Tensor) -> Tensor:\n","        B, m = delta.shape\n","        eps = self.eps\n","        tol = self.tol\n","\n","        z_sq_sum = (z_sq * mask).sum(dim=-1, keepdim=True)\n","        rho_abs = rho.abs().unsqueeze(-1)\n","        upper_bound = delta[:, -1:] + z_sq_sum * rho_abs + 1.0\n","\n","        lo = delta + eps\n","        hi = torch.cat([delta[:, 1:], upper_bound], dim=-1) - eps\n","        lam = 0.5 * (lo + hi)\n","        rho_exp = rho.unsqueeze(-1)\n","\n","        for _step in range(self.max_newton):\n","            delta_exp = delta.unsqueeze(-1)\n","            lam_exp = lam.unsqueeze(-2)\n","            denom = delta_exp - lam_exp\n","\n","            denom_safe = torch.where(\n","                denom.abs() < eps,\n","                torch.full_like(denom, eps) * denom.sign().clamp(min=0.5),\n","                denom\n","            )\n","\n","            z_sq_exp = z_sq.unsqueeze(-1)\n","            mask_exp = mask.unsqueeze(-1)\n","            masked_z = z_sq_exp * mask_exp\n","\n","            terms = masked_z / denom_safe\n","            f = 1.0 + rho_exp * terms.sum(dim=-2)\n","            f_prime = rho_exp * (masked_z / (denom_safe * denom_safe)).sum(dim=-2)\n","\n","            f_prime_safe = torch.where(\n","                f_prime.abs() < eps,\n","                torch.full_like(f_prime, eps),\n","                f_prime\n","            )\n","            delta_lam = -f / f_prime_safe\n","            lam_new = torch.clamp(lam + delta_lam, lo, hi)\n","\n","            f_pos = f > 0\n","            lo = torch.where(f_pos & mask.bool(), lam, lo)\n","            hi = torch.where(~f_pos & mask.bool(), lam, hi)\n","\n","            converged = (f.abs() < tol) | ~mask.bool()\n","            lam = torch.where(converged, lam, lam_new)\n","\n","        return lam\n","\n","\n","# =============================================================================\n","# Phase 2b: Eigenvectors from Secular Equation\n","# =============================================================================\n","\n","def secular_eigenvectors(delta: Tensor, lam: Tensor, z: Tensor,\n","                         mask: Tensor, eps: float = 1e-30) -> Tensor:\n","    delta_exp = delta.unsqueeze(-1)\n","    lam_exp = lam.unsqueeze(-2)\n","    denom = delta_exp - lam_exp\n","\n","    denom_safe = torch.where(\n","        denom.abs() < eps,\n","        torch.full_like(denom, eps) * denom.sign().clamp(min=0.5),\n","        denom\n","    )\n","\n","    z_exp = z.unsqueeze(-1)\n","    mask_exp = mask.unsqueeze(-1)\n","    V = (z_exp * mask_exp) / denom_safe\n","    col_norms = torch.sqrt((V * V).sum(dim=-2, keepdim=True) + eps)\n","    V = V / col_norms\n","    return V\n","\n","\n","# =============================================================================\n","# Phase 2c: Fixed-Depth Tensor Tree D&C\n","# =============================================================================\n","\n","class TensorTreeDC(nn.Module):\n","\n","    def __init__(self, max_n: int,\n","                 max_newton: int = DEFAULT_MAX_NEWTON,\n","                 eps: float = 1e-30, tol: float = 1e-7):\n","        super().__init__()\n","        self.padded_n = 1 << math.ceil(math.log2(max(max_n, 2)))\n","        self.depth = int(math.log2(self.padded_n))\n","        self.max_n = max_n\n","        self.eps = eps\n","        self.secular_solver = SecularNewtonSolver(\n","            max_newton=max_newton, eps=eps, tol=tol\n","        )\n","\n","    def forward(self, d: Tensor, e: Tensor) -> Tuple[Tensor, Tensor]:\n","        B, n = d.shape\n","        pn = self.padded_n\n","        eps = self.eps\n","        device = d.device\n","        dtype = d.dtype\n","\n","        if n < pn:\n","            d_max = d.abs().max(dim=-1, keepdim=True).values + 1.0\n","            pad_diag = d_max + torch.arange(1, pn - n + 1, device=device, dtype=dtype).unsqueeze(0)\n","            d_padded = torch.cat([d, pad_diag], dim=-1)\n","            e_padded = torch.zeros(B, pn - 1, device=device, dtype=dtype)\n","            e_padded[:, :n - 1] = e\n","        else:\n","            d_padded = d.clone()\n","            e_padded = e.clone()\n","\n","        # DOWNWARD PASS\n","        coupling_rho = []\n","        current_d = d_padded.clone()\n","        current_e = e_padded.clone()\n","\n","        for level in range(self.depth):\n","            num_sub = 2 ** level\n","            sub_size = pn // num_sub\n","            half = sub_size // 2\n","\n","            cd = current_d.reshape(B, num_sub, sub_size)\n","            ce = current_e.reshape(B, num_sub, sub_size - 1)\n","\n","            rho = ce[:, :, half - 1].clone()\n","            coupling_rho.append(rho)\n","\n","            cd[:, :, half - 1] = cd[:, :, half - 1] - rho.abs()\n","            cd[:, :, half] = cd[:, :, half] - rho.abs()\n","            ce[:, :, half - 1] = 0.0\n","\n","            left_d = cd[:, :, :half].reshape(B, num_sub * half)\n","            right_d = cd[:, :, half:].reshape(B, num_sub * half)\n","            current_d = torch.stack([left_d.reshape(B, num_sub, half),\n","                                     right_d.reshape(B, num_sub, half)],\n","                                    dim=2).reshape(B, pn)\n","\n","            left_e = ce[:, :, :half - 1].reshape(B, num_sub, half - 1)\n","            right_e = ce[:, :, half:].reshape(B, num_sub, half - 1)\n","            current_e = torch.stack([left_e, right_e], dim=2).reshape(\n","                B, num_sub * 2 * (half - 1))\n","            expected_e_len = pn - 1\n","            if current_e.shape[-1] < expected_e_len:\n","                current_e = torch.nn.functional.pad(\n","                    current_e, (0, expected_e_len - current_e.shape[-1]))\n","\n","        # BASE\n","        base_evals = current_d\n","        V_current = torch.ones(B, pn, 1, 1, device=device, dtype=dtype)\n","\n","        # UPWARD PASS\n","        current_evals = base_evals\n","\n","        for level in range(self.depth - 1, -1, -1):\n","            num_sub = 2 ** level\n","            sub_size = pn // num_sub\n","            half = sub_size // 2\n","            child_size = half\n","\n","            evals_grouped = current_evals.reshape(B, num_sub, 2, child_size)\n","            left_evals = evals_grouped[:, :, 0, :]\n","            right_evals = evals_grouped[:, :, 1, :]\n","            delta = torch.cat([left_evals, right_evals], dim=-1)\n","\n","            V_grouped = V_current.reshape(B, num_sub, 2, child_size, child_size)\n","            V_left = V_grouped[:, :, 0, :, :]\n","            V_right = V_grouped[:, :, 1, :, :]\n","\n","            z_left = V_left[:, :, -1, :]\n","            z_right = V_right[:, :, 0, :]\n","            z_cat = torch.cat([z_left, z_right], dim=-1)\n","\n","            delta_sorted, sort_idx = delta.sort(dim=-1)\n","            z_sorted = z_cat.gather(-1, sort_idx)\n","\n","            rho = coupling_rho[level]\n","            mask = torch.ones(B, num_sub, sub_size, device=device, dtype=dtype)\n","\n","            gaps = (delta_sorted[:, :, 1:] - delta_sorted[:, :, :-1]).abs()\n","            degenerate = gaps < (eps * 100)\n","            avg = 0.5 * (delta_sorted[:, :, :-1] + delta_sorted[:, :, 1:])\n","\n","            delta_defl = delta_sorted.clone()\n","            delta_defl[:, :, :-1] = torch.where(degenerate, avg, delta_sorted[:, :, :-1])\n","            delta_defl[:, :, 1:] = torch.where(degenerate, avg, delta_sorted[:, :, 1:])\n","\n","            z_defl = z_sorted.clone()\n","            defl_kill = torch.ones_like(z_sorted)\n","            defl_kill[:, :, 1:] = torch.where(\n","                degenerate, torch.zeros_like(gaps), torch.ones_like(gaps))\n","            z_defl = z_defl * defl_kill\n","\n","            z_sq = z_defl * z_defl\n","            Bns = B * num_sub\n","            new_evals_flat = self.secular_solver(\n","                delta_defl.reshape(Bns, sub_size),\n","                z_sq.reshape(Bns, sub_size),\n","                rho.reshape(Bns),\n","                mask.reshape(Bns, sub_size),\n","            )\n","            new_evals = new_evals_flat.reshape(B, num_sub, sub_size)\n","\n","            V_secular_flat = secular_eigenvectors(\n","                delta_defl.reshape(Bns, sub_size),\n","                new_evals_flat,\n","                z_defl.reshape(Bns, sub_size),\n","                mask.reshape(Bns, sub_size),\n","                eps=eps\n","            )\n","            V_secular = V_secular_flat.reshape(B, num_sub, sub_size, sub_size)\n","\n","            inv_sort = sort_idx.argsort(dim=-1)\n","            inv_exp = inv_sort.unsqueeze(-1).expand_as(V_secular)\n","            V_unsorted = V_secular.gather(-2, inv_exp)\n","\n","            V_block = torch.zeros(B, num_sub, sub_size, sub_size,\n","                                  device=device, dtype=dtype)\n","            V_block[:, :, :half, :half] = V_left\n","            V_block[:, :, half:, half:] = V_right\n","\n","            V_merged = torch.bmm(\n","                V_block.reshape(Bns, sub_size, sub_size),\n","                V_unsorted.reshape(Bns, sub_size, sub_size)\n","            ).reshape(B, num_sub, sub_size, sub_size)\n","\n","            current_evals = new_evals.reshape(B, pn)\n","            V_current = V_merged\n","\n","        eigenvalues = current_evals\n","        eigenvectors = V_current.squeeze(1)\n","\n","        sorted_evals, sort_perm = eigenvalues.sort(dim=-1)\n","        sort_exp = sort_perm.unsqueeze(-2).expand_as(eigenvectors)\n","        sorted_evecs = eigenvectors.gather(-1, sort_exp)\n","\n","        if n < pn:\n","            sorted_evals = sorted_evals[:, :n]\n","            sorted_evecs = sorted_evecs[:, :n, :n]\n","\n","        return sorted_evals, sorted_evecs\n","\n","\n","# =============================================================================\n","# Phase 2 (alternate): Jacobi for small n\n","# =============================================================================\n","\n","class JacobiEigh(nn.Module):\n","    \"\"\"\n","    Jacobi eigenvalue algorithm for small symmetric matrices.\n","    Fixed sweep count, fully vectorized, zero branches.\n","\n","    COMPILE FIX: Pair indices stored as plain Python lists (not tensors).\n","    Dynamo sees these as constants — no SymInt issues.\n","    \"\"\"\n","\n","    def __init__(self, max_n: int,\n","                 max_sweeps: int = DEFAULT_MAX_JACOBI_SWEEPS,\n","                 eps: float = 1e-30):\n","        super().__init__()\n","        self.max_n = max_n\n","        self.max_sweeps = max_sweeps\n","        self.eps = eps\n","\n","        # CRITICAL: plain Python lists, NOT registered buffers.\n","        # Dynamo traces these as compile-time constants.\n","        pairs = []\n","        for p in range(max_n):\n","            for q in range(p + 1, max_n):\n","                pairs.append((p, q))\n","        self._pairs_p: list[int] = [p for p, q in pairs]\n","        self._pairs_q: list[int] = [q for p, q in pairs]\n","        self._n_pairs: int = len(pairs)\n","\n","    def forward(self, A: Tensor) -> Tuple[Tensor, Tensor]:\n","        \"\"\"\n","        A: [B, n, n] symmetric\n","        Returns: (eigenvalues [B, n] ascending, eigenvectors [B, n, n])\n","        \"\"\"\n","        B, n, _ = A.shape\n","        eps = self.eps\n","\n","        W = A.clone()\n","        V = torch.eye(n, device=A.device, dtype=A.dtype).unsqueeze(0).expand(B, -1, -1).clone()\n","\n","        for _sweep in range(self.max_sweeps):\n","            for idx in range(self._n_pairs):\n","                # Plain Python ints — Dynamo sees these as constants\n","                p: int = self._pairs_p[idx]\n","                q: int = self._pairs_q[idx]\n","\n","                app = W[:, p, p]\n","                aqq = W[:, q, q]\n","                apq = W[:, p, q]\n","\n","                # Givens rotation angle\n","                two_apq = 2.0 * apq\n","                diff = aqq - app\n","\n","                # Safe division: sign-preserving eps guard\n","                abs_two_apq = two_apq.abs().clamp(min=eps)\n","                sign_two_apq = torch.where(two_apq >= 0,\n","                                           torch.ones_like(two_apq),\n","                                           -torch.ones_like(two_apq))\n","                tau = diff / (abs_two_apq * sign_two_apq)\n","\n","                tau_sign = torch.where(tau >= 0,\n","                                       torch.ones_like(tau),\n","                                       -torch.ones_like(tau))\n","                t = tau_sign / (tau.abs() + torch.sqrt(1.0 + tau * tau))\n","\n","                # Zero rotation when off-diagonal is already negligible\n","                skip = (apq.abs() < eps).float()\n","                t = t * (1.0 - skip)\n","\n","                c = 1.0 / torch.sqrt(1.0 + t * t)\n","                s = t * c\n","\n","                # ── Rotate W columns p, q ──\n","                Wp = W[:, :, p].clone()\n","                Wq = W[:, :, q].clone()\n","                c_col = c.unsqueeze(-1)\n","                s_col = s.unsqueeze(-1)\n","                W[:, :, p] = c_col * Wp - s_col * Wq\n","                W[:, :, q] = s_col * Wp + c_col * Wq\n","\n","                # ── Rotate W rows p, q ──\n","                Wp = W[:, p, :].clone()\n","                Wq = W[:, q, :].clone()\n","                W[:, p, :] = c_col * Wp - s_col * Wq\n","                W[:, q, :] = s_col * Wp + c_col * Wq\n","\n","                # ── Exact diagonal repair (prevents accumulation drift) ──\n","                W[:, p, q] = 0.0\n","                W[:, q, p] = 0.0\n","                W[:, p, p] = app - t * apq\n","                W[:, q, q] = aqq + t * apq\n","\n","                # ── Accumulate eigenvectors ──\n","                Vp = V[:, :, p].clone()\n","                Vq = V[:, :, q].clone()\n","                V[:, :, p] = c_col * Vp - s_col * Vq\n","                V[:, :, q] = s_col * Vp + c_col * Vq\n","\n","        # ── Newton-Schulz re-orthogonalization ──\n","        # 2 iterations: orth error 1e-3 → ~1e-12 via bmm (GPU-native)\n","        V = orthogonalize_ns(V, n_iter=2)\n","\n","        # ── Extract and sort ──\n","        eigenvalues = torch.diagonal(W, dim1=-2, dim2=-1)\n","        sorted_evals, sort_perm = eigenvalues.sort(dim=-1)\n","        sort_exp = sort_perm.unsqueeze(-2).expand_as(V)\n","        sorted_evecs = V.gather(-1, sort_exp)\n","\n","        return sorted_evals, sorted_evecs\n","\n","\n","# =============================================================================\n","# Phase 3: Householder Back-Accumulation\n","# =============================================================================\n","\n","class HouseholderBackAccumulate(nn.Module):\n","\n","    def __init__(self, max_n: int, eps: float = 1e-30):\n","        super().__init__()\n","        self.max_n = max_n\n","        self.eps = eps\n","\n","    def forward(self, reflectors: Tensor, Z: Tensor, n: int) -> Tensor:\n","        V = Z.clone()\n","        for k in range(n - 3, -1, -1):\n","            tail_len = n - k - 1\n","            v = reflectors[k, :, :tail_len]\n","            v_col = v.unsqueeze(-1)\n","            V_sub = V[:, k + 1:, :]\n","            vtV = torch.bmm(v_col.transpose(-2, -1), V_sub)\n","            V[:, k + 1:, :] = V_sub - 2.0 * v_col @ vtV\n","        return V\n","\n","\n","# =============================================================================\n","# Validation\n","# =============================================================================\n","\n","class EighValidator(nn.Module):\n","\n","    def forward(self, A: Tensor, eigenvalues: Tensor,\n","                eigenvectors: Tensor) -> Tuple[Tensor, Tensor, Tensor]:\n","        B, n, _ = A.shape\n","        AV = torch.bmm(A, eigenvectors)\n","        VL = eigenvectors * eigenvalues.unsqueeze(-2)\n","        residual = AV - VL\n","        A_norm = torch.linalg.norm(A.reshape(B, -1), dim=-1).clamp(min=1e-30)\n","        residual_norm = torch.linalg.norm(residual.reshape(B, -1), dim=-1) / A_norm\n","\n","        VtV = torch.bmm(eigenvectors.transpose(-2, -1), eigenvectors)\n","        I = torch.eye(n, device=A.device, dtype=A.dtype).unsqueeze(0)\n","        orth_err = torch.linalg.norm((VtV - I).reshape(B, -1), dim=-1)\n","\n","        return residual_norm, orth_err, residual_norm.max()\n","\n","\n","# =============================================================================\n","# Top-Level: CompiledEigh\n","# =============================================================================\n","\n","class CompiledEigh(nn.Module):\n","    \"\"\"\n","    Drop-in replacement for torch.linalg.eigh.\n","\n","    Usage:\n","        solver = CompiledEigh(max_n=6)\n","        solver = torch.compile(solver, fullgraph=True)\n","        eigenvalues, eigenvectors = solver(A)\n","    \"\"\"\n","\n","    def __init__(self, max_n: int,\n","                 use_jacobi: Optional[bool] = None,\n","                 max_newton: int = DEFAULT_MAX_NEWTON,\n","                 max_jacobi_sweeps: int = DEFAULT_MAX_JACOBI_SWEEPS,\n","                 eps: float = 1e-30, tol: float = 1e-7):\n","        super().__init__()\n","        self.max_n = max_n\n","        self.eps = eps\n","\n","        if use_jacobi is None:\n","            use_jacobi = (max_n <= JACOBI_THRESHOLD)\n","        self.use_jacobi = use_jacobi\n","\n","        if use_jacobi:\n","            self.jacobi = JacobiEigh(\n","                max_n=max_n, max_sweeps=max_jacobi_sweeps, eps=eps)\n","        else:\n","            self.tridiag = HouseholderTridiagonalizer(max_n=max_n, eps=eps)\n","            self.dc = TensorTreeDC(\n","                max_n=max_n, max_newton=max_newton, eps=eps, tol=tol)\n","            self.back_accum = HouseholderBackAccumulate(max_n=max_n, eps=eps)\n","\n","        self.validator = EighValidator()\n","\n","    def forward(self, A: Tensor, validate: bool = False\n","                ) -> Tuple[Tensor, Tensor]:\n","        B, n, _ = A.shape\n","\n","        if self.use_jacobi:\n","            eigenvalues, eigenvectors = self.jacobi(A)\n","        else:\n","            A_work = A.clone()\n","            d = torch.empty(B, n, device=A.device, dtype=A.dtype)\n","            e = torch.empty(B, n - 1, device=A.device, dtype=A.dtype)\n","            reflectors = torch.zeros(max(n - 2, 1), B, n,\n","                                     device=A.device, dtype=A.dtype)\n","            self.tridiag(A_work, d, e, reflectors)\n","            eigenvalues, Z = self.dc(d, e)\n","            eigenvectors = self.back_accum(reflectors, Z, n)\n","            # Newton-Schulz re-orthogonalization for D&C path\n","            eigenvectors = orthogonalize_ns(eigenvectors, n_iter=2)\n","\n","        if validate:\n","            res_norm, orth_err, max_err = self.validator(A, eigenvalues, eigenvectors)\n","            print(f\"[CompiledEigh] max residual: {max_err.item():.2e}, \"\n","                  f\"mean orth err: {orth_err.mean().item():.2e}\")\n","\n","        return eigenvalues, eigenvectors\n","\n","\n","# =============================================================================\n","# Functional API\n","# =============================================================================\n","\n","_cached_solvers = {}\n","\n","def compiled_eigh(A: Tensor, validate: bool = False) -> Tuple[Tensor, Tensor]:\n","    B, n, _ = A.shape\n","    key = (n, A.device, A.dtype)\n","    if key not in _cached_solvers:\n","        _cached_solvers[key] = CompiledEigh(max_n=n).to(A.device)\n","    return _cached_solvers[key](A, validate=validate)\n","\n","\n","\"\"\"\n","CompiledEigh — Colab GPU Benchmark v3\n","Fixes:\n","  v2: Jacobi pairs as plain Python lists (Dynamo compile fix), sweeps 6→10\n","  v3: Replaced Gram-Schmidt with Newton-Schulz orthogonalization (all-bmm),\n","      disabled TF32 to ensure fp32 precision on Blackwell\n","\"\"\"\n","\n","import torch\n","import time\n","import gc\n","import sys\n","\n","# ── Ensure full fp32 precision on Ampere/Hopper/Blackwell ──\n","# TF32 uses 10-bit mantissa for matmul which can degrade orthogonality\n","torch.backends.cuda.matmul.allow_tf32 = False\n","torch.backends.cudnn.allow_tf32 = False\n","torch.set_float32_matmul_precision('highest')\n","\n","\n","def sync():\n","    if torch.cuda.is_available():\n","        torch.cuda.synchronize()\n","\n","\n","def gpu_timer(fn, warmup=10, repeats=200):\n","    for _ in range(warmup):\n","        fn()\n","    sync()\n","    start = time.perf_counter()\n","    for _ in range(repeats):\n","        fn()\n","    sync()\n","    return (time.perf_counter() - start) / repeats\n","\n","\n","def make_symmetric_batch(B, n, device, dtype=torch.float32):\n","    R = torch.randn(B, n, n, device=device, dtype=dtype)\n","    return (R + R.transpose(-2, -1)) / 2.0\n","\n","\n","def make_cm_like_batch(B, n, device, dtype=torch.float32):\n","    points = torch.randn(B, n, n, device=device, dtype=dtype)\n","    points = points / (points.norm(dim=-1, keepdim=True) + 1e-8)\n","    return torch.bmm(points, points.transpose(-2, -1)) * 0.3\n","\n","\n","def fmt_time(seconds):\n","    if seconds < 1e-3:\n","        return f\"{seconds*1e6:.1f} us\"\n","    elif seconds < 1.0:\n","        return f\"{seconds*1e3:.2f} ms\"\n","    return f\"{seconds:.3f} s\"\n","\n","\n","# ─── Test 0: Newton-Schulz Diagnostic ───\n","\n","def test_ns_diagnostic(device):\n","    \"\"\"Verify Newton-Schulz orthogonalization works on GPU independently.\"\"\"\n","    print(\"\\n\" + \"=\" * 70)\n","    print(\"  TEST 0: NEWTON-SCHULZ DIAGNOSTIC\")\n","    print(\"=\" * 70)\n","\n","    for n in [5, 6, 8]:\n","        B = 1024\n","        # Create nearly-orthogonal matrix (simulating Jacobi output)\n","        Q, _ = torch.linalg.qr(torch.randn(B, n, n, device=device))\n","        # Perturb to ~1e-3 orthogonality error\n","        noise = torch.randn(B, n, n, device=device) * 1e-3\n","        V_dirty = Q + noise\n","\n","        I_n = torch.eye(n, device=device).unsqueeze(0)\n","\n","        # Before NS\n","        VtV_before = torch.bmm(V_dirty.transpose(-2, -1), V_dirty)\n","        orth_before = torch.linalg.norm((VtV_before - I_n).reshape(B, -1), dim=-1).max().item()\n","\n","        # After NS (2 iterations)\n","        V_clean = orthogonalize_ns(V_dirty, n_iter=2)\n","        VtV_after = torch.bmm(V_clean.transpose(-2, -1), V_clean)\n","        orth_after = torch.linalg.norm((VtV_after - I_n).reshape(B, -1), dim=-1).max().item()\n","\n","        # After NS (3 iterations for comparison)\n","        V_clean3 = orthogonalize_ns(V_dirty, n_iter=3)\n","        VtV_after3 = torch.bmm(V_clean3.transpose(-2, -1), V_clean3)\n","        orth_after3 = torch.linalg.norm((VtV_after3 - I_n).reshape(B, -1), dim=-1).max().item()\n","\n","        print(f\"  n={n}: before={orth_before:.2e}  \"\n","              f\"after(2iter)={orth_after:.2e}  \"\n","              f\"after(3iter)={orth_after3:.2e}\")\n","\n","    # Also test with actual Jacobi output\n","    print(f\"\\n  --- With actual Jacobi output ---\")\n","    for n in [5, 6]:\n","        B = 2048\n","        A = make_symmetric_batch(B, n, device)\n","        solver = JacobiEigh(max_n=n, max_sweeps=10).to(device)\n","\n","        # Run Jacobi WITHOUT the NS cleanup\n","        W = A.clone()\n","        V = torch.eye(n, device=device).unsqueeze(0).expand(B, -1, -1).clone()\n","        for _sweep in range(solver.max_sweeps):\n","            for idx in range(solver._n_pairs):\n","                p, q = solver._pairs_p[idx], solver._pairs_q[idx]\n","                app, aqq, apq = W[:, p, p], W[:, q, q], W[:, p, q]\n","                two_apq = 2.0 * apq\n","                diff = aqq - app\n","                abs_2apq = two_apq.abs().clamp(min=1e-30)\n","                sign_2apq = torch.where(two_apq >= 0,\n","                    torch.ones_like(two_apq), -torch.ones_like(two_apq))\n","                tau = diff / (abs_2apq * sign_2apq)\n","                tau_sign = torch.where(tau >= 0,\n","                    torch.ones_like(tau), -torch.ones_like(tau))\n","                t = tau_sign / (tau.abs() + torch.sqrt(1.0 + tau * tau))\n","                skip = (apq.abs() < 1e-30).float()\n","                t = t * (1.0 - skip)\n","                c = 1.0 / torch.sqrt(1.0 + t * t)\n","                s = t * c\n","                c_col, s_col = c.unsqueeze(-1), s.unsqueeze(-1)\n","                Wp = W[:, :, p].clone(); Wq = W[:, :, q].clone()\n","                W[:, :, p] = c_col * Wp - s_col * Wq\n","                W[:, :, q] = s_col * Wp + c_col * Wq\n","                Wp = W[:, p, :].clone(); Wq = W[:, q, :].clone()\n","                W[:, p, :] = c_col * Wp - s_col * Wq\n","                W[:, q, :] = s_col * Wp + c_col * Wq\n","                W[:, p, q] = 0.0; W[:, q, p] = 0.0\n","                W[:, p, p] = app - t * apq\n","                W[:, q, q] = aqq + t * apq\n","                Vp = V[:, :, p].clone(); Vq = V[:, :, q].clone()\n","                V[:, :, p] = c_col * Vp - s_col * Vq\n","                V[:, :, q] = s_col * Vp + c_col * Vq\n","\n","        I_n = torch.eye(n, device=device).unsqueeze(0)\n","        VtV = torch.bmm(V.transpose(-2, -1), V)\n","        orth_raw = torch.linalg.norm((VtV - I_n).reshape(B, -1), dim=-1).max().item()\n","\n","        V_ns = orthogonalize_ns(V, n_iter=2)\n","        VtV_ns = torch.bmm(V_ns.transpose(-2, -1), V_ns)\n","        orth_ns = torch.linalg.norm((VtV_ns - I_n).reshape(B, -1), dim=-1).max().item()\n","\n","        print(f\"  Jacobi raw n={n}: orth={orth_raw:.2e}  after NS(2)={orth_ns:.2e}\")\n","\n","\n","# ─── Test 1: Accuracy ───\n","\n","def test_accuracy(device):\n","    print(\"\\n\" + \"=\" * 70)\n","    print(\"  TEST 1: ACCURACY vs torch.linalg.eigh\")\n","    print(\"=\" * 70)\n","\n","    validator = EighValidator()\n","    configs = [\n","        (3,   4096, \"3x3   small\"),\n","        (5,   4096, \"5x5   CM matrix size\"),\n","        (6,   4096, \"6x6   pentachoron bordered\"),\n","        (8,   2048, \"8x8   padded CM\"),\n","        (12,  1024, \"12x12 medium\"),\n","        (16,  512,  \"16x16 Jacobi boundary\"),\n","    ]\n","\n","    all_pass = True\n","    for n, B, label in configs:\n","        A = make_symmetric_batch(B, n, device)\n","        ref_vals, ref_vecs = torch.linalg.eigh(A)\n","\n","        solver = CompiledEigh(max_n=n).to(device)\n","        our_vals, our_vecs = solver(A)\n","\n","        val_err = (our_vals - ref_vals).abs().max().item()\n","        val_mean = (our_vals - ref_vals).abs().mean().item()\n","\n","        dots = torch.bmm(ref_vecs.transpose(-2, -1), our_vecs)\n","        alignment = dots.abs().max(dim=-1).values.min().item()\n","\n","        res_norm, orth_err, max_res = validator(A, our_vals, our_vecs)\n","        max_orth = orth_err.max().item()\n","\n","        # Thresholds: eigenval 1e-3, alignment 0.999, orth 1e-4\n","        ok = val_err < 1e-3 and alignment > 0.999 and max_orth < 1e-4\n","        if not ok:\n","            all_pass = False\n","\n","        print(f\"\\n  [{'PASS' if ok else 'FAIL'}] {label}  (n={n}, B={B})\")\n","        print(f\"    eigenvalue err   max={val_err:.2e}  mean={val_mean:.2e}\")\n","        print(f\"    eigvec alignment min={alignment:.8f}\")\n","        print(f\"    residual norm    max={max_res.item():.2e}\")\n","        print(f\"    orthogonality    max={max_orth:.2e}\")\n","\n","    print(f\"\\n  --- CM-like spectral distribution ---\")\n","    for n in [5, 6]:\n","        A = make_cm_like_batch(2048, n, device)\n","        ref_vals, _ = torch.linalg.eigh(A)\n","        solver = CompiledEigh(max_n=n).to(device)\n","        our_vals, our_vecs = solver(A)\n","        val_err = (our_vals - ref_vals).abs().max().item()\n","        res_norm, orth_err, max_res = validator(A, our_vals, our_vecs)\n","        print(f\"    CM-like n={n}: val_err={val_err:.2e}  \"\n","              f\"res={max_res.item():.2e}  orth={orth_err.max().item():.2e}\")\n","\n","    return all_pass\n","\n","\n","# ─── Test 2: torch.compile fullgraph ───\n","\n","def test_compile(device):\n","    print(\"\\n\" + \"=\" * 70)\n","    print(\"  TEST 2: torch.compile(fullgraph=True)\")\n","    print(\"=\" * 70)\n","\n","    results = {}\n","    for n, B, label in [(5, 1024, \"5x5\"), (6, 1024, \"6x6\"), (8, 512, \"8x8\")]:\n","        A = make_symmetric_batch(B, n, device)\n","        solver = CompiledEigh(max_n=n).to(device)\n","\n","        try:\n","            compiled_solver = torch.compile(solver, fullgraph=True)\n","            vals, vecs = compiled_solver(A)\n","            sync()\n","            ref_vals, _ = torch.linalg.eigh(A)\n","            err = (vals - ref_vals).abs().max().item()\n","            results[label] = (\"PASS\", err)\n","            print(f\"  [{label}] fullgraph=True SUCCESS  (val_err={err:.2e})\")\n","        except Exception as e:\n","            results[label] = (\"FAIL\", str(e)[:200])\n","            print(f\"  [{label}] COMPILE FAILED: {str(e)[:200]}\")\n","\n","    return all(v[0] == \"PASS\" for v in results.values())\n","\n","\n","# ─── Test 3: Throughput ───\n","\n","def test_benchmark(device):\n","    print(\"\\n\" + \"=\" * 70)\n","    print(\"  TEST 3: GPU THROUGHPUT BENCHMARK\")\n","    print(\"=\" * 70)\n","    print(f\"  Device: {torch.cuda.get_device_name(0)}\")\n","    print(f\"  Timing: 10 warmup + 200 repeats\\n\")\n","\n","    configs = [\n","        (5,  1024,  \"CM 5x5  B=1024\"),\n","        (5,  4096,  \"CM 5x5  B=4096\"),\n","        (5,  8192,  \"CM 5x5  B=8192\"),\n","        (6,  1024,  \"CM 6x6  B=1024\"),\n","        (6,  4096,  \"CM 6x6  B=4096\"),\n","        (6,  8192,  \"CM 6x6  B=8192\"),\n","        (8,  2048,  \"8x8     B=2048\"),\n","        (16, 1024,  \"16x16   B=1024\"),\n","    ]\n","\n","    print(f\"  {'Config':<22} {'eigh ref':>10} {'ours eager':>12} \"\n","          f\"{'ours compiled':>14} {'vs ref':>8}\")\n","    print(f\"  {'-'*22} {'-'*10} {'-'*12} {'-'*14} {'-'*8}\")\n","\n","    for n, B, label in configs:\n","        A = make_symmetric_batch(B, n, device)\n","\n","        ref_time = gpu_timer(lambda: torch.linalg.eigh(A))\n","\n","        solver = CompiledEigh(max_n=n).to(device)\n","        eager_time = gpu_timer(lambda: solver(A))\n","\n","        try:\n","            compiled_solver = torch.compile(solver, fullgraph=True)\n","            for _ in range(5):\n","                compiled_solver(A)\n","            sync()\n","            compiled_time = gpu_timer(lambda: compiled_solver(A))\n","            compiled_str = fmt_time(compiled_time)\n","            speedup = ref_time / compiled_time\n","            speedup_str = f\"{speedup:.2f}x\"\n","        except Exception:\n","            compiled_str = \"FAIL\"\n","            speedup_str = \"N/A\"\n","\n","        print(f\"  {label:<22} {fmt_time(ref_time):>10} \"\n","              f\"{fmt_time(eager_time):>12} {compiled_str:>14} {speedup_str:>8}\")\n","\n","    print(f\"\\n  --- High batch stress test ---\")\n","    for n in [5, 6]:\n","        for B in [16384, 32768]:\n","            try:\n","                A = make_symmetric_batch(B, n, device)\n","                solver = CompiledEigh(max_n=n).to(device)\n","                compiled_solver = torch.compile(solver, fullgraph=True)\n","                for _ in range(3):\n","                    compiled_solver(A)\n","                sync()\n","                t = gpu_timer(lambda: compiled_solver(A), warmup=5, repeats=100)\n","                ref_t = gpu_timer(lambda: torch.linalg.eigh(A), warmup=5, repeats=100)\n","                print(f\"    n={n} B={B}: compiled={fmt_time(t)}  ref={fmt_time(ref_t)}  \"\n","                      f\"ratio={ref_t/t:.2f}x  throughput={B/t:.0f}/sec\")\n","            except RuntimeError as e:\n","                if \"out of memory\" in str(e).lower():\n","                    print(f\"    n={n} B={B}: OOM\")\n","                    torch.cuda.empty_cache()\n","                else:\n","                    raise\n","\n","\n","# ─── Test 4: Autograd ───\n","\n","def test_autograd(device):\n","    print(\"\\n\" + \"=\" * 70)\n","    print(\"  TEST 4: AUTOGRAD BACKWARD\")\n","    print(\"=\" * 70)\n","\n","    for n, B in [(5, 512), (6, 512)]:\n","        A_ref = make_symmetric_batch(B, n, device).requires_grad_(True)\n","        vals_ref, vecs_ref = torch.linalg.eigh(A_ref)\n","        (vals_ref.sum() + (vecs_ref ** 2).sum()).backward()\n","        grad_ref = A_ref.grad.clone()\n","\n","        # Eager backward\n","        A_e = A_ref.detach().clone().requires_grad_(True)\n","        solver = CompiledEigh(max_n=n).to(device)\n","        try:\n","            vals_e, vecs_e = solver(A_e)\n","            (vals_e.sum() + (vecs_e ** 2).sum()).backward()\n","            err_e = (A_e.grad - grad_ref).abs().max().item()\n","            rel_e = err_e / (grad_ref.abs().max().item() + 1e-30)\n","            print(f\"  [{'PASS' if rel_e < 0.1 else 'WARN'}] n={n} eager backward: \"\n","                  f\"grad_err={err_e:.2e}  rel={rel_e:.2e}\")\n","        except Exception as e:\n","            print(f\"  [FAIL] n={n} eager backward: {e}\")\n","\n","        # Compiled backward (may break — forward fullgraph is the key win)\n","        A_c = A_ref.detach().clone().requires_grad_(True)\n","        try:\n","            compiled_solver = torch.compile(solver)\n","            vals_c, vecs_c = compiled_solver(A_c)\n","            (vals_c.sum() + (vecs_c ** 2).sum()).backward()\n","            err_c = (A_c.grad - grad_ref).abs().max().item()\n","            rel_c = err_c / (grad_ref.abs().max().item() + 1e-30)\n","            print(f\"  [{'PASS' if rel_c < 0.1 else 'WARN'}] n={n} compiled backward: \"\n","                  f\"grad_err={err_c:.2e}  rel={rel_c:.2e}\")\n","        except Exception as e:\n","            print(f\"  [INFO] n={n} compiled backward: {str(e)[:150]}\")\n","            print(f\"         (forward fullgraph is the main win)\")\n","\n","\n","# ─── Test 5: VRAM ───\n","\n","def test_vram(device):\n","    print(\"\\n\" + \"=\" * 70)\n","    print(\"  TEST 5: VRAM USAGE\")\n","    print(\"=\" * 70)\n","\n","    for n, B in [(5, 4096), (6, 4096), (6, 8192), (5, 8192)]:\n","        torch.cuda.empty_cache()\n","        gc.collect()\n","        torch.cuda.reset_peak_memory_stats()\n","        base_mem = torch.cuda.memory_allocated()\n","\n","        A = make_symmetric_batch(B, n, device)\n","        solver = CompiledEigh(max_n=n).to(device)\n","        vals, vecs = solver(A)\n","\n","        peak_mem = torch.cuda.max_memory_allocated()\n","        delta_mb = (peak_mem - base_mem) / (1024 ** 2)\n","        print(f\"  n={n} B={B}: peak delta = {delta_mb:.1f} MB\")\n","\n","        del A, solver, vals, vecs\n","        torch.cuda.empty_cache()\n","        gc.collect()\n","\n","\n","# ─── Main ───\n","\n","def main():\n","    print(\"=\" * 70)\n","    print(\"  CompiledEigh v3 — GPU Benchmark Suite\")\n","    print(\"=\" * 70)\n","\n","    if not torch.cuda.is_available():\n","        print(\"\\n  No CUDA. Run on Colab with A100/H100.\")\n","        sys.exit(1)\n","\n","    device = torch.device('cuda')\n","    print(f\"\\n  GPU:     {torch.cuda.get_device_name(0)}\")\n","    print(f\"  CUDA:    {torch.version.cuda}\")\n","    print(f\"  PyTorch: {torch.__version__}\")\n","    mem_gb = torch.cuda.get_device_properties(0).total_mem / (1024**3)\n","    print(f\"  VRAM:    {mem_gb:.1f} GB\")\n","    print(f\"  TF32 matmul: {torch.backends.cuda.matmul.allow_tf32}\")\n","    print(f\"  float32 precision: {torch.get_float32_matmul_precision()}\")\n","\n","    test_ns_diagnostic(device)\n","    acc_ok = test_accuracy(device)\n","    compile_ok = test_compile(device)\n","    test_benchmark(device)\n","    test_autograd(device)\n","    test_vram(device)\n","\n","    print(\"\\n\" + \"=\" * 70)\n","    print(\"  SUMMARY\")\n","    print(\"=\" * 70)\n","    print(f\"  Accuracy:  {'PASS' if acc_ok else 'FAIL'}\")\n","    print(f\"  Compile:   {'PASS' if compile_ok else 'FAIL'}\")\n","    print(\"=\" * 70)\n","\n","\n","if __name__ == '__main__':\n","    main()"],"metadata":{"colab":{"base_uri":"https://localhost:8080/","height":419},"id":"P0UzSGt3LRX7","executionInfo":{"status":"error","timestamp":1775016653868,"user_tz":420,"elapsed":40,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"9ff58deb-5e96-4816-fb29-f30c1cff6201"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["======================================================================\n","  CompiledEigh v3 — GPU Benchmark Suite\n","======================================================================\n","\n","  GPU:     NVIDIA RTX PRO 6000 Blackwell Server Edition\n","  CUDA:    12.8\n","  PyTorch: 2.10.0+cu128\n"]},{"output_type":"error","ename":"AttributeError","evalue":"'torch._C._CudaDeviceProperties' object has no attribute 'total_mem'","traceback":["\u001b[0;31m---------------------------------------------------------------------------\u001b[0m","\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)","\u001b[0;32m/tmp/ipykernel_61613/2312139352.py\u001b[0m in \u001b[0;36m<cell line: 0>\u001b[0;34m()\u001b[0m\n\u001b[1;32m   1047\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1048\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0m__name__\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m'__main__'\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1049\u001b[0;31m     \u001b[0mmain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m","\u001b[0;32m/tmp/ipykernel_61613/2312139352.py\u001b[0m in \u001b[0;36mmain\u001b[0;34m()\u001b[0m\n\u001b[1;32m   1026\u001b[0m     \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf\"  CUDA:    {torch.version.cuda}\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1027\u001b[0m     \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf\"  PyTorch: {torch.__version__}\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1028\u001b[0;31m     \u001b[0mmem_gb\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcuda\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget_device_properties\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtotal_mem\u001b[0m \u001b[0;34m/\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0;36m1024\u001b[0m\u001b[0;34m**\u001b[0m\u001b[0;36m3\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1029\u001b[0m     \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf\"  VRAM:    {mem_gb:.1f} GB\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1030\u001b[0m     \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf\"  TF32 matmul: {torch.backends.cuda.matmul.allow_tf32}\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;31mAttributeError\u001b[0m: 'torch._C._CudaDeviceProperties' object has no attribute 'total_mem'"]}]},{"cell_type":"code","source":["\"\"\"\n","CompiledEigh: torch.compile(fullgraph=True) drop-in for torch.linalg.eigh.\n","\n","Eliminates all graph breaks, device-host syncs, and dynamic allocation.\n","Output contract matches torch.linalg.eigh:\n","  eigenvalues:  [*, n] real, ascending\n","  eigenvectors: [*, n, n] orthonormal columns\n","\n","Author: AbstractPhil / GeoLIP project\n","\"\"\"\n","\n","import math\n","import torch\n","import torch.nn as nn\n","from torch import Tensor\n","from typing import Tuple, Optional\n","\n","\n","# =============================================================================\n","# Constants\n","# =============================================================================\n","\n","DEFAULT_MAX_NEWTON: int = 8\n","DEFAULT_MAX_JACOBI_SWEEPS: int = 10  # 10 sweeps saturates for n <= 16\n","JACOBI_THRESHOLD: int = 16\n","\n","\n","# =============================================================================\n","# Atom: 2x2 Symmetric Eigenproblem\n","# =============================================================================\n","\n","def eigh_2x2(a: Tensor, b: Tensor, c: Tensor, eps: float = 1e-30\n","             ) -> Tuple[Tensor, Tensor, Tensor, Tensor]:\n","    \"\"\"\n","    Closed-form eigendecomposition of batched 2x2 symmetric matrices.\n","    Returns: (lambda1, lambda2, cos_theta, sin_theta), lambda1 <= lambda2.\n","    \"\"\"\n","    trace = a + c\n","    diff = a - c\n","    two_b = 2.0 * b\n","    hyp = torch.sqrt(diff * diff + two_b * two_b + eps)\n","\n","    lambda1 = 0.5 * (trace - hyp)\n","    lambda2 = 0.5 * (trace + hyp)\n","\n","    vx = two_b\n","    vy = lambda2 - a\n","    norm_v = torch.sqrt(vx * vx + vy * vy + eps)\n","    cos_theta = vy / norm_v\n","    sin_theta = vx / norm_v\n","\n","    return lambda1, lambda2, cos_theta, sin_theta\n","\n","\n","# =============================================================================\n","# Utility: Newton-Schulz Orthogonalization (all-bmm, GPU-native)\n","# =============================================================================\n","\n","def orthogonalize_ns(V: Tensor, n_iter: int = 2) -> Tensor:\n","    \"\"\"\n","    Re-orthogonalize columns of V via Newton-Schulz iteration.\n","\n","    Computes V @ (V^T V)^{-1/2} using the coupled iteration:\n","        X_0 = I,  Y_0 = V^T V\n","        X_{k+1} = 0.5 * X_k @ (3I - Y_k)\n","        Y_{k+1} = 0.5 * (3I - Y_k) @ Y_k\n","\n","    X converges to (V^T V)^{-1/2}, Y converges to I.\n","    Cubically convergent when V^T V ≈ I.\n","\n","    Convergence from ||V^T V - I|| = ε:\n","        1 iteration: error → O(ε²)  ≈ 1e-6 from 1e-3\n","        2 iterations: error → O(ε⁴) ≈ 1e-12 from 1e-3\n","\n","    All ops are bmm — fully compiled, no sequential column processing.\n","\n","    V: [B, n, n] (square, columns are approximate eigenvectors)\n","    Returns: [B, n, n] with orthonormal columns\n","    \"\"\"\n","    B, n, m = V.shape\n","    I_n = torch.eye(m, device=V.device, dtype=V.dtype).unsqueeze(0).expand(B, -1, -1)\n","\n","    # Z = V^T V ≈ I\n","    Y = torch.bmm(V.transpose(-2, -1), V)\n","    X = I_n.clone()\n","\n","    for _ in range(n_iter):\n","        T = 3.0 * I_n - Y           # (3I - Y_k)\n","        X = 0.5 * torch.bmm(X, T)   # X_{k+1}\n","        Y = 0.5 * torch.bmm(T, Y)   # Y_{k+1}\n","\n","    return torch.bmm(V, X)\n","\n","\n","# =============================================================================\n","# Phase 1: Householder Tridiagonalization\n","# =============================================================================\n","\n","class HouseholderTridiagonalizer(nn.Module):\n","    \"\"\"\n","    Reduces batched symmetric A to tridiagonal T = Q^T A Q via\n","    Householder reflections. Fixed loop bounds, compilable.\n","    \"\"\"\n","\n","    def __init__(self, max_n: int, eps: float = 1e-30):\n","        super().__init__()\n","        self.max_n = max_n\n","        self.eps = eps\n","\n","    def forward(self, A: Tensor, d: Tensor, e: Tensor,\n","                reflectors: Tensor) -> None:\n","        B, n, _ = A.shape\n","        eps = self.eps\n","\n","        for k in range(n - 2):\n","            tail_len = n - k - 1\n","            x = A[:, k + 1:, k].clone()\n","\n","            sigma = torch.sqrt((x * x).sum(dim=-1, keepdim=True) + eps)\n","            sign_x0 = torch.where(x[:, 0:1] >= 0,\n","                                  torch.ones_like(sigma),\n","                                  -torch.ones_like(sigma))\n","            alpha = -sign_x0 * sigma\n","\n","            v = x.clone()\n","            v[:, 0:1] = v[:, 0:1] - alpha\n","            v_norm = torch.sqrt((v * v).sum(dim=-1, keepdim=True) + eps)\n","            v = v / v_norm\n","\n","            reflectors[k, :, :tail_len] = v\n","            if tail_len < n:\n","                reflectors[k, :, tail_len:] = 0.0\n","\n","            sub_A = A[:, k + 1:, k + 1:]\n","            v_col = v.unsqueeze(-1)\n","\n","            p = torch.bmm(sub_A, v_col).squeeze(-1)\n","            vtp = (v * p).sum(dim=-1, keepdim=True)\n","            q = p - vtp * v\n","\n","            q_col = q.unsqueeze(-1)\n","            q_row = q.unsqueeze(-2)\n","            v_row = v.unsqueeze(-2)\n","\n","            A[:, k + 1:, k + 1:] -= 2.0 * (v_col @ q_row + q_col @ v_row)\n","            A[:, k, k + 1] = alpha.squeeze(-1)\n","            A[:, k + 1, k] = alpha.squeeze(-1)\n","\n","        for i in range(n):\n","            d[:, i] = A[:, i, i]\n","        for i in range(n - 1):\n","            e[:, i] = A[:, i, i + 1]\n","\n","\n","# =============================================================================\n","# Phase 2a: Secular Equation Newton Solver (Fixed Budget)\n","# =============================================================================\n","\n","class SecularNewtonSolver(nn.Module):\n","\n","    def __init__(self, max_newton: int = DEFAULT_MAX_NEWTON,\n","                 eps: float = 1e-30, tol: float = 1e-7):\n","        super().__init__()\n","        self.max_newton = max_newton\n","        self.eps = eps\n","        self.tol = tol\n","\n","    def forward(self, delta: Tensor, z_sq: Tensor,\n","                rho: Tensor, mask: Tensor) -> Tensor:\n","        B, m = delta.shape\n","        eps = self.eps\n","        tol = self.tol\n","\n","        z_sq_sum = (z_sq * mask).sum(dim=-1, keepdim=True)\n","        rho_abs = rho.abs().unsqueeze(-1)\n","        upper_bound = delta[:, -1:] + z_sq_sum * rho_abs + 1.0\n","\n","        lo = delta + eps\n","        hi = torch.cat([delta[:, 1:], upper_bound], dim=-1) - eps\n","        lam = 0.5 * (lo + hi)\n","        rho_exp = rho.unsqueeze(-1)\n","\n","        for _step in range(self.max_newton):\n","            delta_exp = delta.unsqueeze(-1)\n","            lam_exp = lam.unsqueeze(-2)\n","            denom = delta_exp - lam_exp\n","\n","            denom_safe = torch.where(\n","                denom.abs() < eps,\n","                torch.full_like(denom, eps) * denom.sign().clamp(min=0.5),\n","                denom\n","            )\n","\n","            z_sq_exp = z_sq.unsqueeze(-1)\n","            mask_exp = mask.unsqueeze(-1)\n","            masked_z = z_sq_exp * mask_exp\n","\n","            terms = masked_z / denom_safe\n","            f = 1.0 + rho_exp * terms.sum(dim=-2)\n","            f_prime = rho_exp * (masked_z / (denom_safe * denom_safe)).sum(dim=-2)\n","\n","            f_prime_safe = torch.where(\n","                f_prime.abs() < eps,\n","                torch.full_like(f_prime, eps),\n","                f_prime\n","            )\n","            delta_lam = -f / f_prime_safe\n","            lam_new = torch.clamp(lam + delta_lam, lo, hi)\n","\n","            f_pos = f > 0\n","            lo = torch.where(f_pos & mask.bool(), lam, lo)\n","            hi = torch.where(~f_pos & mask.bool(), lam, hi)\n","\n","            converged = (f.abs() < tol) | ~mask.bool()\n","            lam = torch.where(converged, lam, lam_new)\n","\n","        return lam\n","\n","\n","# =============================================================================\n","# Phase 2b: Eigenvectors from Secular Equation\n","# =============================================================================\n","\n","def secular_eigenvectors(delta: Tensor, lam: Tensor, z: Tensor,\n","                         mask: Tensor, eps: float = 1e-30) -> Tensor:\n","    delta_exp = delta.unsqueeze(-1)\n","    lam_exp = lam.unsqueeze(-2)\n","    denom = delta_exp - lam_exp\n","\n","    denom_safe = torch.where(\n","        denom.abs() < eps,\n","        torch.full_like(denom, eps) * denom.sign().clamp(min=0.5),\n","        denom\n","    )\n","\n","    z_exp = z.unsqueeze(-1)\n","    mask_exp = mask.unsqueeze(-1)\n","    V = (z_exp * mask_exp) / denom_safe\n","    col_norms = torch.sqrt((V * V).sum(dim=-2, keepdim=True) + eps)\n","    V = V / col_norms\n","    return V\n","\n","\n","# =============================================================================\n","# Phase 2c: Fixed-Depth Tensor Tree D&C\n","# =============================================================================\n","\n","class TensorTreeDC(nn.Module):\n","\n","    def __init__(self, max_n: int,\n","                 max_newton: int = DEFAULT_MAX_NEWTON,\n","                 eps: float = 1e-30, tol: float = 1e-7):\n","        super().__init__()\n","        self.padded_n = 1 << math.ceil(math.log2(max(max_n, 2)))\n","        self.depth = int(math.log2(self.padded_n))\n","        self.max_n = max_n\n","        self.eps = eps\n","        self.secular_solver = SecularNewtonSolver(\n","            max_newton=max_newton, eps=eps, tol=tol\n","        )\n","\n","    def forward(self, d: Tensor, e: Tensor) -> Tuple[Tensor, Tensor]:\n","        B, n = d.shape\n","        pn = self.padded_n\n","        eps = self.eps\n","        device = d.device\n","        dtype = d.dtype\n","\n","        if n < pn:\n","            d_max = d.abs().max(dim=-1, keepdim=True).values + 1.0\n","            pad_diag = d_max + torch.arange(1, pn - n + 1, device=device, dtype=dtype).unsqueeze(0)\n","            d_padded = torch.cat([d, pad_diag], dim=-1)\n","            e_padded = torch.zeros(B, pn - 1, device=device, dtype=dtype)\n","            e_padded[:, :n - 1] = e\n","        else:\n","            d_padded = d.clone()\n","            e_padded = e.clone()\n","\n","        # DOWNWARD PASS\n","        coupling_rho = []\n","        current_d = d_padded.clone()\n","        current_e = e_padded.clone()\n","\n","        for level in range(self.depth):\n","            num_sub = 2 ** level\n","            sub_size = pn // num_sub\n","            half = sub_size // 2\n","\n","            cd = current_d.reshape(B, num_sub, sub_size)\n","            ce = current_e.reshape(B, num_sub, sub_size - 1)\n","\n","            rho = ce[:, :, half - 1].clone()\n","            coupling_rho.append(rho)\n","\n","            cd[:, :, half - 1] = cd[:, :, half - 1] - rho.abs()\n","            cd[:, :, half] = cd[:, :, half] - rho.abs()\n","            ce[:, :, half - 1] = 0.0\n","\n","            left_d = cd[:, :, :half].reshape(B, num_sub * half)\n","            right_d = cd[:, :, half:].reshape(B, num_sub * half)\n","            current_d = torch.stack([left_d.reshape(B, num_sub, half),\n","                                     right_d.reshape(B, num_sub, half)],\n","                                    dim=2).reshape(B, pn)\n","\n","            left_e = ce[:, :, :half - 1].reshape(B, num_sub, half - 1)\n","            right_e = ce[:, :, half:].reshape(B, num_sub, half - 1)\n","            current_e = torch.stack([left_e, right_e], dim=2).reshape(\n","                B, num_sub * 2 * (half - 1))\n","            expected_e_len = pn - 1\n","            if current_e.shape[-1] < expected_e_len:\n","                current_e = torch.nn.functional.pad(\n","                    current_e, (0, expected_e_len - current_e.shape[-1]))\n","\n","        # BASE\n","        base_evals = current_d\n","        V_current = torch.ones(B, pn, 1, 1, device=device, dtype=dtype)\n","\n","        # UPWARD PASS\n","        current_evals = base_evals\n","\n","        for level in range(self.depth - 1, -1, -1):\n","            num_sub = 2 ** level\n","            sub_size = pn // num_sub\n","            half = sub_size // 2\n","            child_size = half\n","\n","            evals_grouped = current_evals.reshape(B, num_sub, 2, child_size)\n","            left_evals = evals_grouped[:, :, 0, :]\n","            right_evals = evals_grouped[:, :, 1, :]\n","            delta = torch.cat([left_evals, right_evals], dim=-1)\n","\n","            V_grouped = V_current.reshape(B, num_sub, 2, child_size, child_size)\n","            V_left = V_grouped[:, :, 0, :, :]\n","            V_right = V_grouped[:, :, 1, :, :]\n","\n","            z_left = V_left[:, :, -1, :]\n","            z_right = V_right[:, :, 0, :]\n","            z_cat = torch.cat([z_left, z_right], dim=-1)\n","\n","            delta_sorted, sort_idx = delta.sort(dim=-1)\n","            z_sorted = z_cat.gather(-1, sort_idx)\n","\n","            rho = coupling_rho[level]\n","            mask = torch.ones(B, num_sub, sub_size, device=device, dtype=dtype)\n","\n","            gaps = (delta_sorted[:, :, 1:] - delta_sorted[:, :, :-1]).abs()\n","            degenerate = gaps < (eps * 100)\n","            avg = 0.5 * (delta_sorted[:, :, :-1] + delta_sorted[:, :, 1:])\n","\n","            delta_defl = delta_sorted.clone()\n","            delta_defl[:, :, :-1] = torch.where(degenerate, avg, delta_sorted[:, :, :-1])\n","            delta_defl[:, :, 1:] = torch.where(degenerate, avg, delta_sorted[:, :, 1:])\n","\n","            z_defl = z_sorted.clone()\n","            defl_kill = torch.ones_like(z_sorted)\n","            defl_kill[:, :, 1:] = torch.where(\n","                degenerate, torch.zeros_like(gaps), torch.ones_like(gaps))\n","            z_defl = z_defl * defl_kill\n","\n","            z_sq = z_defl * z_defl\n","            Bns = B * num_sub\n","            new_evals_flat = self.secular_solver(\n","                delta_defl.reshape(Bns, sub_size),\n","                z_sq.reshape(Bns, sub_size),\n","                rho.reshape(Bns),\n","                mask.reshape(Bns, sub_size),\n","            )\n","            new_evals = new_evals_flat.reshape(B, num_sub, sub_size)\n","\n","            V_secular_flat = secular_eigenvectors(\n","                delta_defl.reshape(Bns, sub_size),\n","                new_evals_flat,\n","                z_defl.reshape(Bns, sub_size),\n","                mask.reshape(Bns, sub_size),\n","                eps=eps\n","            )\n","            V_secular = V_secular_flat.reshape(B, num_sub, sub_size, sub_size)\n","\n","            inv_sort = sort_idx.argsort(dim=-1)\n","            inv_exp = inv_sort.unsqueeze(-1).expand_as(V_secular)\n","            V_unsorted = V_secular.gather(-2, inv_exp)\n","\n","            V_block = torch.zeros(B, num_sub, sub_size, sub_size,\n","                                  device=device, dtype=dtype)\n","            V_block[:, :, :half, :half] = V_left\n","            V_block[:, :, half:, half:] = V_right\n","\n","            V_merged = torch.bmm(\n","                V_block.reshape(Bns, sub_size, sub_size),\n","                V_unsorted.reshape(Bns, sub_size, sub_size)\n","            ).reshape(B, num_sub, sub_size, sub_size)\n","\n","            current_evals = new_evals.reshape(B, pn)\n","            V_current = V_merged\n","\n","        eigenvalues = current_evals\n","        eigenvectors = V_current.squeeze(1)\n","\n","        sorted_evals, sort_perm = eigenvalues.sort(dim=-1)\n","        sort_exp = sort_perm.unsqueeze(-2).expand_as(eigenvectors)\n","        sorted_evecs = eigenvectors.gather(-1, sort_exp)\n","\n","        if n < pn:\n","            sorted_evals = sorted_evals[:, :n]\n","            sorted_evecs = sorted_evecs[:, :n, :n]\n","\n","        return sorted_evals, sorted_evecs\n","\n","\n","# =============================================================================\n","# Phase 2 (alternate): Jacobi for small n\n","# =============================================================================\n","\n","class JacobiEigh(nn.Module):\n","    \"\"\"\n","    Jacobi eigenvalue algorithm for small symmetric matrices.\n","    Fixed sweep count, fully vectorized, zero branches.\n","\n","    COMPILE FIX: Pair indices stored as plain Python lists (not tensors).\n","    Dynamo sees these as constants — no SymInt issues.\n","    \"\"\"\n","\n","    def __init__(self, max_n: int,\n","                 max_sweeps: int = DEFAULT_MAX_JACOBI_SWEEPS,\n","                 eps: float = 1e-30):\n","        super().__init__()\n","        self.max_n = max_n\n","        self.max_sweeps = max_sweeps\n","        self.eps = eps\n","\n","        # CRITICAL: plain Python lists, NOT registered buffers.\n","        # Dynamo traces these as compile-time constants.\n","        pairs = []\n","        for p in range(max_n):\n","            for q in range(p + 1, max_n):\n","                pairs.append((p, q))\n","        self._pairs_p: list[int] = [p for p, q in pairs]\n","        self._pairs_q: list[int] = [q for p, q in pairs]\n","        self._n_pairs: int = len(pairs)\n","\n","    def forward(self, A: Tensor) -> Tuple[Tensor, Tensor]:\n","        \"\"\"\n","        A: [B, n, n] symmetric\n","        Returns: (eigenvalues [B, n] ascending, eigenvectors [B, n, n])\n","        \"\"\"\n","        B, n, _ = A.shape\n","        eps = self.eps\n","\n","        W = A.clone()\n","        V = torch.eye(n, device=A.device, dtype=A.dtype).unsqueeze(0).expand(B, -1, -1).clone()\n","\n","        for _sweep in range(self.max_sweeps):\n","            for idx in range(self._n_pairs):\n","                # Plain Python ints — Dynamo sees these as constants\n","                p: int = self._pairs_p[idx]\n","                q: int = self._pairs_q[idx]\n","\n","                app = W[:, p, p]\n","                aqq = W[:, q, q]\n","                apq = W[:, p, q]\n","\n","                # Givens rotation angle\n","                two_apq = 2.0 * apq\n","                diff = aqq - app\n","\n","                # Safe division: sign-preserving eps guard\n","                abs_two_apq = two_apq.abs().clamp(min=eps)\n","                sign_two_apq = torch.where(two_apq >= 0,\n","                                           torch.ones_like(two_apq),\n","                                           -torch.ones_like(two_apq))\n","                tau = diff / (abs_two_apq * sign_two_apq)\n","\n","                tau_sign = torch.where(tau >= 0,\n","                                       torch.ones_like(tau),\n","                                       -torch.ones_like(tau))\n","                t = tau_sign / (tau.abs() + torch.sqrt(1.0 + tau * tau))\n","\n","                # Zero rotation when off-diagonal is already negligible\n","                skip = (apq.abs() < eps).float()\n","                t = t * (1.0 - skip)\n","\n","                c = 1.0 / torch.sqrt(1.0 + t * t)\n","                s = t * c\n","\n","                # ── Rotate W columns p, q ──\n","                Wp = W[:, :, p].clone()\n","                Wq = W[:, :, q].clone()\n","                c_col = c.unsqueeze(-1)\n","                s_col = s.unsqueeze(-1)\n","                W[:, :, p] = c_col * Wp - s_col * Wq\n","                W[:, :, q] = s_col * Wp + c_col * Wq\n","\n","                # ── Rotate W rows p, q ──\n","                Wp = W[:, p, :].clone()\n","                Wq = W[:, q, :].clone()\n","                W[:, p, :] = c_col * Wp - s_col * Wq\n","                W[:, q, :] = s_col * Wp + c_col * Wq\n","\n","                # ── Exact diagonal repair (prevents accumulation drift) ──\n","                W[:, p, q] = 0.0\n","                W[:, q, p] = 0.0\n","                W[:, p, p] = app - t * apq\n","                W[:, q, q] = aqq + t * apq\n","\n","                # ── Accumulate eigenvectors ──\n","                Vp = V[:, :, p].clone()\n","                Vq = V[:, :, q].clone()\n","                V[:, :, p] = c_col * Vp - s_col * Vq\n","                V[:, :, q] = s_col * Vp + c_col * Vq\n","\n","        # ── Newton-Schulz re-orthogonalization ──\n","        # 2 iterations: orth error 1e-3 → ~1e-12 via bmm (GPU-native)\n","        V = orthogonalize_ns(V, n_iter=2)\n","\n","        # ── Extract and sort ──\n","        eigenvalues = torch.diagonal(W, dim1=-2, dim2=-1)\n","        sorted_evals, sort_perm = eigenvalues.sort(dim=-1)\n","        sort_exp = sort_perm.unsqueeze(-2).expand_as(V)\n","        sorted_evecs = V.gather(-1, sort_exp)\n","\n","        return sorted_evals, sorted_evecs\n","\n","\n","# =============================================================================\n","# Phase 3: Householder Back-Accumulation\n","# =============================================================================\n","\n","class HouseholderBackAccumulate(nn.Module):\n","\n","    def __init__(self, max_n: int, eps: float = 1e-30):\n","        super().__init__()\n","        self.max_n = max_n\n","        self.eps = eps\n","\n","    def forward(self, reflectors: Tensor, Z: Tensor, n: int) -> Tensor:\n","        V = Z.clone()\n","        for k in range(n - 3, -1, -1):\n","            tail_len = n - k - 1\n","            v = reflectors[k, :, :tail_len]\n","            v_col = v.unsqueeze(-1)\n","            V_sub = V[:, k + 1:, :]\n","            vtV = torch.bmm(v_col.transpose(-2, -1), V_sub)\n","            V[:, k + 1:, :] = V_sub - 2.0 * v_col @ vtV\n","        return V\n","\n","\n","# =============================================================================\n","# Validation\n","# =============================================================================\n","\n","class EighValidator(nn.Module):\n","\n","    def forward(self, A: Tensor, eigenvalues: Tensor,\n","                eigenvectors: Tensor) -> Tuple[Tensor, Tensor, Tensor]:\n","        B, n, _ = A.shape\n","        AV = torch.bmm(A, eigenvectors)\n","        VL = eigenvectors * eigenvalues.unsqueeze(-2)\n","        residual = AV - VL\n","        A_norm = torch.linalg.norm(A.reshape(B, -1), dim=-1).clamp(min=1e-30)\n","        residual_norm = torch.linalg.norm(residual.reshape(B, -1), dim=-1) / A_norm\n","\n","        VtV = torch.bmm(eigenvectors.transpose(-2, -1), eigenvectors)\n","        I = torch.eye(n, device=A.device, dtype=A.dtype).unsqueeze(0)\n","        orth_err = torch.linalg.norm((VtV - I).reshape(B, -1), dim=-1)\n","\n","        return residual_norm, orth_err, residual_norm.max()\n","\n","\n","# =============================================================================\n","# Top-Level: CompiledEigh\n","# =============================================================================\n","\n","class CompiledEigh(nn.Module):\n","    \"\"\"\n","    Drop-in replacement for torch.linalg.eigh.\n","\n","    Usage:\n","        solver = CompiledEigh(max_n=6)\n","        solver = torch.compile(solver, fullgraph=True)\n","        eigenvalues, eigenvectors = solver(A)\n","    \"\"\"\n","\n","    def __init__(self, max_n: int,\n","                 use_jacobi: Optional[bool] = None,\n","                 max_newton: int = DEFAULT_MAX_NEWTON,\n","                 max_jacobi_sweeps: int = DEFAULT_MAX_JACOBI_SWEEPS,\n","                 eps: float = 1e-30, tol: float = 1e-7):\n","        super().__init__()\n","        self.max_n = max_n\n","        self.eps = eps\n","\n","        if use_jacobi is None:\n","            use_jacobi = (max_n <= JACOBI_THRESHOLD)\n","        self.use_jacobi = use_jacobi\n","\n","        if use_jacobi:\n","            self.jacobi = JacobiEigh(\n","                max_n=max_n, max_sweeps=max_jacobi_sweeps, eps=eps)\n","        else:\n","            self.tridiag = HouseholderTridiagonalizer(max_n=max_n, eps=eps)\n","            self.dc = TensorTreeDC(\n","                max_n=max_n, max_newton=max_newton, eps=eps, tol=tol)\n","            self.back_accum = HouseholderBackAccumulate(max_n=max_n, eps=eps)\n","\n","        self.validator = EighValidator()\n","\n","    def forward(self, A: Tensor, validate: bool = False\n","                ) -> Tuple[Tensor, Tensor]:\n","        B, n, _ = A.shape\n","\n","        if self.use_jacobi:\n","            eigenvalues, eigenvectors = self.jacobi(A)\n","        else:\n","            A_work = A.clone()\n","            d = torch.empty(B, n, device=A.device, dtype=A.dtype)\n","            e = torch.empty(B, n - 1, device=A.device, dtype=A.dtype)\n","            reflectors = torch.zeros(max(n - 2, 1), B, n,\n","                                     device=A.device, dtype=A.dtype)\n","            self.tridiag(A_work, d, e, reflectors)\n","            eigenvalues, Z = self.dc(d, e)\n","            eigenvectors = self.back_accum(reflectors, Z, n)\n","            # Newton-Schulz re-orthogonalization for D&C path\n","            eigenvectors = orthogonalize_ns(eigenvectors, n_iter=2)\n","\n","        if validate:\n","            res_norm, orth_err, max_err = self.validator(A, eigenvalues, eigenvectors)\n","            print(f\"[CompiledEigh] max residual: {max_err.item():.2e}, \"\n","                  f\"mean orth err: {orth_err.mean().item():.2e}\")\n","\n","        return eigenvalues, eigenvectors\n","\n","\n","# =============================================================================\n","# Functional API\n","# =============================================================================\n","\n","_cached_solvers = {}\n","\n","def compiled_eigh(A: Tensor, validate: bool = False) -> Tuple[Tensor, Tensor]:\n","    B, n, _ = A.shape\n","    key = (n, A.device, A.dtype)\n","    if key not in _cached_solvers:\n","        _cached_solvers[key] = CompiledEigh(max_n=n).to(A.device)\n","    return _cached_solvers[key](A, validate=validate)\n","\n","\n","\"\"\"\n","CompiledEigh — Colab GPU Benchmark v3\n","Fixes:\n","  v2: Jacobi pairs as plain Python lists (Dynamo compile fix), sweeps 6→10\n","  v3: Replaced Gram-Schmidt with Newton-Schulz orthogonalization (all-bmm),\n","      disabled TF32 to ensure fp32 precision on Blackwell\n","\"\"\"\n","\n","import torch\n","import time\n","import gc\n","import sys\n","\n","# ── Ensure full fp32 precision on Ampere/Hopper/Blackwell ──\n","# TF32 uses 10-bit mantissa for matmul which can degrade orthogonality\n","torch.backends.cuda.matmul.allow_tf32 = False\n","torch.backends.cudnn.allow_tf32 = False\n","torch.set_float32_matmul_precision('highest')\n","\n","\n","def sync():\n","    if torch.cuda.is_available():\n","        torch.cuda.synchronize()\n","\n","\n","def gpu_timer(fn, warmup=10, repeats=200):\n","    for _ in range(warmup):\n","        fn()\n","    sync()\n","    start = time.perf_counter()\n","    for _ in range(repeats):\n","        fn()\n","    sync()\n","    return (time.perf_counter() - start) / repeats\n","\n","\n","def make_symmetric_batch(B, n, device, dtype=torch.float32):\n","    R = torch.randn(B, n, n, device=device, dtype=dtype)\n","    return (R + R.transpose(-2, -1)) / 2.0\n","\n","\n","def make_cm_like_batch(B, n, device, dtype=torch.float32):\n","    points = torch.randn(B, n, n, device=device, dtype=dtype)\n","    points = points / (points.norm(dim=-1, keepdim=True) + 1e-8)\n","    return torch.bmm(points, points.transpose(-2, -1)) * 0.3\n","\n","\n","def fmt_time(seconds):\n","    if seconds < 1e-3:\n","        return f\"{seconds*1e6:.1f} us\"\n","    elif seconds < 1.0:\n","        return f\"{seconds*1e3:.2f} ms\"\n","    return f\"{seconds:.3f} s\"\n","\n","\n","# ─── Test 0: Newton-Schulz Diagnostic ───\n","\n","def test_ns_diagnostic(device):\n","    \"\"\"Verify Newton-Schulz orthogonalization works on GPU independently.\"\"\"\n","    print(\"\\n\" + \"=\" * 70)\n","    print(\"  TEST 0: NEWTON-SCHULZ DIAGNOSTIC\")\n","    print(\"=\" * 70)\n","\n","    for n in [5, 6, 8]:\n","        B = 1024\n","        # Create nearly-orthogonal matrix (simulating Jacobi output)\n","        Q, _ = torch.linalg.qr(torch.randn(B, n, n, device=device))\n","        # Perturb to ~1e-3 orthogonality error\n","        noise = torch.randn(B, n, n, device=device) * 1e-3\n","        V_dirty = Q + noise\n","\n","        I_n = torch.eye(n, device=device).unsqueeze(0)\n","\n","        # Before NS\n","        VtV_before = torch.bmm(V_dirty.transpose(-2, -1), V_dirty)\n","        orth_before = torch.linalg.norm((VtV_before - I_n).reshape(B, -1), dim=-1).max().item()\n","\n","        # After NS (2 iterations)\n","        V_clean = orthogonalize_ns(V_dirty, n_iter=2)\n","        VtV_after = torch.bmm(V_clean.transpose(-2, -1), V_clean)\n","        orth_after = torch.linalg.norm((VtV_after - I_n).reshape(B, -1), dim=-1).max().item()\n","\n","        # After NS (3 iterations for comparison)\n","        V_clean3 = orthogonalize_ns(V_dirty, n_iter=3)\n","        VtV_after3 = torch.bmm(V_clean3.transpose(-2, -1), V_clean3)\n","        orth_after3 = torch.linalg.norm((VtV_after3 - I_n).reshape(B, -1), dim=-1).max().item()\n","\n","        print(f\"  n={n}: before={orth_before:.2e}  \"\n","              f\"after(2iter)={orth_after:.2e}  \"\n","              f\"after(3iter)={orth_after3:.2e}\")\n","\n","    # Also test with actual Jacobi output\n","    print(f\"\\n  --- With actual Jacobi output ---\")\n","    for n in [5, 6]:\n","        B = 2048\n","        A = make_symmetric_batch(B, n, device)\n","        solver = JacobiEigh(max_n=n, max_sweeps=10).to(device)\n","\n","        # Run Jacobi WITHOUT the NS cleanup\n","        W = A.clone()\n","        V = torch.eye(n, device=device).unsqueeze(0).expand(B, -1, -1).clone()\n","        for _sweep in range(solver.max_sweeps):\n","            for idx in range(solver._n_pairs):\n","                p, q = solver._pairs_p[idx], solver._pairs_q[idx]\n","                app, aqq, apq = W[:, p, p], W[:, q, q], W[:, p, q]\n","                two_apq = 2.0 * apq\n","                diff = aqq - app\n","                abs_2apq = two_apq.abs().clamp(min=1e-30)\n","                sign_2apq = torch.where(two_apq >= 0,\n","                    torch.ones_like(two_apq), -torch.ones_like(two_apq))\n","                tau = diff / (abs_2apq * sign_2apq)\n","                tau_sign = torch.where(tau >= 0,\n","                    torch.ones_like(tau), -torch.ones_like(tau))\n","                t = tau_sign / (tau.abs() + torch.sqrt(1.0 + tau * tau))\n","                skip = (apq.abs() < 1e-30).float()\n","                t = t * (1.0 - skip)\n","                c = 1.0 / torch.sqrt(1.0 + t * t)\n","                s = t * c\n","                c_col, s_col = c.unsqueeze(-1), s.unsqueeze(-1)\n","                Wp = W[:, :, p].clone(); Wq = W[:, :, q].clone()\n","                W[:, :, p] = c_col * Wp - s_col * Wq\n","                W[:, :, q] = s_col * Wp + c_col * Wq\n","                Wp = W[:, p, :].clone(); Wq = W[:, q, :].clone()\n","                W[:, p, :] = c_col * Wp - s_col * Wq\n","                W[:, q, :] = s_col * Wp + c_col * Wq\n","                W[:, p, q] = 0.0; W[:, q, p] = 0.0\n","                W[:, p, p] = app - t * apq\n","                W[:, q, q] = aqq + t * apq\n","                Vp = V[:, :, p].clone(); Vq = V[:, :, q].clone()\n","                V[:, :, p] = c_col * Vp - s_col * Vq\n","                V[:, :, q] = s_col * Vp + c_col * Vq\n","\n","        I_n = torch.eye(n, device=device).unsqueeze(0)\n","        VtV = torch.bmm(V.transpose(-2, -1), V)\n","        orth_raw = torch.linalg.norm((VtV - I_n).reshape(B, -1), dim=-1).max().item()\n","\n","        V_ns = orthogonalize_ns(V, n_iter=2)\n","        VtV_ns = torch.bmm(V_ns.transpose(-2, -1), V_ns)\n","        orth_ns = torch.linalg.norm((VtV_ns - I_n).reshape(B, -1), dim=-1).max().item()\n","\n","        print(f\"  Jacobi raw n={n}: orth={orth_raw:.2e}  after NS(2)={orth_ns:.2e}\")\n","\n","\n","# ─── Test 1: Accuracy ───\n","\n","def test_accuracy(device):\n","    print(\"\\n\" + \"=\" * 70)\n","    print(\"  TEST 1: ACCURACY vs torch.linalg.eigh\")\n","    print(\"=\" * 70)\n","\n","    validator = EighValidator()\n","    configs = [\n","        (3,   4096, \"3x3   small\"),\n","        (5,   4096, \"5x5   CM matrix size\"),\n","        (6,   4096, \"6x6   pentachoron bordered\"),\n","        (8,   2048, \"8x8   padded CM\"),\n","        (12,  1024, \"12x12 medium\"),\n","        (16,  512,  \"16x16 Jacobi boundary\"),\n","    ]\n","\n","    all_pass = True\n","    for n, B, label in configs:\n","        A = make_symmetric_batch(B, n, device)\n","        ref_vals, ref_vecs = torch.linalg.eigh(A)\n","\n","        solver = CompiledEigh(max_n=n).to(device)\n","        our_vals, our_vecs = solver(A)\n","\n","        val_err = (our_vals - ref_vals).abs().max().item()\n","        val_mean = (our_vals - ref_vals).abs().mean().item()\n","\n","        dots = torch.bmm(ref_vecs.transpose(-2, -1), our_vecs)\n","        alignment = dots.abs().max(dim=-1).values.min().item()\n","\n","        res_norm, orth_err, max_res = validator(A, our_vals, our_vecs)\n","        max_orth = orth_err.max().item()\n","\n","        # Thresholds: eigenval 1e-3, alignment 0.999, orth 1e-4\n","        ok = val_err < 1e-3 and alignment > 0.999 and max_orth < 1e-4\n","        if not ok:\n","            all_pass = False\n","\n","        print(f\"\\n  [{'PASS' if ok else 'FAIL'}] {label}  (n={n}, B={B})\")\n","        print(f\"    eigenvalue err   max={val_err:.2e}  mean={val_mean:.2e}\")\n","        print(f\"    eigvec alignment min={alignment:.8f}\")\n","        print(f\"    residual norm    max={max_res.item():.2e}\")\n","        print(f\"    orthogonality    max={max_orth:.2e}\")\n","\n","    print(f\"\\n  --- CM-like spectral distribution ---\")\n","    for n in [5, 6]:\n","        A = make_cm_like_batch(2048, n, device)\n","        ref_vals, _ = torch.linalg.eigh(A)\n","        solver = CompiledEigh(max_n=n).to(device)\n","        our_vals, our_vecs = solver(A)\n","        val_err = (our_vals - ref_vals).abs().max().item()\n","        res_norm, orth_err, max_res = validator(A, our_vals, our_vecs)\n","        print(f\"    CM-like n={n}: val_err={val_err:.2e}  \"\n","              f\"res={max_res.item():.2e}  orth={orth_err.max().item():.2e}\")\n","\n","    return all_pass\n","\n","\n","# ─── Test 2: torch.compile fullgraph ───\n","\n","def test_compile(device):\n","    print(\"\\n\" + \"=\" * 70)\n","    print(\"  TEST 2: torch.compile(fullgraph=True)\")\n","    print(\"=\" * 70)\n","\n","    results = {}\n","    for n, B, label in [(5, 1024, \"5x5\"), (6, 1024, \"6x6\"), (8, 512, \"8x8\")]:\n","        A = make_symmetric_batch(B, n, device)\n","        solver = CompiledEigh(max_n=n).to(device)\n","\n","        try:\n","            compiled_solver = torch.compile(solver, fullgraph=True)\n","            vals, vecs = compiled_solver(A)\n","            sync()\n","            ref_vals, _ = torch.linalg.eigh(A)\n","            err = (vals - ref_vals).abs().max().item()\n","            results[label] = (\"PASS\", err)\n","            print(f\"  [{label}] fullgraph=True SUCCESS  (val_err={err:.2e})\")\n","        except Exception as e:\n","            results[label] = (\"FAIL\", str(e)[:200])\n","            print(f\"  [{label}] COMPILE FAILED: {str(e)[:200]}\")\n","\n","    return all(v[0] == \"PASS\" for v in results.values())\n","\n","\n","# ─── Test 3: Throughput ───\n","\n","def test_benchmark(device):\n","    print(\"\\n\" + \"=\" * 70)\n","    print(\"  TEST 3: GPU THROUGHPUT BENCHMARK\")\n","    print(\"=\" * 70)\n","    print(f\"  Device: {torch.cuda.get_device_name(0)}\")\n","    print(f\"  Timing: 10 warmup + 200 repeats\\n\")\n","\n","    configs = [\n","        (5,  1024,  \"CM 5x5  B=1024\"),\n","        (5,  4096,  \"CM 5x5  B=4096\"),\n","        (5,  8192,  \"CM 5x5  B=8192\"),\n","        (6,  1024,  \"CM 6x6  B=1024\"),\n","        (6,  4096,  \"CM 6x6  B=4096\"),\n","        (6,  8192,  \"CM 6x6  B=8192\"),\n","        (8,  2048,  \"8x8     B=2048\"),\n","        (16, 1024,  \"16x16   B=1024\"),\n","    ]\n","\n","    print(f\"  {'Config':<22} {'eigh ref':>10} {'ours eager':>12} \"\n","          f\"{'ours compiled':>14} {'vs ref':>8}\")\n","    print(f\"  {'-'*22} {'-'*10} {'-'*12} {'-'*14} {'-'*8}\")\n","\n","    for n, B, label in configs:\n","        A = make_symmetric_batch(B, n, device)\n","\n","        ref_time = gpu_timer(lambda: torch.linalg.eigh(A))\n","\n","        solver = CompiledEigh(max_n=n).to(device)\n","        eager_time = gpu_timer(lambda: solver(A))\n","\n","        try:\n","            compiled_solver = torch.compile(solver, fullgraph=True)\n","            for _ in range(5):\n","                compiled_solver(A)\n","            sync()\n","            compiled_time = gpu_timer(lambda: compiled_solver(A))\n","            compiled_str = fmt_time(compiled_time)\n","            speedup = ref_time / compiled_time\n","            speedup_str = f\"{speedup:.2f}x\"\n","        except Exception:\n","            compiled_str = \"FAIL\"\n","            speedup_str = \"N/A\"\n","\n","        print(f\"  {label:<22} {fmt_time(ref_time):>10} \"\n","              f\"{fmt_time(eager_time):>12} {compiled_str:>14} {speedup_str:>8}\")\n","\n","    print(f\"\\n  --- High batch stress test ---\")\n","    for n in [5, 6]:\n","        for B in [16384, 32768]:\n","            try:\n","                A = make_symmetric_batch(B, n, device)\n","                solver = CompiledEigh(max_n=n).to(device)\n","                compiled_solver = torch.compile(solver, fullgraph=True)\n","                for _ in range(3):\n","                    compiled_solver(A)\n","                sync()\n","                t = gpu_timer(lambda: compiled_solver(A), warmup=5, repeats=100)\n","                ref_t = gpu_timer(lambda: torch.linalg.eigh(A), warmup=5, repeats=100)\n","                print(f\"    n={n} B={B}: compiled={fmt_time(t)}  ref={fmt_time(ref_t)}  \"\n","                      f\"ratio={ref_t/t:.2f}x  throughput={B/t:.0f}/sec\")\n","            except RuntimeError as e:\n","                if \"out of memory\" in str(e).lower():\n","                    print(f\"    n={n} B={B}: OOM\")\n","                    torch.cuda.empty_cache()\n","                else:\n","                    raise\n","\n","\n","# ─── Test 4: Autograd ───\n","\n","def test_autograd(device):\n","    print(\"\\n\" + \"=\" * 70)\n","    print(\"  TEST 4: AUTOGRAD BACKWARD\")\n","    print(\"=\" * 70)\n","\n","    for n, B in [(5, 512), (6, 512)]:\n","        A_ref = make_symmetric_batch(B, n, device).requires_grad_(True)\n","        vals_ref, vecs_ref = torch.linalg.eigh(A_ref)\n","        (vals_ref.sum() + (vecs_ref ** 2).sum()).backward()\n","        grad_ref = A_ref.grad.clone()\n","\n","        # Eager backward\n","        A_e = A_ref.detach().clone().requires_grad_(True)\n","        solver = CompiledEigh(max_n=n).to(device)\n","        try:\n","            vals_e, vecs_e = solver(A_e)\n","            (vals_e.sum() + (vecs_e ** 2).sum()).backward()\n","            err_e = (A_e.grad - grad_ref).abs().max().item()\n","            rel_e = err_e / (grad_ref.abs().max().item() + 1e-30)\n","            print(f\"  [{'PASS' if rel_e < 0.1 else 'WARN'}] n={n} eager backward: \"\n","                  f\"grad_err={err_e:.2e}  rel={rel_e:.2e}\")\n","        except Exception as e:\n","            print(f\"  [FAIL] n={n} eager backward: {e}\")\n","\n","        # Compiled backward (may break — forward fullgraph is the key win)\n","        A_c = A_ref.detach().clone().requires_grad_(True)\n","        try:\n","            compiled_solver = torch.compile(solver)\n","            vals_c, vecs_c = compiled_solver(A_c)\n","            (vals_c.sum() + (vecs_c ** 2).sum()).backward()\n","            err_c = (A_c.grad - grad_ref).abs().max().item()\n","            rel_c = err_c / (grad_ref.abs().max().item() + 1e-30)\n","            print(f\"  [{'PASS' if rel_c < 0.1 else 'WARN'}] n={n} compiled backward: \"\n","                  f\"grad_err={err_c:.2e}  rel={rel_c:.2e}\")\n","        except Exception as e:\n","            print(f\"  [INFO] n={n} compiled backward: {str(e)[:150]}\")\n","            print(f\"         (forward fullgraph is the main win)\")\n","\n","\n","# ─── Test 5: VRAM ───\n","\n","def test_vram(device):\n","    print(\"\\n\" + \"=\" * 70)\n","    print(\"  TEST 5: VRAM USAGE\")\n","    print(\"=\" * 70)\n","\n","    for n, B in [(5, 4096), (6, 4096), (6, 8192), (5, 8192)]:\n","        torch.cuda.empty_cache()\n","        gc.collect()\n","        torch.cuda.reset_peak_memory_stats()\n","        base_mem = torch.cuda.memory_allocated()\n","\n","        A = make_symmetric_batch(B, n, device)\n","        solver = CompiledEigh(max_n=n).to(device)\n","        vals, vecs = solver(A)\n","\n","        peak_mem = torch.cuda.max_memory_allocated()\n","        delta_mb = (peak_mem - base_mem) / (1024 ** 2)\n","        print(f\"  n={n} B={B}: peak delta = {delta_mb:.1f} MB\")\n","\n","        del A, solver, vals, vecs\n","        torch.cuda.empty_cache()\n","        gc.collect()\n","\n","\n","# ─── Main ───\n","\n","def main():\n","    print(\"=\" * 70)\n","    print(\"  CompiledEigh v3 — GPU Benchmark Suite\")\n","    print(\"=\" * 70)\n","\n","    if not torch.cuda.is_available():\n","        print(\"\\n  No CUDA. Run on Colab with A100/H100.\")\n","        sys.exit(1)\n","\n","    device = torch.device('cuda')\n","    print(f\"\\n  GPU:     {torch.cuda.get_device_name(0)}\")\n","    print(f\"  CUDA:    {torch.version.cuda}\")\n","    print(f\"  PyTorch: {torch.__version__}\")\n","    mem_gb = torch.cuda.get_device_properties(0).total_memory / (1024**3)\n","    print(f\"  VRAM:    {mem_gb:.1f} GB\")\n","    print(f\"  TF32 matmul: {torch.backends.cuda.matmul.allow_tf32}\")\n","    print(f\"  float32 precision: {torch.get_float32_matmul_precision()}\")\n","\n","    test_ns_diagnostic(device)\n","    acc_ok = test_accuracy(device)\n","    compile_ok = test_compile(device)\n","    test_benchmark(device)\n","    test_autograd(device)\n","    test_vram(device)\n","\n","    print(\"\\n\" + \"=\" * 70)\n","    print(\"  SUMMARY\")\n","    print(\"=\" * 70)\n","    print(f\"  Accuracy:  {'PASS' if acc_ok else 'FAIL'}\")\n","    print(f\"  Compile:   {'PASS' if compile_ok else 'FAIL'}\")\n","    print(\"=\" * 70)\n","\n","\n","if __name__ == '__main__':\n","    main()"],"metadata":{"colab":{"base_uri":"https://localhost:8080/","height":1000},"id":"crYe5tzkJLvs","executionInfo":{"status":"error","timestamp":1775017485514,"user_tz":420,"elapsed":728060,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"8ffab1a4-f441-44c1-a169-3bcf6d4ae17e"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["======================================================================\n","  CompiledEigh v3 — GPU Benchmark Suite\n","======================================================================\n","\n","  GPU:     NVIDIA RTX PRO 6000 Blackwell Server Edition\n","  CUDA:    12.8\n","  PyTorch: 2.10.0+cu128\n","  VRAM:    95.0 GB\n","  TF32 matmul: False\n","  float32 precision: highest\n","\n","======================================================================\n","  TEST 0: NEWTON-SCHULZ DIAGNOSTIC\n","======================================================================\n","  n=5: before=1.23e-02  after(2iter)=6.16e-03  after(3iter)=9.24e-03\n","  n=6: before=1.33e-02  after(2iter)=6.63e-03  after(3iter)=9.99e-03\n","  n=8: before=1.66e-02  after(2iter)=8.31e-03  after(3iter)=1.25e-02\n","\n","  --- With actual Jacobi output ---\n","  Jacobi raw n=5: orth=1.86e-06  after NS(2)=1.04e-06\n","  Jacobi raw n=6: orth=2.45e-06  after NS(2)=1.27e-06\n","\n","======================================================================\n","  TEST 1: ACCURACY vs torch.linalg.eigh\n","======================================================================\n","\n","  [PASS] 3x3   small  (n=3, B=4096)\n","    eigenvalue err   max=2.15e-06  mean=2.25e-07\n","    eigvec alignment min=0.99999946\n","    residual norm    max=6.80e-07\n","    orthogonality    max=6.13e-07\n","\n","  [PASS] 5x5   CM matrix size  (n=5, B=4096)\n","    eigenvalue err   max=5.01e-06  mean=4.69e-07\n","    eigvec alignment min=0.99999928\n","    residual norm    max=1.09e-06\n","    orthogonality    max=9.46e-07\n","\n","  [PASS] 6x6   pentachoron bordered  (n=6, B=4096)\n","    eigenvalue err   max=5.72e-06  mean=6.18e-07\n","    eigvec alignment min=0.99999905\n","    residual norm    max=1.33e-06\n","    orthogonality    max=1.28e-06\n","\n","  [PASS] 8x8   padded CM  (n=8, B=2048)\n","    eigenvalue err   max=9.06e-06  mean=9.37e-07\n","    eigvec alignment min=0.99999905\n","    residual norm    max=1.37e-06\n","    orthogonality    max=1.68e-06\n","\n","  [PASS] 12x12 medium  (n=12, B=1024)\n","    eigenvalue err   max=1.29e-05  mean=1.87e-06\n","    eigvec alignment min=0.99999893\n","    residual norm    max=1.89e-06\n","    orthogonality    max=2.78e-06\n","\n","  [PASS] 16x16 Jacobi boundary  (n=16, B=512)\n","    eigenvalue err   max=1.91e-05  mean=3.12e-06\n","    eigvec alignment min=0.99999869\n","    residual norm    max=2.45e-06\n","    orthogonality    max=4.25e-06\n","\n","  --- CM-like spectral distribution ---\n","    CM-like n=5: val_err=1.07e-06  res=1.02e-06  orth=1.05e-06\n","    CM-like n=6: val_err=1.31e-06  res=1.26e-06  orth=1.24e-06\n","\n","======================================================================\n","  TEST 2: torch.compile(fullgraph=True)\n","======================================================================\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["  [5x5] fullgraph=True SUCCESS  (val_err=3.34e-06)\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["  [6x6] fullgraph=True SUCCESS  (val_err=7.15e-06)\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["  [8x8] fullgraph=True SUCCESS  (val_err=5.25e-06)\n","\n","======================================================================\n","  TEST 3: GPU THROUGHPUT BENCHMARK\n","======================================================================\n","  Device: NVIDIA RTX PRO 6000 Blackwell Server Edition\n","  Timing: 10 warmup + 200 repeats\n","\n","  Config                   eigh ref   ours eager  ours compiled   vs ref\n","  ---------------------- ---------- ------------ -------------- --------\n","  CM 5x5  B=1024           114.1 us     30.88 ms        4.54 ms    0.03x\n"]},{"output_type":"error","ename":"KeyboardInterrupt","evalue":"","traceback":["\u001b[0;31m---------------------------------------------------------------------------\u001b[0m","\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)","\u001b[0;32m/tmp/ipykernel_61613/3247892068.py\u001b[0m in \u001b[0;36m<cell line: 0>\u001b[0;34m()\u001b[0m\n\u001b[1;32m   1047\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1048\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0m__name__\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m'__main__'\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1049\u001b[0;31m     \u001b[0mmain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m","\u001b[0;32m/tmp/ipykernel_61613/3247892068.py\u001b[0m in \u001b[0;36mmain\u001b[0;34m()\u001b[0m\n\u001b[1;32m   1034\u001b[0m     \u001b[0macc_ok\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtest_accuracy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1035\u001b[0m     \u001b[0mcompile_ok\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtest_compile\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1036\u001b[0;31m     \u001b[0mtest_benchmark\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1037\u001b[0m     \u001b[0mtest_autograd\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1038\u001b[0m     \u001b[0mtest_vram\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/tmp/ipykernel_61613/3247892068.py\u001b[0m in \u001b[0;36mtest_benchmark\u001b[0;34m(device)\u001b[0m\n\u001b[1;32m    909\u001b[0m             \u001b[0mcompiled_solver\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcompile\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msolver\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfullgraph\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    910\u001b[0m             \u001b[0;32mfor\u001b[0m \u001b[0m_\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m5\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 911\u001b[0;31m                 \u001b[0mcompiled_solver\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mA\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    912\u001b[0m             \u001b[0msync\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    913\u001b[0m             \u001b[0mcompiled_time\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mgpu_timer\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;32mlambda\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mcompiled_solver\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mA\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_dynamo/eval_frame.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    463\u001b[0m             )\n\u001b[1;32m    464\u001b[0m         \u001b[0;32mwith\u001b[0m \u001b[0m_set_in_optimized_module\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 465\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0msuper\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__call__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    466\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    467\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_aot_compile\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minputs\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mlist\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_dynamo\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0maot_compile\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mModelInput\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1774\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_compiled_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[misc]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1775\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1776\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1777\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1778\u001b[0m     \u001b[0;31m# torchrec tests the code consistency with the following code\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1785\u001b[0m                 \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_pre_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_hooks\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1786\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1787\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1788\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1789\u001b[0m         \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_dynamo/eval_frame.py\u001b[0m in \u001b[0;36mcompile_wrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    951\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    952\u001b[0m                 \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 953\u001b[0;31m                     \u001b[0;32mreturn\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    954\u001b[0m                 \u001b[0;32mexcept\u001b[0m \u001b[0mUnsupported\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    955\u001b[0m                     \u001b[0;32mif\u001b[0m \u001b[0mconfig\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mverbose\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1774\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_compiled_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[misc]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1775\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1776\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1777\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1778\u001b[0m     \u001b[0;31m# torchrec tests the code consistency with the following code\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1785\u001b[0m                 \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_pre_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_hooks\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1786\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1787\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1788\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1789\u001b[0m         \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_dynamo/convert_frame.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, frame, cache_entry, frame_state)\u001b[0m\n\u001b[1;32m   2200\u001b[0m         \u001b[0;32mwith\u001b[0m \u001b[0mcompile_lock\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_disable_current_modes\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2201\u001b[0m             \u001b[0;31m# skip=1: skip this frame\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2202\u001b[0;31m             result = self._torchdynamo_orig_backend(\n\u001b[0m\u001b[1;32m   2203\u001b[0m                 \u001b[0mframe\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcache_entry\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mhooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mframe_state\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mskip\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2204\u001b[0m             )\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_dynamo/convert_frame.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, frame, cache_entry, hooks, frame_state, skip)\u001b[0m\n\u001b[1;32m    705\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    706\u001b[0m         \u001b[0;32mwith\u001b[0m \u001b[0mcompile_context\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mCompileContext\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcompile_id\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 707\u001b[0;31m             result = _compile(\n\u001b[0m\u001b[1;32m    708\u001b[0m                 \u001b[0mframe\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mf_code\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    709\u001b[0m                 \u001b[0mframe\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mf_globals\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_dynamo/convert_frame.py\u001b[0m in \u001b[0;36m_compile\u001b[0;34m(code, globals, locals, builtins, closure, compiler_fn, one_graph, export, export_constraints, hooks, cache_entry, cache_size, frame, frame_state, compile_id, skip, package, convert_frame_box)\u001b[0m\n\u001b[1;32m   1750\u001b[0m         \u001b[0mtracer_output\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1751\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1752\u001b[0;31m             \u001b[0mguarded_code\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtracer_output\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcompile_inner\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcode\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mone_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhooks\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1753\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1754\u001b[0m             \u001b[0;31m# NB: We only put_code_state in success case.  Success case here\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_utils_internal.py\u001b[0m in \u001b[0;36mwrapper_function\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m     95\u001b[0m             \u001b[0;31m# in stack traces when profiling is not enabled.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     96\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mStrobelightCompileTimeProfiler\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0menabled\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 97\u001b[0;31m                 \u001b[0;32mreturn\u001b[0m \u001b[0mfunction\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     98\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     99\u001b[0m             return StrobelightCompileTimeProfiler.profile_compile_time(\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_dynamo/convert_frame.py\u001b[0m in \u001b[0;36mcompile_inner\u001b[0;34m(code, one_graph, hooks)\u001b[0m\n\u001b[1;32m   1431\u001b[0m             )\n\u001b[1;32m   1432\u001b[0m             \u001b[0mstack\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0menter_context\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mCompileTimeInstructionCounter\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrecord\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1433\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0m_compile_inner\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcode\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mone_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhooks\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1434\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1435\u001b[0m         return (\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_dynamo/convert_frame.py\u001b[0m in \u001b[0;36m_compile_inner\u001b[0;34m(code, one_graph, hooks)\u001b[0m\n\u001b[1;32m   1465\u001b[0m         \u001b[0mout_code\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1466\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1467\u001b[0;31m             dynamo_output = compile_frame(\n\u001b[0m\u001b[1;32m   1468\u001b[0m                 \u001b[0mcode\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1469\u001b[0m                 \u001b[0mglobals\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_dynamo/convert_frame.py\u001b[0m in \u001b[0;36mcompile_frame\u001b[0;34m(code, globals, locals, builtins, closure, compiler_fn, one_graph, restart_reasons, export, export_constraints, frame_state, distributed_state, package)\u001b[0m\n\u001b[1;32m   1339\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1340\u001b[0m             \u001b[0;32mwith\u001b[0m \u001b[0mdynamo_timed\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf\"compile_attempt_{attempt}\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlog_pt2_compile_event\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1341\u001b[0;31m                 \u001b[0mbytecode\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtracer_output\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtransform_code_object\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcode\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtransform\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1342\u001b[0m                 \u001b[0;32massert\u001b[0m \u001b[0mtracer_output\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1343\u001b[0m                 return DynamoOutput(\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_dynamo/bytecode_transformation.py\u001b[0m in \u001b[0;36mtransform_code_object\u001b[0;34m(code, transformations, safe)\u001b[0m\n\u001b[1;32m   1598\u001b[0m     \u001b[0mpropagate_line_nums\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minstructions\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1599\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1600\u001b[0;31m     \u001b[0mtracer_output\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtransformations\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minstructions\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcode_options\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1601\u001b[0m     \u001b[0m_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbytecode\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mclean_and_assemble_instructions\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minstructions\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkeys\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcode_options\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1602\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mbytecode\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtracer_output\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_dynamo/convert_frame.py\u001b[0m in \u001b[0;36mtransform\u001b[0;34m(instructions, code_options)\u001b[0m\n\u001b[1;32m   1311\u001b[0m             \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0moverrides\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_get_current_function_mode_stack\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1312\u001b[0m         )\n\u001b[0;32m-> 1313\u001b[0;31m         tracer_output = trace_frame(\n\u001b[0m\u001b[1;32m   1314\u001b[0m             \u001b[0mcode\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1315\u001b[0m             \u001b[0mglobals\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_dynamo/convert_frame.py\u001b[0m in \u001b[0;36m_fn\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    326\u001b[0m             \u001b[0mreset_user_object_tracking\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    327\u001b[0m             \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 328\u001b[0;31m                 \u001b[0;32mreturn\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    329\u001b[0m             \u001b[0;32mfinally\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    330\u001b[0m                 \u001b[0mcleanup\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mclose\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_dynamo/convert_frame.py\u001b[0m in \u001b[0;36mtrace_frame\u001b[0;34m(code, globals, locals, builtins, closure, compiler_fn, tf_mode_stack, one_graph, speculation_log, instructions, code_options, export, export_constraints, frame_state, distributed_state, package)\u001b[0m\n\u001b[1;32m    836\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    837\u001b[0m     \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 838\u001b[0;31m         \u001b[0mrun_tracer\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    839\u001b[0m         \u001b[0mtracer_output\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mDynamoTracerOutput\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtracer\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    840\u001b[0m         \u001b[0moutput\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtracer_output\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0moutput_graph\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_dynamo/convert_frame.py\u001b[0m in \u001b[0;36mrun_tracer\u001b[0;34m()\u001b[0m\n\u001b[1;32m    817\u001b[0m             \u001b[0mtracer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0moutput\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmark_bytecode_tracing_start\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    818\u001b[0m             \u001b[0;32mwith\u001b[0m \u001b[0mtracing\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtracer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0moutput\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtracing_context\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtracer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mset_current_tx\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 819\u001b[0;31m                 \u001b[0mtracer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    820\u001b[0m         \u001b[0;32mexcept\u001b[0m \u001b[0mexc\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mUnspecializeRestartAnalysis\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    821\u001b[0m             \u001b[0mspeculation_log\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mclear\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[has-type]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_dynamo/symbolic_convert.py\u001b[0m in \u001b[0;36mrun\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m   1652\u001b[0m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstart_point\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minstruction_pointer\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1653\u001b[0m                 \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1654\u001b[0;31m                     \u001b[0;32mwhile\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1655\u001b[0m                         \u001b[0;32mpass\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1656\u001b[0m                 \u001b[0;32mexcept\u001b[0m \u001b[0mException\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_dynamo/symbolic_convert.py\u001b[0m in \u001b[0;36mstep\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m   1332\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1333\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1334\u001b[0;31m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdispatch_table\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0minst\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mopcode\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minst\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1335\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0moutput\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshould_exit\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1336\u001b[0m         \u001b[0;32mexcept\u001b[0m \u001b[0mTensorifyScalarRestartAnalysis\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_dynamo/symbolic_convert.py\u001b[0m in \u001b[0;36mRETURN_VALUE\u001b[0;34m(self, inst)\u001b[0m\n\u001b[1;32m   4690\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   4691\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mRETURN_VALUE\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minst\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mInstruction\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 4692\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_return\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minst\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   4693\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   4694\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mRETURN_CONST\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minst\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mInstruction\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_dynamo/symbolic_convert.py\u001b[0m in \u001b[0;36m_return\u001b[0;34m(self, inst)\u001b[0m\n\u001b[1;32m   4672\u001b[0m         )\n\u001b[1;32m   4673\u001b[0m         \u001b[0mlog\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdebug\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"return triggered compile\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 4674\u001b[0;31m         all_stack_locals_metadata = self.output.compile_subgraph(\n\u001b[0m\u001b[1;32m   4675\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   4676\u001b[0m             reason=GraphCompileReason(\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_dynamo/output_graph.py\u001b[0m in \u001b[0;36mcompile_subgraph\u001b[0;34m(self, tx, reason, partial_convert, stack_pops)\u001b[0m\n\u001b[1;32m   1698\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mcount_calls\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgraph\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m!=\u001b[0m \u001b[0;36m0\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpass2\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgraph_outputs\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m!=\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1699\u001b[0m                 output.extend(\n\u001b[0;32m-> 1700\u001b[0;31m                     \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcompile_and_call_fx_graph\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpass2\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgraph_output_vars\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mroot\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1701\u001b[0m                 )\n\u001b[1;32m   1702\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_dynamo/output_graph.py\u001b[0m in \u001b[0;36mcompile_and_call_fx_graph\u001b[0;34m(self, tx, rv, root)\u001b[0m\n\u001b[1;32m   2223\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2224\u001b[0m             \u001b[0;32mwith\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrestore_global_state\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2225\u001b[0;31m                 \u001b[0mcompiled_fn\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcall_user_compiler\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mgm\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexample_inputs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   2226\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2227\u001b[0m             \u001b[0;32mfrom\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_lazy_graph_module\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0m_LazyGraphModule\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_dynamo/output_graph.py\u001b[0m in \u001b[0;36mcall_user_compiler\u001b[0;34m(self, gm, example_inputs)\u001b[0m\n\u001b[1;32m   2375\u001b[0m             \u001b[0mdynamo_compile_column_us\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m\"aot_autograd_cumulative_compile_time_us\"\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2376\u001b[0m         ):\n\u001b[0;32m-> 2377\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_user_compiler\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mgm\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mexample_inputs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   2378\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2379\u001b[0m     def _call_user_compiler(\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_dynamo/output_graph.py\u001b[0m in \u001b[0;36m_call_user_compiler\u001b[0;34m(self, gm, example_inputs)\u001b[0m\n\u001b[1;32m   2410\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mconfig\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mverify_correctness\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2411\u001b[0m                 \u001b[0mcompiler_fn\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mWrapperBackend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcompiler_fn\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2412\u001b[0;31m             \u001b[0mcompiled_fn\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcompiler_fn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mgm\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mexample_inputs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   2413\u001b[0m             \u001b[0m_step_logger\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlogging\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mINFO\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34mf\"done compiler function {name}\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2414\u001b[0m             \u001b[0;32massert\u001b[0m \u001b[0mcallable\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcompiled_fn\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"compiler_fn did not return callable\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_dynamo/repro/after_dynamo.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, gm, example_inputs, **kwargs)\u001b[0m\n\u001b[1;32m    154\u001b[0m                     \u001b[0;32mraise\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    155\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 156\u001b[0;31m             \u001b[0mcompiled_gm\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcompiler_fn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mgm\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mexample_inputs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    157\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    158\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mcompiled_gm\u001b[0m  \u001b[0;31m# type: ignore[return-value]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/__init__.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, model_, inputs_)\u001b[0m\n\u001b[1;32m   2433\u001b[0m         \u001b[0;32mfrom\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_inductor\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcompile_fx\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mcompile_fx\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2434\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2435\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mcompile_fx\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minputs_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mconfig_patches\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mconfig\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   2436\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2437\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mget_compiler_config\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_inductor/compile_fx.py\u001b[0m in \u001b[0;36mcompile_fx\u001b[0;34m(model_, example_inputs_, inner_compile, config_patches, decompositions, ignore_shape_env)\u001b[0m\n\u001b[1;32m   2535\u001b[0m                 )\n\u001b[1;32m   2536\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2537\u001b[0;31m     return _maybe_wrap_and_compile_fx_main(\n\u001b[0m\u001b[1;32m   2538\u001b[0m         \u001b[0mmodel_\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2539\u001b[0m         \u001b[0mexample_inputs_\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_inductor/compile_fx.py\u001b[0m in \u001b[0;36m_maybe_wrap_and_compile_fx_main\u001b[0;34m(model_, example_inputs_, inner_compile, decompositions, ignore_shape_env)\u001b[0m\n\u001b[1;32m   2612\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2613\u001b[0m     \u001b[0;31m# Finally do the actual work!\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2614\u001b[0;31m     return _compile_fx_main(\n\u001b[0m\u001b[1;32m   2615\u001b[0m         \u001b[0mmodel_\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2616\u001b[0m         \u001b[0mexample_inputs_\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_inductor/compile_fx.py\u001b[0m in \u001b[0;36m_compile_fx_main\u001b[0;34m(model_, example_inputs_, inner_compile, decompositions, ignore_shape_env)\u001b[0m\n\u001b[1;32m   2807\u001b[0m         ):\n\u001b[1;32m   2808\u001b[0m             \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2809\u001b[0;31m                 return aot_autograd(\n\u001b[0m\u001b[1;32m   2810\u001b[0m                     \u001b[0mfw_compiler\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mfw_compiler\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2811\u001b[0m                     \u001b[0mbw_compiler\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mbw_compiler\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_dynamo/backends/common.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, gm, example_inputs, **kwargs)\u001b[0m\n\u001b[1;32m    121\u001b[0m             \u001b[0;31m# NB: NOT cloned!\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    122\u001b[0m             \u001b[0;32mwith\u001b[0m \u001b[0menable_aot_logging\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpatch_config\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 123\u001b[0;31m                 \u001b[0mcg\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0maot_module_simplified\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mgm\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mexample_inputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    124\u001b[0m                 \u001b[0mcounters\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"aot_autograd\"\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"ok\"\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m+=\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    125\u001b[0m                 \u001b[0;32mreturn\u001b[0m \u001b[0mdisable\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcg\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mreason\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m\"do not trace AOT-compiled graph\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_functorch/aot_autograd.py\u001b[0m in \u001b[0;36maot_module_simplified\u001b[0;34m(mod, args, fw_compiler, bw_compiler, partition_fn, decompositions, keep_inference_input_mutations, inference_compiler, cudagraphs, boxed_forward_device_index, ignore_shape_env, disable_functionalization)\u001b[0m\n\u001b[1;32m   1112\u001b[0m                 \u001b[0mshape_env\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1113\u001b[0m             )\n\u001b[0;32m-> 1114\u001b[0;31m             \u001b[0maot_graph_capture\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0maot_stage1_graph_capture\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0maot_state\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfunctional_call\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1115\u001b[0m             compiled_fn, _ = aot_stage2_compile(\n\u001b[1;32m   1116\u001b[0m                 \u001b[0maot_state\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_functorch/_aot_autograd/graph_compile.py\u001b[0m in \u001b[0;36maot_stage1_graph_capture\u001b[0;34m(aot_state, orig_flat_fn)\u001b[0m\n\u001b[1;32m    229\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    230\u001b[0m             graph, updated_flat_args, updated_flat_args_descs, maybe_subclass_meta = (\n\u001b[0;32m--> 231\u001b[0;31m                 aot_dispatch_base_graph(\n\u001b[0m\u001b[1;32m    232\u001b[0m                     \u001b[0mflat_fn\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    233\u001b[0m                     \u001b[0maot_state\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mflat_args\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_functorch/_aot_autograd/graph_capture.py\u001b[0m in \u001b[0;36maot_dispatch_base_graph\u001b[0;34m(flat_fn, flat_args, flat_args_descs, aot_config, fw_metadata)\u001b[0m\n\u001b[1;32m    250\u001b[0m     )\n\u001b[1;32m    251\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 252\u001b[0;31m     fw_module = _create_graph(\n\u001b[0m\u001b[1;32m    253\u001b[0m         \u001b[0mfn_to_trace\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    254\u001b[0m         \u001b[0mupdated_flat_args_subclasses_desugared\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_functorch/_aot_autograd/graph_capture.py\u001b[0m in \u001b[0;36m_create_graph\u001b[0;34m(f, args, args_descs, aot_config)\u001b[0m\n\u001b[1;32m     87\u001b[0m         \u001b[0mctx\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     88\u001b[0m     ):\n\u001b[0;32m---> 89\u001b[0;31m         fx_g = make_fx(\n\u001b[0m\u001b[1;32m     90\u001b[0m             \u001b[0minner_f\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     91\u001b[0m             \u001b[0mdecomposition_table\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0maot_config\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdecompositions\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py\u001b[0m in \u001b[0;36mwrapped\u001b[0;34m(*args)\u001b[0m\n\u001b[1;32m   2699\u001b[0m     \u001b[0;34m@\u001b[0m\u001b[0mfunctools\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mwraps\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2700\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mwrapped\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mobject\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mGraphModule\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2701\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mmake_fx_tracer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrace\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   2702\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2703\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mwrapped\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py\u001b[0m in \u001b[0;36mtrace\u001b[0;34m(self, f, *args)\u001b[0m\n\u001b[1;32m   2624\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mtrace\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mCallable\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mobject\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mfx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mGraphModule\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2625\u001b[0m         \u001b[0;32mwith\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_init_modes_from_inputs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2626\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_trace_inner\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   2627\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2628\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mis_hop_subgraph_tracer\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mbool\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py\u001b[0m in \u001b[0;36m_trace_inner\u001b[0;34m(self, f, *args)\u001b[0m\n\u001b[1;32m   2586\u001b[0m             \u001b[0;32massert\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfx_tracer\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2587\u001b[0m             \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2588\u001b[0;31m                 t = dispatch_trace(\n\u001b[0m\u001b[1;32m   2589\u001b[0m                     \u001b[0mwrap_key\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfunc\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfx_tracer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpre_dispatch\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2590\u001b[0m                     \u001b[0mtracer\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfx_tracer\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_compile.py\u001b[0m in \u001b[0;36minner\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m     52\u001b[0m                 \u001b[0mfn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__dynamo_disable\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdisable_fn\u001b[0m  \u001b[0;31m# type: ignore[attr-defined]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     53\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 54\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mdisable_fn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     55\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     56\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0minner\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_dynamo/eval_frame.py\u001b[0m in \u001b[0;36m_fn\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m   1179\u001b[0m                         ):\n\u001b[1;32m   1180\u001b[0m                             \u001b[0;32mreturn\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1181\u001b[0;31m                     \u001b[0;32mreturn\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1182\u001b[0m                 \u001b[0;32mfinally\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1183\u001b[0m                     \u001b[0mset_eval_frame\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py\u001b[0m in \u001b[0;36mdispatch_trace\u001b[0;34m(root, tracer, concrete_args)\u001b[0m\n\u001b[1;32m   1458\u001b[0m     \u001b[0mconcrete_args\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mOptional\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mtuple\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mAny\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m...\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1459\u001b[0m ) -> GraphModule:\n\u001b[0;32m-> 1460\u001b[0;31m     \u001b[0mgraph\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtracer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrace\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mroot\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mconcrete_args\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[arg-type]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1461\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1462\u001b[0m     \u001b[0;31m# NB: be careful not to DCE .item() calls\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_dynamo/eval_frame.py\u001b[0m in \u001b[0;36m_fn\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m   1179\u001b[0m                         ):\n\u001b[1;32m   1180\u001b[0m                             \u001b[0;32mreturn\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1181\u001b[0;31m                     \u001b[0;32mreturn\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1182\u001b[0m                 \u001b[0;32mfinally\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1183\u001b[0m                     \u001b[0mset_eval_frame\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/fx/_symbolic_trace.py\u001b[0m in \u001b[0;36mtrace\u001b[0;34m(self, root, concrete_args)\u001b[0m\n\u001b[1;32m    877\u001b[0m                     \u001b[0;34m\"output\"\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    878\u001b[0m                     \u001b[0;34m\"output\"\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 879\u001b[0;31m                     \u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcreate_arg\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    880\u001b[0m                     \u001b[0;34m{\u001b[0m\u001b[0;34m}\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    881\u001b[0m                     \u001b[0mtype_expr\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mann\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"return\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py\u001b[0m in \u001b[0;36mwrapped\u001b[0;34m(*proxies, **_unused)\u001b[0m\n\u001b[1;32m   1524\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mget_proxy_slot\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mt\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtracer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mt\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;32mlambda\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mproxy\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[attr-defined]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1525\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1526\u001b[0;31m         \u001b[0mout\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mtensors\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type:ignore[call-arg]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1527\u001b[0m         \u001b[0mout\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpytree\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtree_map_only\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mTensor\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mget_tensor_proxy_slot\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mout\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1528\u001b[0m         out = pytree.tree_map_only(\n","\u001b[0;32m<string>\u001b[0m in \u001b[0;36m<lambda>\u001b[0;34m(arg0, arg1, arg2, arg3)\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_functorch/_aot_autograd/graph_capture.py\u001b[0m in \u001b[0;36minner_f\u001b[0;34m(*args)\u001b[0m\n\u001b[1;32m     70\u001b[0m             \u001b[0;32mnonlocal\u001b[0m \u001b[0mout_descs\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     71\u001b[0m             \u001b[0;32massert\u001b[0m \u001b[0mout_descs\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 72\u001b[0;31m             \u001b[0mout\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mout_descs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcall_and_expect_output_descs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     73\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mout\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     74\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_functorch/_aot_autograd/utils.py\u001b[0m in \u001b[0;36mcall_and_expect_output_descs\u001b[0;34m(fn, args)\u001b[0m\n\u001b[1;32m    736\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    737\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mcall_and_expect_output_descs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 738\u001b[0;31m     \u001b[0mouts_pair\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    739\u001b[0m     \u001b[0;32massert\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mouts_pair\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtuple\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mouts_pair\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;36m2\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mfn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mouts_pair\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    740\u001b[0m     \u001b[0mouts\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mouts_descs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mouts_pair\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_functorch/_aot_autograd/graph_capture_wrappers.py\u001b[0m in \u001b[0;36minner_fn\u001b[0;34m(*args)\u001b[0m\n\u001b[1;32m   1104\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1105\u001b[0m             \u001b[0;31m# Run the joint\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1106\u001b[0;31m             \u001b[0mouts\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mouts_descs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcall_and_expect_output_descs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1107\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1108\u001b[0m         \u001b[0;31m# Return both the tokens and the outputs\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_functorch/_aot_autograd/utils.py\u001b[0m in \u001b[0;36mcall_and_expect_output_descs\u001b[0;34m(fn, args)\u001b[0m\n\u001b[1;32m    736\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    737\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mcall_and_expect_output_descs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 738\u001b[0;31m     \u001b[0mouts_pair\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    739\u001b[0m     \u001b[0;32massert\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mouts_pair\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtuple\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mouts_pair\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;36m2\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mfn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mouts_pair\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    740\u001b[0m     \u001b[0mouts\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mouts_descs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mouts_pair\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_functorch/_aot_autograd/graph_capture_wrappers.py\u001b[0m in \u001b[0;36m_functionalized_f_helper\u001b[0;34m(*args)\u001b[0m\n\u001b[1;32m    810\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    811\u001b[0m                 \u001b[0;31m# Run the joint\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 812\u001b[0;31m                 \u001b[0mf_outs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mf_outs_descs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcall_and_expect_output_descs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mf_args\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    813\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    814\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mtrace_joint\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_functorch/_aot_autograd/utils.py\u001b[0m in \u001b[0;36mcall_and_expect_output_descs\u001b[0;34m(fn, args)\u001b[0m\n\u001b[1;32m    736\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    737\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mcall_and_expect_output_descs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 738\u001b[0;31m     \u001b[0mouts_pair\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    739\u001b[0m     \u001b[0;32massert\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mouts_pair\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtuple\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mouts_pair\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;36m2\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mfn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mouts_pair\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    740\u001b[0m     \u001b[0mouts\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mouts_descs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mouts_pair\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_functorch/_aot_autograd/graph_capture_wrappers.py\u001b[0m in \u001b[0;36minner_fn\u001b[0;34m(*args)\u001b[0m\n\u001b[1;32m    110\u001b[0m     \u001b[0;34m@\u001b[0m\u001b[0msimple_wraps\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfn\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    111\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0minner_fn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 112\u001b[0;31m         \u001b[0mouts\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mouts_descs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcall_and_expect_output_descs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    113\u001b[0m         \u001b[0;32massert\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmeta\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0moutput_info\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mouts\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    114\u001b[0m         \u001b[0;31m# The compiled fw will return mutated input tensors, *including* metadata-only mutation.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_functorch/_aot_autograd/utils.py\u001b[0m in \u001b[0;36mcall_and_expect_output_descs\u001b[0;34m(fn, args)\u001b[0m\n\u001b[1;32m    736\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    737\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mcall_and_expect_output_descs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 738\u001b[0;31m     \u001b[0mouts_pair\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    739\u001b[0m     \u001b[0;32massert\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mouts_pair\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtuple\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mouts_pair\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;36m2\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mfn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mouts_pair\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    740\u001b[0m     \u001b[0mouts\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mouts_descs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mouts_pair\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_functorch/_aot_autograd/graph_compile.py\u001b[0m in \u001b[0;36morig_flat_fn2\u001b[0;34m(*args)\u001b[0m\n\u001b[1;32m    183\u001b[0m     \u001b[0;34m@\u001b[0m\u001b[0msimple_wraps\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0morig_flat_fn\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    184\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0morig_flat_fn2\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mFxValue\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mtuple\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mlist\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mFxValue\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlist\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mAOTOutput\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 185\u001b[0;31m         \u001b[0mout\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0morig_flat_fn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    186\u001b[0m         out_descs: list[AOTOutput] = type(out)(  # type: ignore[assignment]\n\u001b[1;32m    187\u001b[0m             \u001b[0mPlainAOTOutput\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[misc]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_functorch/_aot_autograd/graph_capture_wrappers.py\u001b[0m in \u001b[0;36mfunctional_call\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m   1373\u001b[0m                         \u001b[0;32massert\u001b[0m \u001b[0mfake_mode\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1374\u001b[0m                         \u001b[0mfake_mode\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mepoch\u001b[0m \u001b[0;34m+=\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1375\u001b[0;31m                         \u001b[0mout\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mPropagateUnbackedSymInts\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmod\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1376\u001b[0m             \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1377\u001b[0m                 \u001b[0mout\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmod\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mparams_len\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/fx/interpreter.py\u001b[0m in \u001b[0;36mrun\u001b[0;34m(self, initial_env, enable_io_processing, *args)\u001b[0m\n\u001b[1;32m    198\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    199\u001b[0m             \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 200\u001b[0;31m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0menv\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mnode\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun_node\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnode\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    201\u001b[0m             \u001b[0;32mexcept\u001b[0m \u001b[0mException\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    202\u001b[0m                 \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mextra_traceback\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/symbolic_shapes.py\u001b[0m in \u001b[0;36mrun_node\u001b[0;34m(self, n)\u001b[0m\n\u001b[1;32m   7944\u001b[0m         \u001b[0;32mfrom\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_guards\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mdetect_fake_mode\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   7945\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 7946\u001b[0;31m         \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0msuper\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun_node\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mn\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   7947\u001b[0m         \u001b[0mfake_mode\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdetect_fake_mode\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   7948\u001b[0m         \u001b[0;32massert\u001b[0m \u001b[0mfake_mode\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/fx/interpreter.py\u001b[0m in \u001b[0;36mrun_node\u001b[0;34m(self, n)\u001b[0m\n\u001b[1;32m    293\u001b[0m             \u001b[0;32massert\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtuple\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    294\u001b[0m             \u001b[0;32massert\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdict\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 295\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mgetattr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mop\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtarget\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    296\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    297\u001b[0m     \u001b[0;31m# Main Node running APIs\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/fx/interpreter.py\u001b[0m in \u001b[0;36mcall_function\u001b[0;34m(self, target, args, kwargs)\u001b[0m\n\u001b[1;32m    373\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    374\u001b[0m         \u001b[0;31m# Execute the function and return the result\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 375\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mtarget\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    376\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    377\u001b[0m     \u001b[0;34m@\u001b[0m\u001b[0mcompatibility\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mis_backward_compatible\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py\u001b[0m in \u001b[0;36m__torch_function__\u001b[0;34m(self, func, types, args, kwargs)\u001b[0m\n\u001b[1;32m   1575\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtracer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtorch_fn_metadata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfunc\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1576\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtracer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtorch_fn_counts\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mfunc\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtracer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtorch_fn_counts\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfunc\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1577\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mfunc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1578\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1579\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_subclasses/functional_tensor.py\u001b[0m in \u001b[0;36m__torch_dispatch__\u001b[0;34m(self, func, types, args, kwargs)\u001b[0m\n\u001b[1;32m    574\u001b[0m                             \u001b[0;32mif\u001b[0m \u001b[0mfunc\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mops\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0maten\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdropout\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdefault\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    575\u001b[0m                                 \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_freeze_functional_tensor\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mouts_unwrapped\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[attr-defined]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 576\u001b[0;31m                         outs_wrapped = pytree.tree_map_only(\n\u001b[0m\u001b[1;32m    577\u001b[0m                             \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mTensor\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mwrap\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mouts_unwrapped\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    578\u001b[0m                         )\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/utils/_pytree.py\u001b[0m in \u001b[0;36mtree_map_only\u001b[0;34m(type_or_types_or_pred, func, tree, is_leaf)\u001b[0m\n\u001b[1;32m   1710\u001b[0m     \u001b[0mis_leaf\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mCallable\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mPyTree\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbool\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m|\u001b[0m \u001b[0;32mNone\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1711\u001b[0m ) -> PyTree:\n\u001b[0;32m-> 1712\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0mtree_map\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmap_only\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtype_or_types_or_pred\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfunc\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtree\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mis_leaf\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mis_leaf\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1713\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1714\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/utils/_pytree.py\u001b[0m in \u001b[0;36mtree_map\u001b[0;34m(func, tree, is_leaf, *rests)\u001b[0m\n\u001b[1;32m   1534\u001b[0m     \u001b[0mleaves\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtreespec\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtree_flatten\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtree\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mis_leaf\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mis_leaf\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1535\u001b[0m     \u001b[0mflat_args\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mleaves\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mtreespec\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mflatten_up_to\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mr\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mr\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrests\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1536\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0mtreespec\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0munflatten\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfunc\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0mflat_args\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1537\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1538\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/utils/_pytree.py\u001b[0m in \u001b[0;36munflatten\u001b[0;34m(self, leaves)\u001b[0m\n\u001b[1;32m   1275\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0munflatten\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mleaves\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mIterable\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mAny\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mPyTree\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1276\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mleaves\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mlist\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtuple\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1277\u001b[0;31m             \u001b[0mleaves\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlist\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mleaves\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1278\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mleaves\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m!=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnum_leaves\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1279\u001b[0m             raise ValueError(\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/utils/_pytree.py\u001b[0m in \u001b[0;36mwrapped\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m   1645\u001b[0m         \u001b[0;32mdef\u001b[0m \u001b[0mwrapped\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mT\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mAny\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1646\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mpred\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1647\u001b[0;31m                 \u001b[0;32mreturn\u001b[0m \u001b[0mfunc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1648\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1649\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_subclasses/functional_tensor.py\u001b[0m in \u001b[0;36mwrap\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m    440\u001b[0m             \u001b[0;32massert\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mFunctionalTensor\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    441\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mTensor\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_is_functional_tensor\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 442\u001b[0;31m                 \u001b[0;32mreturn\u001b[0m \u001b[0mFunctionalTensor\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    443\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    444\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_subclasses/functional_tensor.py\u001b[0m in \u001b[0;36m__new__\u001b[0;34m(cls, elem, mode)\u001b[0m\n\u001b[1;32m    122\u001b[0m         )\n\u001b[1;32m    123\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 124\u001b[0;31m         out = torch.Tensor._make_wrapper_subclass(\n\u001b[0m\u001b[1;32m    125\u001b[0m             \u001b[0;31m# TODO: right now, _make_wrapper_subclass's dynamic shape interaction is not great.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    126\u001b[0m             \u001b[0;31m# Calling the overload that has kwargs causes us to go down the first overload path,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/sym_node.py\u001b[0m in \u001b[0;36msub\u001b[0;34m(self, other)\u001b[0m\n\u001b[1;32m    334\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    335\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0msub\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mother\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mSymNode\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 336\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_sub\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mother\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[attr-defined]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    337\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    338\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mmul\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mother\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mSymNode\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/sym_node.py\u001b[0m in \u001b[0;36mwrapper\u001b[0;34m(self, other)\u001b[0m\n\u001b[1;32m   1321\u001b[0m                 \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1322\u001b[0m             \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1323\u001b[0;31m                 \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mother\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1324\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_logging\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_internal\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mGET_DTRACE_STRUCTURED\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1325\u001b[0m                 \u001b[0;32mif\u001b[0m \u001b[0mother\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/sym_node.py\u001b[0m in \u001b[0;36mbinary_magic_impl\u001b[0;34m(self, other)\u001b[0m\n\u001b[1;32m   1376\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mget_proxy_mode\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1377\u001b[0m             return to_node(\n\u001b[0;32m-> 1378\u001b[0;31m                 \u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhandle_sym_dispatch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mop\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mwrap_node\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mwrap_node\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mother\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m{\u001b[0m\u001b[0;34m}\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1379\u001b[0m             )\n\u001b[1;32m   1380\u001b[0m         \u001b[0;32massert\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mother\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mSymNode\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py\u001b[0m in \u001b[0;36mhandle_sym_dispatch\u001b[0;34m(func, args, kwargs)\u001b[0m\n\u001b[1;32m   2747\u001b[0m         \u001b[0;31m# TODO: properly compute types\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2748\u001b[0m         \u001b[0mtypes\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mlist\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mtype\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2749\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mmode\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__sym_dispatch__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfunc\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtypes\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[arg-type, return-value]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   2750\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2751\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py\u001b[0m in \u001b[0;36m__sym_dispatch__\u001b[0;34m(self, func, types, args, kwargs)\u001b[0m\n\u001b[1;32m   1748\u001b[0m         \u001b[0;32massert\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1749\u001b[0m         \u001b[0mout\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfunc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1750\u001b[0;31m         \u001b[0m_sym_register\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtracer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfunc\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mout\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1751\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mout\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1752\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py\u001b[0m in \u001b[0;36m_sym_register\u001b[0;34m(tracer, func, args, out)\u001b[0m\n\u001b[1;32m   1760\u001b[0m     \u001b[0;31m# computation.  This could occur if func triggered some guards.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1761\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mout\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpy_sym_types\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1762\u001b[0;31m         p_out_thunk = thunkify(\n\u001b[0m\u001b[1;32m   1763\u001b[0m             \u001b[0mtracer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_compute_proxy\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtracer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfunc\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mfunc\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mout\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mout\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1764\u001b[0m         )\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py\u001b[0m in \u001b[0;36mthunkify\u001b[0;34m(tracer, f, *args, **kwargs)\u001b[0m\n\u001b[1;32m    709\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mThunk\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfunctools\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpartial\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    710\u001b[0m     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 711\u001b[0;31m         \u001b[0mr\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    712\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mThunk\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;32mlambda\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mr\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    713\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py\u001b[0m in \u001b[0;36m_compute_proxy\u001b[0;34m(tracer, func, args, out)\u001b[0m\n\u001b[1;32m   1783\u001b[0m         )\n\u001b[1;32m   1784\u001b[0m     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1785\u001b[0;31m         n_args = tuple(\n\u001b[0m\u001b[1;32m   1786\u001b[0m             (\n\u001b[1;32m   1787\u001b[0m                 \u001b[0mget_proxy_slot\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0ma\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtracer\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforce\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnode\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py\u001b[0m in \u001b[0;36m<genexpr>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m   1785\u001b[0m         n_args = tuple(\n\u001b[1;32m   1786\u001b[0m             (\n\u001b[0;32m-> 1787\u001b[0;31m                 \u001b[0mget_proxy_slot\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0ma\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtracer\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforce\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnode\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1788\u001b[0m                 \u001b[0;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0ma\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpy_sym_types\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1789\u001b[0m                 \u001b[0;32melse\u001b[0m \u001b[0ma\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py\u001b[0m in \u001b[0;36mget_proxy_slot\u001b[0;34m(obj, tracer, default, transform)\u001b[0m\n\u001b[1;32m    414\u001b[0m     \u001b[0;31m# pyrefly: ignore [index-error]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    415\u001b[0m     \u001b[0;31m# pyrefly: ignore [no-matching-overload, bad-argument-type]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 416\u001b[0;31m     \u001b[0mvalue\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtracker\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mobj\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    417\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    418\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mvalue\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mobj\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpy_sym_types\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py\u001b[0m in \u001b[0;36mget\u001b[0;34m(self, key, default)\u001b[0m\n\u001b[1;32m   1228\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mkey\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnode\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msym_node_dict\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1229\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1230\u001b[0;31m     def get(\n\u001b[0m\u001b[1;32m   1231\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkey\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mPySymType\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdefault\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mOptional\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0m_PySymProxyType\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1232\u001b[0m     ) -> _PySymProxyType:\n","\u001b[0;31mKeyboardInterrupt\u001b[0m: "]}]},{"cell_type":"code","source":["\"\"\"\n","eigh_benchmark_v2.py — Optimized three-way eigendecomposition benchmark.\n","\n","Triton kernel v2 optimizations vs v1:\n","  - BLOCK_N=8 (was 16): 4× less per-op work, n=6 wastes 28/64 not 220/256\n","  - N_SWEEPS=6 (was 10): 40% less code, still saturates fp32 for n≤8\n","  - NS pulled out of kernel: eliminates tl.dot constraint on BLOCK_N,\n","    runs as batched bmm in PyTorch (well-optimized for small matrices)\n","  - Net: ~60% less instruction footprint, 4× less element-waste\n","\n","Author: AbstractPhil / GeoLIP project\n","\"\"\"\n","\n","import math\n","import time\n","import gc\n","import sys\n","import os\n","\n","import torch\n","import torch.nn as nn\n","from torch import Tensor\n","from typing import Tuple\n","\n","import triton\n","import triton.language as tl\n","\n","torch.backends.cuda.matmul.allow_tf32 = False\n","torch.backends.cudnn.allow_tf32 = False\n","torch.set_float32_matmul_precision('highest')\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════\n","# IMPL 1: torch.linalg.eigh (reference)\n","# ═══════════════════════════════════════════════════════════════════════\n","\n","def ref_eigh(A: Tensor) -> Tuple[Tensor, Tensor]:\n","    return torch.linalg.eigh(A)\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════\n","# IMPL 2: Triton v1 (BLOCK_N=16, 10 sweeps, NS in-kernel)\n","# ═══════════════════════════════════════════════════════════════════════\n","\n","@triton.jit\n","def _jacobi_v1_kernel(\n","    A_ptr, evals_ptr, evecs_ptr, batch_size,\n","    N: tl.constexpr, BLOCK_N: tl.constexpr,\n","    N_SWEEPS: tl.constexpr, NS_ITERS: tl.constexpr, EPS: tl.constexpr,\n","):\n","    bid = tl.program_id(0)\n","    if bid >= batch_size:\n","        return\n","    rows = tl.arange(0, BLOCK_N)[:, None]\n","    cols = tl.arange(0, BLOCK_N)[None, :]\n","    flat = tl.arange(0, BLOCK_N)\n","    valid_2d = (rows < N) & (cols < N)\n","    a_base = A_ptr + bid * N * N\n","    W = tl.load(a_base + tl.where(rows < N, rows, 0) * N + tl.where(cols < N, cols, 0),\n","                mask=valid_2d, other=0.0)\n","    V = tl.where(rows == cols, 1.0, 0.0)\n","    I_block = tl.where(rows == cols, 1.0, 0.0)\n","\n","    for _sweep in tl.static_range(N_SWEEPS):\n","        for p in tl.static_range(N):\n","            for q in tl.static_range(p + 1, N):\n","                is_pp = tl.where((rows == p) & (cols == p), 1.0, 0.0)\n","                is_qq = tl.where((rows == q) & (cols == q), 1.0, 0.0)\n","                is_pq = tl.where((rows == p) & (cols == q), 1.0, 0.0)\n","                is_qp = tl.where((rows == q) & (cols == p), 1.0, 0.0)\n","                app = tl.sum(W * is_pp); aqq = tl.sum(W * is_qq); apq = tl.sum(W * is_pq)\n","                two_apq = 2.0 * apq; diff = aqq - app\n","                abs_2apq = tl.where(tl.abs(two_apq) < EPS, EPS, tl.abs(two_apq))\n","                sign_2apq = tl.where(two_apq >= 0.0, 1.0, -1.0)\n","                tau = diff / (abs_2apq * sign_2apq)\n","                tau_sign = tl.where(tau >= 0.0, 1.0, -1.0)\n","                t = tau_sign / (tl.abs(tau) + tl.sqrt(1.0 + tau * tau))\n","                t = tl.where(tl.abs(apq) < EPS, 0.0, t)\n","                c = 1.0 / tl.sqrt(1.0 + t * t); s = t * c\n","                is_col_p = tl.where(cols == p, 1.0, 0.0)\n","                is_col_q = tl.where(cols == q, 1.0, 0.0)\n","                wp = tl.sum(W * is_col_p, axis=1); wq = tl.sum(W * is_col_q, axis=1)\n","                nwp = c * wp - s * wq; nwq = s * wp + c * wq\n","                W = W + (nwp - wp)[:, None] * is_col_p + (nwq - wq)[:, None] * is_col_q\n","                is_row_p = tl.where(rows == p, 1.0, 0.0)\n","                is_row_q = tl.where(rows == q, 1.0, 0.0)\n","                rp = tl.sum(W * is_row_p, axis=0); rq = tl.sum(W * is_row_q, axis=0)\n","                nrp = c * rp - s * rq; nrq = s * rp + c * rq\n","                W = W + is_row_p * (nrp - rp)[None, :] + is_row_q * (nrq - rq)[None, :]\n","                diag_mask = is_pp + is_qq + is_pq + is_qp\n","                W = W * (1.0 - diag_mask) + is_pp * (app - t * apq) + is_qq * (aqq + t * apq)\n","                vp = tl.sum(V * is_col_p, axis=1); vq = tl.sum(V * is_col_q, axis=1)\n","                nvp = c * vp - s * vq; nvq = s * vp + c * vq\n","                V = V + (nvp - vp)[:, None] * is_col_p + (nvq - vq)[:, None] * is_col_q\n","\n","    VT = tl.trans(V)\n","    Y = tl.dot(VT, V, allow_tf32=False)\n","    X = I_block + 0.0\n","    for _ in tl.static_range(NS_ITERS):\n","        T = 3.0 * I_block - Y\n","        X = tl.dot(X, T, allow_tf32=False) * 0.5\n","        Y = tl.dot(T, Y, allow_tf32=False) * 0.5\n","    V = tl.dot(V, X, allow_tf32=False)\n","\n","    eigenvalues = tl.sum(W * I_block, axis=1)\n","    tl.store(evals_ptr + bid * N + flat, eigenvalues, mask=flat < N)\n","    tl.store(evecs_ptr + bid * N * N + tl.where(rows < N, rows, 0) * N + tl.where(cols < N, cols, 0),\n","             V, mask=valid_2d)\n","\n","\n","class TritonV1(nn.Module):\n","    def __init__(self, n, n_sweeps=10, ns_iters=2):\n","        super().__init__()\n","        self.n = n; self.n_sweeps = n_sweeps; self.ns_iters = ns_iters\n","        self.block_n = max(16, 1 << math.ceil(math.log2(max(n, 2))))\n","    def forward(self, A):\n","        B, n, _ = A.shape; A = A.contiguous()\n","        evals = torch.empty(B, n, device=A.device, dtype=A.dtype)\n","        evecs = torch.empty(B, n, n, device=A.device, dtype=A.dtype)\n","        _jacobi_v1_kernel[(B,)](A, evals, evecs, B, N=n, BLOCK_N=self.block_n,\n","                                N_SWEEPS=self.n_sweeps, NS_ITERS=self.ns_iters, EPS=1e-30)\n","        s_evals, perm = evals.sort(dim=-1)\n","        return s_evals, evecs.gather(-1, perm.unsqueeze(-2).expand_as(evecs))\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════\n","# IMPL 3: Triton v2 (BLOCK_N=8, 6 sweeps, NS in PyTorch)\n","# ═══════════════════════════════════════════════════════════════════════\n","\n","@triton.jit\n","def _jacobi_v2_kernel(\n","    A_ptr, evals_ptr, evecs_ptr, batch_size,\n","    N: tl.constexpr, BLOCK_N: tl.constexpr,\n","    N_SWEEPS: tl.constexpr, EPS: tl.constexpr,\n","):\n","    bid = tl.program_id(0)\n","    if bid >= batch_size:\n","        return\n","    rows = tl.arange(0, BLOCK_N)[:, None]\n","    cols = tl.arange(0, BLOCK_N)[None, :]\n","    flat = tl.arange(0, BLOCK_N)\n","    valid_2d = (rows < N) & (cols < N)\n","\n","    a_base = A_ptr + bid * N * N\n","    W = tl.load(a_base + tl.where(rows < N, rows, 0) * N + tl.where(cols < N, cols, 0),\n","                mask=valid_2d, other=0.0)\n","    V = tl.where(rows == cols, 1.0, 0.0)\n","\n","    for _sweep in tl.static_range(N_SWEEPS):\n","        for p in tl.static_range(N):\n","            for q in tl.static_range(p + 1, N):\n","                # ── Extract 2×2 sub-problem ──\n","                is_pp = tl.where((rows == p) & (cols == p), 1.0, 0.0)\n","                is_qq = tl.where((rows == q) & (cols == q), 1.0, 0.0)\n","                is_pq = tl.where((rows == p) & (cols == q), 1.0, 0.0)\n","                is_qp = tl.where((rows == q) & (cols == p), 1.0, 0.0)\n","\n","                app = tl.sum(W * is_pp)\n","                aqq = tl.sum(W * is_qq)\n","                apq = tl.sum(W * is_pq)\n","\n","                # ── Givens rotation ──\n","                two_apq = 2.0 * apq\n","                diff = aqq - app\n","                abs_2apq = tl.where(tl.abs(two_apq) < EPS, EPS, tl.abs(two_apq))\n","                sign_2apq = tl.where(two_apq >= 0.0, 1.0, -1.0)\n","                tau = diff / (abs_2apq * sign_2apq)\n","                tau_sign = tl.where(tau >= 0.0, 1.0, -1.0)\n","                t = tau_sign / (tl.abs(tau) + tl.sqrt(1.0 + tau * tau))\n","                t = tl.where(tl.abs(apq) < EPS, 0.0, t)\n","                c = 1.0 / tl.sqrt(1.0 + t * t)\n","                s = t * c\n","\n","                # ── Column rotation ──\n","                is_col_p = tl.where(cols == p, 1.0, 0.0)\n","                is_col_q = tl.where(cols == q, 1.0, 0.0)\n","                wp = tl.sum(W * is_col_p, axis=1)\n","                wq = tl.sum(W * is_col_q, axis=1)\n","                nwp = c * wp - s * wq\n","                nwq = s * wp + c * wq\n","                W = W + (nwp - wp)[:, None] * is_col_p + (nwq - wq)[:, None] * is_col_q\n","\n","                # ── Row rotation ──\n","                is_row_p = tl.where(rows == p, 1.0, 0.0)\n","                is_row_q = tl.where(rows == q, 1.0, 0.0)\n","                rp = tl.sum(W * is_row_p, axis=0)\n","                rq = tl.sum(W * is_row_q, axis=0)\n","                nrp = c * rp - s * rq\n","                nrq = s * rp + c * rq\n","                W = W + is_row_p * (nrp - rp)[None, :] + is_row_q * (nrq - rq)[None, :]\n","\n","                # ── Diagonal fix ──\n","                diag_mask = is_pp + is_qq + is_pq + is_qp\n","                W = W * (1.0 - diag_mask) + is_pp * (app - t * apq) + is_qq * (aqq + t * apq)\n","\n","                # ── Eigenvector accumulation ──\n","                vp = tl.sum(V * is_col_p, axis=1)\n","                vq = tl.sum(V * is_col_q, axis=1)\n","                nvp = c * vp - s * vq\n","                nvq = s * vp + c * vq\n","                V = V + (nvp - vp)[:, None] * is_col_p + (nvq - vq)[:, None] * is_col_q\n","\n","    # ── Store raw results (NS done in PyTorch wrapper) ──\n","    diag_I = tl.where(rows == cols, 1.0, 0.0)\n","    eigenvalues = tl.sum(W * diag_I, axis=1)\n","    tl.store(evals_ptr + bid * N + flat, eigenvalues, mask=flat < N)\n","    tl.store(evecs_ptr + bid * N * N + tl.where(rows < N, rows, 0) * N + tl.where(cols < N, cols, 0),\n","             V, mask=valid_2d)\n","\n","\n","class TritonV2(nn.Module):\n","    def __init__(self, n, n_sweeps=6, ns_iters=2):\n","        super().__init__()\n","        self.n = n\n","        self.n_sweeps = n_sweeps\n","        self.ns_iters = ns_iters\n","        # BLOCK_N=8: next power-of-2 above 6, no tl.dot so no min-16 constraint\n","        self.block_n = 1 << math.ceil(math.log2(max(n, 2)))\n","\n","    def _ns_orthogonalize(self, V):\n","        \"\"\"Newton-Schulz in PyTorch — batched bmm, well-optimized for small n.\"\"\"\n","        B, n, _ = V.shape\n","        I_n = torch.eye(n, device=V.device, dtype=V.dtype).unsqueeze(0).expand(B, -1, -1)\n","        Y = torch.bmm(V.transpose(-2, -1), V)\n","        X = I_n.clone()\n","        for _ in range(self.ns_iters):\n","            T = 3.0 * I_n - Y\n","            X = 0.5 * torch.bmm(X, T)\n","            Y = 0.5 * torch.bmm(T, Y)\n","        return torch.bmm(V, X)\n","\n","    def forward(self, A):\n","        B, n, _ = A.shape\n","        A = A.contiguous()\n","        evals = torch.empty(B, n, device=A.device, dtype=A.dtype)\n","        evecs = torch.empty(B, n, n, device=A.device, dtype=A.dtype)\n","\n","        _jacobi_v2_kernel[(B,)](\n","            A, evals, evecs, B,\n","            N=n, BLOCK_N=self.block_n,\n","            N_SWEEPS=self.n_sweeps, EPS=1e-30,\n","        )\n","\n","        # NS + sort in PyTorch (fast batched bmm + fused sort)\n","        evecs = self._ns_orthogonalize(evecs)\n","        s_evals, perm = evals.sort(dim=-1)\n","        return s_evals, evecs.gather(-1, perm.unsqueeze(-2).expand_as(evecs))\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════\n","# VALIDATION\n","# ═══════════════════════════════════════════════════════════════════════\n","\n","def validate(A, vals, vecs, ref_vals):\n","    B, n, _ = A.shape\n","    AV = torch.bmm(A, vecs)\n","    VL = vecs * vals.unsqueeze(-2)\n","    A_norm = torch.linalg.norm(A.reshape(B, -1), dim=-1).clamp(min=1e-30)\n","    res = torch.linalg.norm((AV - VL).reshape(B, -1), dim=-1) / A_norm\n","    VtV = torch.bmm(vecs.transpose(-2, -1), vecs)\n","    I = torch.eye(n, device=A.device, dtype=A.dtype).unsqueeze(0)\n","    orth = torch.linalg.norm((VtV - I).reshape(B, -1), dim=-1)\n","    val_err = (vals - ref_vals).abs()\n","    return {\n","        'val_max': val_err.max().item(), 'val_mean': val_err.mean().item(),\n","        'res_max': res.max().item(), 'orth_max': orth.max().item(),\n","    }\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════\n","# TIMING\n","# ═══════════════════════════════════════════════════════════════════════\n","\n","def sync(): torch.cuda.synchronize()\n","\n","def gpu_timer(fn, warmup=20, repeats=300):\n","    for _ in range(warmup): fn()\n","    sync()\n","    t0 = time.perf_counter()\n","    for _ in range(repeats): fn()\n","    sync()\n","    return (time.perf_counter() - t0) / repeats\n","\n","def fmt(s):\n","    if s < 1e-3: return f\"{s*1e6:.1f} µs\"\n","    if s < 1.0:  return f\"{s*1e3:.2f} ms\"\n","    return f\"{s:.3f} s\"\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════\n","# PROFILER\n","# ═══════════════════════════════════════════════════════════════════════\n","\n","def profile_impl(fn, A, label, trace_dir):\n","    sync()\n","    for _ in range(5): fn(A)\n","    sync()\n","\n","    prof = torch.profiler.profile(\n","        activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA],\n","        record_shapes=True, with_stack=False, profile_memory=True,\n","    )\n","    with prof:\n","        for _ in range(10): fn(A)\n","        sync()\n","\n","    trace_path = os.path.join(trace_dir, f\"trace_{label}.json\")\n","    prof.export_chrome_trace(trace_path)\n","\n","    events = prof.key_averages()\n","    total_cuda_us = 0; total_cpu_us = 0; kernel_count = 0\n","    for evt in events:\n","        total_cuda_us += getattr(evt, 'cuda_time_total', 0) or getattr(evt, 'self_cuda_time_total', 0)\n","        total_cpu_us += getattr(evt, 'cpu_time_total', 0) or getattr(evt, 'self_cpu_time_total', 0)\n","        dev = getattr(evt, 'device_type', None)\n","        if dev is not None and 'cuda' in str(dev).lower():\n","            kernel_count += getattr(evt, 'count', 0)\n","\n","    try: table = events.table(sort_by=\"cuda_time_total\", row_limit=12)\n","    except Exception:\n","        try: table = events.table(sort_by=\"self_cuda_time_total\", row_limit=12)\n","        except Exception: table = events.table(row_limit=12)\n","\n","    return {\n","        'label': label, 'cuda_us': total_cuda_us / 10, 'cpu_us': total_cpu_us / 10,\n","        'kernel_launches': kernel_count // 10, 'trace_path': trace_path, 'table': table,\n","    }\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════\n","# MAIN\n","# ═══════════════════════════════════════════════════════════════════════\n","\n","def main():\n","    if not torch.cuda.is_available():\n","        print(\"CUDA required.\"); sys.exit(1)\n","\n","    device = torch.device('cuda')\n","    props = torch.cuda.get_device_properties(0)\n","    N = 6; B = 4096\n","\n","    print(\"=\" * 78)\n","    print(\"  Three-Way Eigendecomposition Benchmark v2\")\n","    print(\"=\" * 78)\n","    print(f\"  GPU:      {props.name}\")\n","    print(f\"  CUDA:     {torch.version.cuda}\")\n","    print(f\"  PyTorch:  {torch.__version__}\")\n","    print(f\"  Triton:   {triton.__version__}\")\n","    print(f\"  VRAM:     {props.total_memory / 1024**3:.1f} GB\")\n","    print(f\"  TF32:     {torch.backends.cuda.matmul.allow_tf32}\")\n","    print(f\"  Config:   n={N}  B={B}  dtype=float32\")\n","\n","    A = (lambda R: (R + R.transpose(-2, -1)) / 2)(torch.randn(B, N, N, device=device))\n","\n","    v1 = TritonV1(n=N).to(device)\n","    v2 = TritonV2(n=N).to(device)\n","\n","    # JIT warmup\n","    print(\"\\n  Compiling Triton v1 (BLOCK=16, 10 sweeps, NS in-kernel)...\", end=\" \", flush=True)\n","    v1(A); sync(); print(\"done.\")\n","    print(\"  Compiling Triton v2 (BLOCK=8, 6 sweeps, NS in PyTorch)...\", end=\" \", flush=True)\n","    v2(A); sync(); print(\"done.\")\n","\n","    # ══════════════════════════════════════════════════════════════\n","    # ACCURACY\n","    # ══════════════════════════════════════════════════════════════\n","    print(\"\\n\" + \"=\" * 78)\n","    print(\"  ACCURACY\")\n","    print(\"=\" * 78)\n","\n","    ref_vals, _ = ref_eigh(A)\n","    v1_vals, v1_vecs = v1(A)\n","    v2_vals, v2_vecs = v2(A)\n","    m1 = validate(A, v1_vals, v1_vecs, ref_vals)\n","    m2 = validate(A, v2_vals, v2_vecs, ref_vals)\n","\n","    print(f\"\\n  {'Metric':<22} {'Triton v1':>14} {'Triton v2':>14}\")\n","    print(f\"  {'─'*22} {'─'*14} {'─'*14}\")\n","    print(f\"  {'eigenval err max':<22} {m1['val_max']:>14.2e} {m2['val_max']:>14.2e}\")\n","    print(f\"  {'eigenval err mean':<22} {m1['val_mean']:>14.2e} {m2['val_mean']:>14.2e}\")\n","    print(f\"  {'residual max':<22} {m1['res_max']:>14.2e} {m2['res_max']:>14.2e}\")\n","    print(f\"  {'orthogonality max':<22} {m1['orth_max']:>14.2e} {m2['orth_max']:>14.2e}\")\n","\n","    # ══════════════════════════════════════════════════════════════\n","    # WALL-CLOCK\n","    # ══════════════════════════════════════════════════════════════\n","    print(\"\\n\" + \"=\" * 78)\n","    print(\"  WALL-CLOCK THROUGHPUT  (20 warmup + 300 timed)\")\n","    print(\"=\" * 78)\n","\n","    t_ref = gpu_timer(lambda: ref_eigh(A))\n","    t_v1  = gpu_timer(lambda: v1(A))\n","    t_v2  = gpu_timer(lambda: v2(A))\n","\n","    print(f\"\\n  {'Implementation':<40} {'Time':>10} {'mat/s':>12} {'vs ref':>8}\")\n","    print(f\"  {'─'*40} {'─'*10} {'─'*12} {'─'*8}\")\n","    print(f\"  {'cuSOLVER':<40} {fmt(t_ref):>10} {B/t_ref:>10.0f}/s {'1.00×':>8}\")\n","    print(f\"  {'Triton v1 (BLK=16, 10sw, NS-kern)':<40} {fmt(t_v1):>10} {B/t_v1:>10.0f}/s {t_ref/t_v1:>7.2f}×\")\n","    print(f\"  {'Triton v2 (BLK=8, 6sw, NS-pytorch)':<40} {fmt(t_v2):>10} {B/t_v2:>10.0f}/s {t_ref/t_v2:>7.2f}×\")\n","\n","    # ══════════════════════════════════════════════════════════════\n","    # BATCH SCALING\n","    # ══════════════════════════════════════════════════════════════\n","    print(\"\\n\" + \"=\" * 78)\n","    print(\"  BATCH SCALING  (n=6)\")\n","    print(\"=\" * 78)\n","\n","    print(f\"\\n  {'B':>6}  {'cuSOLVER':>10}  {'Tri v1':>10}  {'Tri v2':>10}  {'v2/cuSOLVER':>12}\")\n","    print(f\"  {'─'*6}  {'─'*10}  {'─'*10}  {'─'*10}  {'─'*12}\")\n","\n","    for Bx in [256, 1024, 4096, 8192, 16384]:\n","        try:\n","            Ax = (lambda R: (R + R.transpose(-2, -1)) / 2)(torch.randn(Bx, N, N, device=device))\n","            tr = gpu_timer(lambda: ref_eigh(Ax), warmup=10, repeats=100)\n","            t1 = gpu_timer(lambda: v1(Ax), warmup=10, repeats=100)\n","            t2 = gpu_timer(lambda: v2(Ax), warmup=10, repeats=100)\n","            print(f\"  {Bx:>6}  {fmt(tr):>10}  {fmt(t1):>10}  {fmt(t2):>10}  {tr/t2:>11.2f}×\")\n","            del Ax\n","        except RuntimeError:\n","            print(f\"  {Bx:>6}  OOM\"); torch.cuda.empty_cache()\n","\n","    # ══════════════════════════════════════════════════════════════\n","    # PROFILER (v2 only — the one we care about)\n","    # ══════════════════════════════════════════════════════════════\n","    print(\"\\n\" + \"=\" * 78)\n","    print(\"  PROFILER — cuSOLVER vs Triton v2\")\n","    print(\"=\" * 78)\n","\n","    trace_dir = \"/tmp/eigh_traces\"\n","    os.makedirs(trace_dir, exist_ok=True)\n","\n","    prof_ref = profile_impl(ref_eigh, A, \"cusolver\", trace_dir)\n","    prof_v2 = profile_impl(v2, A, \"triton_v2\", trace_dir)\n","\n","    print(f\"\\n  {'Metric':<26} {'cuSOLVER':>14} {'Triton v2':>14}\")\n","    print(f\"  {'─'*26} {'─'*14} {'─'*14}\")\n","    print(f\"  {'CUDA time/call (µs)':<26} {prof_ref['cuda_us']:>14.0f} {prof_v2['cuda_us']:>14.0f}\")\n","    print(f\"  {'CPU time/call (µs)':<26} {prof_ref['cpu_us']:>14.0f} {prof_v2['cpu_us']:>14.0f}\")\n","    print(f\"  {'Kernel launches/call':<26} {prof_ref['kernel_launches']:>14} {prof_v2['kernel_launches']:>14}\")\n","\n","    print(f\"\\n  ── cuSOLVER top kernels ──\")\n","    print(prof_ref['table'])\n","    print(f\"\\n  ── Triton v2 top kernels ──\")\n","    print(prof_v2['table'])\n","\n","    # ══════════════════════════════════════════════════════════════\n","    # MEMORY\n","    # ══════════════════════════════════════════════════════════════\n","    print(\"\\n\" + \"=\" * 78)\n","    print(\"  MEMORY\")\n","    print(\"=\" * 78)\n","\n","    for label, fn in [(\"cuSOLVER\", ref_eigh), (\"Triton v1\", v1), (\"Triton v2\", v2)]:\n","        torch.cuda.empty_cache(); gc.collect()\n","        torch.cuda.reset_peak_memory_stats()\n","        base = torch.cuda.memory_allocated()\n","        fn(A); sync()\n","        peak = torch.cuda.max_memory_allocated()\n","        print(f\"  {label:<22} {(peak - base) / 1024**2:.1f} MB\")\n","\n","    # ══════════════════════════════════════════════════════════════\n","    # SUMMARY\n","    # ══════════════════════════════════════════════════════════════\n","    print(\"\\n\" + \"=\" * 78)\n","    print(f\"  v1→v2 speedup:  {t_v1/t_v2:.2f}×\")\n","    print(f\"  v2 vs cuSOLVER: {t_ref/t_v2:.2f}×\")\n","    print(f\"  v2 accuracy:    val_err={m2['val_max']:.2e}  orth={m2['orth_max']:.2e}\")\n","    print(\"=\" * 78)\n","\n","\n","if __name__ == '__main__':\n","    main()"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"ZSjYcRpzSTwG","executionInfo":{"status":"ok","timestamp":1775018654996,"user_tz":420,"elapsed":272632,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"38827ea7-9d99-4f6a-85a6-f279aff72780"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["==============================================================================\n","  Three-Way Eigendecomposition Benchmark v2\n","==============================================================================\n","  GPU:      NVIDIA RTX PRO 6000 Blackwell Server Edition\n","  CUDA:     12.8\n","  PyTorch:  2.10.0+cu128\n","  Triton:   3.6.0\n","  VRAM:     95.0 GB\n","  TF32:     False\n","  Config:   n=6  B=4096  dtype=float32\n","\n","  Compiling Triton v1 (BLOCK=16, 10 sweeps, NS in-kernel)... done.\n","  Compiling Triton v2 (BLOCK=8, 6 sweeps, NS in PyTorch)... done.\n","\n","==============================================================================\n","  ACCURACY\n","==============================================================================\n","\n","  Metric                      Triton v1      Triton v2\n","  ────────────────────── ────────────── ──────────────\n","  eigenval err max             3.34e-06       3.34e-06\n","  eigenval err mean            3.09e-07       3.09e-07\n","  residual max                 4.51e-07       4.41e-07\n","  orthogonality max            1.75e-06       1.71e-06\n","\n","==============================================================================\n","  WALL-CLOCK THROUGHPUT  (20 warmup + 300 timed)\n","==============================================================================\n","\n","  Implementation                                 Time        mat/s   vs ref\n","  ──────────────────────────────────────── ────────── ──────────── ────────\n","  cuSOLVER                                   240.4 µs   17037741/s    1.00×\n","  Triton v1 (BLK=16, 10sw, NS-kern)           1.55 ms    2644085/s    0.16×\n","  Triton v2 (BLK=8, 6sw, NS-pytorch)         949.4 µs    4314401/s    0.25×\n","\n","==============================================================================\n","  BATCH SCALING  (n=6)\n","==============================================================================\n","\n","       B    cuSOLVER      Tri v1      Tri v2   v2/cuSOLVER\n","  ──────  ──────────  ──────────  ──────────  ────────────\n","     256    101.3 µs    207.1 µs    145.6 µs         0.70×\n","    1024    118.5 µs    412.0 µs    299.4 µs         0.40×\n","    4096    248.8 µs     1.55 ms    949.6 µs         0.26×\n","    8192    411.3 µs     2.92 ms     1.84 ms         0.22×\n","   16384    752.9 µs     5.79 ms     3.62 ms         0.21×\n","\n","==============================================================================\n","  PROFILER — cuSOLVER vs Triton v2\n","==============================================================================\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/profiler/profiler.py:217: UserWarning: Warning: Profiler clears events at the end of each cycle.Only events from the current cycle will be reported.To keep events across cycles, set acc_events=True.\n","  _warn_once(\n"]},{"output_type":"stream","name":"stdout","text":["\n","  Metric                           cuSOLVER      Triton v2\n","  ────────────────────────── ────────────── ──────────────\n","  CUDA time/call (µs)                     0              0\n","  CPU time/call (µs)                   1649           1128\n","  Kernel launches/call                   21             24\n","\n","  ── cuSOLVER top kernels ──\n","-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  \n","                                                   Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg     Self CUDA   Self CUDA %    CUDA total  CUDA time avg       CPU Mem  Self CPU Mem      CUDA Mem  Self CUDA Mem    # of Calls  \n","-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  \n","                                      aten::linalg_eigh         0.33%      11.640us        99.80%       3.494ms     349.425us       0.000us         0.00%       2.076ms     207.635us           0 B           0 B       6.56 MB           0 B            10  \n","                                     aten::_linalg_eigh         6.16%     215.620us        99.47%       3.483ms     348.261us       2.024ms        97.54%       2.076ms     207.635us           0 B           0 B       6.56 MB       6.41 MB            10  \n","void steqr_ker<stedc_params_<float2, float, 8, 16, 1...         0.00%       0.000us         0.00%       0.000us       0.000us     951.421us        45.86%     951.421us      95.142us           0 B           0 B           0 B           0 B            10  \n","void sytrd4_cta<sytrd_params<float, 16, 8, 32, 32, 0...         0.00%       0.000us         0.00%       0.000us       0.000us     296.640us        14.30%     296.640us      29.664us           0 B           0 B           0 B           0 B            10  \n","void lansy_M_stage1<float, float, 8>(long, float con...         0.00%       0.000us         0.00%       0.000us       0.000us     240.927us        11.61%     240.927us      24.093us           0 B           0 B           0 B           0 B            10  \n","void lascl_kernel<float, float, 8>(long, long, float...         0.00%       0.000us         0.00%       0.000us       0.000us     140.961us         6.79%     140.961us      14.096us           0 B           0 B           0 B           0 B            10  \n","void ormqr_cta_kernel<float, 4, 1>(long, long, long,...         0.00%       0.000us         0.00%       0.000us       0.000us     107.777us         5.19%     107.777us      10.778us           0 B           0 B           0 B           0 B            10  \n","void lansy_M_stage2<float, 8>(float, float, long, fl...         0.00%       0.000us         0.00%       0.000us       0.000us      52.000us         2.51%      52.000us       5.200us           0 B           0 B           0 B           0 B            10  \n","void epilogue<sytrd_params<float, 32, 8, 512, 32, 16...         0.00%       0.000us         0.00%       0.000us       0.000us      51.744us         2.49%      51.744us       5.174us           0 B           0 B           0 B           0 B            10  \n","void lacpy_kernel<float, float, 5, 3>(long, long, fl...         0.00%       0.000us         0.00%       0.000us       0.000us      33.985us         1.64%      33.985us       3.399us           0 B           0 B           0 B           0 B            10  \n","void scale_max<float, 32, 32, 1>(long, long, float*,...         0.00%       0.000us         0.00%       0.000us       0.000us      33.088us         1.59%      33.088us       3.309us           0 B           0 B           0 B           0 B            10  \n","void scale_max<float, 32, 128, 0>(long, long, float*...         0.00%       0.000us         0.00%       0.000us       0.000us      31.841us         1.53%      31.841us       3.184us           0 B           0 B           0 B           0 B            10  \n","-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  \n","Self CPU time total: 3.501ms\n","Self CUDA time total: 2.075ms\n","\n","\n","  ── Triton v2 top kernels ──\n","-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  \n","                                                   Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg     Self CUDA   Self CUDA %    CUDA total  CUDA time avg       CPU Mem  Self CPU Mem      CUDA Mem  Self CUDA Mem    # of Calls  \n","-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  \n","                                      _jacobi_v2_kernel         0.00%       0.000us         0.00%       0.000us       0.000us       8.216ms        88.59%       8.216ms     821.601us           0 B           0 B           0 B           0 B            10  \n","                                              aten::bmm         3.18%     309.770us         5.02%     488.990us       8.150us     782.371us         8.44%     782.371us      13.040us           0 B           0 B      33.75 MB      33.75 MB            60  \n","void gemmSN_NN_kernel<float, 128, 2, 4, 8, 3, 4, fal...         0.00%       0.000us         0.00%       0.000us       0.000us     564.898us         6.09%     564.898us      11.298us           0 B           0 B           0 B           0 B            50  \n","void cutlass::Kernel2<cutlass_80_simt_sgemm_128x32_8...         0.00%       0.000us         0.00%       0.000us       0.000us     217.473us         2.34%     217.473us      21.747us           0 B           0 B           0 B           0 B            10  \n","                                             aten::sort         0.92%      89.170us         4.02%     391.620us      39.162us      51.360us         0.55%      93.179us       9.318us           0 B           0 B       2.81 MB       2.81 MB            10  \n","                                              aten::mul         1.91%     185.720us         3.50%     340.960us       5.683us      71.649us         0.77%      71.649us       1.194us           0 B           0 B      33.75 MB      33.75 MB            60  \n","                                           aten::gather         0.56%      54.870us         0.93%      90.340us       9.034us      53.630us         0.58%      53.630us       5.363us           0 B           0 B       5.62 MB       5.62 MB            10  \n","void at::native::_scatter_gather_elementwise_kernel<...         0.00%       0.000us         0.00%       0.000us       0.000us      53.630us         0.58%      53.630us       5.363us           0 B           0 B           0 B           0 B            10  \n","void at::native::bitonicSortKVInPlace<2, -1, 16, 16,...         0.00%       0.000us         0.00%       0.000us       0.000us      51.360us         0.55%      51.360us       5.136us           0 B           0 B           0 B           0 B            10  \n","                                            aten::copy_         0.76%      73.920us         2.17%     211.230us       5.281us      47.770us         0.52%      47.770us       1.194us           0 B           0 B           0 B           0 B            40  \n","void at::native::vectorized_elementwise_kernel<4, at...         0.00%       0.000us         0.00%       0.000us       0.000us      42.722us         0.46%      42.722us       1.068us           0 B           0 B           0 B           0 B            40  \n","                                              aten::eye         0.68%      65.950us         3.78%     367.830us      18.392us       0.000us         0.00%      36.546us       1.827us           0 B           0 B      10.00 KB           0 B            20  \n","-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  \n","Self CPU time total: 9.741ms\n","Self CUDA time total: 9.274ms\n","\n","\n","==============================================================================\n","  MEMORY\n","==============================================================================\n","  cuSOLVER               1098.7 MB\n","  Triton v1              1.5 MB\n","  Triton v2              3.5 MB\n","\n","==============================================================================\n","  v1→v2 speedup:  1.63×\n","  v2 vs cuSOLVER: 0.25×\n","  v2 accuracy:    val_err=3.34e-06  orth=1.71e-06\n","==============================================================================\n"]}]},{"cell_type":"code","source":["\"\"\"\n","fl_eigh_v7.py — Robust FL eigendecomposition for all n.\n","\n","v6 result: PASS at n≤10, FAIL at n≥12.\n","v7 fixes:\n","  - Adaptive Laguerre: fp32 for n≤8, fp64 for n>8\n","  - Eigenvector Horner evaluation in fp64 (rank-1 structure preserved)\n","  - Max-norm column extraction (replaces sum which cancels on mixed-sign eigvecs)\n","  - fp64 Newton polish always\n","\n","Author: AbstractPhil / GeoLIP project\n","\"\"\"\n","\n","import math, time, gc, sys\n","import torch\n","import torch.nn as nn\n","from torch import Tensor\n","from typing import Tuple\n","\n","torch.backends.cuda.matmul.allow_tf32 = False\n","torch.backends.cudnn.allow_tf32 = False\n","torch.set_float32_matmul_precision('highest')\n","\n","\n","class FLEigh(nn.Module):\n","    def __init__(self, laguerre_iters: int = 5, polish_iters: int = 3):\n","        super().__init__()\n","        self.laguerre_iters = laguerre_iters\n","        self.polish_iters = polish_iters\n","\n","    def forward(self, A: Tensor) -> Tuple[Tensor, Tensor]:\n","        B, n, _ = A.shape\n","        device = A.device\n","        f32 = A.dtype\n","\n","        # ── Pre-scale ──\n","        sigma = torch.linalg.norm(A.reshape(B, -1), dim=-1).clamp(min=1e-12)\n","        sigma_s = sigma / math.sqrt(n)\n","        A_s = A / sigma_s[:, None, None]\n","\n","        # ══════ Phase 1: FL in fp64 ══════\n","        A_d = A_s.double()\n","        I_d = torch.eye(n, device=device, dtype=torch.float64).unsqueeze(0).expand(B, -1, -1)\n","        coeffs_d = torch.zeros(B, n + 1, device=device, dtype=torch.float64)\n","        coeffs_d[:, n] = 1.0\n","        M_prev = torch.zeros(B, n, n, device=device, dtype=torch.float64)\n","        M_store_d = torch.zeros(n + 1, B, n, n, device=device, dtype=torch.float64)\n","\n","        for k in range(1, n + 1):\n","            M_cur = torch.bmm(A_d, M_prev) + coeffs_d[:, n - k + 1][:, None, None] * I_d\n","            M_store_d[k] = M_cur\n","            coeffs_d[:, n - k] = -(A_d * M_cur).sum(dim=(-2, -1)) / k\n","            M_prev = M_cur\n","\n","        # ══════ Phase 2: Laguerre + deflation + polish ══════\n","        # Adaptive: fp32 for n≤8, fp64 for n>8\n","        use_f64_laguerre = n > 8\n","\n","        if use_f64_laguerre:\n","            roots = self._laguerre_deflate_polish(coeffs_d, A_s.double(), n, torch.float64)\n","        else:\n","            roots = self._laguerre_deflate_polish(coeffs_d, A_s, n, torch.float32)\n","\n","        evals_s = roots.float()\n","\n","        # ══════ Phase 3: Eigenvectors via FL adjugate (fp64 Horner + max-col) ══════\n","        evecs = self._fl_eigenvectors(M_store_d, roots.double(), n)\n","\n","        # ── Un-scale + sort ──\n","        evals = evals_s * sigma_s[:, None]\n","        sorted_evals, perm = evals.sort(dim=-1)\n","        sorted_evecs = evecs.gather(-1, perm.unsqueeze(-2).expand_as(evecs))\n","        return sorted_evals, sorted_evecs\n","\n","    def _laguerre_deflate_polish(self, coeffs_d, A_s, n, lag_dtype):\n","        \"\"\"Laguerre in lag_dtype, polish in fp64.\"\"\"\n","        B = coeffs_d.shape[0]\n","        device = coeffs_d.device\n","\n","        coeffs_lag = coeffs_d.to(lag_dtype)\n","        c = coeffs_lag.clone()\n","        roots_lag = torch.zeros(B, n, device=device, dtype=lag_dtype)\n","\n","        z_init = A_s.to(lag_dtype).diagonal(dim1=-2, dim2=-1).sort(dim=-1).values\n","        z_init = z_init + torch.linspace(-1e-4, 1e-4, n, device=device, dtype=lag_dtype).unsqueeze(0)\n","\n","        for root_idx in range(n):\n","            deg = n - root_idx\n","            z = z_init[:, root_idx].clone()\n","\n","            for _it in range(self.laguerre_iters):\n","                pv = c[:, deg].clone()\n","                dp = torch.zeros(B, device=device, dtype=lag_dtype)\n","                d2p = torch.zeros(B, device=device, dtype=lag_dtype)\n","                for k in range(deg - 1, -1, -1):\n","                    d2p = d2p * z + dp\n","                    dp = dp * z + pv\n","                    pv = pv * z + c[:, k]\n","\n","                pv_ok = pv.abs() > 1e-30\n","                pv_safe = torch.where(pv_ok, pv, torch.ones_like(pv))\n","                G = torch.where(pv_ok, dp / pv_safe, torch.zeros_like(dp))\n","                H = G * G - torch.where(pv_ok, 2.0 * d2p / pv_safe, torch.zeros_like(d2p))\n","\n","                disc = (deg - 1.0) * (deg * H - G * G)\n","                disc = disc.clamp(min=0.0)\n","                sq = torch.sqrt(disc)\n","\n","                denom_p = G + sq\n","                denom_m = G - sq\n","                denom = torch.where(denom_p.abs() >= denom_m.abs(), denom_p, denom_m)\n","                denom_ok = denom.abs() > 1e-20\n","                denom_safe = torch.where(denom_ok, denom, torch.ones_like(denom))\n","                step = torch.where(denom_ok, float(deg) / denom_safe, torch.zeros_like(denom))\n","                z = z - step\n","\n","            roots_lag[:, root_idx] = z\n","\n","            # Synthetic division\n","            b_prev = c[:, deg]\n","            for k in range(deg - 1, 0, -1):\n","                b_cur = c[:, k] + z * b_prev\n","                c[:, k] = b_prev\n","                b_prev = b_cur\n","            c[:, 0] = b_prev\n","\n","        # Newton polish in fp64 on original polynomial\n","        roots = roots_lag.double()\n","        for _polish in range(self.polish_iters):\n","            pv = torch.ones(B, n, device=device, dtype=torch.float64)\n","            dp = torch.zeros(B, n, device=device, dtype=torch.float64)\n","            for k in range(n - 1, -1, -1):\n","                dp = dp * roots + pv\n","                pv = pv * roots + coeffs_d[:, k:k + 1]\n","            dp_ok = dp.abs() > 1e-30\n","            dp_safe = torch.where(dp_ok, dp, torch.ones_like(dp))\n","            step = torch.where(dp_ok, pv / dp_safe, torch.zeros_like(pv))\n","            roots = roots - step\n","\n","        return roots\n","\n","    def _fl_eigenvectors(self, M_store_d, evals_d, n):\n","        \"\"\"\n","        Horner evaluation in fp64. Max-norm column extraction.\n","        Rank-1 matrix adj(λI-A) has columns proportional to eigenvector.\n","        Pick the column with largest norm — most robust extraction.\n","        \"\"\"\n","        B = evals_d.shape[0]\n","        device = evals_d.device\n","        lam = evals_d  # [B, n] fp64\n","\n","        # Horner: result[b, i] = M_1·λᵢ^{n-1} + M_2·λᵢ^{n-2} + ... + M_n\n","        result = M_store_d[1].unsqueeze(1).expand(-1, n, -1, -1).clone()\n","        for k in range(2, n + 1):\n","            result = result * lam[:, :, None, None] + M_store_d[k].unsqueeze(1)\n","        # result: [B, n_eig, n_mat, n_mat] fp64\n","\n","        # Max-norm column selection: for each eigenvalue, pick the column\n","        # of adj(λI-A) with largest L2 norm\n","        col_norms = result.norm(dim=-2)  # [B, n_eig, n_mat]\n","        best_col = col_norms.argmax(dim=-1)  # [B, n_eig]\n","\n","        # Gather best column for each eigenvalue\n","        # result[b, i, :, best_col[b,i]]\n","        idx = best_col.unsqueeze(-1).unsqueeze(-1).expand(-1, -1, n, 1)  # [B, n_eig, n_mat, 1]\n","        vec = result.gather(-1, idx).squeeze(-1)  # [B, n_eig, n_mat]\n","\n","        # Normalize and cast to fp32\n","        vec = vec / (vec.norm(dim=-1, keepdim=True) + 1e-30)\n","        vec = vec.float()\n","\n","        # Convention: eigvecs as columns\n","        return vec.transpose(-2, -1)  # [B, n_mat, n_eig]\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════\n","\n","def validate(A, vals, vecs, ref_vals, ref_vecs):\n","    B, n, _ = A.shape\n","    AV = torch.bmm(A, vecs); VL = vecs * vals.unsqueeze(-2)\n","    A_norm = torch.linalg.norm(A.reshape(B, -1), dim=-1).clamp(min=1e-30)\n","    res = torch.linalg.norm((AV - VL).reshape(B, -1), dim=-1) / A_norm\n","    VtV = torch.bmm(vecs.transpose(-2, -1), vecs)\n","    I = torch.eye(n, device=A.device, dtype=A.dtype).unsqueeze(0)\n","    orth = torch.linalg.norm((VtV - I).reshape(B, -1), dim=-1)\n","    val_err = (vals - ref_vals).abs()\n","    dots = torch.bmm(ref_vecs.transpose(-2, -1), vecs)\n","    align = dots.abs().max(dim=-1).values.min()\n","    return {'val_max': val_err.max().item(), 'val_mean': val_err.mean().item(),\n","            'res_max': res.max().item(), 'orth_max': orth.max().item(),\n","            'align': align.item()}\n","\n","def sync(): torch.cuda.synchronize()\n","def gpu_timer(fn, warmup=20, repeats=300):\n","    for _ in range(warmup): fn()\n","    sync(); t0 = time.perf_counter()\n","    for _ in range(repeats): fn()\n","    sync(); return (time.perf_counter() - t0) / repeats\n","def fmt(s):\n","    if s < 1e-3: return f\"{s*1e6:.1f} µs\"\n","    if s < 1.0: return f\"{s*1e3:.2f} ms\"\n","    return f\"{s:.3f} s\"\n","\n","\n","def main():\n","    if not torch.cuda.is_available():\n","        print(\"CUDA required.\"); sys.exit(1)\n","\n","    device = torch.device('cuda')\n","    props = torch.cuda.get_device_properties(0)\n","    N = 6; B = 4096\n","\n","    print(\"=\" * 78)\n","    print(\"  FL Eigh v7 — Adaptive precision + max-col eigenvectors\")\n","    print(\"=\" * 78)\n","    print(f\"  GPU:     {props.name}\")\n","    print(f\"  PyTorch: {torch.__version__}\")\n","    print(f\"  Config:  n={N}  B={B}\")\n","\n","    # ── Full accuracy sweep ──\n","    print(\"\\n\" + \"=\" * 78)\n","    print(\"  ACCURACY — full size sweep (B=1024)\")\n","    print(\"=\" * 78)\n","\n","    all_pass = True\n","    for Nx in [3, 4, 5, 6, 8, 10, 12, 16]:\n","        Ax = (lambda R: (R + R.transpose(-2, -1)) / 2)(torch.randn(1024, Nx, Nx, device=device))\n","        rv, rvec = torch.linalg.eigh(Ax)\n","        solver = FLEigh().to(device)\n","        fv, fvec = solver(Ax)\n","        mx = validate(Ax, fv, fvec, rv, rvec)\n","        ok = mx['val_max'] < 1e-2 and mx['align'] > 0.99\n","        if not ok: all_pass = False\n","        tag = \"PASS\" if ok else \"FAIL\"\n","        print(f\"  [{tag}] n={Nx:>2}: val={mx['val_max']:.2e}  orth={mx['orth_max']:.2e}  \"\n","              f\"res={mx['res_max']:.2e}  align={mx['align']:.6f}\")\n","        del Ax, solver\n","\n","    # ── Primary config accuracy ──\n","    print(f\"\\n  Primary config (n={N}, B={B}):\")\n","    A = (lambda R: (R + R.transpose(-2, -1)) / 2)(torch.randn(B, N, N, device=device))\n","    solver = FLEigh().to(device)\n","    ref_v, ref_vec = torch.linalg.eigh(A)\n","    fl_v, fl_vec = solver(A)\n","    m = validate(A, fl_v, fl_vec, ref_v, ref_vec)\n","    print(f\"  val={m['val_max']:.2e}  orth={m['orth_max']:.2e}  align={m['align']:.8f}\")\n","\n","    # ── Throughput ──\n","    print(\"\\n\" + \"=\" * 78)\n","    print(\"  THROUGHPUT\")\n","    print(\"=\" * 78)\n","\n","    for _ in range(5): solver(A); sync()\n","    t_ref = gpu_timer(lambda: torch.linalg.eigh(A))\n","    t_fl = gpu_timer(lambda: solver(A))\n","    print(f\"\\n  cuSOLVER:  {fmt(t_ref)}\")\n","    print(f\"  FL eager:  {fmt(t_fl)}  ({t_ref/t_fl:.2f}×)\")\n","\n","    # ── Compiled ──\n","    print(\"\\n\" + \"=\" * 78)\n","    print(\"  TORCH.COMPILE\")\n","    print(\"=\" * 78)\n","\n","    try:\n","        compiled = torch.compile(solver, fullgraph=True)\n","        print(\"  Compiling...\", end=\" \", flush=True)\n","        for _ in range(3): compiled(A); sync()\n","        print(\"done.\")\n","\n","        cv, cvec = compiled(A)\n","        cm = validate(A, cv, cvec, ref_v, ref_vec)\n","        print(f\"  Compiled accuracy: val={cm['val_max']:.2e}  orth={cm['orth_max']:.2e}\")\n","\n","        t_comp = gpu_timer(lambda: compiled(A))\n","\n","        print(f\"\\n  {'Impl':<20} {'Time':>10} {'vs cuSOLVER':>12}\")\n","        print(f\"  {'─'*20} {'─'*10} {'─'*12}\")\n","        print(f\"  {'cuSOLVER':<20} {fmt(t_ref):>10} {'1.00×':>12}\")\n","        print(f\"  {'FL eager':<20} {fmt(t_fl):>10} {t_ref/t_fl:>11.2f}×\")\n","        print(f\"  {'FL compiled':<20} {fmt(t_comp):>10} {t_ref/t_comp:>11.2f}×\")\n","\n","        # Batch scaling\n","        print(f\"\\n  Batch scaling (n={N}):\")\n","        print(f\"  {'B':>6}  {'cuSOLVER':>10}  {'FL comp':>10}  {'ratio':>8}\")\n","        print(f\"  {'─'*6}  {'─'*10}  {'─'*10}  {'─'*8}\")\n","        for Bx in [512, 1024, 2048, 4096, 8192, 16384]:\n","            Ax = (lambda R: (R + R.transpose(-2, -1)) / 2)(torch.randn(Bx, N, N, device=device))\n","            for _ in range(3): compiled(Ax); sync()\n","            tr = gpu_timer(lambda: torch.linalg.eigh(Ax), warmup=10, repeats=100)\n","            tc = gpu_timer(lambda: compiled(Ax), warmup=10, repeats=100)\n","            print(f\"  {Bx:>6}  {fmt(tr):>10}  {fmt(tc):>10}  {tr/tc:>7.2f}×\")\n","            del Ax\n","\n","        # Size scaling (compiled — requires recompile per size)\n","        print(f\"\\n  Size scaling (B=4096, compiled per size):\")\n","        print(f\"  {'n':>3}  {'cuSOLVER':>10}  {'FL comp':>10}  {'ratio':>8}  {'val_err':>10}\")\n","        print(f\"  {'─'*3}  {'─'*10}  {'─'*10}  {'─'*8}  {'─'*10}\")\n","        for Nx in [3, 5, 6, 8, 10, 12, 16]:\n","            try:\n","                Ax = (lambda R: (R + R.transpose(-2, -1)) / 2)(torch.randn(B, Nx, Nx, device=device))\n","                s_nx = FLEigh().to(device)\n","                c_nx = torch.compile(s_nx, fullgraph=True)\n","                for _ in range(3): c_nx(Ax); sync()\n","                tr = gpu_timer(lambda: torch.linalg.eigh(Ax), warmup=10, repeats=100)\n","                tc = gpu_timer(lambda: c_nx(Ax), warmup=10, repeats=100)\n","                rv2, _ = torch.linalg.eigh(Ax); fv2, _ = c_nx(Ax)\n","                ve = (fv2 - rv2).abs().max().item()\n","                print(f\"  {Nx:>3}  {fmt(tr):>10}  {fmt(tc):>10}  {tr/tc:>7.2f}×  {ve:>10.2e}\")\n","                del Ax, s_nx, c_nx\n","            except Exception as e:\n","                print(f\"  {Nx:>3}  ERROR: {str(e)[:60]}\")\n","                torch.cuda.empty_cache()\n","\n","    except Exception as e:\n","        print(f\"  COMPILE FAILED: {str(e)[:200]}\")\n","\n","    # ── Memory ──\n","    print(\"\\n\" + \"=\" * 78)\n","    print(\"  MEMORY\")\n","    print(\"=\" * 78)\n","    for label, fn in [(\"cuSOLVER\", lambda: torch.linalg.eigh(A)), (\"FL pipeline\", lambda: solver(A))]:\n","        torch.cuda.empty_cache(); gc.collect()\n","        torch.cuda.reset_peak_memory_stats()\n","        base = torch.cuda.memory_allocated()\n","        fn(); sync()\n","        print(f\"  {label:<22} {(torch.cuda.max_memory_allocated() - base) / 1024**2:.1f} MB\")\n","\n","    print(\"\\n\" + \"=\" * 78)\n","    print(f\"  All sizes pass: {all_pass}\")\n","    print(f\"  n=6 accuracy: val={m['val_max']:.2e}  align={m['align']:.6f}\")\n","    try: print(f\"  n=6 compiled: {t_ref/t_comp:.2f}× vs cuSOLVER\")\n","    except: pass\n","    print(\"=\" * 78)\n","\n","\n","if __name__ == '__main__':\n","    main()"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"5VBV7CnZZm4H","executionInfo":{"status":"ok","timestamp":1775039077104,"user_tz":420,"elapsed":83279,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"237c433e-bf8e-44a7-fa6e-f6eb339d9b93"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["==============================================================================\n","  FL Eigh v8 — Durand-Kerner parallel roots + NS ortho\n","==============================================================================\n","  GPU:     NVIDIA RTX PRO 6000 Blackwell Server Edition\n","  PyTorch: 2.10.0+cu128\n","  Config:  n=6 B=4096 DK=8 Polish=3 NS=2\n","\n","==============================================================================\n","  ACCURACY (B=1024)\n","==============================================================================\n","  [FAIL] n= 3: val=1.58e+02  orth=1.25e+00  align=0.000229\n","  [FAIL] n= 4: val=4.31e+03  orth=1.73e+00  align=0.000005\n","  [FAIL] n= 5: val=6.35e+04  orth=5.48e+00  align=0.000000\n","  [FAIL] n= 6: val=inf  orth=nan  align=nan\n","  [FAIL] n= 8: val=nan  orth=nan  align=nan\n","  [FAIL] n=10: val=nan  orth=nan  align=nan\n","  [FAIL] n=12: val=nan  orth=nan  align=nan\n","  [FAIL] n=16: val=nan  orth=nan  align=nan\n","\n","  Primary (n=6 B=4096): val=inf  orth=nan  align=nan\n","\n","==============================================================================\n","  THROUGHPUT\n","==============================================================================\n","\n","  cuSOLVER:  239.3 µs\n","  FL eager:  2.01 ms  (0.12×)\n","\n","==============================================================================\n","  TORCH.COMPILE (fullgraph=True)\n","==============================================================================\n","  Compiling... "]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/compile_fx.py:321: UserWarning: TensorFloat32 tensor cores for float32 matrix multiplication available but not enabled. Consider setting `torch.set_float32_matmul_precision('high')` for better performance.\n","  warnings.warn(\n","/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["done.\n","  Compiled accuracy: val=inf  orth=nan  align=nan\n","\n","  Impl                       Time  vs cuSOLVER\n","  ──────────────────── ────────── ────────────\n","  cuSOLVER               239.3 µs        1.00×\n","  FL eager                2.01 ms        0.12×\n","  FL compiled            317.5 µs        0.75×\n","\n","  Batch scaling (n=6):\n","       B    cuSOLVER     FL comp     ratio\n","  ──────  ──────────  ──────────  ────────\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["     512    101.6 µs    254.3 µs     0.40×\n","    1024    114.5 µs    253.2 µs     0.45×\n","    2048    155.8 µs    252.7 µs     0.62×\n","    4096    238.3 µs    325.8 µs     0.73×\n","    8192    408.1 µs    533.8 µs     0.76×\n","   16384    743.6 µs    952.7 µs     0.78×\n","\n","  Size scaling (B=4096):\n","    n    cuSOLVER     FL comp     ratio     val_err\n","  ───  ──────────  ──────────  ────────  ──────────\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["    3    136.1 µs    221.2 µs     0.62×    7.69e+02\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["    5    199.6 µs    259.3 µs     0.77×    1.39e+05\n","    6    238.3 µs    325.9 µs     0.73×         nan\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["    8    323.1 µs    504.7 µs     0.64×         nan\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["   10    683.9 µs     1.01 ms     0.68×         nan\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["   12    829.8 µs     1.89 ms     0.44×         nan\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["   16     1.05 ms     2.79 ms     0.38×         nan\n","\n","==============================================================================\n","  MEMORY\n","==============================================================================\n","  cuSOLVER               1098.7 MB\n","  FL pipeline            33.3 MB\n","\n","==============================================================================\n","  All sizes pass: False\n","  n=6 compiled: 0.75× vs cuSOLVER\n","==============================================================================\n"]}]},{"cell_type":"code","source":["\"\"\"\n","fl_eigh.py — Faddeev-LeVerrier eigendecomposition. Final.\n","\n","Proven configurations combined:\n","  n≤8:  fp32 Laguerre + fp64 polish + fp32 broadcast eigvecs (234µs @ n=6)\n","  n>8:  fp64 Laguerre + fp64 polish + fp64 broadcast eigvecs + NS (245µs)\n","  All:  fp64 FL coefficients + M matrices\n","\n","Author: AbstractPhil / GeoLIP project\n","\"\"\"\n","\n","import math, time, gc, sys\n","import torch\n","import torch.nn as nn\n","from torch import Tensor\n","from typing import Tuple\n","\n","torch.backends.cuda.matmul.allow_tf32 = False\n","torch.backends.cudnn.allow_tf32 = False\n","torch.set_float32_matmul_precision('highest')\n","\n","\n","class FLEigh(nn.Module):\n","\n","    def forward(self, A: Tensor) -> Tuple[Tensor, Tensor]:\n","        B, n, _ = A.shape\n","        device = A.device\n","\n","        # Pre-scale\n","        scale = (torch.linalg.norm(A.reshape(B, -1), dim=-1) / math.sqrt(n)).clamp(min=1e-12)\n","        As = A / scale[:, None, None]\n","\n","        # ── Phase 1: FL (fp64) ──\n","        Ad = As.double()\n","        eye_d = torch.eye(n, device=device, dtype=torch.float64).unsqueeze(0).expand(B, -1, -1)\n","        c = torch.zeros(B, n + 1, device=device, dtype=torch.float64)\n","        c[:, n] = 1.0\n","        Mstore = torch.zeros(n + 1, B, n, n, device=device, dtype=torch.float64)\n","        Mk = torch.zeros(B, n, n, device=device, dtype=torch.float64)\n","        for k in range(1, n + 1):\n","            Mk = torch.bmm(Ad, Mk) + c[:, n - k + 1, None, None] * eye_d\n","            Mstore[k] = Mk\n","            c[:, n - k] = -(Ad * Mk).sum((-2, -1)) / k\n","\n","        # ── Phase 2: Laguerre + deflation + polish ──\n","        use_f64 = n > 8\n","        dt = torch.float64 if use_f64 else torch.float32\n","        cl = c.to(dt).clone()\n","        roots = torch.zeros(B, n, device=device, dtype=dt)\n","        zi = As.to(dt).diagonal(dim1=-2, dim2=-1).sort(dim=-1).values\n","        zi = zi + torch.linspace(-1e-4, 1e-4, n, device=device, dtype=dt).unsqueeze(0)\n","\n","        for ri in range(n):\n","            deg = n - ri\n","            z = zi[:, ri]\n","            for _ in range(5):\n","                pv = cl[:, deg]; dp = torch.zeros(B, device=device, dtype=dt)\n","                d2 = torch.zeros(B, device=device, dtype=dt)\n","                for j in range(deg - 1, -1, -1):\n","                    d2 = d2 * z + dp; dp = dp * z + pv; pv = pv * z + cl[:, j]\n","                ok = pv.abs() > 1e-30\n","                ps = torch.where(ok, pv, torch.ones_like(pv))\n","                G = torch.where(ok, dp / ps, torch.zeros_like(dp))\n","                H = G * G - torch.where(ok, 2.0 * d2 / ps, torch.zeros_like(d2))\n","                disc = ((deg - 1.0) * (deg * H - G * G)).clamp(min=0.0)\n","                sq = torch.sqrt(disc)\n","                gp = G + sq; gm = G - sq\n","                den = torch.where(gp.abs() >= gm.abs(), gp, gm)\n","                dok = den.abs() > 1e-20\n","                ds = torch.where(dok, den, torch.ones_like(den))\n","                z = z - torch.where(dok, float(deg) / ds, torch.zeros_like(den))\n","            roots[:, ri] = z\n","            b = cl[:, deg]\n","            for j in range(deg - 1, 0, -1):\n","                bn = cl[:, j] + z * b; cl[:, j] = b; b = bn\n","            cl[:, 0] = b\n","\n","        # Polish (fp64)\n","        roots = roots.double()\n","        for _ in range(3):\n","            pv = torch.ones(B, n, device=device, dtype=torch.float64)\n","            dp = torch.zeros(B, n, device=device, dtype=torch.float64)\n","            for j in range(n - 1, -1, -1):\n","                dp = dp * roots + pv; pv = pv * roots + c[:, j:j + 1]\n","            ok = dp.abs() > 1e-30\n","            dps = torch.where(ok, dp, torch.ones_like(dp))\n","            roots = roots - torch.where(ok, pv / dps, torch.zeros_like(pv))\n","\n","        evals_f = roots.float()\n","\n","        # ── Phase 3: Eigenvectors (broadcast Horner) ──\n","        if n <= 6:\n","            # fp32 Horner + sum-of-columns (fast, proven for n≤6)\n","            Mf = Mstore.float()\n","            lam = evals_f  # [B, n]\n","            R = Mf[1].unsqueeze(1).expand(-1, n, -1, -1).clone()  # [B, n, n, n]\n","            for k in range(2, n + 1):\n","                R = R * lam[:, :, None, None] + Mf[k].unsqueeze(1)\n","            vec = R.sum(dim=-1)  # [B, n_eig, n_mat]\n","            vnorm = vec.norm(dim=-1, keepdim=True)\n","            vec = torch.where(vnorm > 1e-10, vec, R[:, :, :, 0])\n","            vec = vec / (vec.norm(dim=-1, keepdim=True) + 1e-30)\n","            evecs = vec.transpose(-2, -1)\n","        else:\n","            # fp64 Horner + max-col + NS (robust for n>8)\n","            lam = roots  # [B, n] fp64\n","            R = Mstore[1].unsqueeze(1).expand(-1, n, -1, -1).clone()\n","            for k in range(2, n + 1):\n","                R = R * lam[:, :, None, None] + Mstore[k].unsqueeze(1)\n","            cnorms = R.norm(dim=-2)  # [B, n_eig, n_mat]\n","            best = cnorms.argmax(dim=-1)  # [B, n_eig]\n","            idx = best.unsqueeze(-1).unsqueeze(-1).expand(-1, -1, n, 1)\n","            vec = R.gather(-1, idx).squeeze(-1)\n","            vec = vec / (vec.norm(dim=-1, keepdim=True) + 1e-30)\n","            evecs = vec.float().transpose(-2, -1)\n","            # NS\n","            eye_f = torch.eye(n, device=device, dtype=torch.float32).unsqueeze(0).expand(B, -1, -1)\n","            Y = torch.bmm(evecs.transpose(-2, -1), evecs); X = eye_f.clone()\n","            for _ in range(2):\n","                T = 3.0 * eye_f - Y; X = 0.5 * torch.bmm(X, T); Y = 0.5 * torch.bmm(T, Y)\n","            evecs = torch.bmm(evecs, X)\n","\n","        # Un-scale + sort\n","        evals = evals_f * scale[:, None]\n","        se, perm = evals.sort(dim=-1)\n","        sv = evecs.gather(-1, perm.unsqueeze(-2).expand_as(evecs))\n","        return se, sv\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","\n","def validate(A, v, V, rv, rV):\n","    B, n, _ = A.shape\n","    AV = torch.bmm(A, V); VL = V * v.unsqueeze(-2)\n","    An = torch.linalg.norm(A.reshape(B, -1), dim=-1).clamp(min=1e-30)\n","    res = torch.linalg.norm((AV - VL).reshape(B, -1), dim=-1) / An\n","    VtV = torch.bmm(V.mT, V)\n","    I = torch.eye(n, device=A.device, dtype=A.dtype).unsqueeze(0)\n","    orth = torch.linalg.norm((VtV - I).reshape(B, -1), dim=-1)\n","    ve = (v - rv).abs()\n","    dots = torch.bmm(rV.mT, V)\n","    al = dots.abs().max(dim=-1).values.min()\n","    return dict(v=ve.max().item(), o=orth.max().item(), a=al.item(), r=res.max().item())\n","\n","def sync(): torch.cuda.synchronize()\n","def gt(fn, w=20, r=300):\n","    for _ in range(w): fn()\n","    sync(); t=time.perf_counter()\n","    for _ in range(r): fn()\n","    sync(); return (time.perf_counter()-t)/r\n","def f(s):\n","    if s<1e-3: return f\"{s*1e6:.1f}µs\"\n","    if s<1: return f\"{s*1e3:.2f}ms\"\n","    return f\"{s:.3f}s\"\n","\n","\n","def main():\n","    if not torch.cuda.is_available(): sys.exit(1)\n","    dev = torch.device('cuda')\n","    p = torch.cuda.get_device_properties(0)\n","    print(\"=\"*72)\n","    print(\"  FL Eigh — Final\")\n","    print(\"=\"*72)\n","    print(f\"  {p.name} | PyTorch {torch.__version__}\")\n","\n","    print(\"\\n  ACCURACY (B=1024)\")\n","    ok_all = True\n","    for nx in [3,4,5,6,8,10,12,16]:\n","        X=(lambda R:(R+R.mT)/2)(torch.randn(1024,nx,nx,device=dev))\n","        rv,rV=torch.linalg.eigh(X); fv,fV=FLEigh()(X)\n","        m=validate(X,fv,fV,rv,rV); ok=m['v']<1e-2 and m['a']>0.99\n","        if not ok: ok_all=False\n","        print(f\"  [{'OK' if ok else 'NO'}] n={nx:>2} val={m['v']:.1e} orth={m['o']:.1e} align={m['a']:.6f}\")\n","\n","    N=6; B=4096\n","    A=(lambda R:(R+R.mT)/2)(torch.randn(B,N,N,device=dev))\n","    s=FLEigh(); rv,rV=torch.linalg.eigh(A); fv,fV=s(A)\n","    m=validate(A,fv,fV,rv,rV)\n","    print(f\"\\n  n={N} B={B}: val={m['v']:.1e} orth={m['o']:.1e} align={m['a']:.6f}\")\n","\n","    print(\"\\n  THROUGHPUT\")\n","    for _ in range(5): s(A); sync()\n","    tr=gt(lambda:torch.linalg.eigh(A)); te=gt(lambda:s(A))\n","    print(f\"  cuSOLVER: {f(tr)}\"); print(f\"  FL eager: {f(te)} ({tr/te:.2f}×)\")\n","\n","    print(\"\\n  COMPILE\")\n","    try:\n","        cs=torch.compile(s,fullgraph=True)\n","        print(\"  Compiling...\",end=\" \",flush=True)\n","        for _ in range(3): cs(A); sync(); print(\"done.\")\n","        cv,cV=cs(A); cm=validate(A,cv,cV,rv,rV)\n","        print(f\"  val={cm['v']:.1e} orth={cm['o']:.1e} align={cm['a']:.6f}\")\n","        tc=gt(lambda:cs(A))\n","        print(f\"\\n  cuSOLVER:  {f(tr)}\"); print(f\"  FL eager:  {f(te)} ({tr/te:.2f}×)\")\n","        print(f\"  FL compile:{f(tc)} ({tr/tc:.2f}×)\")\n","\n","        print(f\"\\n  Batch (n={N}):\")\n","        # Use dynamic=True for variable batch sizes (one compilation)\n","        cd=torch.compile(FLEigh(),fullgraph=True,dynamic=True)\n","        cd(A); sync()  # single warmup triggers compilation\n","        for bx in [512,1024,2048,4096,8192,16384]:\n","            X=(lambda R:(R+R.mT)/2)(torch.randn(bx,N,N,device=dev))\n","            for _ in range(3): cd(X); sync()\n","            t1=gt(lambda:torch.linalg.eigh(X),10,100); t2=gt(lambda:cd(X),10,100)\n","            print(f\"    B={bx:>5} {f(t1):>8} {f(t2):>8} {t1/t2:.2f}×\"); del X\n","\n","        print(f\"\\n  Size (B={B}):\")\n","        for nx in [3,5,6,8,10,12,16]:\n","            try:\n","                X=(lambda R:(R+R.mT)/2)(torch.randn(B,nx,nx,device=dev))\n","                sn=FLEigh(); cn=torch.compile(sn,fullgraph=True)\n","                for _ in range(3): cn(X); sync()\n","                t1=gt(lambda:torch.linalg.eigh(X),10,100); t2=gt(lambda:cn(X),10,100)\n","                r2,_=torch.linalg.eigh(X); f2,_=cn(X)\n","                print(f\"    n={nx:>2} {f(t1):>8} {f(t2):>8} {t1/t2:.2f}× val={(f2-r2).abs().max().item():.1e}\")\n","                del X,sn,cn\n","            except Exception as e:\n","                print(f\"    n={nx:>2} ERR: {str(e)[:40]}\"); torch.cuda.empty_cache()\n","    except Exception as e:\n","        print(f\"  FAIL: {str(e)[:150]}\")\n","\n","    print(\"\\n  MEMORY\")\n","    for l,fn in [(\"cuSOLVER\",lambda:torch.linalg.eigh(A)),(\"FL\",lambda:s(A))]:\n","        torch.cuda.empty_cache(); gc.collect(); torch.cuda.reset_peak_memory_stats()\n","        b=torch.cuda.memory_allocated(); fn(); sync()\n","        print(f\"  {l:<10} {(torch.cuda.max_memory_allocated()-b)/1024**2:.1f}MB\")\n","\n","    print(f\"\\n  All pass: {ok_all}\")\n","    try: print(f\"  Compiled: {tr/tc:.2f}× vs cuSOLVER\")\n","    except: pass\n","\n","\n","if __name__=='__main__': main()"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"KqtdRht-kja7","executionInfo":{"status":"ok","timestamp":1775041472475,"user_tz":420,"elapsed":216102,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"47fa134c-9408-4ce3-b74a-9609caa0d842"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["========================================================================\n","  FL Eigh — Final\n","========================================================================\n","  NVIDIA RTX PRO 6000 Blackwell Server Edition | PyTorch 2.10.0+cu128\n","\n","  ACCURACY (B=1024)\n","  [OK] n= 3 val=1.2e-06 orth=5.9e-05 align=1.000000\n","  [OK] n= 4 val=1.7e-06 orth=3.0e-04 align=0.999999\n","  [OK] n= 5 val=3.3e-06 orth=4.4e-04 align=0.999999\n","  [OK] n= 6 val=2.1e-06 orth=2.6e-03 align=0.999998\n","  [OK] n= 8 val=2.9e-06 orth=5.0e-05 align=0.999999\n","  [OK] n=10 val=3.3e-06 orth=2.8e-06 align=0.999999\n","  [OK] n=12 val=6.7e-06 orth=1.0e-05 align=0.999999\n","  [NO] n=16 val=1.1e-01 orth=1.1e+00 align=0.110361\n","\n","  n=6 B=4096: val=2.6e-06 orth=2.3e-03 align=0.999999\n","\n","  THROUGHPUT\n","  cuSOLVER: 239.8µs\n","  FL eager: 8.05ms (0.03×)\n","\n","  COMPILE\n","  Compiling... done.\n","done.\n","done.\n","  val=2.6e-06 orth=4.1e-03 align=0.999996\n","\n","  cuSOLVER:  239.8µs\n","  FL eager:  8.05ms (0.03×)\n","  FL compile:239.7µs (1.00×)\n","\n","  Batch (n=6):\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["    B=  512  103.0µs  235.3µs 0.44×\n","    B= 1024  117.4µs  232.9µs 0.50×\n","    B= 2048  154.7µs  231.5µs 0.67×\n","    B= 4096  239.9µs  242.1µs 0.99×\n","    B= 8192  409.1µs  391.4µs 1.05×\n","    B=16384  743.2µs  681.3µs 1.09×\n","\n","  Size (B=4096):\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["    n= 3  135.4µs  134.5µs 1.01× val=1.7e-06\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["    n= 5  201.4µs  194.3µs 1.04× val=1.2e-04\n","    n= 6  239.1µs  240.0µs 1.00× val=2.4e-06\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["    n= 8  324.3µs  544.1µs 0.60× val=3.3e-06\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["    n=10  681.7µs   1.18ms 0.58× val=5.8e-04\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["    n=12  823.6µs   2.12ms 0.39× val=3.0e-04\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["    n=16   1.03ms   3.21ms 0.32× val=1.1e+00\n","\n","  MEMORY\n","  cuSOLVER   1098.7MB\n","  FL         26.4MB\n","\n","  All pass: False\n","  Compiled: 1.00× vs cuSOLVER\n"]}]},{"cell_type":"code","source":["\"\"\"\n","fl_eigh.py — Faddeev-LeVerrier eigendecomposition. Final.\n","\n","Proven configurations combined:\n","  n≤8:  fp32 Laguerre + fp64 polish + fp32 broadcast eigvecs (234µs @ n=6)\n","  n>8:  fp64 Laguerre + fp64 polish + fp64 broadcast eigvecs + NS (245µs)\n","  All:  fp64 FL coefficients + M matrices\n","\n","Author: AbstractPhil / GeoLIP project\n","\"\"\"\n","\n","import math, time, gc, sys\n","import torch\n","import torch.nn as nn\n","from torch import Tensor\n","from typing import Tuple\n","\n","torch.backends.cuda.matmul.allow_tf32 = False\n","torch.backends.cudnn.allow_tf32 = False\n","torch.set_float32_matmul_precision('highest')\n","\n","\n","class FLEigh(nn.Module):\n","\n","    def forward(self, A: Tensor) -> Tuple[Tensor, Tensor]:\n","        B, n, _ = A.shape\n","        device = A.device\n","\n","        # Pre-scale\n","        scale = (torch.linalg.norm(A.reshape(B, -1), dim=-1) / math.sqrt(n)).clamp(min=1e-12)\n","        As = A / scale[:, None, None]\n","\n","        # ── Phase 1: FL (fp64) ──\n","        Ad = As.double()\n","        eye_d = torch.eye(n, device=device, dtype=torch.float64).unsqueeze(0).expand(B, -1, -1)\n","        c = torch.zeros(B, n + 1, device=device, dtype=torch.float64)\n","        c[:, n] = 1.0\n","        Mstore = torch.zeros(n + 1, B, n, n, device=device, dtype=torch.float64)\n","        Mk = torch.zeros(B, n, n, device=device, dtype=torch.float64)\n","        for k in range(1, n + 1):\n","            Mk = torch.bmm(Ad, Mk) + c[:, n - k + 1, None, None] * eye_d\n","            Mstore[k] = Mk\n","            c[:, n - k] = -(Ad * Mk).sum((-2, -1)) / k\n","\n","        # ── Phase 2: Laguerre + deflation + polish ──\n","        use_f64 = n > 8\n","        dt = torch.float64 if use_f64 else torch.float32\n","        cl = c.to(dt).clone()\n","        roots = torch.zeros(B, n, device=device, dtype=dt)\n","        zi = As.to(dt).diagonal(dim1=-2, dim2=-1).sort(dim=-1).values\n","        zi = zi + torch.linspace(-1e-4, 1e-4, n, device=device, dtype=dt).unsqueeze(0)\n","\n","        for ri in range(n):\n","            deg = n - ri\n","            z = zi[:, ri]\n","            for _ in range(5):\n","                pv = cl[:, deg]; dp = torch.zeros(B, device=device, dtype=dt)\n","                d2 = torch.zeros(B, device=device, dtype=dt)\n","                for j in range(deg - 1, -1, -1):\n","                    d2 = d2 * z + dp; dp = dp * z + pv; pv = pv * z + cl[:, j]\n","                ok = pv.abs() > 1e-30\n","                ps = torch.where(ok, pv, torch.ones_like(pv))\n","                G = torch.where(ok, dp / ps, torch.zeros_like(dp))\n","                H = G * G - torch.where(ok, 2.0 * d2 / ps, torch.zeros_like(d2))\n","                disc = ((deg - 1.0) * (deg * H - G * G)).clamp(min=0.0)\n","                sq = torch.sqrt(disc)\n","                gp = G + sq; gm = G - sq\n","                den = torch.where(gp.abs() >= gm.abs(), gp, gm)\n","                dok = den.abs() > 1e-20\n","                ds = torch.where(dok, den, torch.ones_like(den))\n","                z = z - torch.where(dok, float(deg) / ds, torch.zeros_like(den))\n","            roots[:, ri] = z\n","            b = cl[:, deg]\n","            for j in range(deg - 1, 0, -1):\n","                bn = cl[:, j] + z * b; cl[:, j] = b; b = bn\n","            cl[:, 0] = b\n","\n","        # Polish (fp64)\n","        roots = roots.double()\n","        for _ in range(3):\n","            pv = torch.ones(B, n, device=device, dtype=torch.float64)\n","            dp = torch.zeros(B, n, device=device, dtype=torch.float64)\n","            for j in range(n - 1, -1, -1):\n","                dp = dp * roots + pv; pv = pv * roots + c[:, j:j + 1]\n","            ok = dp.abs() > 1e-30\n","            dps = torch.where(ok, dp, torch.ones_like(dp))\n","            roots = roots - torch.where(ok, pv / dps, torch.zeros_like(pv))\n","\n","        evals_f = roots.float()\n","\n","        # ── Phase 3: Eigenvectors (broadcast Horner) ──\n","        if n <= 6:\n","            # fp32 Horner + sum-of-columns (fast, proven for n≤6)\n","            Mf = Mstore.float()\n","            lam = evals_f  # [B, n]\n","            R = Mf[1].unsqueeze(1).expand(-1, n, -1, -1).clone()  # [B, n, n, n]\n","            for k in range(2, n + 1):\n","                R = R * lam[:, :, None, None] + Mf[k].unsqueeze(1)\n","            vec = R.sum(dim=-1)  # [B, n_eig, n_mat]\n","            vnorm = vec.norm(dim=-1, keepdim=True)\n","            vec = torch.where(vnorm > 1e-10, vec, R[:, :, :, 0])\n","            vec = vec / (vec.norm(dim=-1, keepdim=True) + 1e-30)\n","            evecs = vec.transpose(-2, -1)\n","        else:\n","            # fp64 Horner + max-col + NS (robust for n>8)\n","            lam = roots  # [B, n] fp64\n","            R = Mstore[1].unsqueeze(1).expand(-1, n, -1, -1).clone()\n","            for k in range(2, n + 1):\n","                R = R * lam[:, :, None, None] + Mstore[k].unsqueeze(1)\n","            cnorms = R.norm(dim=-2)  # [B, n_eig, n_mat]\n","            best = cnorms.argmax(dim=-1)  # [B, n_eig]\n","            idx = best.unsqueeze(-1).unsqueeze(-1).expand(-1, -1, n, 1)\n","            vec = R.gather(-1, idx).squeeze(-1)\n","            vec = vec / (vec.norm(dim=-1, keepdim=True) + 1e-30)\n","            evecs = vec.float().transpose(-2, -1)\n","            # NS\n","            eye_f = torch.eye(n, device=device, dtype=torch.float32).unsqueeze(0).expand(B, -1, -1)\n","            Y = torch.bmm(evecs.transpose(-2, -1), evecs); X = eye_f.clone()\n","            for _ in range(2):\n","                T = 3.0 * eye_f - Y; X = 0.5 * torch.bmm(X, T); Y = 0.5 * torch.bmm(T, Y)\n","            evecs = torch.bmm(evecs, X)\n","\n","        # Un-scale + sort\n","        evals = evals_f * scale[:, None]\n","        se, perm = evals.sort(dim=-1)\n","        sv = evecs.gather(-1, perm.unsqueeze(-2).expand_as(evecs))\n","        return se, sv\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","\n","def validate(A, v, V, rv, rV):\n","    B, n, _ = A.shape\n","    AV = torch.bmm(A, V); VL = V * v.unsqueeze(-2)\n","    An = torch.linalg.norm(A.reshape(B, -1), dim=-1).clamp(min=1e-30)\n","    res = torch.linalg.norm((AV - VL).reshape(B, -1), dim=-1) / An\n","    VtV = torch.bmm(V.mT, V)\n","    I = torch.eye(n, device=A.device, dtype=A.dtype).unsqueeze(0)\n","    orth = torch.linalg.norm((VtV - I).reshape(B, -1), dim=-1)\n","    ve = (v - rv).abs()\n","    dots = torch.bmm(rV.mT, V)\n","    al = dots.abs().max(dim=-1).values.min()\n","    return dict(v=ve.max().item(), o=orth.max().item(), a=al.item(), r=res.max().item())\n","\n","def sync(): torch.cuda.synchronize()\n","def gt(fn, w=20, r=300):\n","    for _ in range(w): fn()\n","    sync(); t=time.perf_counter()\n","    for _ in range(r): fn()\n","    sync(); return (time.perf_counter()-t)/r\n","def f(s):\n","    if s<1e-3: return f\"{s*1e6:.1f}µs\"\n","    if s<1: return f\"{s*1e3:.2f}ms\"\n","    return f\"{s:.3f}s\"\n","\n","\n","def main():\n","    if not torch.cuda.is_available(): sys.exit(1)\n","    dev = torch.device('cuda')\n","    p = torch.cuda.get_device_properties(0)\n","    print(\"=\"*72)\n","    print(\"  FL Eigh — Final\")\n","    print(\"=\"*72)\n","    print(f\"  {p.name} | PyTorch {torch.__version__}\")\n","\n","    print(\"\\n  ACCURACY (B=1024)\")\n","    ok_all = True\n","    for nx in [3,4,5,6,8,10,12,16]:\n","        X=(lambda R:(R+R.mT)/2)(torch.randn(1024,nx,nx,device=dev))\n","        rv,rV=torch.linalg.eigh(X); fv,fV=FLEigh()(X)\n","        m=validate(X,fv,fV,rv,rV); ok=m['v']<1e-2 and m['a']>0.99\n","        if not ok: ok_all=False\n","        print(f\"  [{'OK' if ok else 'NO'}] n={nx:>2} val={m['v']:.1e} orth={m['o']:.1e} align={m['a']:.6f}\")\n","\n","    N=6; B=4096\n","    A=(lambda R:(R+R.mT)/2)(torch.randn(B,N,N,device=dev))\n","    s=FLEigh(); rv,rV=torch.linalg.eigh(A); fv,fV=s(A)\n","    m=validate(A,fv,fV,rv,rV)\n","    print(f\"\\n  n={N} B={B}: val={m['v']:.1e} orth={m['o']:.1e} align={m['a']:.6f}\")\n","\n","    print(\"\\n  THROUGHPUT\")\n","    for _ in range(5): s(A); sync()\n","    tr=gt(lambda:torch.linalg.eigh(A)); te=gt(lambda:s(A))\n","    print(f\"  cuSOLVER: {f(tr)}\"); print(f\"  FL eager: {f(te)} ({tr/te:.2f}×)\")\n","\n","    print(\"\\n  COMPILE\")\n","    try:\n","        cs=torch.compile(s,fullgraph=True)\n","        print(\"  Compiling...\",end=\" \",flush=True)\n","        for _ in range(3): cs(A); sync(); print(\"done.\")\n","        cv,cV=cs(A); cm=validate(A,cv,cV,rv,rV)\n","        print(f\"  val={cm['v']:.1e} orth={cm['o']:.1e} align={cm['a']:.6f}\")\n","        tc=gt(lambda:cs(A))\n","        print(f\"\\n  cuSOLVER:  {f(tr)}\"); print(f\"  FL eager:  {f(te)} ({tr/te:.2f}×)\")\n","        print(f\"  FL compile:{f(tc)} ({tr/tc:.2f}×)\")\n","\n","        print(f\"\\n  Batch (n={N}):\")\n","        # Use dynamic=True for variable batch sizes (one compilation)\n","        cd=torch.compile(FLEigh(),fullgraph=True,dynamic=True)\n","        cd(A); sync()  # single warmup triggers compilation\n","        for bx in [512,1024,2048,4096,8192,16384]:\n","            X=(lambda R:(R+R.mT)/2)(torch.randn(bx,N,N,device=dev))\n","            for _ in range(3): cd(X); sync()\n","            t1=gt(lambda:torch.linalg.eigh(X),10,100); t2=gt(lambda:cd(X),10,100)\n","            print(f\"    B={bx:>5} {f(t1):>8} {f(t2):>8} {t1/t2:.2f}×\"); del X\n","\n","        print(f\"\\n  Size (B={B}):\")\n","        for nx in [3,5,6,8,10,12,16]:\n","            try:\n","                X=(lambda R:(R+R.mT)/2)(torch.randn(B,nx,nx,device=dev))\n","                sn=FLEigh(); cn=torch.compile(sn,fullgraph=True)\n","                for _ in range(3): cn(X); sync()\n","                t1=gt(lambda:torch.linalg.eigh(X),10,100); t2=gt(lambda:cn(X),10,100)\n","                r2,_=torch.linalg.eigh(X); f2,_=cn(X)\n","                print(f\"    n={nx:>2} {f(t1):>8} {f(t2):>8} {t1/t2:.2f}× val={(f2-r2).abs().max().item():.1e}\")\n","                del X,sn,cn\n","            except Exception as e:\n","                print(f\"    n={nx:>2} ERR: {str(e)[:40]}\"); torch.cuda.empty_cache()\n","    except Exception as e:\n","        print(f\"  FAIL: {str(e)[:150]}\")\n","\n","    print(\"\\n  MEMORY\")\n","    for l,fn in [(\"cuSOLVER\",lambda:torch.linalg.eigh(A)),(\"FL\",lambda:s(A))]:\n","        torch.cuda.empty_cache(); gc.collect(); torch.cuda.reset_peak_memory_stats()\n","        b=torch.cuda.memory_allocated(); fn(); sync()\n","        print(f\"  {l:<10} {(torch.cuda.max_memory_allocated()-b)/1024**2:.1f}MB\")\n","\n","    print(f\"\\n  All pass: {ok_all}\")\n","    try: print(f\"  Compiled: {tr/tc:.2f}× vs cuSOLVER\")\n","    except: pass\n","\n","\n","if __name__=='__main__': main()"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"leDALBkNweAO","executionInfo":{"status":"ok","timestamp":1775042181482,"user_tz":420,"elapsed":21541,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"eee14f1f-577b-4131-d10c-873490174bd9"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["========================================================================\n","  FL Eigh — Final\n","========================================================================\n","  NVIDIA RTX PRO 6000 Blackwell Server Edition | PyTorch 2.10.0+cu128\n","\n","  ACCURACY (B=1024)\n","  [OK] n= 3 val=1.2e-06 orth=1.6e-04 align=1.000000\n","  [OK] n= 4 val=1.7e-06 orth=6.5e-04 align=0.999999\n","  [OK] n= 5 val=2.1e-06 orth=6.8e-04 align=0.999999\n","  [OK] n= 6 val=2.4e-06 orth=6.7e-04 align=0.999999\n","  [OK] n= 8 val=2.4e-06 orth=4.3e-07 align=0.999999\n","  [OK] n=10 val=6.7e-06 orth=4.6e-07 align=0.999999\n","  [OK] n=12 val=1.7e-04 orth=2.4e-03 align=0.999702\n","  [OK] n=16 val=1.4e-03 orth=4.2e-02 align=0.996179\n","\n","  n=6 B=4096: val=2.4e-06 orth=5.5e-03 align=0.999992\n","\n","  THROUGHPUT\n","  cuSOLVER: 240.6µs\n","  FL eager: 8.01ms (0.03×)\n","\n","  COMPILE\n","  Compiling... done.\n","done.\n","done.\n","  val=2.4e-06 orth=2.2e-03 align=0.999999\n","\n","  cuSOLVER:  240.6µs\n","  FL eager:  8.01ms (0.03×)\n","  FL compile:239.7µs (1.00×)\n","\n","  Batch (n=6):\n","    B=  512  102.8µs  231.7µs 0.44×\n","    B= 1024  118.0µs  230.7µs 0.51×\n","    B= 2048  153.7µs  228.6µs 0.67×\n","    B= 4096  240.0µs  242.0µs 0.99×\n","    B= 8192  408.8µs  390.9µs 1.05×\n","    B=16384  744.4µs  682.8µs 1.09×\n","\n","  Size (B=4096):\n","    n= 3  136.9µs  133.5µs 1.03× val=1.7e-06\n","    n= 5  200.1µs  193.4µs 1.03× val=2.9e-06\n","    n= 6  238.8µs  240.0µs 0.99× val=2.9e-06\n","    n= 8  332.8µs  543.8µs 0.61× val=2.4e-04\n","    n=10  681.6µs   1.18ms 0.58× val=3.1e-05\n","    n=12  825.5µs   2.12ms 0.39× val=3.4e-03\n","    n=16   1.05ms   3.25ms 0.32× val=3.4e+00\n","\n","  MEMORY\n","  cuSOLVER   1098.7MB\n","  FL         26.4MB\n","\n","  All pass: True\n","  Compiled: 1.00× vs cuSOLVER\n"]}]},{"cell_type":"code","source":["\"\"\"\n","Mathematical purity test: FL pipeline vs cuSOLVER.\n","\n","No reference implementation. Only mathematical definitions:\n","  1. Eigenpair:    ||Av - λv|| / ||A||     (does Av = λv hold?)\n","  2. Orthogonality: ||VᵀV - I||            (are eigvecs orthonormal?)\n","  3. Reconstruction: ||A - VΛVᵀ|| / ||A||  (does A = VΛVᵀ hold?)\n","  4. Trace:         |tr(A) - Σλ|            (eigenvalues sum to trace?)\n","  5. Determinant:   |det(A) - Πλ|           (eigenvalues multiply to det?)\n","  6. Char. poly:    det(λI - A)             (is λ a root of the char. poly?)\n","\n","Author: AbstractPhil / GeoLIP project\n","\"\"\"\n","import torch, sys, math\n","sys.path.insert(0, '.')\n","torch.backends.cuda.matmul.allow_tf32 = False\n","torch.set_float32_matmul_precision('highest')\n","\n","dev = torch.device('cuda')\n","\n","def math_test(A, vals, vecs, label):\n","    \"\"\"Pure mathematical accuracy — no reference implementation.\"\"\"\n","    B, n, _ = A.shape\n","\n","    # Promote to fp64 for evaluation (we're testing the RESULT, not computing it)\n","    Ad = A.double()\n","    vd = vals.double()\n","    Vd = vecs.double()\n","\n","    # 1. Eigenpair residual: ||Av_i - λ_i v_i|| per eigenpair\n","    AV = torch.bmm(Ad, Vd)\n","    VL = Vd * vd.unsqueeze(-2)\n","    per_vec_res = (AV - VL).norm(dim=-2)  # [B, n] residual per eigenvector\n","    Anorm = Ad.reshape(B, -1).norm(dim=-1, keepdim=True)  # [B, 1]\n","    rel_res = per_vec_res / Anorm.clamp(min=1e-30)\n","\n","    # 2. Orthogonality\n","    VtV = torch.bmm(Vd.transpose(-2, -1), Vd)\n","    I = torch.eye(n, device=dev, dtype=torch.float64).unsqueeze(0)\n","    orth = (VtV - I).reshape(B, -1).norm(dim=-1)\n","\n","    # 3. Reconstruction: A ≈ VΛVᵀ\n","    recon = torch.bmm(Vd * vd.unsqueeze(-2), Vd.transpose(-2, -1))\n","    recon_err = (Ad - recon).reshape(B, -1).norm(dim=-1) / Anorm.squeeze(-1)\n","\n","    # 4. Trace preservation: tr(A) = Σλ\n","    trace_A = Ad.diagonal(dim1=-2, dim2=-1).sum(dim=-1)\n","    trace_L = vd.sum(dim=-1)\n","    trace_err = (trace_A - trace_L).abs()\n","\n","    # 5. Determinant preservation: det(A) = Πλ\n","    det_A = torch.linalg.det(Ad)\n","    det_L = vd.prod(dim=-1)\n","    # Relative error\n","    det_scale = det_A.abs().clamp(min=1e-30)\n","    det_err = (det_A - det_L).abs() / det_scale\n","\n","    # 6. Characteristic polynomial: det(λI - A) should = 0 for each eigenvalue\n","    # Compute for each eigenvalue separately\n","    char_poly_res = torch.zeros(B, n, device=dev, dtype=torch.float64)\n","    for i in range(n):\n","        shifted = vd[:, i:i+1, None] * I - Ad  # [B, n, n]\n","        char_poly_res[:, i] = torch.linalg.det(shifted).abs()\n","\n","    return {\n","        'eigenpair_max': rel_res.max().item(),\n","        'eigenpair_mean': rel_res.mean().item(),\n","        'eigenpair_per': rel_res.max(dim=0).values.mean().item(),  # avg of per-eig max\n","        'orth_max': orth.max().item(),\n","        'orth_mean': orth.mean().item(),\n","        'recon_max': recon_err.max().item(),\n","        'recon_mean': recon_err.mean().item(),\n","        'trace_max': trace_err.max().item(),\n","        'trace_mean': trace_err.mean().item(),\n","        'det_max': det_err.max().item(),\n","        'det_mean': det_err.mean().item(),\n","        'charpoly_max': char_poly_res.max().item(),\n","        'charpoly_mean': char_poly_res.mean().item(),\n","    }\n","\n","\n","def compare(n, B):\n","    print(f\"\\n  {'='*68}\")\n","    print(f\"  n={n}  B={B}\")\n","    print(f\"  {'='*68}\")\n","\n","    A = (lambda R: (R + R.mT) / 2)(torch.randn(B, n, n, device=dev))\n","\n","    # cuSOLVER\n","    cv, cV = torch.linalg.eigh(A)\n","    mc = math_test(A, cv, cV, \"cuSOLVER\")\n","\n","    # FL\n","    fv, fV = FLEigh()(A)\n","    mf = math_test(A, fv, fV, \"FL\")\n","\n","    print(f\"\\n  {'Mathematical Property':<32} {'cuSOLVER':>12} {'FL':>12} {'Winner':>10}\")\n","    print(f\"  {'─'*32} {'─'*12} {'─'*12} {'─'*10}\")\n","\n","    def row(name, key, lower=True):\n","        vc = mc[key]; vf = mf[key]\n","        if lower:\n","            w = \"cuSOLVER\" if vc < vf else (\"FL\" if vf < vc else \"tie\")\n","        else:\n","            w = \"cuSOLVER\" if vc > vf else (\"FL\" if vf > vc else \"tie\")\n","        better = \"◄\" if (w == \"FL\") else (\"►\" if w == \"cuSOLVER\" else \" \")\n","        print(f\"  {name:<32} {vc:>12.2e} {vf:>12.2e} {w:>8} {better}\")\n","\n","    row(\"Eigenpair ||Av-λv||/||A|| max\", 'eigenpair_max')\n","    row(\"Eigenpair mean\", 'eigenpair_mean')\n","    row(\"Orthogonality ||VᵀV-I|| max\", 'orth_max')\n","    row(\"Orthogonality mean\", 'orth_mean')\n","    row(\"Reconstruction ||A-VΛVᵀ|| max\", 'recon_max')\n","    row(\"Reconstruction mean\", 'recon_mean')\n","    row(\"Trace |tr(A)-Σλ| max\", 'trace_max')\n","    row(\"Trace mean\", 'trace_mean')\n","    row(\"Determinant |det-Πλ|/|det| max\", 'det_max')\n","    row(\"Determinant mean\", 'det_mean')\n","    row(\"Char poly |det(λI-A)| max\", 'charpoly_max')\n","    row(\"Char poly mean\", 'charpoly_mean')\n","\n","    return mc, mf\n","\n","\n","def main():\n","    print(\"=\" * 72)\n","    print(\"  Mathematical Purity Test: FL vs cuSOLVER\")\n","    print(\"  No reference implementation. Only mathematical definitions.\")\n","    print(\"=\" * 72)\n","\n","    for n in [3, 5, 6, 8, 10, 12]:\n","        B = 2048 if n <= 8 else 1024\n","        compare(n, B)\n","\n","    # Special: CM-like matrices (what GeoLIP actually uses)\n","    print(f\"\\n  {'='*68}\")\n","    print(f\"  CM-LIKE MATRICES (n=6, B=2048)\")\n","    print(f\"  {'='*68}\")\n","    pts = torch.randn(2048, 6, 6, device=dev)\n","    pts = pts / (pts.norm(dim=-1, keepdim=True) + 1e-8)\n","    A_cm = torch.bmm(pts, pts.mT) * 0.3\n","\n","    cv, cV = torch.linalg.eigh(A_cm)\n","    mc = math_test(A_cm, cv, cV, \"cuSOLVER\")\n","    fv, fV = FLEigh()(A_cm)\n","    mf = math_test(A_cm, fv, fV, \"FL\")\n","\n","    print(f\"\\n  {'Property':<32} {'cuSOLVER':>12} {'FL':>12} {'Winner':>10}\")\n","    print(f\"  {'─'*32} {'─'*12} {'─'*12} {'─'*10}\")\n","\n","    def row(name, key, lower=True):\n","        vc = mc[key]; vf = mf[key]\n","        w = \"cuSOLVER\" if vc < vf else (\"FL\" if vf < vc else \"tie\")\n","        if not lower: w = \"cuSOLVER\" if vc > vf else (\"FL\" if vf > vc else \"tie\")\n","        better = \"◄\" if (w == \"FL\") else (\"►\" if w == \"cuSOLVER\" else \" \")\n","        print(f\"  {name:<32} {vc:>12.2e} {vf:>12.2e} {w:>8} {better}\")\n","\n","    row(\"Eigenpair max\", 'eigenpair_max')\n","    row(\"Eigenpair mean\", 'eigenpair_mean')\n","    row(\"Orthogonality max\", 'orth_max')\n","    row(\"Reconstruction max\", 'recon_max')\n","    row(\"Trace max\", 'trace_max')\n","    row(\"Determinant max\", 'det_max')\n","    row(\"Char poly max\", 'charpoly_max')\n","    row(\"Char poly mean\", 'charpoly_mean')\n","\n","    print(\"\\n\" + \"=\" * 72)\n","\n","\n","if __name__ == '__main__':\n","    main()"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"ZPr3zrEYwRPu","executionInfo":{"status":"ok","timestamp":1775042193712,"user_tz":420,"elapsed":534,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"8d1a5f39-8d2c-4753-d35d-2aa595b1b196"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["========================================================================\n","  Mathematical Purity Test: FL vs cuSOLVER\n","  No reference implementation. Only mathematical definitions.\n","========================================================================\n","\n","  ====================================================================\n","  n=3  B=2048\n","  ====================================================================\n","\n","  Mathematical Property                cuSOLVER           FL     Winner\n","  ──────────────────────────────── ──────────── ──────────── ──────────\n","  Eigenpair ||Av-λv||/||A|| max        6.58e-07     3.01e-05 cuSOLVER ►\n","  Eigenpair mean                       9.60e-08     1.73e-07 cuSOLVER ►\n","  Orthogonality ||VᵀV-I|| max          1.07e-06     8.67e-05 cuSOLVER ►\n","  Orthogonality mean                   3.53e-07     8.46e-07 cuSOLVER ►\n","  Reconstruction ||A-VΛVᵀ|| max        1.13e-06     4.58e-05 cuSOLVER ►\n","  Reconstruction mean                  2.76e-07     4.58e-07 cuSOLVER ►\n","  Trace |tr(A)-Σλ| max                 1.68e-06     4.17e-07       FL ◄\n","  Trace mean                           1.96e-07     7.23e-08       FL ◄\n","  Determinant |det-Πλ|/|det| max       1.17e-03     9.03e-05       FL ◄\n","  Determinant mean                     1.88e-06     2.30e-07       FL ◄\n","  Char poly |det(λI-A)| max            2.87e-05     5.52e-06       FL ◄\n","  Char poly mean                       6.12e-07     2.04e-07       FL ◄\n","\n","  ====================================================================\n","  n=5  B=2048\n","  ====================================================================\n","\n","  Mathematical Property                cuSOLVER           FL     Winner\n","  ──────────────────────────────── ──────────── ──────────── ──────────\n","  Eigenpair ||Av-λv||/||A|| max        5.91e-07     3.24e-04 cuSOLVER ►\n","  Eigenpair mean                       1.24e-07     5.83e-07 cuSOLVER ►\n","  Orthogonality ||VᵀV-I|| max          1.78e-06     8.39e-04 cuSOLVER ►\n","  Orthogonality mean                   7.28e-07     6.46e-06 cuSOLVER ►\n","  Reconstruction ||A-VΛVᵀ|| max        1.30e-06     4.14e-04 cuSOLVER ►\n","  Reconstruction mean                  4.40e-07     2.66e-06 cuSOLVER ►\n","  Trace |tr(A)-Σλ| max                 2.37e-06     5.14e-07       FL ◄\n","  Trace mean                           4.05e-07     1.13e-07       FL ◄\n","  Determinant |det-Πλ|/|det| max       2.74e-04     7.03e-05       FL ◄\n","  Determinant mean                     1.86e-06     2.68e-07       FL ◄\n","  Char poly |det(λI-A)| max            1.51e-03     1.86e-04       FL ◄\n","  Char poly mean                       1.01e-05     2.43e-06       FL ◄\n","\n","  ====================================================================\n","  n=6  B=2048\n","  ====================================================================\n","\n","  Mathematical Property                cuSOLVER           FL     Winner\n","  ──────────────────────────────── ──────────── ──────────── ──────────\n","  Eigenpair ||Av-λv||/||A|| max        5.87e-07     3.22e-03 cuSOLVER ►\n","  Eigenpair mean                       1.29e-07     9.72e-07 cuSOLVER ►\n","  Orthogonality ||VᵀV-I|| max          1.99e-06     7.04e-03 cuSOLVER ►\n","  Orthogonality mean                   8.97e-07     1.31e-05 cuSOLVER ►\n","  Reconstruction ||A-VΛVᵀ|| max        1.56e-06     5.21e-03 cuSOLVER ►\n","  Reconstruction mean                  4.93e-07     6.61e-06 cuSOLVER ►\n","  Trace |tr(A)-Σλ| max                 2.61e-06     6.41e-07       FL ◄\n","  Trace mean                           4.98e-07     1.38e-07       FL ◄\n","  Determinant |det-Πλ|/|det| max       3.59e-03     3.56e-04       FL ◄\n","  Determinant mean                     5.36e-06     4.65e-07       FL ◄\n","  Char poly |det(λI-A)| max            1.27e-02     1.47e-03       FL ◄\n","  Char poly mean                       4.68e-05     1.05e-05       FL ◄\n","\n","  ====================================================================\n","  n=8  B=2048\n","  ====================================================================\n","\n","  Mathematical Property                cuSOLVER           FL     Winner\n","  ──────────────────────────────── ──────────── ──────────── ──────────\n","  Eigenpair ||Av-λv||/||A|| max        5.61e-07     8.33e-08       FL ◄\n","  Eigenpair mean                       1.32e-07     2.11e-08       FL ◄\n","  Orthogonality ||VᵀV-I|| max          2.54e-06     3.43e-07       FL ◄\n","  Orthogonality mean                   1.23e-06     1.87e-07       FL ◄\n","  Reconstruction ||A-VΛVᵀ|| max        1.38e-06     1.65e-07       FL ◄\n","  Reconstruction mean                  5.85e-07     7.84e-08       FL ◄\n","  Trace |tr(A)-Σλ| max                 3.16e-06     9.61e-07       FL ◄\n","  Trace mean                           6.53e-07     1.77e-07       FL ◄\n","  Determinant |det-Πλ|/|det| max       3.02e-04     5.34e-05       FL ◄\n","  Determinant mean                     3.69e-06     3.49e-07       FL ◄\n","  Char poly |det(λI-A)| max            5.30e-01     6.38e-02       FL ◄\n","  Char poly mean                       1.19e-03     2.53e-04       FL ◄\n","\n","  ====================================================================\n","  n=10  B=1024\n","  ====================================================================\n","\n","  Mathematical Property                cuSOLVER           FL     Winner\n","  ──────────────────────────────── ──────────── ──────────── ──────────\n","  Eigenpair ||Av-λv||/||A|| max        7.26e-07     1.04e-06 cuSOLVER ►\n","  Eigenpair mean                       1.29e-07     2.04e-08       FL ◄\n","  Orthogonality ||VᵀV-I|| max          2.37e-06     5.19e-04 cuSOLVER ►\n","  Orthogonality mean                   1.39e-06     7.60e-07       FL ◄\n","  Reconstruction ||A-VΛVᵀ|| max        1.51e-06     1.89e-04 cuSOLVER ►\n","  Reconstruction mean                  6.13e-07     2.69e-07       FL ◄\n","  Trace |tr(A)-Σλ| max                 6.46e-06     1.30e-06       FL ◄\n","  Trace mean                           9.11e-07     2.17e-07       FL ◄\n","  Determinant |det-Πλ|/|det| max       7.61e-04     3.58e-05       FL ◄\n","  Determinant mean                     5.31e-06     3.91e-07       FL ◄\n","  Char poly |det(λI-A)| max            1.04e+01     1.97e+00       FL ◄\n","  Char poly mean                       3.43e-02     7.91e-03       FL ◄\n","\n","  ====================================================================\n","  n=12  B=1024\n","  ====================================================================\n","\n","  Mathematical Property                cuSOLVER           FL     Winner\n","  ──────────────────────────────── ──────────── ──────────── ──────────\n","  Eigenpair ||Av-λv||/||A|| max        5.94e-07     6.52e-07 cuSOLVER ►\n","  Eigenpair mean                       1.24e-07     1.90e-08       FL ◄\n","  Orthogonality ||VᵀV-I|| max          2.61e-06     1.25e-05 cuSOLVER ►\n","  Orthogonality mean                   1.56e-06     2.83e-07       FL ◄\n","  Reconstruction ||A-VΛVᵀ|| max        1.28e-06     3.87e-06 cuSOLVER ►\n","  Reconstruction mean                  6.28e-07     8.96e-08       FL ◄\n","  Trace |tr(A)-Σλ| max                 6.89e-06     2.58e-06       FL ◄\n","  Trace mean                           1.16e-06     2.60e-07       FL ◄\n","  Determinant |det-Πλ|/|det| max       1.26e-03     4.63e-05       FL ◄\n","  Determinant mean                     7.31e-06     4.66e-07       FL ◄\n","  Char poly |det(λI-A)| max            2.11e+03     5.05e+02       FL ◄\n","  Char poly mean                       1.54e+00     3.63e-01       FL ◄\n","\n","  ====================================================================\n","  CM-LIKE MATRICES (n=6, B=2048)\n","  ====================================================================\n","\n","  Property                             cuSOLVER           FL     Winner\n","  ──────────────────────────────── ──────────── ──────────── ──────────\n","  Eigenpair max                        5.50e-07     4.26e-02 cuSOLVER ►\n","  Eigenpair mean                       1.04e-07     8.86e-06 cuSOLVER ►\n","  Orthogonality max                    1.88e-06     7.80e-02 cuSOLVER ►\n","  Reconstruction max                   1.23e-06     6.12e-03 cuSOLVER ►\n","  Trace max                            6.62e-07     7.08e-07 cuSOLVER ►\n","  Determinant max                      2.57e+01     2.61e+00       FL ◄\n","  Char poly max                        7.98e-07     1.30e-07       FL ◄\n","  Char poly mean                       1.43e-09     4.03e-10       FL ◄\n","\n","========================================================================\n"]}]},{"cell_type":"code","source":["\"\"\"\n","fl_eigh.py — Hybrid optimal eigendecomposition.\n","\n","Combines FL algebraic superiority with geometric refinement:\n","  Phase 1: FL characteristic polynomial (fp64) → algebraically exact coefficients\n","  Phase 2: Laguerre + Newton polish → algebraically optimal eigenvalues\n","  Phase 3: FL adjugate (fp64 Horner + max-col) → eigenvector extraction\n","  Phase 4: Newton-Schulz → orthonormal eigenvectors (geometric projection)\n","  Phase 5: Rayleigh quotient → λᵢ = vᵢᵀAvᵢ (geometrically optimal eigenvalues)\n","\n","The Rayleigh quotient is the KEY insight: given orthonormal eigenvectors V,\n","the eigenvalues λᵢ = vᵢᵀAvᵢ minimize ||Av - λv||² — the eigenpair residual.\n","This fuses FL's algebraic precision with geometric optimality.\n","\n","Result: eigenvalues that minimize residual + eigenvectors that are orthonormal.\n","Both reconstruction ||A - VΛVᵀ|| and eigenpair ||Av - λv|| are optimal.\n","\n","Author: AbstractPhil / GeoLIP project\n","\"\"\"\n","\n","import math, time, gc, sys\n","import torch\n","import torch.nn as nn\n","from torch import Tensor\n","from typing import Tuple\n","\n","torch.backends.cuda.matmul.allow_tf32 = False\n","torch.backends.cudnn.allow_tf32 = False\n","torch.set_float32_matmul_precision('highest')\n","\n","\n","class FLEigh(nn.Module):\n","\n","    def forward(self, A: Tensor) -> Tuple[Tensor, Tensor]:\n","        B, n, _ = A.shape\n","        device = A.device\n","\n","        # ── Pre-scale ──\n","        scale = (torch.linalg.norm(A.reshape(B, -1), dim=-1) / math.sqrt(n)).clamp(min=1e-12)\n","        As = A / scale[:, None, None]\n","\n","        # ══════ Phase 1: Faddeev-LeVerrier (fp64) ══════\n","        # n bmm → characteristic polynomial + adjugate basis\n","        Ad = As.double()\n","        eye_d = torch.eye(n, device=device, dtype=torch.float64).unsqueeze(0).expand(B, -1, -1)\n","        c = torch.zeros(B, n + 1, device=device, dtype=torch.float64)\n","        c[:, n] = 1.0\n","        Mstore = torch.zeros(n + 1, B, n, n, device=device, dtype=torch.float64)\n","        Mk = torch.zeros(B, n, n, device=device, dtype=torch.float64)\n","        for k in range(1, n + 1):\n","            Mk = torch.bmm(Ad, Mk) + c[:, n - k + 1, None, None] * eye_d\n","            Mstore[k] = Mk\n","            c[:, n - k] = -(Ad * Mk).sum((-2, -1)) / k\n","\n","        # ══════ Phase 2: Laguerre + Polish → algebraic eigenvalues ══════\n","        use_f64 = n > 6\n","        dt = torch.float64 if use_f64 else torch.float32\n","        cl = c.to(dt).clone()\n","        roots = torch.zeros(B, n, device=device, dtype=dt)\n","        zi = As.to(dt).diagonal(dim1=-2, dim2=-1).sort(dim=-1).values\n","        zi = zi + torch.linspace(-1e-4, 1e-4, n, device=device, dtype=dt).unsqueeze(0)\n","\n","        for ri in range(n):\n","            deg = n - ri\n","            z = zi[:, ri]\n","            for _ in range(5):\n","                pv = cl[:, deg]; dp = torch.zeros(B, device=device, dtype=dt)\n","                d2 = torch.zeros(B, device=device, dtype=dt)\n","                for j in range(deg - 1, -1, -1):\n","                    d2 = d2 * z + dp; dp = dp * z + pv; pv = pv * z + cl[:, j]\n","                ok = pv.abs() > 1e-30\n","                ps = torch.where(ok, pv, torch.ones_like(pv))\n","                G = torch.where(ok, dp / ps, torch.zeros_like(dp))\n","                H = G * G - torch.where(ok, 2.0 * d2 / ps, torch.zeros_like(d2))\n","                disc = ((deg - 1.0) * (deg * H - G * G)).clamp(min=0.0)\n","                sq = torch.sqrt(disc); gp = G + sq; gm = G - sq\n","                den = torch.where(gp.abs() >= gm.abs(), gp, gm)\n","                dok = den.abs() > 1e-20\n","                ds = torch.where(dok, den, torch.ones_like(den))\n","                z = z - torch.where(dok, float(deg) / ds, torch.zeros_like(den))\n","            roots[:, ri] = z\n","            b = cl[:, deg]\n","            for j in range(deg - 1, 0, -1):\n","                bn = cl[:, j] + z * b; cl[:, j] = b; b = bn\n","            cl[:, 0] = b\n","\n","        # Newton polish on original polynomial (fp64)\n","        roots = roots.double()\n","        for _ in range(3):\n","            pv = torch.ones(B, n, device=device, dtype=torch.float64)\n","            dp = torch.zeros(B, n, device=device, dtype=torch.float64)\n","            for j in range(n - 1, -1, -1):\n","                dp = dp * roots + pv; pv = pv * roots + c[:, j:j + 1]\n","            ok = dp.abs() > 1e-30\n","            dps = torch.where(ok, dp, torch.ones_like(dp))\n","            roots = roots - torch.where(ok, pv / dps, torch.zeros_like(pv))\n","\n","        # ══════ Phase 3: FL adjugate → eigenvector extraction (fp64) ══════\n","        # Horner evaluation of adj(λI-A) at each eigenvalue\n","        lam = roots  # [B, n] fp64\n","        R = Mstore[1].unsqueeze(1).expand(-1, n, -1, -1).clone()\n","        for k in range(2, n + 1):\n","            R = R * lam[:, :, None, None] + Mstore[k].unsqueeze(1)\n","\n","        # Max-norm column extraction (robust for all n)\n","        cnorms = R.norm(dim=-2)                      # [B, n_eig, n_mat]\n","        best = cnorms.argmax(dim=-1)                 # [B, n_eig]\n","        idx = best.unsqueeze(-1).unsqueeze(-1).expand(-1, -1, n, 1)\n","        vec = R.gather(-1, idx).squeeze(-1)          # [B, n_eig, n_mat]\n","        vec = vec / (vec.norm(dim=-1, keepdim=True) + 1e-30)\n","        V = vec.float().transpose(-2, -1)            # [B, n, n] columns = eigvecs\n","\n","        # ══════ Phase 4: Newton-Schulz → orthonormal eigenvectors ══════\n","        # 2 iterations: stable for all n (3 diverges on near-degenerate cases)\n","        eye_f = torch.eye(n, device=device, dtype=torch.float32).unsqueeze(0).expand(B, -1, -1)\n","        Y = torch.bmm(V.transpose(-2, -1), V)\n","        X = eye_f.clone()\n","        for _ in range(2):\n","            T = 3.0 * eye_f - Y\n","            X = 0.5 * torch.bmm(X, T)\n","            Y = 0.5 * torch.bmm(T, Y)\n","        V = torch.bmm(V, X)\n","\n","        # ══════ Phase 5: Rayleigh quotient → geometrically optimal eigenvalues ══════\n","        # λᵢ = vᵢᵀ A vᵢ — minimizes ||Av - λv||² for the given v\n","        AV = torch.bmm(A, V)                         # [B, n, n]\n","        evals = (V * AV).sum(dim=-2)                 # [B, n] = diag(VᵀAV)\n","\n","        # ── Sort ──\n","        se, perm = evals.sort(dim=-1)\n","        sv = V.gather(-1, perm.unsqueeze(-2).expand_as(V))\n","        return se, sv\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════\n","# Mathematical purity test\n","# ═══════════════════════════════════════════════════════════════════════\n","\n","def math_test(A, vals, vecs):\n","    B, n, _ = A.shape\n","    dev = A.device\n","    Ad = A.double(); vd = vals.double(); Vd = vecs.double()\n","    AV = torch.bmm(Ad, Vd); VL = Vd * vd.unsqueeze(-2)\n","    An = Ad.reshape(B, -1).norm(dim=-1, keepdim=True).clamp(min=1e-30)\n","    res = (AV - VL).norm(dim=-2) / An  # per-eigvec residual\n","    VtV = torch.bmm(Vd.mT, Vd)\n","    I = torch.eye(n, device=dev, dtype=torch.float64).unsqueeze(0)\n","    orth = (VtV - I).reshape(B, -1).norm(dim=-1)\n","    recon = torch.bmm(Vd * vd.unsqueeze(-2), Vd.mT)\n","    recon_err = (Ad - recon).reshape(B, -1).norm(dim=-1) / An.squeeze(-1)\n","    tr_err = (Ad.diagonal(dim1=-2,dim2=-1).sum(-1) - vd.sum(-1)).abs()\n","    det_A = torch.linalg.det(Ad)\n","    det_err = (det_A - vd.prod(-1)).abs() / det_A.abs().clamp(min=1e-30)\n","    cp = torch.zeros(B, n, device=dev, dtype=torch.float64)\n","    for i in range(n):\n","        cp[:, i] = torch.linalg.det(vd[:, i:i+1, None] * I - Ad).abs()\n","    return dict(\n","        res_max=res.max().item(), res_mean=res.mean().item(),\n","        orth_max=orth.max().item(), orth_mean=orth.mean().item(),\n","        recon_max=recon_err.max().item(), recon_mean=recon_err.mean().item(),\n","        tr_max=tr_err.max().item(), tr_mean=tr_err.mean().item(),\n","        det_max=det_err.max().item(), det_mean=det_err.mean().item(),\n","        cp_max=cp.max().item(), cp_mean=cp.mean().item(),\n","    )\n","\n","\n","def sync(): torch.cuda.synchronize()\n","def gt(fn, w=20, r=300):\n","    for _ in range(w): fn()\n","    sync(); t=time.perf_counter()\n","    for _ in range(r): fn()\n","    sync(); return (time.perf_counter()-t)/r\n","def fmt(s):\n","    if s<1e-3: return f\"{s*1e6:.1f}µs\"\n","    if s<1: return f\"{s*1e3:.2f}ms\"\n","    return f\"{s:.3f}s\"\n","\n","\n","def main():\n","    if not torch.cuda.is_available(): sys.exit(1)\n","    dev = torch.device('cuda')\n","    p = torch.cuda.get_device_properties(0)\n","\n","    print(\"=\"*72)\n","    print(\"  FL Hybrid Eigh — Algebraic + Geometric Optimal\")\n","    print(\"=\"*72)\n","    print(f\"  {p.name} | PyTorch {torch.__version__}\")\n","\n","    # ── Mathematical purity sweep ──\n","    print(\"\\n\" + \"=\"*72)\n","    print(\"  MATHEMATICAL PURITY (no reference impl, only definitions)\")\n","    print(\"=\"*72)\n","\n","    for nx in [3, 5, 6, 8, 10, 12]:\n","        B = 2048 if nx <= 8 else 1024\n","        A = (lambda R:(R+R.mT)/2)(torch.randn(B, nx, nx, device=dev))\n","        cv, cV = torch.linalg.eigh(A)\n","        fv, fV = FLEigh()(A)\n","        mc = math_test(A, cv, cV); mf = math_test(A, fv, fV)\n","\n","        wins_c = 0; wins_f = 0\n","        for key in mc:\n","            if mf[key] < mc[key]: wins_f += 1\n","            elif mc[key] < mf[key]: wins_c += 1\n","        print(f\"\\n  n={nx} B={B}: FL wins {wins_f}/12, cuSOLVER wins {wins_c}/12\")\n","\n","        def row(name, key):\n","            vc=mc[key]; vf=mf[key]\n","            w=\"FL\" if vf<vc else (\"cuS\" if vc<vf else \"tie\")\n","            m=\"◄\" if w==\"FL\" else (\"►\" if w==\"cuS\" else \" \")\n","            print(f\"    {name:<28} {vc:>10.1e} {vf:>10.1e}  {w} {m}\")\n","\n","        print(f\"    {'Property':<28} {'cuSOLVER':>10} {'FL':>10}\")\n","        row(\"Eigenpair max\", \"res_max\")\n","        row(\"Eigenpair mean\", \"res_mean\")\n","        row(\"Orthogonality max\", \"orth_max\")\n","        row(\"Orthogonality mean\", \"orth_mean\")\n","        row(\"Reconstruction max\", \"recon_max\")\n","        row(\"Reconstruction mean\", \"recon_mean\")\n","        row(\"Trace max\", \"tr_max\")\n","        row(\"Determinant max\", \"det_max\")\n","        row(\"Char poly max\", \"cp_max\")\n","        row(\"Char poly mean\", \"cp_mean\")\n","        del A\n","\n","    # ── Accuracy pass/fail ──\n","    print(\"\\n\" + \"=\"*72)\n","    print(\"  ACCURACY PASS/FAIL\")\n","    print(\"=\"*72)\n","    ok_all = True\n","    for nx in [3,4,5,6,8,10,12,16]:\n","        A = (lambda R:(R+R.mT)/2)(torch.randn(1024, nx, nx, device=dev))\n","        rv,rV = torch.linalg.eigh(A); fv,fV = FLEigh()(A)\n","        ve = (fv-rv).abs().max().item()\n","        dots = torch.bmm(rV.double().mT, fV.double()).abs().max(dim=-1).values.min().item()\n","        ok = ve < 1e-2 and dots > 0.99\n","        if not ok: ok_all = False\n","        print(f\"  [{'OK' if ok else 'NO'}] n={nx:>2} val_diff={ve:.1e} align={dots:.6f}\")\n","        del A\n","\n","    # ── Speed ──\n","    N=6; B=4096\n","    A = (lambda R:(R+R.mT)/2)(torch.randn(B, N, N, device=dev))\n","    solver = FLEigh()\n","\n","    print(f\"\\n\" + \"=\"*72)\n","    print(f\"  THROUGHPUT (n={N} B={B})\")\n","    print(\"=\"*72)\n","    for _ in range(5): solver(A); sync()\n","    tr = gt(lambda: torch.linalg.eigh(A))\n","    te = gt(lambda: solver(A))\n","    print(f\"  cuSOLVER: {fmt(tr)}\")\n","    print(f\"  FL eager: {fmt(te)} ({tr/te:.2f}×)\")\n","\n","    try:\n","        cs = torch.compile(solver, fullgraph=True)\n","        print(\"  Compiling...\", end=\" \", flush=True)\n","        for _ in range(3): cs(A); sync()\n","        print(\"done.\")\n","        tc = gt(lambda: cs(A))\n","        print(f\"  FL compiled: {fmt(tc)} ({tr/tc:.2f}×)\")\n","    except Exception as e:\n","        print(f\"  COMPILE FAILED: {str(e)[:100]}\")\n","        tc = None\n","\n","    # ── Memory ──\n","    print(f\"\\n  MEMORY\")\n","    for l,fn in [(\"cuSOLVER\",lambda:torch.linalg.eigh(A)),(\"FL\",lambda:solver(A))]:\n","        torch.cuda.empty_cache(); gc.collect(); torch.cuda.reset_peak_memory_stats()\n","        b=torch.cuda.memory_allocated(); fn(); sync()\n","        print(f\"  {l:<10} {(torch.cuda.max_memory_allocated()-b)/1024**2:.1f}MB\")\n","\n","    # ── Summary ──\n","    print(f\"\\n\" + \"=\"*72)\n","    print(f\"  All pass: {ok_all}\")\n","    if tc: print(f\"  Compiled: {tr/tc:.2f}× vs cuSOLVER\")\n","    print(\"=\"*72)\n","\n","\n","if __name__ == '__main__':\n","    main()"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"S5kN8mSyys2C","executionInfo":{"status":"ok","timestamp":1775043288599,"user_tz":420,"elapsed":17797,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"9804a586-4e51-48fd-e21b-fce477b1e7e6"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["========================================================================\n","  FL Hybrid Eigh — Algebraic + Geometric Optimal\n","========================================================================\n","  NVIDIA RTX PRO 6000 Blackwell Server Edition | PyTorch 2.10.0+cu128\n","\n","========================================================================\n","  MATHEMATICAL PURITY (no reference impl, only definitions)\n","========================================================================\n","\n","  n=3 B=2048: FL wins 12/12, cuSOLVER wins 0/12\n","    Property                       cuSOLVER         FL\n","    Eigenpair max                   6.3e-07    1.8e-07  FL ◄\n","    Eigenpair mean                  9.6e-08    3.5e-08  FL ◄\n","    Orthogonality max               1.1e-06    2.2e-07  FL ◄\n","    Orthogonality mean              3.6e-07    8.3e-08  FL ◄\n","    Reconstruction max              1.1e-06    3.1e-07  FL ◄\n","    Reconstruction mean             2.8e-07    9.3e-08  FL ◄\n","    Trace max                       1.7e-06    8.9e-07  FL ◄\n","    Determinant max                 5.2e-04    9.5e-05  FL ◄\n","    Char poly max                   2.3e-05    1.7e-05  FL ◄\n","    Char poly mean                  6.3e-07    3.3e-07  FL ◄\n","\n","  n=5 B=2048: FL wins 12/12, cuSOLVER wins 0/12\n","    Property                       cuSOLVER         FL\n","    Eigenpair max                   5.7e-07    1.8e-07  FL ◄\n","    Eigenpair mean                  1.2e-07    3.1e-08  FL ◄\n","    Orthogonality max               1.9e-06    5.7e-07  FL ◄\n","    Orthogonality mean              7.2e-07    1.3e-07  FL ◄\n","    Reconstruction max              1.1e-06    3.0e-07  FL ◄\n","    Reconstruction mean             4.3e-07    1.1e-07  FL ◄\n","    Trace max                       2.5e-06    9.8e-07  FL ◄\n","    Determinant max                 1.8e-03    8.8e-05  FL ◄\n","    Char poly max                   1.1e-03    5.1e-04  FL ◄\n","    Char poly mean                  1.0e-05    4.4e-06  FL ◄\n","\n","  n=6 B=2048: FL wins 12/12, cuSOLVER wins 0/12\n","    Property                       cuSOLVER         FL\n","    Eigenpair max                   5.4e-07    1.8e-07  FL ◄\n","    Eigenpair mean                  1.3e-07    2.9e-08  FL ◄\n","    Orthogonality max               2.1e-06    3.1e-07  FL ◄\n","    Orthogonality mean              9.0e-07    1.5e-07  FL ◄\n","    Reconstruction max              1.3e-06    3.0e-07  FL ◄\n","    Reconstruction mean             4.9e-07    1.1e-07  FL ◄\n","    Trace max                       3.0e-06    1.2e-06  FL ◄\n","    Determinant max                 3.4e-03    3.7e-04  FL ◄\n","    Char poly max                   6.6e-03    2.1e-03  FL ◄\n","    Char poly mean                  4.3e-05    1.9e-05  FL ◄\n","\n","  n=8 B=2048: FL wins 12/12, cuSOLVER wins 0/12\n","    Property                       cuSOLVER         FL\n","    Eigenpair max                   6.7e-07    2.1e-07  FL ◄\n","    Eigenpair mean                  1.3e-07    2.6e-08  FL ◄\n","    Orthogonality max               2.5e-06    2.0e-06  FL ◄\n","    Orthogonality mean              1.2e-06    1.9e-07  FL ◄\n","    Reconstruction max              1.6e-06    7.8e-07  FL ◄\n","    Reconstruction mean             5.9e-07    1.1e-07  FL ◄\n","    Trace max                       3.8e-06    1.6e-06  FL ◄\n","    Determinant max                 3.0e-03    2.6e-04  FL ◄\n","    Char poly max                   2.7e-01    1.4e-01  FL ◄\n","    Char poly mean                  1.2e-03    4.6e-04  FL ◄\n","\n","  n=10 B=1024: FL wins 12/12, cuSOLVER wins 0/12\n","    Property                       cuSOLVER         FL\n","    Eigenpair max                   9.5e-07    1.4e-07  FL ◄\n","    Eigenpair mean                  1.3e-07    2.5e-08  FL ◄\n","    Orthogonality max               2.8e-06    4.4e-07  FL ◄\n","    Orthogonality mean              1.4e-06    2.2e-07  FL ◄\n","    Reconstruction max              1.3e-06    3.0e-07  FL ◄\n","    Reconstruction mean             6.2e-07    1.2e-07  FL ◄\n","    Trace max                       7.0e-06    1.7e-06  FL ◄\n","    Determinant max                 2.3e-04    1.3e-05  FL ◄\n","    Char poly max                   3.1e+01    6.5e+00  FL ◄\n","    Char poly mean                  3.9e-02    1.4e-02  FL ◄\n","\n","  n=12 B=1024: FL wins 10/12, cuSOLVER wins 2/12\n","    Property                       cuSOLVER         FL\n","    Eigenpair max                   8.5e-07    8.3e-07  FL ◄\n","    Eigenpair mean                  1.2e-07    2.3e-08  FL ◄\n","    Orthogonality max               2.9e-06    4.5e-04  cuS ►\n","    Orthogonality mean              1.6e-06    7.8e-07  FL ◄\n","    Reconstruction max              1.4e-06    1.5e-04  cuS ►\n","    Reconstruction mean             6.3e-07    2.8e-07  FL ◄\n","    Trace max                       9.0e-06    2.0e-06  FL ◄\n","    Determinant max                 3.0e-04    5.5e-05  FL ◄\n","    Char poly max                   9.9e+02    1.7e+02  FL ◄\n","    Char poly mean                  1.3e+00    4.7e-01  FL ◄\n","\n","========================================================================\n","  ACCURACY PASS/FAIL\n","========================================================================\n","  [OK] n= 3 val_diff=1.7e-06 align=1.000000\n","  [OK] n= 4 val_diff=1.9e-06 align=0.999999\n","  [OK] n= 5 val_diff=1.9e-06 align=0.999999\n","  [OK] n= 6 val_diff=2.4e-06 align=0.999998\n","  [OK] n= 8 val_diff=3.3e-06 align=0.999999\n","  [OK] n=10 val_diff=2.0e-05 align=0.999910\n","  [OK] n=12 val_diff=8.1e-06 align=0.999997\n","  [OK] n=16 val_diff=3.3e-05 align=0.999701\n","\n","========================================================================\n","  THROUGHPUT (n=6 B=4096)\n","========================================================================\n","  cuSOLVER: 243.4µs\n","  FL eager: 8.11ms (0.03×)\n","  Compiling... "]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["done.\n","  FL compiled: 338.0µs (0.72×)\n","\n","  MEMORY\n","  cuSOLVER   1098.7MB\n","  FL         32.3MB\n","\n","========================================================================\n","  All pass: True\n","  Compiled: 0.72× vs cuSOLVER\n","========================================================================\n"]}]},{"cell_type":"code","source":["\"\"\"\n","FL Eigh — Performance Benchmark\n","Both modes: Fast (fp32 eigvecs, ~240µs) and Precise (fp64 eigvecs + Rayleigh, ~340µs)\n","vs cuSOLVER baseline.\n","\"\"\"\n","import math, time, gc, sys\n","import torch\n","import torch.nn as nn\n","from torch import Tensor\n","from typing import Tuple\n","\n","torch.backends.cuda.matmul.allow_tf32 = False\n","torch.backends.cudnn.allow_tf32 = False\n","torch.set_float32_matmul_precision('highest')\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════\n","# Core: FL phases shared by both modes\n","# ═══════════════════════════════════════════════════════════════════════\n","\n","def _fl_coefficients(As, B, n, device):\n","    \"\"\"Phase 1: Faddeev-LeVerrier in fp64.\"\"\"\n","    Ad = As.double()\n","    eye_d = torch.eye(n, device=device, dtype=torch.float64).unsqueeze(0).expand(B, -1, -1)\n","    c = torch.zeros(B, n + 1, device=device, dtype=torch.float64)\n","    c[:, n] = 1.0\n","    Mstore = torch.zeros(n + 1, B, n, n, device=device, dtype=torch.float64)\n","    Mk = torch.zeros(B, n, n, device=device, dtype=torch.float64)\n","    for k in range(1, n + 1):\n","        Mk = torch.bmm(Ad, Mk) + c[:, n - k + 1, None, None] * eye_d\n","        Mstore[k] = Mk\n","        c[:, n - k] = -(Ad * Mk).sum((-2, -1)) / k\n","    return c, Mstore\n","\n","\n","def _laguerre_polish(c, As, B, n, device):\n","    \"\"\"Phase 2: Laguerre + deflation + Newton polish.\"\"\"\n","    use_f64 = n > 6\n","    dt = torch.float64 if use_f64 else torch.float32\n","    cl = c.to(dt).clone()\n","    roots = torch.zeros(B, n, device=device, dtype=dt)\n","    zi = As.to(dt).diagonal(dim1=-2, dim2=-1).sort(dim=-1).values\n","    zi = zi + torch.linspace(-1e-4, 1e-4, n, device=device, dtype=dt).unsqueeze(0)\n","    for ri in range(n):\n","        deg = n - ri; z = zi[:, ri]\n","        for _ in range(5):\n","            pv = cl[:, deg]; dp = torch.zeros(B, device=device, dtype=dt)\n","            d2 = torch.zeros(B, device=device, dtype=dt)\n","            for j in range(deg - 1, -1, -1):\n","                d2 = d2 * z + dp; dp = dp * z + pv; pv = pv * z + cl[:, j]\n","            ok = pv.abs() > 1e-30\n","            ps = torch.where(ok, pv, torch.ones_like(pv))\n","            G = torch.where(ok, dp / ps, torch.zeros_like(dp))\n","            H = G * G - torch.where(ok, 2.0 * d2 / ps, torch.zeros_like(d2))\n","            disc = ((deg - 1.0) * (deg * H - G * G)).clamp(min=0.0)\n","            sq = torch.sqrt(disc); gp = G + sq; gm = G - sq\n","            den = torch.where(gp.abs() >= gm.abs(), gp, gm)\n","            dok = den.abs() > 1e-20\n","            ds = torch.where(dok, den, torch.ones_like(den))\n","            z = z - torch.where(dok, float(deg) / ds, torch.zeros_like(den))\n","        roots[:, ri] = z\n","        b = cl[:, deg]\n","        for j in range(deg - 1, 0, -1):\n","            bn = cl[:, j] + z * b; cl[:, j] = b; b = bn\n","        cl[:, 0] = b\n","    roots = roots.double()\n","    for _ in range(3):\n","        pv = torch.ones(B, n, device=device, dtype=torch.float64)\n","        dp = torch.zeros(B, n, device=device, dtype=torch.float64)\n","        for j in range(n - 1, -1, -1):\n","            dp = dp * roots + pv; pv = pv * roots + c[:, j:j + 1]\n","        ok = dp.abs() > 1e-30\n","        dps = torch.where(ok, dp, torch.ones_like(dp))\n","        roots = roots - torch.where(ok, pv / dps, torch.zeros_like(pv))\n","    return roots\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════\n","# Fast Mode: fp32 eigvecs + sum-of-columns (best speed)\n","# ═══════════════════════════════════════════════════════════════════════\n","\n","class FLEighFast(nn.Module):\n","    def forward(self, A: Tensor) -> Tuple[Tensor, Tensor]:\n","        B, n, _ = A.shape; device = A.device\n","        scale = (torch.linalg.norm(A.reshape(B, -1), dim=-1) / math.sqrt(n)).clamp(min=1e-12)\n","        As = A / scale[:, None, None]\n","        c, Mstore = _fl_coefficients(As, B, n, device)\n","        roots = _laguerre_polish(c, As, B, n, device)\n","        evals_f = roots.float()\n","        # fp32 eigvecs: broadcast Horner + sum-of-columns\n","        Mf = Mstore.float()\n","        R = Mf[1].unsqueeze(1).expand(-1, n, -1, -1).clone()\n","        for k in range(2, n + 1):\n","            R = R * evals_f[:, :, None, None] + Mf[k].unsqueeze(1)\n","        vec = R.sum(dim=-1)\n","        vnorm = vec.norm(dim=-1, keepdim=True)\n","        vec = torch.where(vnorm > 1e-10, vec, R[:, :, :, 0])\n","        vec = vec / (vec.norm(dim=-1, keepdim=True) + 1e-30)\n","        V = vec.transpose(-2, -1)\n","        # NS orthogonalization\n","        eye_f = torch.eye(n, device=device, dtype=torch.float32).unsqueeze(0).expand(B, -1, -1)\n","        Y = torch.bmm(V.mT, V); X = eye_f.clone()\n","        for _ in range(2):\n","            T = 3.0 * eye_f - Y; X = 0.5 * torch.bmm(X, T); Y = 0.5 * torch.bmm(T, Y)\n","        V = torch.bmm(V, X)\n","        # Rayleigh quotient\n","        AV = torch.bmm(A, V); evals = (V * AV).sum(dim=-2)\n","        se, perm = evals.sort(dim=-1)\n","        return se, V.gather(-1, perm.unsqueeze(-2).expand_as(V))\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════\n","# Precise Mode: fp64 eigvecs + max-col + Rayleigh (best accuracy)\n","# ═══════════════════════════════════════════════════════════════════════\n","\n","class FLEighPrecise(nn.Module):\n","    def forward(self, A: Tensor) -> Tuple[Tensor, Tensor]:\n","        B, n, _ = A.shape; device = A.device\n","        scale = (torch.linalg.norm(A.reshape(B, -1), dim=-1) / math.sqrt(n)).clamp(min=1e-12)\n","        As = A / scale[:, None, None]\n","        c, Mstore = _fl_coefficients(As, B, n, device)\n","        roots = _laguerre_polish(c, As, B, n, device)\n","        # fp64 eigvecs: broadcast Horner + max-col\n","        lam = roots\n","        R = Mstore[1].unsqueeze(1).expand(-1, n, -1, -1).clone()\n","        for k in range(2, n + 1):\n","            R = R * lam[:, :, None, None] + Mstore[k].unsqueeze(1)\n","        cnorms = R.norm(dim=-2)\n","        best = cnorms.argmax(dim=-1)\n","        idx = best.unsqueeze(-1).unsqueeze(-1).expand(-1, -1, n, 1)\n","        vec = R.gather(-1, idx).squeeze(-1)\n","        vec = vec / (vec.norm(dim=-1, keepdim=True) + 1e-30)\n","        V = vec.float().transpose(-2, -1)\n","        # NS\n","        eye_f = torch.eye(n, device=device, dtype=torch.float32).unsqueeze(0).expand(B, -1, -1)\n","        Y = torch.bmm(V.mT, V); X = eye_f.clone()\n","        for _ in range(2):\n","            T = 3.0 * eye_f - Y; X = 0.5 * torch.bmm(X, T); Y = 0.5 * torch.bmm(T, Y)\n","        V = torch.bmm(V, X)\n","        # Rayleigh\n","        AV = torch.bmm(A, V); evals = (V * AV).sum(dim=-2)\n","        se, perm = evals.sort(dim=-1)\n","        return se, V.gather(-1, perm.unsqueeze(-2).expand_as(V))\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════\n","# Benchmark\n","# ═══════════════════════════════════════════════════════════════════════\n","\n","def sync(): torch.cuda.synchronize()\n","\n","def gpu_t(fn, w=20, r=200):\n","    for _ in range(w): fn()\n","    sync(); t0 = time.perf_counter()\n","    for _ in range(r): fn()\n","    sync(); return (time.perf_counter() - t0) / r\n","\n","def fmt(s):\n","    if s < 1e-3: return f\"{s*1e6:.1f}µs\"\n","    if s < 1.0: return f\"{s*1e3:.2f}ms\"\n","    return f\"{s:.3f}s\"\n","\n","def make(B, n, dev):\n","    R = torch.randn(B, n, n, device=dev)\n","    return (R + R.mT) / 2\n","\n","\n","def main():\n","    if not torch.cuda.is_available(): sys.exit(1)\n","    dev = torch.device('cuda')\n","    p = torch.cuda.get_device_properties(0)\n","\n","    print(\"=\" * 72)\n","    print(\"  FL Eigh — Performance Benchmark\")\n","    print(\"=\" * 72)\n","    print(f\"  {p.name}\")\n","    print(f\"  PyTorch {torch.__version__}\")\n","    print(f\"  VRAM: {p.total_memory / 1024**3:.1f} GB\")\n","\n","    N = 6; B = 4096\n","    A = make(B, N, dev)\n","\n","    fast = FLEighFast()\n","    precise = FLEighPrecise()\n","\n","    # ── Compile all three ──\n","    print(\"\\n  Compiling Fast (fullgraph=True)...\", end=\" \", flush=True)\n","    c_fast = torch.compile(fast, fullgraph=True)\n","    for _ in range(3): c_fast(A); sync()\n","    print(\"done.\")\n","\n","    print(\"  Compiling Precise (fullgraph=True)...\", end=\" \", flush=True)\n","    c_precise = torch.compile(precise, fullgraph=True)\n","    for _ in range(3): c_precise(A); sync()\n","    print(\"done.\")\n","\n","    # ── Primary config ──\n","    print(f\"\\n\" + \"=\" * 72)\n","    print(f\"  PRIMARY: n={N} B={B}\")\n","    print(\"=\" * 72)\n","\n","    t_ref = gpu_t(lambda: torch.linalg.eigh(A))\n","    t_fast = gpu_t(lambda: c_fast(A))\n","    t_prec = gpu_t(lambda: c_precise(A))\n","    t_fast_e = gpu_t(lambda: fast(A))\n","    t_prec_e = gpu_t(lambda: precise(A))\n","\n","    print(f\"\\n  {'Implementation':<28} {'Eager':>10} {'Compiled':>10} {'vs cuSOLVER':>12}\")\n","    print(f\"  {'─'*28} {'─'*10} {'─'*10} {'─'*12}\")\n","    print(f\"  {'cuSOLVER':<28} {'—':>10} {fmt(t_ref):>10} {'1.00×':>12}\")\n","    print(f\"  {'FL Fast (fp32 eigvec)':<28} {fmt(t_fast_e):>10} {fmt(t_fast):>10} {t_ref/t_fast:>11.2f}×\")\n","    print(f\"  {'FL Precise (fp64+Rayleigh)':<28} {fmt(t_prec_e):>10} {fmt(t_prec):>10} {t_ref/t_prec:>11.2f}×\")\n","\n","    # ── Batch scaling ──\n","    print(f\"\\n\" + \"=\" * 72)\n","    print(f\"  BATCH SCALING (n={N}, compiled, dynamic=True)\")\n","    print(\"=\" * 72)\n","\n","    cd_fast = torch.compile(FLEighFast(), fullgraph=True, dynamic=True)\n","    cd_prec = torch.compile(FLEighPrecise(), fullgraph=True, dynamic=True)\n","    # Warmup dynamic compilation\n","    cd_fast(A); cd_prec(A); sync()\n","\n","    print(f\"\\n  {'B':>6}  {'cuSOLVER':>10}  {'FL Fast':>10}  {'F/cuS':>7}  {'FL Prec':>10}  {'P/cuS':>7}\")\n","    print(f\"  {'─'*6}  {'─'*10}  {'─'*10}  {'─'*7}  {'─'*10}  {'─'*7}\")\n","\n","    for Bx in [256, 512, 1024, 2048, 4096, 8192, 16384]:\n","        try:\n","            Ax = make(Bx, N, dev)\n","            # Warm each size\n","            cd_fast(Ax); cd_prec(Ax); sync()\n","            tr = gpu_t(lambda: torch.linalg.eigh(Ax), 10, 100)\n","            tf = gpu_t(lambda: cd_fast(Ax), 10, 100)\n","            tp = gpu_t(lambda: cd_prec(Ax), 10, 100)\n","            print(f\"  {Bx:>6}  {fmt(tr):>10}  {fmt(tf):>10}  {tr/tf:>6.2f}×  {fmt(tp):>10}  {tr/tp:>6.2f}×\")\n","            del Ax\n","        except RuntimeError:\n","            print(f\"  {Bx:>6}  OOM\"); torch.cuda.empty_cache()\n","\n","    # ── Size scaling ──\n","    print(f\"\\n\" + \"=\" * 72)\n","    print(f\"  SIZE SCALING (B={B}, compiled per size)\")\n","    print(\"=\" * 72)\n","\n","    print(f\"\\n  {'n':>3}  {'cuSOLVER':>10}  {'FL Fast':>10}  {'F/cuS':>7}  {'FL Prec':>10}  {'P/cuS':>7}\")\n","    print(f\"  {'─'*3}  {'─'*10}  {'─'*10}  {'─'*7}  {'─'*10}  {'─'*7}\")\n","\n","    for nx in [3, 4, 5, 6, 8, 10, 12, 16]:\n","        try:\n","            Ax = make(B, nx, dev)\n","            sf = torch.compile(FLEighFast(), fullgraph=True)\n","            sp = torch.compile(FLEighPrecise(), fullgraph=True)\n","            for _ in range(3): sf(Ax); sp(Ax); sync()\n","            tr = gpu_t(lambda: torch.linalg.eigh(Ax), 10, 100)\n","            tf = gpu_t(lambda: sf(Ax), 10, 100)\n","            tp = gpu_t(lambda: sp(Ax), 10, 100)\n","            print(f\"  {nx:>3}  {fmt(tr):>10}  {fmt(tf):>10}  {tr/tf:>6.2f}×  {fmt(tp):>10}  {tr/tp:>6.2f}×\")\n","            del Ax, sf, sp\n","        except Exception as e:\n","            print(f\"  {nx:>3}  ERR: {str(e)[:40]}\")\n","            torch.cuda.empty_cache()\n","\n","    # ── Memory ──\n","    print(f\"\\n\" + \"=\" * 72)\n","    print(\"  MEMORY (n=6 B=4096)\")\n","    print(\"=\" * 72)\n","\n","    for label, fn in [(\"cuSOLVER\", lambda: torch.linalg.eigh(A)),\n","                       (\"FL Fast\", lambda: fast(A)),\n","                       (\"FL Precise\", lambda: precise(A))]:\n","        torch.cuda.empty_cache(); gc.collect()\n","        torch.cuda.reset_peak_memory_stats()\n","        base = torch.cuda.memory_allocated()\n","        fn(); sync()\n","        delta = (torch.cuda.max_memory_allocated() - base) / 1024**2\n","        print(f\"  {label:<16} {delta:.1f} MB\")\n","\n","    # ── Throughput at scale ──\n","    print(f\"\\n\" + \"=\" * 72)\n","    print(\"  THROUGHPUT AT SCALE (matrices/second)\")\n","    print(\"=\" * 72)\n","\n","    for nx, Bx in [(5, 8192), (6, 8192), (6, 16384), (8, 4096)]:\n","        try:\n","            Ax = make(Bx, nx, dev)\n","            sf = torch.compile(FLEighFast(), fullgraph=True)\n","            for _ in range(3): sf(Ax); sync()\n","            tr = gpu_t(lambda: torch.linalg.eigh(Ax), 10, 100)\n","            tf = gpu_t(lambda: sf(Ax), 10, 100)\n","            thr_r = Bx / tr; thr_f = Bx / tf\n","            print(f\"  n={nx} B={Bx:>5}: cuSOLVER {thr_r/1e6:.2f}M/s  FL {thr_f/1e6:.2f}M/s  ({tf/tr:.2f}× time)\")\n","            del Ax, sf\n","        except Exception as e:\n","            print(f\"  n={nx} B={Bx:>5}: {str(e)[:40]}\")\n","            torch.cuda.empty_cache()\n","\n","    print(f\"\\n\" + \"=\" * 72)\n","    print(f\"  Fast compiled (n=6 B=4096): {fmt(t_fast)} ({t_ref/t_fast:.2f}× vs cuSOLVER)\")\n","    print(f\"  Precise compiled:           {fmt(t_prec)} ({t_ref/t_prec:.2f}× vs cuSOLVER)\")\n","    print(f\"  cuSOLVER:                   {fmt(t_ref)}\")\n","    print(\"=\" * 72)\n","\n","\n","if __name__ == '__main__':\n","    main()"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"trGWXuyO1seL","executionInfo":{"status":"ok","timestamp":1775044941873,"user_tz":420,"elapsed":33984,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"deac3f44-234d-4137-ab2a-7464b38c72e8"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["========================================================================\n","  FL Eigh — Performance Benchmark\n","========================================================================\n","  NVIDIA RTX PRO 6000 Blackwell Server Edition\n","  PyTorch 2.10.0+cu128\n","  VRAM: 95.0 GB\n","\n","  Compiling Fast (fullgraph=True)... done.\n","  Compiling Precise (fullgraph=True)... done.\n","\n","========================================================================\n","  PRIMARY: n=6 B=4096\n","========================================================================\n","\n","  Implementation                    Eager   Compiled  vs cuSOLVER\n","  ──────────────────────────── ────────── ────────── ────────────\n","  cuSOLVER                              —    241.8µs        1.00×\n","  FL Fast (fp32 eigvec)            8.09ms    334.0µs        0.72×\n","  FL Precise (fp64+Rayleigh)       8.09ms    338.1µs        0.72×\n","\n","========================================================================\n","  BATCH SCALING (n=6, compiled, dynamic=True)\n","========================================================================\n","\n","       B    cuSOLVER     FL Fast    F/cuS     FL Prec    P/cuS\n","  ──────  ──────────  ──────────  ───────  ──────────  ───────\n","     256     101.5µs     299.6µs    0.34×     289.3µs    0.35×\n","     512     103.6µs     294.8µs    0.35×     290.1µs    0.36×\n","    1024     116.3µs     298.4µs    0.39×     291.6µs    0.40×\n","    2048     154.4µs     295.8µs    0.52×     288.9µs    0.53×\n","    4096     245.1µs     336.3µs    0.73×     338.3µs    0.72×\n","    8192     408.6µs     547.3µs    0.75×     558.2µs    0.73×\n","   16384     744.4µs     970.0µs    0.77×     995.8µs    0.75×\n","\n","========================================================================\n","  SIZE SCALING (B=4096, compiled per size)\n","========================================================================\n","\n","    n    cuSOLVER     FL Fast    F/cuS     FL Prec    P/cuS\n","  ───  ──────────  ──────────  ───────  ──────────  ───────\n","    3     136.7µs     208.9µs    0.65×     212.6µs    0.64×\n","    4     163.8µs     220.0µs    0.74×     216.8µs    0.76×\n","    5     199.8µs     271.2µs    0.74×     275.4µs    0.73×\n","    6     238.7µs     334.3µs    0.71×     338.3µs    0.71×\n","    8     323.2µs     621.5µs    0.52×     649.2µs    0.50×\n","   10     681.0µs      1.18ms    0.58×      1.20ms    0.57×\n"]},{"output_type":"stream","name":"stderr","text":["W0401 12:02:39.384000 5191 torch/_dynamo/convert_frame.py:1676] [12/8] torch._dynamo hit config.recompile_limit (8)\n","W0401 12:02:39.384000 5191 torch/_dynamo/convert_frame.py:1676] [12/8]    function: 'forward' (/tmp/ipykernel_5191/3914342972.py:83)\n","W0401 12:02:39.384000 5191 torch/_dynamo/convert_frame.py:1676] [12/8]    last reason: 12/7: tensor 'A' size mismatch at index 1. expected 12, actual 16\n","W0401 12:02:39.384000 5191 torch/_dynamo/convert_frame.py:1676] [12/8] To log all recompilation reasons, use TORCH_LOGS=\"recompiles\".\n","W0401 12:02:39.384000 5191 torch/_dynamo/convert_frame.py:1676] [12/8] To diagnose recompilation issues, see https://pytorch.org/docs/main/compile/programming_model.recompilation.html\n"]},{"output_type":"stream","name":"stdout","text":["   12     826.4µs      2.15ms    0.39×      2.16ms    0.38×\n","   16  ERR: recompile_limit reached with fullgraph=T\n","\n","========================================================================\n","  MEMORY (n=6 B=4096)\n","========================================================================\n","  cuSOLVER         1098.7 MB\n","  FL Fast          23.0 MB\n"]},{"output_type":"stream","name":"stderr","text":["W0401 12:02:47.690000 5191 torch/_dynamo/convert_frame.py:1676] [12/9] torch._dynamo hit config.recompile_limit (8)\n","W0401 12:02:47.690000 5191 torch/_dynamo/convert_frame.py:1676] [12/9]    function: 'forward' (/tmp/ipykernel_5191/3914342972.py:83)\n","W0401 12:02:47.690000 5191 torch/_dynamo/convert_frame.py:1676] [12/9]    last reason: 12/7: tensor 'A' size mismatch at index 0. expected 4096, actual 8192\n","W0401 12:02:47.690000 5191 torch/_dynamo/convert_frame.py:1676] [12/9] To log all recompilation reasons, use TORCH_LOGS=\"recompiles\".\n","W0401 12:02:47.690000 5191 torch/_dynamo/convert_frame.py:1676] [12/9] To diagnose recompilation issues, see https://pytorch.org/docs/main/compile/programming_model.recompilation.html\n","W0401 12:02:47.693000 5191 torch/_dynamo/convert_frame.py:1676] [12/10] torch._dynamo hit config.recompile_limit (8)\n","W0401 12:02:47.693000 5191 torch/_dynamo/convert_frame.py:1676] [12/10]    function: 'forward' (/tmp/ipykernel_5191/3914342972.py:83)\n","W0401 12:02:47.693000 5191 torch/_dynamo/convert_frame.py:1676] [12/10]    last reason: 12/7: tensor 'A' size mismatch at index 0. expected 4096, actual 8192\n","W0401 12:02:47.693000 5191 torch/_dynamo/convert_frame.py:1676] [12/10] To log all recompilation reasons, use TORCH_LOGS=\"recompiles\".\n","W0401 12:02:47.693000 5191 torch/_dynamo/convert_frame.py:1676] [12/10] To diagnose recompilation issues, see https://pytorch.org/docs/main/compile/programming_model.recompilation.html\n","W0401 12:02:47.695000 5191 torch/_dynamo/convert_frame.py:1676] [12/11] torch._dynamo hit config.recompile_limit (8)\n","W0401 12:02:47.695000 5191 torch/_dynamo/convert_frame.py:1676] [12/11]    function: 'forward' (/tmp/ipykernel_5191/3914342972.py:83)\n","W0401 12:02:47.695000 5191 torch/_dynamo/convert_frame.py:1676] [12/11]    last reason: 12/7: tensor 'A' size mismatch at index 0. expected 4096, actual 16384\n","W0401 12:02:47.695000 5191 torch/_dynamo/convert_frame.py:1676] [12/11] To log all recompilation reasons, use TORCH_LOGS=\"recompiles\".\n","W0401 12:02:47.695000 5191 torch/_dynamo/convert_frame.py:1676] [12/11] To diagnose recompilation issues, see https://pytorch.org/docs/main/compile/programming_model.recompilation.html\n"]},{"output_type":"stream","name":"stdout","text":["  FL Precise       29.1 MB\n","\n","========================================================================\n","  THROUGHPUT AT SCALE (matrices/second)\n","========================================================================\n","  n=5 B= 8192: recompile_limit reached with fullgraph=T\n","  n=6 B= 8192: recompile_limit reached with fullgraph=T\n","  n=6 B=16384: recompile_limit reached with fullgraph=T\n","  n=8 B= 4096: cuSOLVER 12.69M/s  FL 6.58M/s  (1.93× time)\n","\n","========================================================================\n","  Fast compiled (n=6 B=4096): 334.0µs (0.72× vs cuSOLVER)\n","  Precise compiled:           338.1µs (0.72× vs cuSOLVER)\n","  cuSOLVER:                   241.8µs\n","========================================================================\n"]}]},{"cell_type":"code","source":["\"\"\"\n","fl_eigh_triton.py — Full FL eigendecomposition as a single Triton kernel.\n","\n","Entire pipeline in registers — zero global memory traffic between phases.\n","One kernel launch for the complete eigendecomposition of B matrices.\n","\n","Key insight: interleaved FL recomputation + Horner accumulation avoids\n","storing N M-matrices (which would exceed register limits). Each eigenvector\n","costs N-1 matmuls but keeps peak registers at ~240/255.\n","\n","Register budget per thread (n=6):\n","  Phase 1 (FL coeffs):   A(72) + M(72) + c(14) = 158  ✓\n","  Phase 2 (Laguerre):    c(14) + roots(12) + work(50) = 76  ✓\n","  Phase 3 (eigvec):      A(72) + M(72) + R(72) + c(14) + λ(2) = 232  ✓\n","  Phase 4 (NS+Rayleigh): V(36) + work(72) = 108  ✓\n","\n","Author: AbstractPhil / GeoLIP project\n","\"\"\"\n","\n","import math, time, gc, sys\n","import torch\n","import triton\n","import triton.language as tl\n","from torch import Tensor\n","from typing import Tuple\n","\n","torch.backends.cuda.matmul.allow_tf32 = False\n","torch.backends.cudnn.allow_tf32 = False\n","torch.set_float32_matmul_precision('highest')\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════\n","# Triton kernel: complete FL eigendecomposition\n","# ═══════════════════════════════════════════════════════════════════════\n","\n","@triton.jit\n","def _fl_eigh_kernel(\n","    A_ptr, evals_ptr, evecs_ptr,\n","    B,\n","    N: tl.constexpr,       # matrix size (6)\n","    BLOCK_B: tl.constexpr,  # batch elements per program (32)\n","):\n","    N2: tl.constexpr = N * N\n","    pid = tl.program_id(0)\n","    b_idx = pid * BLOCK_B + tl.arange(0, BLOCK_B)\n","    b_mask = b_idx < B\n","\n","    # ════════════════════════════════════════════════════════════\n","    # Load A + pre-scale\n","    # ════════════════════════════════════════════════════════════\n","    # Flatten: A[b,i,j] at offset b*N2 + i*N + j\n","    # We store matrix elements as separate register variables\n","    # a_00..a_55 each of shape [BLOCK_B]\n","\n","    # Compute Frobenius norm while loading\n","    frob_sq = tl.zeros((BLOCK_B,), dtype=tl.float64)\n","\n","    # Load all N2 elements — unrolled by compiler since N is constexpr\n","    a = tl.zeros((BLOCK_B,), dtype=tl.float64)  # placeholder\n","    # We'll use a flat storage approach with explicit indexing\n","    # For n=6: 36 elements loaded individually\n","    # Using a helper array pattern that Triton JIT handles:\n","\n","    # Actually, we need 36 separate register variables.\n","    # Triton handles this via tl.zeros + store into offset patterns.\n","    # Let's use the pointer arithmetic approach:\n","\n","    # Load A into local fp64 registers\n","    # We process each element (i,j) explicitly\n","    # Store in a flat \"register file\" via global memory staging\n","\n","    # SIMPLER APPROACH: use a [BLOCK_B, N, N] tensor in registers\n","    # Triton represents this as a 3D block\n","\n","    # Load A as [BLOCK_B, N*N] then reshape mentally\n","    a_flat = tl.zeros((BLOCK_B, N2), dtype=tl.float64)\n","\n","    # Hmm — Triton doesn't support 2D+ indexing easily for non-power-of-2.\n","    # For N=6, N2=36. We need BLOCK dimensions to be present.\n","    # Let's use the approach of loading element-by-element.\n","\n","    # PRACTICAL APPROACH: use pointers with explicit loops\n","    # Each loop iteration loads one matrix element for all BLOCK_B matrices\n","\n","    base_offs = b_idx * N2  # [BLOCK_B] starting offset for each matrix\n","\n","    # We'll store the matrix as N2 separate [BLOCK_B] vectors\n","    # Using a compile-time list isn't possible in Triton directly.\n","    # Instead, we'll use a workspace buffer or tl.store/load to scratchpad.\n","\n","    # ACTUALLY: the cleanest Triton pattern for small fixed-size matrices\n","    # is to process everything through a global scratch buffer.\n","    # But that defeats the \"all in registers\" goal.\n","\n","    # Let me use the STRIDE approach: keep everything in a [BLOCK_B, N2] 2D tensor.\n","    # Triton CAN handle 2D tensors.\n","\n","    # Load A as [BLOCK_B, N2]\n","    # offs[b, ij] = b * N2 + ij\n","    ij_range = tl.arange(0, N2)  # [N2] — must be constexpr-sized\n","    # 2D offset: [BLOCK_B, N2]\n","    offs_a = base_offs[:, None] + ij_range[None, :]  # [BLOCK_B, N2]\n","    mask_a = b_mask[:, None] & (ij_range[None, :] < N2)\n","\n","    a = tl.load(A_ptr + offs_a, mask=mask_a, other=0.0).to(tl.float64)\n","    # a is [BLOCK_B, N2] fp64\n","\n","    # Frobenius norm\n","    frob_sq = tl.sum(a * a, axis=1)  # [BLOCK_B]\n","    scale = tl.sqrt(frob_sq / N)\n","    scale = tl.where(scale > 1e-12, scale, 1e-12)\n","    inv_s = 1.0 / scale\n","\n","    # Pre-scale A\n","    a = a * inv_s[:, None]\n","\n","    # ════════════════════════════════════════════════════════════\n","    # Phase 1: Faddeev-LeVerrier (fp64)\n","    # Compute coefficients c[0..N] only. M matrices recomputed in Phase 3.\n","    # ════════════════════════════════════════════════════════════\n","\n","    # Coefficients: c[k] for k=0..N. c[N] = 1.\n","    # Store as [BLOCK_B, N+1]\n","    c_range = tl.arange(0, N + 1)  # [N+1] = [7] for N=6\n","    c = tl.zeros((BLOCK_B, N + 1), dtype=tl.float64)\n","    # c[:, N] = 1.0\n","    c = tl.where(c_range[None, :] == N, 1.0, c)\n","\n","    # M starts as zero matrix [BLOCK_B, N2]\n","    m = tl.zeros((BLOCK_B, N2), dtype=tl.float64)\n","\n","    # N iterations of FL recursion\n","    for k_iter in tl.static_range(1, N + 1):\n","        # m_new = A @ m + c[:, N-k_iter+1] * I\n","        # Batched matmul using sum-reduction:\n","        # m_new[b, i*N+j] = sum_l a[b, i*N+l] * m[b, l*N+j]\n","\n","        m_new = tl.zeros((BLOCK_B, N2), dtype=tl.float64)\n","        for i in tl.static_range(N):\n","            for j in tl.static_range(N):\n","                acc = tl.zeros((BLOCK_B,), dtype=tl.float64)\n","                for l in tl.static_range(N):\n","                    acc += a[:, i * N + l] * m[:, l * N + j]\n","                # Add c[N-k_iter+1] * I[i,j]\n","                if i == j:\n","                    acc += c[:, N - k_iter + 1]\n","                m_new_ij = acc\n","                # Store into m_new — but we can't index-assign into 2D tl tensor\n","                # We need to build m_new column by column\n","                # Actually, tl doesn't support m_new[:, idx] = val assignment\n","\n","                # WORKAROUND: accumulate via where mask\n","                ij_idx = i * N + j\n","                m_new = tl.where(ij_range[None, :] == ij_idx, acc[:, None], m_new)\n","\n","        # trace(A @ m_new) = sum_{i,l} a[b, i*N+l] * m_new[b, l*N+i]\n","        tr = tl.zeros((BLOCK_B,), dtype=tl.float64)\n","        for i in tl.static_range(N):\n","            for l in tl.static_range(N):\n","                tr += a[:, i * N + l] * m_new[:, l * N + i]\n","\n","        # c[:, N-k_iter] = -tr / k_iter\n","        c_val = -tr / float(k_iter)\n","        c = tl.where(c_range[None, :] == (N - k_iter), c_val[:, None], c)\n","\n","        m = m_new\n","\n","    # ════════════════════════════════════════════════════════════\n","    # Phase 2: Laguerre root-finding + Newton polish (fp64)\n","    # ════════════════════════════════════════════════════════════\n","\n","    # Initialize roots from sorted diagonal of A_scaled\n","    # Diagonal: a[:, 0], a[:, 7], a[:, 14], a[:, 21], a[:, 28], a[:, 35] for N=6\n","    diag = tl.zeros((BLOCK_B, N), dtype=tl.float64)\n","    n_range = tl.arange(0, N)\n","    for i in tl.static_range(N):\n","        diag_val = a[:, i * N + i]\n","        diag = tl.where(n_range[None, :] == i, diag_val[:, None], diag)\n","\n","    # Sort diagonal — simple bubble sort (N is small)\n","    for pass_i in tl.static_range(N - 1):\n","        for j in tl.static_range(N - 1):\n","            d_j = diag[:, j]\n","            d_j1 = diag[:, j + 1]\n","            should_swap = d_j > d_j1\n","            new_j = tl.where(should_swap, d_j1, d_j)\n","            new_j1 = tl.where(should_swap, d_j, d_j1)\n","            diag = tl.where(n_range[None, :] == j, new_j[:, None], diag)\n","            diag = tl.where(n_range[None, :] == (j + 1), new_j1[:, None], diag)\n","\n","    # Perturbation\n","    pert = tl.zeros((N,), dtype=tl.float64)\n","    for i in tl.static_range(N):\n","        pert_val = -1e-4 + 2e-4 * float(i) / float(N - 1) if N > 1 else 0.0\n","        pert = tl.where(n_range == i, pert_val, pert)\n","    z_init = diag + pert[None, :]\n","\n","    # Working copy of coefficients for deflation\n","    cl = tl.zeros((BLOCK_B, N + 1), dtype=tl.float64) + c\n","\n","    roots = tl.zeros((BLOCK_B, N), dtype=tl.float64)\n","\n","    # Sequential Laguerre + deflation\n","    for ri in tl.static_range(N):\n","        deg = N - ri\n","        z = z_init[:, ri]  # [BLOCK_B]\n","\n","        # 5 Laguerre iterations\n","        for _lag in tl.static_range(5):\n","            # Horner: p, p', p''/2\n","            pv = tl.zeros((BLOCK_B,), dtype=tl.float64)\n","            # pv = cl[:, deg]\n","            for d in tl.static_range(N + 1):\n","                if d == deg:\n","                    pv = cl[:, d]\n","\n","            dp = tl.zeros((BLOCK_B,), dtype=tl.float64)\n","            d2 = tl.zeros((BLOCK_B,), dtype=tl.float64)\n","\n","            for j_rev in tl.static_range(N):\n","                j_idx = N - 1 - j_rev  # descending\n","                # Only process if j_idx < deg\n","                # For constexpr N=6 and varying deg, we need conditional\n","                # Since deg = N - ri and both are constexpr-derivable...\n","                # Actually ri is a static_range variable, so deg is known\n","                if j_idx < deg:\n","                    d2 = d2 * z + dp\n","                    dp = dp * z + pv\n","                    # pv = pv * z + cl[:, j_idx]\n","                    pv = pv * z + cl[:, j_idx]\n","\n","            # Laguerre step\n","            ok = tl.abs(pv) > 1e-30\n","            ps = tl.where(ok, pv, 1.0)\n","            G = tl.where(ok, dp / ps, 0.0)\n","            H = G * G - tl.where(ok, 2.0 * d2 / ps, 0.0)\n","            disc = tl.maximum((float(deg) - 1.0) * (float(deg) * H - G * G), 0.0)\n","            sq = tl.sqrt(disc)\n","            gp = G + sq\n","            gm = G - sq\n","            den = tl.where(tl.abs(gp) >= tl.abs(gm), gp, gm)\n","            dok = tl.abs(den) > 1e-20\n","            ds = tl.where(dok, den, 1.0)\n","            z = z - tl.where(dok, float(deg) / ds, 0.0)\n","\n","        # Store root\n","        roots = tl.where(n_range[None, :] == ri, z[:, None], roots)\n","\n","        # Synthetic division: cl /= (x - z)\n","        # b = cl[:, deg], then for j = deg-1..1: b_new = cl[:,j] + z*b, cl[:,j] = b\n","        b_div = tl.zeros((BLOCK_B,), dtype=tl.float64)\n","        for d in tl.static_range(N + 1):\n","            if d == deg:\n","                b_div = cl[:, d]\n","\n","        for j_rev in tl.static_range(N):\n","            j_idx = N - 1 - j_rev  # descending\n","            if j_idx > 0 and j_idx < deg:\n","                bn = cl[:, j_idx] + z * b_div\n","                cl = tl.where(c_range[None, :] == j_idx, b_div[:, None], cl)\n","                b_div = bn\n","        cl = tl.where(c_range[None, :] == 0, b_div[:, None], cl)\n","\n","    # Newton polish on ORIGINAL polynomial (3 iterations)\n","    for _pol in tl.static_range(3):\n","        pv = tl.zeros((BLOCK_B, N), dtype=tl.float64) + 1.0\n","        dp = tl.zeros((BLOCK_B, N), dtype=tl.float64)\n","        for j_rev in tl.static_range(N):\n","            j_idx = N - 1 - j_rev\n","            dp = dp * roots + pv\n","            pv = pv * roots + c[:, j_idx:j_idx + 1]\n","        ok = tl.abs(dp) > 1e-30\n","        dps = tl.where(ok, dp, 1.0)\n","        roots = roots - tl.where(ok, pv / dps, 0.0)\n","\n","    # ════════════════════════════════════════════════════════════\n","    # Phase 3: Eigenvectors via interleaved FL+Horner (fp64)\n","    # For each eigenvalue, recompute M[k] on the fly and accumulate Horner\n","    # Peak registers: A(36×2) + M(36×2) + R(36×2) + misc ≈ 240\n","    # ════════════════════════════════════════════════════════════\n","\n","    # Output eigenvectors as [BLOCK_B, N2] fp32\n","    evecs = tl.zeros((BLOCK_B, N2), dtype=tl.float32)\n","\n","    for ei in tl.static_range(N):\n","        lam = roots[:, ei]  # [BLOCK_B] fp64\n","\n","        # Recompute FL + Horner interleaved:\n","        # M starts at 0. R starts at 0.\n","        # Step k=1: M[1] = A@0 + c[N]*I = I. R = I (first Horner term)\n","        # Step k=2: M[2] = A@I + c[N-1]*I. R = R*λ + M[2]\n","        # ...\n","\n","        m_loc = tl.zeros((BLOCK_B, N2), dtype=tl.float64)\n","        r_loc = tl.zeros((BLOCK_B, N2), dtype=tl.float64)\n","\n","        for k_step in tl.static_range(1, N + 1):\n","            # M[k] = A @ m_loc + c[N-k+1] * I\n","            m_new_loc = tl.zeros((BLOCK_B, N2), dtype=tl.float64)\n","            for i in tl.static_range(N):\n","                for j in tl.static_range(N):\n","                    acc = tl.zeros((BLOCK_B,), dtype=tl.float64)\n","                    for l in tl.static_range(N):\n","                        acc += a[:, i * N + l] * m_loc[:, l * N + j]\n","                    if i == j:\n","                        acc += c[:, N - k_step + 1]\n","                    ij_idx = i * N + j\n","                    m_new_loc = tl.where(ij_range[None, :] == ij_idx, acc[:, None], m_new_loc)\n","\n","            # Horner: R = R * λ + M[k] (for k>=2), R = M[1] (for k=1)\n","            if k_step == 1:\n","                r_loc = m_new_loc + 0.0  # copy\n","            else:\n","                r_loc = r_loc * lam[:, None] + m_new_loc\n","\n","            m_loc = m_new_loc\n","\n","        # r_loc is [BLOCK_B, N2] fp64 — rank-1 adjugate at eigenvalue\n","        # Max-norm column extraction\n","        best_col = tl.zeros((BLOCK_B,), dtype=tl.int32)\n","        best_norm = tl.zeros((BLOCK_B,), dtype=tl.float64) - 1.0\n","\n","        for j in tl.static_range(N):\n","            col_sq = tl.zeros((BLOCK_B,), dtype=tl.float64)\n","            for i in tl.static_range(N):\n","                val = r_loc[:, i * N + j]\n","                col_sq += val * val\n","            is_better = col_sq > best_norm\n","            best_norm = tl.where(is_better, col_sq, best_norm)\n","            best_col = tl.where(is_better, j, best_col)\n","\n","        # Extract best column and normalize\n","        vec = tl.zeros((BLOCK_B, N), dtype=tl.float64)\n","        for j in tl.static_range(N):\n","            for i in tl.static_range(N):\n","                val = r_loc[:, i * N + j]\n","                vec = tl.where((n_range[None, :] == i) & (best_col[:, None] == j),\n","                               val[:, None], vec)\n","\n","        vnorm = tl.sqrt(tl.sum(vec * vec, axis=1) + 1e-60)\n","        vec = vec / vnorm[:, None]\n","\n","        # Store as column ei of evecs: evecs[:, i*N + ei] = vec[:, i]\n","        for i in tl.static_range(N):\n","            evecs = tl.where(ij_range[None, :] == (i * N + ei),\n","                             vec[:, i].to(tl.float32)[:, None], evecs)\n","\n","    # ════════════════════════════════════════════════════════════\n","    # Phase 4: Newton-Schulz (fp32) — 2 iterations\n","    # ════════════════════════════════════════════════════════════\n","\n","    # Y = VᵀV, X = I\n","    y = tl.zeros((BLOCK_B, N2), dtype=tl.float32)\n","    for i in tl.static_range(N):\n","        for j in tl.static_range(N):\n","            acc = tl.zeros((BLOCK_B,), dtype=tl.float32)\n","            for l in tl.static_range(N):\n","                # VᵀV[i,j] = sum_l V[l,i]*V[l,j] = sum_l evecs[:,l*N+i]*evecs[:,l*N+j]\n","                acc += evecs[:, l * N + i] * evecs[:, l * N + j]\n","            y = tl.where(ij_range[None, :] == (i * N + j), acc[:, None], y)\n","\n","    x = tl.zeros((BLOCK_B, N2), dtype=tl.float32)\n","    for i in tl.static_range(N):\n","        x = tl.where(ij_range[None, :] == (i * N + i), 1.0, x)\n","\n","    for _ns in tl.static_range(2):\n","        # T = 3I - Y\n","        t_mat = -y\n","        for i in tl.static_range(N):\n","            t_mat = tl.where(ij_range[None, :] == (i * N + i),\n","                             3.0 - y[:, i * N + i:i * N + i + 1], t_mat)\n","\n","        # X_new = 0.5 * X @ T\n","        x_new = tl.zeros((BLOCK_B, N2), dtype=tl.float32)\n","        for i in tl.static_range(N):\n","            for j in tl.static_range(N):\n","                acc = tl.zeros((BLOCK_B,), dtype=tl.float32)\n","                for l in tl.static_range(N):\n","                    acc += x[:, i * N + l] * t_mat[:, l * N + j]\n","                x_new = tl.where(ij_range[None, :] == (i * N + j), (0.5 * acc)[:, None], x_new)\n","\n","        # Y_new = 0.5 * T @ Y\n","        y_new = tl.zeros((BLOCK_B, N2), dtype=tl.float32)\n","        for i in tl.static_range(N):\n","            for j in tl.static_range(N):\n","                acc = tl.zeros((BLOCK_B,), dtype=tl.float32)\n","                for l in tl.static_range(N):\n","                    acc += t_mat[:, i * N + l] * y[:, l * N + j]\n","                y_new = tl.where(ij_range[None, :] == (i * N + j), (0.5 * acc)[:, None], y_new)\n","\n","        x = x_new\n","        y = y_new\n","\n","    # V_orth = V @ X\n","    v_orth = tl.zeros((BLOCK_B, N2), dtype=tl.float32)\n","    for i in tl.static_range(N):\n","        for j in tl.static_range(N):\n","            acc = tl.zeros((BLOCK_B,), dtype=tl.float32)\n","            for l in tl.static_range(N):\n","                acc += evecs[:, i * N + l] * x[:, l * N + j]\n","            v_orth = tl.where(ij_range[None, :] == (i * N + j), acc[:, None], v_orth)\n","    evecs = v_orth\n","\n","    # ════════════════════════════════════════════════════════════\n","    # Phase 5: Rayleigh quotient (fp32)\n","    # λᵢ = vᵢᵀ A vᵢ = sum_l,m V[l,i] * A[l,m] * V[m,i]\n","    # ════════════════════════════════════════════════════════════\n","\n","    a_f32 = a.to(tl.float32)  # scaled A in fp32\n","    evals_out = tl.zeros((BLOCK_B, N), dtype=tl.float32)\n","\n","    for ei in tl.static_range(N):\n","        # AV_i[l] = sum_m A[l,m] * V[m,i]\n","        # λ_i = sum_l V[l,i] * AV_i[l]\n","        lam_i = tl.zeros((BLOCK_B,), dtype=tl.float32)\n","        for l in tl.static_range(N):\n","            av_l = tl.zeros((BLOCK_B,), dtype=tl.float32)\n","            for mm in tl.static_range(N):\n","                av_l += a_f32[:, l * N + mm] * evecs[:, mm * N + ei]\n","            lam_i += evecs[:, l * N + ei] * av_l\n","        # Un-scale eigenvalue\n","        evals_out = tl.where(n_range[None, :] == ei,\n","                             (lam_i * scale.to(tl.float32))[:, None], evals_out)\n","\n","    # ════════════════════════════════════════════════════════════\n","    # Sort eigenvalues ascending + permute eigenvectors\n","    # Simple bubble sort (N is small)\n","    # ════════════════════════════════════════════════════════════\n","\n","    perm = tl.zeros((BLOCK_B, N), dtype=tl.int32)\n","    for i in tl.static_range(N):\n","        perm = tl.where(n_range[None, :] == i, i, perm)\n","\n","    for pass_i in tl.static_range(N - 1):\n","        for j in tl.static_range(N - 1):\n","            e_j = evals_out[:, j]\n","            e_j1 = evals_out[:, j + 1]\n","            should_swap = e_j > e_j1\n","            new_j = tl.where(should_swap, e_j1, e_j)\n","            new_j1 = tl.where(should_swap, e_j, e_j1)\n","            evals_out = tl.where(n_range[None, :] == j, new_j[:, None], evals_out)\n","            evals_out = tl.where(n_range[None, :] == (j + 1), new_j1[:, None], evals_out)\n","\n","            p_j = perm[:, j]\n","            p_j1 = perm[:, j + 1]\n","            new_pj = tl.where(should_swap, p_j1, p_j)\n","            new_pj1 = tl.where(should_swap, p_j, p_j1)\n","            perm = tl.where(n_range[None, :] == j, new_pj[:, None], perm)\n","            perm = tl.where(n_range[None, :] == (j + 1), new_pj1[:, None], perm)\n","\n","    # Permute eigenvectors according to sorted order\n","    evecs_sorted = tl.zeros((BLOCK_B, N2), dtype=tl.float32)\n","    for col_out in tl.static_range(N):\n","        for col_src in tl.static_range(N):\n","            is_match = (perm[:, col_out] == col_src)\n","            for row in tl.static_range(N):\n","                src_val = evecs[:, row * N + col_src]\n","                evecs_sorted = tl.where(\n","                    (ij_range[None, :] == (row * N + col_out)) & is_match[:, None],\n","                    src_val[:, None], evecs_sorted)\n","\n","    # ════════════════════════════════════════════════════════════\n","    # Store results\n","    # ════════════════════════════════════════════════════════════\n","\n","    # Eigenvalues: [B, N]\n","    offs_ev = b_idx[:, None] * N + n_range[None, :]\n","    mask_ev = b_mask[:, None] & (n_range[None, :] < N)\n","    tl.store(evals_ptr + offs_ev, evals_out, mask=mask_ev)\n","\n","    # Eigenvectors: [B, N, N]\n","    offs_vec = b_idx[:, None] * N2 + ij_range[None, :]\n","    mask_vec = b_mask[:, None] & (ij_range[None, :] < N2)\n","    tl.store(evecs_ptr + offs_vec, evecs_sorted, mask=mask_vec)\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════\n","# Python wrapper\n","# ═══════════════════════════════════════════════════════════════════════\n","\n","class FLEighTriton:\n","    \"\"\"Triton-accelerated FL eigendecomposition.\"\"\"\n","\n","    @staticmethod\n","    def apply(A: Tensor) -> Tuple[Tensor, Tensor]:\n","        B, n, _ = A.shape\n","        assert n <= 16, f\"FLEighTriton supports n<=16, got {n}\"\n","        device = A.device\n","\n","        evals = torch.empty(B, n, device=device, dtype=torch.float32)\n","        evecs = torch.empty(B, n, n, device=device, dtype=torch.float32)\n","\n","        BLOCK_B = 32\n","        grid = ((B + BLOCK_B - 1) // BLOCK_B,)\n","\n","        _fl_eigh_kernel[grid](\n","            A.contiguous(), evals, evecs,\n","            B, n, BLOCK_B,\n","        )\n","\n","        return evals, evecs\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════\n","# Torch.compile baseline (for comparison)\n","# ═══════════════════════════════════════════════════════════════════════\n","\n","import torch.nn as nn\n","\n","class FLEighCompiled(nn.Module):\n","    def forward(self, A):\n","        B, n, _ = A.shape; device = A.device\n","        scale = (torch.linalg.norm(A.reshape(B,-1),dim=-1)/math.sqrt(n)).clamp(min=1e-12)\n","        As = A / scale[:,None,None]\n","        Ad = As.double()\n","        eye_d = torch.eye(n,device=device,dtype=torch.float64).unsqueeze(0).expand(B,-1,-1)\n","        c = torch.zeros(B,n+1,device=device,dtype=torch.float64); c[:,n]=1.0\n","        Mstore = torch.zeros(n+1,B,n,n,device=device,dtype=torch.float64)\n","        Mk = torch.zeros(B,n,n,device=device,dtype=torch.float64)\n","        for k in range(1,n+1):\n","            Mk = torch.bmm(Ad,Mk) + c[:,n-k+1,None,None]*eye_d\n","            Mstore[k] = Mk; c[:,n-k] = -(Ad*Mk).sum((-2,-1))/k\n","        dt = torch.float64 if n>6 else torch.float32\n","        cl = c.to(dt).clone(); roots = torch.zeros(B,n,device=device,dtype=dt)\n","        zi = As.to(dt).diagonal(dim1=-2,dim2=-1).sort(dim=-1).values\n","        zi = zi + torch.linspace(-1e-4,1e-4,n,device=device,dtype=dt).unsqueeze(0)\n","        for ri in range(n):\n","            deg=n-ri; z=zi[:,ri]\n","            for _ in range(5):\n","                pv=cl[:,deg]; dp=torch.zeros(B,device=device,dtype=dt); d2=torch.zeros(B,device=device,dtype=dt)\n","                for j in range(deg-1,-1,-1):\n","                    d2=d2*z+dp; dp=dp*z+pv; pv=pv*z+cl[:,j]\n","                ok=pv.abs()>1e-30; ps=torch.where(ok,pv,torch.ones_like(pv))\n","                G=torch.where(ok,dp/ps,torch.zeros_like(dp))\n","                H=G*G-torch.where(ok,2.0*d2/ps,torch.zeros_like(d2))\n","                disc=((deg-1.0)*(deg*H-G*G)).clamp(min=0.0); sq=torch.sqrt(disc)\n","                gp=G+sq; gm=G-sq; den=torch.where(gp.abs()>=gm.abs(),gp,gm)\n","                dok=den.abs()>1e-20; ds=torch.where(dok,den,torch.ones_like(den))\n","                z=z-torch.where(dok,float(deg)/ds,torch.zeros_like(den))\n","            roots[:,ri]=z; b_d=cl[:,deg]\n","            for j in range(deg-1,0,-1):\n","                bn=cl[:,j]+z*b_d; cl[:,j]=b_d; b_d=bn\n","            cl[:,0]=b_d\n","        roots=roots.double()\n","        for _ in range(3):\n","            pv=torch.ones(B,n,device=device,dtype=torch.float64); dp=torch.zeros(B,n,device=device,dtype=torch.float64)\n","            for j in range(n-1,-1,-1):\n","                dp=dp*roots+pv; pv=pv*roots+c[:,j:j+1]\n","            ok=dp.abs()>1e-30; dps=torch.where(ok,dp,torch.ones_like(dp))\n","            roots=roots-torch.where(ok,pv/dps,torch.zeros_like(pv))\n","        lam=roots; R=Mstore[1].unsqueeze(1).expand(-1,n,-1,-1).clone()\n","        for k in range(2,n+1): R=R*lam[:,:,None,None]+Mstore[k].unsqueeze(1)\n","        cn=R.norm(dim=-2); best=cn.argmax(dim=-1)\n","        idx=best.unsqueeze(-1).unsqueeze(-1).expand(-1,-1,n,1)\n","        vec=R.gather(-1,idx).squeeze(-1); vec=vec/(vec.norm(dim=-1,keepdim=True)+1e-30)\n","        V=vec.float().transpose(-2,-1)\n","        eye_f=torch.eye(n,device=device,dtype=torch.float32).unsqueeze(0).expand(B,-1,-1)\n","        Y=torch.bmm(V.mT,V); X=eye_f.clone()\n","        for _ in range(2):\n","            T=3.0*eye_f-Y; X=0.5*torch.bmm(X,T); Y=0.5*torch.bmm(T,Y)\n","        V=torch.bmm(V,X); AV=torch.bmm(A,V); evals=(V*AV).sum(dim=-2)\n","        se,perm=evals.sort(dim=-1)\n","        return se,V.gather(-1,perm.unsqueeze(-2).expand_as(V))\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════\n","# Benchmark\n","# ═══════════════════════════════════════════════════════════════════════\n","\n","def sync(): torch.cuda.synchronize()\n","def gt(fn,w=20,r=200):\n","    for _ in range(w): fn()\n","    sync(); t=time.perf_counter()\n","    for _ in range(r): fn()\n","    sync(); return (time.perf_counter()-t)/r\n","def fmt(s):\n","    if s<1e-3: return f\"{s*1e6:.1f}µs\"\n","    if s<1: return f\"{s*1e3:.2f}ms\"\n","    return f\"{s:.3f}s\"\n","\n","def validate(A, v, V, rv, rV):\n","    B,n,_=A.shape\n","    ve=(v-rv).abs().max().item()\n","    dots=torch.bmm(rV.double().mT, V.double()).abs().max(dim=-1).values.min().item()\n","    AV=torch.bmm(A.double(),V.double()); VL=V.double()*v.double().unsqueeze(-2)\n","    An=A.double().reshape(B,-1).norm(dim=-1).clamp(min=1e-30)\n","    res=(AV-VL).reshape(B,-1).norm(dim=-1)/An\n","    return dict(val=ve, align=dots, res=res.max().item())\n","\n","\n","def main():\n","    if not torch.cuda.is_available(): sys.exit(1)\n","    dev=torch.device('cuda')\n","    p=torch.cuda.get_device_properties(0)\n","\n","    print(\"=\"*72)\n","    print(\"  FL Eigh — Triton Kernel Benchmark\")\n","    print(\"=\"*72)\n","    print(f\"  {p.name}\")\n","    print(f\"  PyTorch {torch.__version__}\")\n","    print(f\"  Triton {triton.__version__}\")\n","\n","    N=6; B=4096\n","    A=(lambda R:(R+R.mT)/2)(torch.randn(B,N,N,device=dev))\n","    rv,rV=torch.linalg.eigh(A)\n","\n","    # ── Accuracy ──\n","    print(f\"\\n  ACCURACY (n={N} B={B})\")\n","    try:\n","        tv,tV = FLEighTriton.apply(A)\n","        mt=validate(A,tv,tV,rv,rV)\n","        print(f\"  Triton:  val={mt['val']:.1e} align={mt['align']:.6f} res={mt['res']:.1e}\")\n","    except Exception as e:\n","        print(f\"  Triton FAILED: {str(e)[:200]}\")\n","        tv=None\n","\n","    # torch.compile baseline\n","    solver=FLEighCompiled()\n","    fv,fV=solver(A); mc=validate(A,fv,fV,rv,rV)\n","    print(f\"  Compile: val={mc['val']:.1e} align={mc['align']:.6f} res={mc['res']:.1e}\")\n","\n","    # ── Throughput ──\n","    print(f\"\\n  THROUGHPUT (n={N} B={B})\")\n","    tr=gt(lambda:torch.linalg.eigh(A))\n","    print(f\"  cuSOLVER:  {fmt(tr)}\")\n","\n","    if tv is not None:\n","        tt=gt(lambda:FLEighTriton.apply(A))\n","        print(f\"  Triton:    {fmt(tt)} ({tr/tt:.2f}× vs cuSOLVER)\")\n","\n","    try:\n","        cs=torch.compile(solver,fullgraph=True)\n","        for _ in range(3): cs(A); sync()\n","        tc=gt(lambda:cs(A))\n","        print(f\"  Compiled:  {fmt(tc)} ({tr/tc:.2f}× vs cuSOLVER)\")\n","    except Exception as e:\n","        print(f\"  Compile FAILED: {str(e)[:100]}\")\n","        tc=None\n","\n","    # ── Batch scaling ──\n","    if tv is not None:\n","        print(f\"\\n  BATCH SCALING (n={N})\")\n","        print(f\"  {'B':>6}  {'cuSOLVER':>10}  {'Triton':>10}  {'T/cuS':>7}\")\n","        print(f\"  {'─'*6}  {'─'*10}  {'─'*10}  {'─'*7}\")\n","        for Bx in [256,512,1024,2048,4096,8192,16384]:\n","            Ax=(lambda R:(R+R.mT)/2)(torch.randn(Bx,N,N,device=dev))\n","            t1=gt(lambda:torch.linalg.eigh(Ax),10,100)\n","            t2=gt(lambda:FLEighTriton.apply(Ax),10,100)\n","            print(f\"  {Bx:>6}  {fmt(t1):>10}  {fmt(t2):>10}  {t1/t2:>6.2f}×\")\n","            del Ax\n","\n","    print(\"=\"*72)\n","\n","\n","if __name__=='__main__': main()"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"pvkAuTi-9s8K","executionInfo":{"status":"ok","timestamp":1775045631824,"user_tz":420,"elapsed":2550,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"326b2c82-9088-4c16-b451-1cc94d5cd041"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["========================================================================\n","  FL Eigh — Triton Kernel Benchmark\n","========================================================================\n","  NVIDIA RTX PRO 6000 Blackwell Server Edition\n","  PyTorch 2.10.0+cu128\n","  Triton 3.6.0\n","\n","  ACCURACY (n=6 B=4096)\n","  Triton FAILED: at 40:13:\n","    # Triton handles this via tl.zeros + store into offset patterns.\n","    # Let's use the pointer arithmetic approach:\n","\n","    # Load A into local fp64 registers\n","    # We process each element (i\n","  Compile: val=2.9e-06 align=0.999999 res=2.4e-07\n","\n","  THROUGHPUT (n=6 B=4096)\n","  cuSOLVER:  238.7µs\n","  Compiled:  338.1µs (0.71× vs cuSOLVER)\n","========================================================================\n"]}]},{"cell_type":"code","source":["\"\"\"\n","FL Eigh Triton — Generated Kernel.\n","Python generates fully-unrolled Triton source with explicit named variables.\n","No lists, no 2D tensors, no tl.where assignment. Just loads, FMAs, stores.\n","\"\"\"\n","import math, time, gc, sys\n","import torch\n","import triton\n","import triton.language as tl\n","from torch import Tensor\n","from typing import Tuple\n","\n","torch.backends.cuda.matmul.allow_tf32 = False\n","torch.backends.cudnn.allow_tf32 = False\n","torch.set_float32_matmul_precision('highest')\n","\n","\n","def _gen_kernel(N):\n","    \"\"\"Generate Triton kernel source for NxN eigendecomposition.\"\"\"\n","    N2 = N * N\n","    L = []  # lines\n","    def E(s): L.append(s)\n","    def var(prefix, i, j=None):\n","        return f\"{prefix}{i}\" if j is None else f\"{prefix}{i}_{j}\"\n","\n","    E(\"@triton.jit\")\n","    E(f\"def _fl_eigh_gen(A_ptr, evals_ptr, evecs_ptr, B, BLOCK_B: tl.constexpr):\")\n","    E(f\"    pid = tl.program_id(0)\")\n","    E(f\"    bid = pid * BLOCK_B + tl.arange(0, BLOCK_B)\")\n","    E(f\"    mask = bid < B\")\n","    E(f\"    off = bid * {N2}\")\n","\n","    # Load A\n","    E(f\"    # Load A\")\n","    E(f\"    frob = tl.zeros((BLOCK_B,), dtype=tl.float64)\")\n","    for i in range(N):\n","        for j in range(N):\n","            v = f\"a{i}_{j}\"\n","            E(f\"    {v} = tl.load(A_ptr + off + {i*N+j}, mask=mask, other=0.0).to(tl.float64)\")\n","            E(f\"    frob = frob + {v} * {v}\")\n","\n","    # Pre-scale\n","    E(f\"    sc = tl.sqrt(frob * {1.0/N})\")\n","    E(f\"    sc = tl.where(sc > 1e-12, sc, 1e-12)\")\n","    E(f\"    inv_sc = 1.0 / sc\")\n","    for i in range(N):\n","        for j in range(N):\n","            E(f\"    a{i}_{j} = a{i}_{j} * inv_sc\")\n","\n","    # Phase 1: FL coefficients only (dont store M)\n","    E(f\"    # Phase 1: FL coefficients\")\n","    for k in range(N+1):\n","        E(f\"    c{k} = tl.zeros((BLOCK_B,), dtype=tl.float64)\" +\n","          (f\" + 1.0\" if k == N else \"\"))\n","\n","    # M starts as zero\n","    for i in range(N):\n","        for j in range(N):\n","            E(f\"    m{i}_{j} = tl.zeros((BLOCK_B,), dtype=tl.float64)\")\n","\n","    for kk in range(1, N+1):\n","        E(f\"    # FL iter k={kk}\")\n","        c_diag = f\"c{N - kk + 1}\"\n","        # mn = A @ m + c_diag * I\n","        for i in range(N):\n","            for j in range(N):\n","                terms = [f\"a{i}_{l} * m{l}_{j}\" for l in range(N)]\n","                expr = \" + \".join(terms)\n","                if i == j:\n","                    expr += f\" + {c_diag}\"\n","                E(f\"    mn{i}_{j} = {expr}\")\n","        # trace(A @ mn)\n","        terms = [f\"a{i}_{l} * mn{l}_{i}\" for i in range(N) for l in range(N)]\n","        E(f\"    tr = \" + \" + \".join(terms))\n","        E(f\"    c{N - kk} = -tr * {1.0/kk}\")\n","        # m = mn\n","        for i in range(N):\n","            for j in range(N):\n","                E(f\"    m{i}_{j} = mn{i}_{j}\")\n","\n","    # Phase 2: Laguerre + deflation\n","    E(f\"    # Phase 2: Laguerre\")\n","    # Get sorted diagonal\n","    diag_vars = [f\"a{i}_{i}\" for i in range(N)]\n","    for i in range(N):\n","        E(f\"    d{i} = {diag_vars[i]}\")\n","    # Bubble sort\n","    for _ in range(N):\n","        for j in range(N-1):\n","            E(f\"    _sw = d{j} > d{j+1}\")\n","            E(f\"    _a, _b = d{j}, d{j+1}\")\n","            E(f\"    d{j} = tl.where(_sw, _b, _a)\")\n","            E(f\"    d{j+1} = tl.where(_sw, _a, _b)\")\n","\n","    # Perturbation\n","    for i in range(N):\n","        p = -1e-4 + 2e-4 * i / max(N-1, 1)\n","        E(f\"    d{i} = d{i} + {p}\")\n","\n","    # Working coefficients\n","    for k in range(N+1):\n","        E(f\"    w{k} = c{k} + 0.0\")\n","\n","    # Sequential roots\n","    for ri in range(N):\n","        deg = N - ri\n","        E(f\"    # Root {ri} (deg={deg})\")\n","        E(f\"    z = d{ri}\")\n","        for _lag in range(5):\n","            E(f\"    pv = w{deg}\")\n","            E(f\"    dp = tl.zeros((BLOCK_B,), dtype=tl.float64)\")\n","            E(f\"    d2 = tl.zeros((BLOCK_B,), dtype=tl.float64)\")\n","            for j in range(deg-1, -1, -1):\n","                E(f\"    d2 = d2 * z + dp\")\n","                E(f\"    dp = dp * z + pv\")\n","                E(f\"    pv = pv * z + w{j}\")\n","            E(f\"    ok = tl.abs(pv) > 1e-30\")\n","            E(f\"    ps = tl.where(ok, pv, 1.0)\")\n","            E(f\"    G = tl.where(ok, dp / ps, 0.0)\")\n","            E(f\"    H = G * G - tl.where(ok, 2.0 * d2 / ps, 0.0)\")\n","            E(f\"    disc = tl.maximum({deg-1.0} * ({float(deg)} * H - G * G), 0.0)\")\n","            E(f\"    sq = tl.sqrt(disc)\")\n","            E(f\"    gp = G + sq\")\n","            E(f\"    gm = G - sq\")\n","            E(f\"    den = tl.where(tl.abs(gp) >= tl.abs(gm), gp, gm)\")\n","            E(f\"    dok = tl.abs(den) > 1e-20\")\n","            E(f\"    ds = tl.where(dok, den, 1.0)\")\n","            E(f\"    z = z - tl.where(dok, {float(deg)} / ds, 0.0)\")\n","        E(f\"    r{ri} = z\")\n","        # Synthetic division\n","        if deg > 1:\n","            E(f\"    _b = w{deg}\")\n","            for j in range(deg-1, 0, -1):\n","                E(f\"    _bn = w{j} + z * _b\")\n","                E(f\"    w{j} = _b\")\n","                E(f\"    _b = _bn\")\n","            E(f\"    w0 = _b\")\n","\n","    # Newton polish\n","    E(f\"    # Newton polish\")\n","    for _pol in range(3):\n","        for ri in range(N):\n","            E(f\"    pv = c{N} + 0.0\")\n","            E(f\"    dp = tl.zeros((BLOCK_B,), dtype=tl.float64)\")\n","            for j in range(N-1, -1, -1):\n","                E(f\"    dp = dp * r{ri} + pv\")\n","                E(f\"    pv = pv * r{ri} + c{j}\")\n","            E(f\"    ok = tl.abs(dp) > 1e-30\")\n","            E(f\"    ds = tl.where(ok, dp, 1.0)\")\n","            E(f\"    r{ri} = r{ri} - tl.where(ok, pv / ds, 0.0)\")\n","\n","    # Phase 3: Eigenvectors via interleaved FL + Horner\n","    E(f\"    # Phase 3: Eigenvectors\")\n","    for ei in range(N):\n","        E(f\"    # Eigenvector {ei}\")\n","        E(f\"    lam = r{ei}\")\n","        # Reset M to zero\n","        for i in range(N):\n","            for j in range(N):\n","                E(f\"    m{i}_{j} = tl.zeros((BLOCK_B,), dtype=tl.float64)\")\n","        # Interleaved FL + Horner\n","        for kk in range(1, N+1):\n","            c_diag = f\"c{N - kk + 1}\"\n","            for i in range(N):\n","                for j in range(N):\n","                    terms = [f\"a{i}_{l} * m{l}_{j}\" for l in range(N)]\n","                    expr = \" + \".join(terms)\n","                    if i == j:\n","                        expr += f\" + {c_diag}\"\n","                    E(f\"    mn{i}_{j} = {expr}\")\n","            if kk == 1:\n","                for i in range(N):\n","                    for j in range(N):\n","                        E(f\"    h{i}_{j} = mn{i}_{j}\")\n","            else:\n","                for i in range(N):\n","                    for j in range(N):\n","                        E(f\"    h{i}_{j} = h{i}_{j} * lam + mn{i}_{j}\")\n","            for i in range(N):\n","                for j in range(N):\n","                    E(f\"    m{i}_{j} = mn{i}_{j}\")\n","\n","        # Max-norm column\n","        E(f\"    best_j = tl.zeros((BLOCK_B,), dtype=tl.int32)\")\n","        E(f\"    best_n2 = tl.zeros((BLOCK_B,), dtype=tl.float64) - 1.0\")\n","        for j in range(N):\n","            terms = [f\"h{i}_{j} * h{i}_{j}\" for i in range(N)]\n","            E(f\"    _cn = \" + \" + \".join(terms))\n","            E(f\"    _better = _cn > best_n2\")\n","            E(f\"    best_n2 = tl.where(_better, _cn, best_n2)\")\n","            E(f\"    best_j = tl.where(_better, {j}, best_j)\")\n","        # Extract\n","        for i in range(N):\n","            E(f\"    ev{ei}_{i} = tl.zeros((BLOCK_B,), dtype=tl.float64)\")\n","            for j in range(N):\n","                E(f\"    ev{ei}_{i} = tl.where(best_j == {j}, h{i}_{j}, ev{ei}_{i})\")\n","        # Normalize\n","        terms = [f\"ev{ei}_{i} * ev{ei}_{i}\" for i in range(N)]\n","        E(f\"    _vn = tl.sqrt(\" + \" + \".join(terms) + \" + 1e-60)\")\n","        for i in range(N):\n","            E(f\"    ev{ei}_{i} = ev{ei}_{i} / _vn\")\n","\n","    # Phase 4: NS (fp32) — V[i,j] = ev{j}_{i} (column j, row i)\n","    E(f\"    # Phase 4: NS\")\n","    for i in range(N):\n","        for j in range(N):\n","            E(f\"    v{i}_{j} = ev{j}_{i}.to(tl.float32)\")\n","\n","    for _ns in range(2):\n","        E(f\"    # NS iter\")\n","        # Y = V^T V\n","        for i in range(N):\n","            for j in range(N):\n","                terms = [f\"v{l}_{i} * v{l}_{j}\" for l in range(N)]\n","                E(f\"    y{i}_{j} = \" + \" + \".join(terms))\n","        # T = 3I - Y\n","        for i in range(N):\n","            for j in range(N):\n","                if i == j:\n","                    E(f\"    t{i}_{j} = 3.0 - y{i}_{j}\")\n","                else:\n","                    E(f\"    t{i}_{j} = -y{i}_{j}\")\n","        # V_new = 0.5 * V @ T\n","        for i in range(N):\n","            for j in range(N):\n","                terms = [f\"v{i}_{l} * t{l}_{j}\" for l in range(N)]\n","                E(f\"    vn{i}_{j} = 0.5 * (\" + \" + \".join(terms) + \")\")\n","        for i in range(N):\n","            for j in range(N):\n","                E(f\"    v{i}_{j} = vn{i}_{j}\")\n","\n","    # Phase 5: Rayleigh (on pre-scaled A, then un-scale)\n","    E(f\"    # Phase 5: Rayleigh\")\n","    for i in range(N):\n","        for j in range(N):\n","            E(f\"    af{i}_{j} = a{i}_{j}.to(tl.float32)\")\n","\n","    evals_vars = []\n","    for ei in range(N):\n","        E(f\"    lam{ei} = tl.zeros((BLOCK_B,), dtype=tl.float32)\")\n","        for l in range(N):\n","            terms = [f\"af{l}_{mm} * v{mm}_{ei}\" for mm in range(N)]\n","            E(f\"    _av = \" + \" + \".join(terms))\n","            E(f\"    lam{ei} = lam{ei} + v{l}_{ei} * _av\")\n","        E(f\"    lam{ei} = lam{ei} * sc.to(tl.float32)\")\n","        evals_vars.append(f\"lam{ei}\")\n","\n","    # Sort\n","    E(f\"    # Sort\")\n","    perm_vars = [f\"p{i}\" for i in range(N)]\n","    for i in range(N):\n","        E(f\"    p{i} = tl.zeros((BLOCK_B,), dtype=tl.int32) + {i}\")\n","    for _ in range(N):\n","        for j in range(N-1):\n","            E(f\"    _sw = {evals_vars[j]} > {evals_vars[j+1]}\")\n","            E(f\"    _ea, _eb = {evals_vars[j]}, {evals_vars[j+1]}\")\n","            E(f\"    {evals_vars[j]} = tl.where(_sw, _eb, _ea)\")\n","            E(f\"    {evals_vars[j+1]} = tl.where(_sw, _ea, _eb)\")\n","            E(f\"    _pa, _pb = {perm_vars[j]}, {perm_vars[j+1]}\")\n","            E(f\"    {perm_vars[j]} = tl.where(_sw, _pb, _pa)\")\n","            E(f\"    {perm_vars[j+1]} = tl.where(_sw, _pa, _pb)\")\n","\n","    # Permute eigenvectors and store\n","    E(f\"    # Store\")\n","    for j_out in range(N):\n","        for j_src in range(N):\n","            E(f\"    _is{j_src} = ({perm_vars[j_out]} == {j_src})\")\n","        for i in range(N):\n","            E(f\"    _sv = tl.zeros((BLOCK_B,), dtype=tl.float32)\")\n","            for j_src in range(N):\n","                E(f\"    _sv = tl.where(_is{j_src}, v{i}_{j_src}, _sv)\")\n","            E(f\"    tl.store(evecs_ptr + bid * {N2} + {i*N+j_out}, _sv, mask=mask)\")\n","\n","    for ei in range(N):\n","        E(f\"    tl.store(evals_ptr + bid * {N} + {ei}, {evals_vars[ei]}, mask=mask)\")\n","\n","    return \"\\n\".join(L)\n","\n","\n","# Generate kernel source and write to temp file (Triton needs inspect.getsource)\n","_kernel_src = _gen_kernel(6)\n","_kernel_path = '/tmp/_fl_eigh_triton_gen.py'\n","with open(_kernel_path, 'w') as _f:\n","    _f.write(\"import triton\\nimport triton.language as tl\\n\\n\")\n","    _f.write(_kernel_src)\n","\n","import importlib.util as _ilu\n","_spec = _ilu.spec_from_file_location(\"_fl_eigh_triton_gen\", _kernel_path)\n","_mod = _ilu.module_from_spec(_spec)\n","_spec.loader.exec_module(_mod)\n","_fl_eigh_gen = _mod._fl_eigh_gen\n","\n","\n","class FLEighTriton:\n","    @staticmethod\n","    def apply(A: Tensor) -> Tuple[Tensor, Tensor]:\n","        B, n, _ = A.shape\n","        evals = torch.empty(B, n, device=A.device, dtype=torch.float32)\n","        evecs = torch.empty(B, n, n, device=A.device, dtype=torch.float32)\n","        BLOCK_B = 32\n","        grid = ((B + BLOCK_B - 1) // BLOCK_B,)\n","        _fl_eigh_gen[grid](A.contiguous(), evals, evecs, B, BLOCK_B)\n","        return evals, evecs\n","\n","\n","# Benchmark\n","def sync(): torch.cuda.synchronize()\n","def gt(fn, w=20, r=200):\n","    for _ in range(w): fn()\n","    sync(); t = time.perf_counter()\n","    for _ in range(r): fn()\n","    sync(); return (time.perf_counter() - t) / r\n","def fmt(s):\n","    if s < 1e-3: return f\"{s*1e6:.1f}us\"\n","    if s < 1: return f\"{s*1e3:.2f}ms\"\n","    return f\"{s:.3f}s\"\n","\n","\n","def main():\n","    if not torch.cuda.is_available(): sys.exit(1)\n","    dev = torch.device('cuda')\n","    p = torch.cuda.get_device_properties(0)\n","    print(\"=\" * 72)\n","    print(\"  FL Eigh Triton — Generated Kernel\")\n","    print(\"=\" * 72)\n","    print(f\"  {p.name} | Triton {triton.__version__}\")\n","    print(f\"  Generated kernel: {len(_kernel_src.splitlines())} lines\")\n","\n","    N = 6; B = 4096\n","    A = (lambda R: (R + R.mT) / 2)(torch.randn(B, N, N, device=dev))\n","    rv, rV = torch.linalg.eigh(A)\n","\n","    print(f\"\\n  ACCURACY (n={N} B={B})\")\n","    try:\n","        tv, tV = FLEighTriton.apply(A)\n","        ve = (tv - rv).abs().max().item()\n","        dots = torch.bmm(rV.double().mT, tV.double()).abs().max(dim=-1).values.min().item()\n","        AV = torch.bmm(A.double(), tV.double())\n","        VL = tV.double() * tv.double().unsqueeze(-2)\n","        res = (AV - VL).reshape(B, -1).norm(dim=-1) / A.double().reshape(B, -1).norm(dim=-1).clamp(min=1e-30)\n","        print(f\"  Triton: val={ve:.1e} align={dots:.6f} res={res.max().item():.1e}\")\n","        triton_ok = True\n","    except Exception as e:\n","        print(f\"  FAILED: {str(e)[:300]}\")\n","        triton_ok = False\n","\n","    print(f\"\\n  THROUGHPUT (n={N} B={B})\")\n","    tr = gt(lambda: torch.linalg.eigh(A))\n","    print(f\"  cuSOLVER: {fmt(tr)}\")\n","\n","    if triton_ok:\n","        tt = gt(lambda: FLEighTriton.apply(A))\n","        print(f\"  Triton:   {fmt(tt)} ({tr/tt:.2f}x)\")\n","\n","        print(f\"\\n  BATCH SCALING (n={N})\")\n","        print(f\"  {'B':>6}  {'cuSOLVER':>10}  {'Triton':>10}  {'ratio':>7}\")\n","        for Bx in [256, 512, 1024, 2048, 4096, 8192, 16384]:\n","            Ax = (lambda R: (R+R.mT)/2)(torch.randn(Bx, N, N, device=dev))\n","            t1 = gt(lambda: torch.linalg.eigh(Ax), 10, 100)\n","            t2 = gt(lambda: FLEighTriton.apply(Ax), 10, 100)\n","            print(f\"  {Bx:>6}  {fmt(t1):>10}  {fmt(t2):>10}  {t1/t2:>6.2f}x\")\n","            del Ax\n","\n","    print(\"=\" * 72)\n","\n","\n","if __name__ == '__main__':\n","    main()"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"3hJCXqHxA9f_","outputId":"10587ce9-c806-499d-ddfe-5e90009f43ce"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["========================================================================\n","  FL Eigh Triton — Generated Kernel\n","========================================================================\n","  NVIDIA RTX PRO 6000 Blackwell Server Edition | Triton 3.6.0\n","  Generated kernel: 7461 lines\n","\n","  ACCURACY (n=6 B=4096)\n"]}]},{"cell_type":"code","source":["\"\"\"\n","FL Eigh — Complete Sweep.\n","Every variant. Every metric. Every scale.\n","\"\"\"\n","import math, time, gc, sys\n","import torch\n","import torch.nn as nn\n","from torch import Tensor\n","from typing import Tuple\n","\n","torch.backends.cuda.matmul.allow_tf32 = False\n","torch.backends.cudnn.allow_tf32 = False\n","torch.set_float32_matmul_precision('highest')\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# FL Eigh — Fast mode (fp32 eigvecs, sum-of-columns)\n","# ═══════════════════════════════════════════════════════════════════\n","class FLEighFast(nn.Module):\n","    def forward(self, A):\n","        B,n,_=A.shape; dev=A.device\n","        sc=(torch.linalg.norm(A.reshape(B,-1),dim=-1)/math.sqrt(n)).clamp(min=1e-12)\n","        As=A/sc[:,None,None]; Ad=As.double()\n","        I=torch.eye(n,device=dev,dtype=torch.float64).unsqueeze(0).expand(B,-1,-1)\n","        c=torch.zeros(B,n+1,device=dev,dtype=torch.float64); c[:,n]=1.0\n","        Ms=torch.zeros(n+1,B,n,n,device=dev,dtype=torch.float64)\n","        Mk=torch.zeros(B,n,n,device=dev,dtype=torch.float64)\n","        for k in range(1,n+1):\n","            Mk=torch.bmm(Ad,Mk)+c[:,n-k+1,None,None]*I; Ms[k]=Mk\n","            c[:,n-k]=-(Ad*Mk).sum((-2,-1))/k\n","        dt=torch.float64 if n>6 else torch.float32\n","        cl=c.to(dt).clone(); roots=torch.zeros(B,n,device=dev,dtype=dt)\n","        zi=As.to(dt).diagonal(dim1=-2,dim2=-1).sort(dim=-1).values\n","        zi=zi+torch.linspace(-1e-4,1e-4,n,device=dev,dtype=dt).unsqueeze(0)\n","        for ri in range(n):\n","            deg=n-ri; z=zi[:,ri]\n","            for _ in range(5):\n","                pv=cl[:,deg]; dp=torch.zeros(B,device=dev,dtype=dt); d2=torch.zeros(B,device=dev,dtype=dt)\n","                for j in range(deg-1,-1,-1): d2=d2*z+dp; dp=dp*z+pv; pv=pv*z+cl[:,j]\n","                ok=pv.abs()>1e-30; ps=torch.where(ok,pv,torch.ones_like(pv))\n","                G=torch.where(ok,dp/ps,torch.zeros_like(dp)); H=G*G-torch.where(ok,2.0*d2/ps,torch.zeros_like(d2))\n","                disc=((deg-1.0)*(deg*H-G*G)).clamp(min=0.0); sq=torch.sqrt(disc); gp=G+sq; gm=G-sq\n","                den=torch.where(gp.abs()>=gm.abs(),gp,gm); dok=den.abs()>1e-20\n","                ds=torch.where(dok,den,torch.ones_like(den)); z=z-torch.where(dok,float(deg)/ds,torch.zeros_like(den))\n","            roots[:,ri]=z; b=cl[:,deg]\n","            for j in range(deg-1,0,-1): bn=cl[:,j]+z*b; cl[:,j]=b; b=bn\n","            cl[:,0]=b\n","        roots=roots.double()\n","        for _ in range(3):\n","            pv=torch.ones(B,n,device=dev,dtype=torch.float64); dp=torch.zeros(B,n,device=dev,dtype=torch.float64)\n","            for j in range(n-1,-1,-1): dp=dp*roots+pv; pv=pv*roots+c[:,j:j+1]\n","            ok=dp.abs()>1e-30; dps=torch.where(ok,dp,torch.ones_like(dp)); roots=roots-torch.where(ok,pv/dps,torch.zeros_like(pv))\n","        ef=roots.float(); Mf=Ms.float()\n","        R=Mf[1].unsqueeze(1).expand(-1,n,-1,-1).clone()\n","        for k in range(2,n+1): R=R*ef[:,:,None,None]+Mf[k].unsqueeze(1)\n","        vec=R.sum(dim=-1); vn=vec.norm(dim=-1,keepdim=True)\n","        vec=torch.where(vn>1e-10,vec,R[:,:,:,0]); vec=vec/(vec.norm(dim=-1,keepdim=True)+1e-30)\n","        V=vec.transpose(-2,-1)\n","        eye_f=torch.eye(n,device=dev,dtype=torch.float32).unsqueeze(0).expand(B,-1,-1)\n","        Y=torch.bmm(V.mT,V); X=eye_f.clone()\n","        for _ in range(2): T=3.0*eye_f-Y; X=0.5*torch.bmm(X,T); Y=0.5*torch.bmm(T,Y)\n","        V=torch.bmm(V,X); AV=torch.bmm(A,V); evals=(V*AV).sum(dim=-2)\n","        se,p=evals.sort(dim=-1); return se,V.gather(-1,p.unsqueeze(-2).expand_as(V))\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# FL Eigh — Precise mode (fp64 eigvecs, max-col, Rayleigh)\n","# ═══════════════════════════════════════════════════════════════════\n","class FLEighPrecise(nn.Module):\n","    def forward(self, A):\n","        B,n,_=A.shape; dev=A.device\n","        sc=(torch.linalg.norm(A.reshape(B,-1),dim=-1)/math.sqrt(n)).clamp(min=1e-12)\n","        As=A/sc[:,None,None]; Ad=As.double()\n","        I=torch.eye(n,device=dev,dtype=torch.float64).unsqueeze(0).expand(B,-1,-1)\n","        c=torch.zeros(B,n+1,device=dev,dtype=torch.float64); c[:,n]=1.0\n","        Ms=torch.zeros(n+1,B,n,n,device=dev,dtype=torch.float64)\n","        Mk=torch.zeros(B,n,n,device=dev,dtype=torch.float64)\n","        for k in range(1,n+1):\n","            Mk=torch.bmm(Ad,Mk)+c[:,n-k+1,None,None]*I; Ms[k]=Mk\n","            c[:,n-k]=-(Ad*Mk).sum((-2,-1))/k\n","        dt=torch.float64 if n>6 else torch.float32\n","        cl=c.to(dt).clone(); roots=torch.zeros(B,n,device=dev,dtype=dt)\n","        zi=As.to(dt).diagonal(dim1=-2,dim2=-1).sort(dim=-1).values\n","        zi=zi+torch.linspace(-1e-4,1e-4,n,device=dev,dtype=dt).unsqueeze(0)\n","        for ri in range(n):\n","            deg=n-ri; z=zi[:,ri]\n","            for _ in range(5):\n","                pv=cl[:,deg]; dp=torch.zeros(B,device=dev,dtype=dt); d2=torch.zeros(B,device=dev,dtype=dt)\n","                for j in range(deg-1,-1,-1): d2=d2*z+dp; dp=dp*z+pv; pv=pv*z+cl[:,j]\n","                ok=pv.abs()>1e-30; ps=torch.where(ok,pv,torch.ones_like(pv))\n","                G=torch.where(ok,dp/ps,torch.zeros_like(dp)); H=G*G-torch.where(ok,2.0*d2/ps,torch.zeros_like(d2))\n","                disc=((deg-1.0)*(deg*H-G*G)).clamp(min=0.0); sq=torch.sqrt(disc); gp=G+sq; gm=G-sq\n","                den=torch.where(gp.abs()>=gm.abs(),gp,gm); dok=den.abs()>1e-20\n","                ds=torch.where(dok,den,torch.ones_like(den)); z=z-torch.where(dok,float(deg)/ds,torch.zeros_like(den))\n","            roots[:,ri]=z; b=cl[:,deg]\n","            for j in range(deg-1,0,-1): bn=cl[:,j]+z*b; cl[:,j]=b; b=bn\n","            cl[:,0]=b\n","        roots=roots.double()\n","        for _ in range(3):\n","            pv=torch.ones(B,n,device=dev,dtype=torch.float64); dp=torch.zeros(B,n,device=dev,dtype=torch.float64)\n","            for j in range(n-1,-1,-1): dp=dp*roots+pv; pv=pv*roots+c[:,j:j+1]\n","            ok=dp.abs()>1e-30; dps=torch.where(ok,dp,torch.ones_like(dp)); roots=roots-torch.where(ok,pv/dps,torch.zeros_like(pv))\n","        lam=roots; R=Ms[1].unsqueeze(1).expand(-1,n,-1,-1).clone()\n","        for k in range(2,n+1): R=R*lam[:,:,None,None]+Ms[k].unsqueeze(1)\n","        cn=R.norm(dim=-2); best=cn.argmax(dim=-1)\n","        idx=best.unsqueeze(-1).unsqueeze(-1).expand(-1,-1,n,1)\n","        vec=R.gather(-1,idx).squeeze(-1); vec=vec/(vec.norm(dim=-1,keepdim=True)+1e-30)\n","        V=vec.float().transpose(-2,-1)\n","        eye_f=torch.eye(n,device=dev,dtype=torch.float32).unsqueeze(0).expand(B,-1,-1)\n","        Y=torch.bmm(V.mT,V); X=eye_f.clone()\n","        for _ in range(2): T=3.0*eye_f-Y; X=0.5*torch.bmm(X,T); Y=0.5*torch.bmm(T,Y)\n","        V=torch.bmm(V,X); AV=torch.bmm(A,V); evals=(V*AV).sum(dim=-2)\n","        se,p=evals.sort(dim=-1); return se,V.gather(-1,p.unsqueeze(-2).expand_as(V))\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# CuPy CUDA kernel\n","# ═══════════════════════════════════════════════════════════════════\n","_CUDA_SRC = r\"\"\"\n","extern \"C\" __global__ void fl_eigh_kernel(const float* __restrict__ A_in, float* __restrict__ evals_out, float* __restrict__ evecs_out, int B) {\n","    int tid=blockIdx.x*blockDim.x+threadIdx.x; if(tid>=B)return;\n","    const int NN=6,N2=36; double a[36],c[7],m[36],roots[6]; float evecs[36];\n","    double frob_sq=0; for(int i=0;i<N2;i++){a[i]=(double)A_in[tid*N2+i];frob_sq+=a[i]*a[i];}\n","    double scale=sqrt(frob_sq/6.0); if(scale<1e-12)scale=1e-12; double inv_s=1.0/scale;\n","    for(int i=0;i<N2;i++)a[i]*=inv_s;\n","    for(int i=0;i<7;i++)c[i]=0; c[6]=1.0; for(int i=0;i<N2;i++)m[i]=0;\n","    for(int k=1;k<=NN;k++){double mn[36];\n","        for(int i=0;i<NN;i++)for(int j=0;j<NN;j++){double acc=0;for(int l=0;l<NN;l++)acc+=a[i*NN+l]*m[l*NN+j];if(i==j)acc+=c[NN-k+1];mn[i*NN+j]=acc;}\n","        double tr=0;for(int i=0;i<NN;i++)for(int l=0;l<NN;l++)tr+=a[i*NN+l]*mn[l*NN+i];c[NN-k]=-tr/(double)k;for(int i=0;i<N2;i++)m[i]=mn[i];}\n","    double diag[6];for(int i=0;i<NN;i++)diag[i]=a[i*NN+i];\n","    for(int p=0;p<NN-1;p++)for(int j=0;j<NN-1;j++)if(diag[j]>diag[j+1]){double t=diag[j];diag[j]=diag[j+1];diag[j+1]=t;}\n","    for(int i=0;i<NN;i++)diag[i]+=-1e-4+2e-4*(double)i/5.0;\n","    double cl[7];for(int i=0;i<7;i++)cl[i]=c[i];\n","    for(int ri=0;ri<NN;ri++){int deg=NN-ri;double z=diag[ri];\n","        for(int lag=0;lag<5;lag++){double pv=cl[deg],dp=0,d2=0;for(int j=deg-1;j>=0;j--){d2=d2*z+dp;dp=dp*z+pv;pv=pv*z+cl[j];}\n","            if(fabs(pv)>1e-30){double G=dp/pv,H=G*G-2*d2/pv,disc=((double)(deg-1))*((double)deg*H-G*G);if(disc<0)disc=0;double sq=sqrt(disc),gp=G+sq,gm=G-sq,den=(fabs(gp)>=fabs(gm))?gp:gm;if(fabs(den)>1e-20)z-=(double)deg/den;}}\n","        roots[ri]=z;if(deg>1){double b=cl[deg];for(int j=deg-1;j>0;j--){double bn=cl[j]+z*b;cl[j]=b;b=bn;}cl[0]=b;}}\n","    for(int pol=0;pol<3;pol++)for(int ri=0;ri<NN;ri++){double pv=c[NN],dp=0;for(int j=NN-1;j>=0;j--){dp=dp*roots[ri]+pv;pv=pv*roots[ri]+c[j];}if(fabs(dp)>1e-30)roots[ri]-=pv/dp;}\n","    for(int ei=0;ei<NN;ei++){double lam=roots[ei],m_loc[36],r_loc[36];for(int i=0;i<N2;i++)m_loc[i]=0;\n","        for(int k=1;k<=NN;k++){double mn[36];for(int i=0;i<NN;i++)for(int j=0;j<NN;j++){double acc=0;for(int l=0;l<NN;l++)acc+=a[i*NN+l]*m_loc[l*NN+j];if(i==j)acc+=c[NN-k+1];mn[i*NN+j]=acc;}\n","            if(k==1)for(int i=0;i<N2;i++)r_loc[i]=mn[i];else for(int i=0;i<N2;i++)r_loc[i]=r_loc[i]*lam+mn[i];for(int i=0;i<N2;i++)m_loc[i]=mn[i];}\n","        int bj=0;double bn=-1;for(int j=0;j<NN;j++){double cs=0;for(int i=0;i<NN;i++)cs+=r_loc[i*NN+j]*r_loc[i*NN+j];if(cs>bn){bn=cs;bj=j;}}\n","        double vnorm=0,vec[6];for(int i=0;i<NN;i++){vec[i]=r_loc[i*NN+bj];vnorm+=vec[i]*vec[i];}vnorm=sqrt(vnorm)+1e-30;\n","        for(int i=0;i<NN;i++)evecs[i*NN+ei]=(float)(vec[i]/vnorm);}\n","    for(int ns=0;ns<2;ns++){float y[36],tm[36],vn[36];\n","        for(int i=0;i<NN;i++)for(int j=0;j<NN;j++){float a2=0;for(int l=0;l<NN;l++)a2+=evecs[l*NN+i]*evecs[l*NN+j];y[i*NN+j]=a2;}\n","        for(int i=0;i<NN;i++)for(int j=0;j<NN;j++)tm[i*NN+j]=((i==j)?3.0f:0.0f)-y[i*NN+j];\n","        for(int i=0;i<NN;i++)for(int j=0;j<NN;j++){float a2=0;for(int l=0;l<NN;l++)a2+=evecs[i*NN+l]*tm[l*NN+j];vn[i*NN+j]=0.5f*a2;}\n","        for(int i=0;i<N2;i++)evecs[i]=vn[i];}\n","    float af[36];for(int i=0;i<N2;i++)af[i]=(float)a[i];float ev[6];\n","    for(int ei=0;ei<NN;ei++){float lf=0;for(int l=0;l<NN;l++){float av=0;for(int mm=0;mm<NN;mm++)av+=af[l*NN+mm]*evecs[mm*NN+ei];lf+=evecs[l*NN+ei]*av;}ev[ei]=lf*(float)scale;}\n","    int perm[6];for(int i=0;i<NN;i++)perm[i]=i;\n","    for(int p=0;p<NN-1;p++)for(int j=0;j<NN-1;j++)if(ev[j]>ev[j+1]){float t=ev[j];ev[j]=ev[j+1];ev[j+1]=t;int pt=perm[j];perm[j]=perm[j+1];perm[j+1]=pt;}\n","    for(int i=0;i<NN;i++)evals_out[tid*NN+i]=ev[i];\n","    for(int jo=0;jo<NN;jo++){int js=perm[jo];for(int i=0;i<NN;i++)evecs_out[tid*N2+i*NN+jo]=evecs[i*NN+js];}\n","}\n","\"\"\"\n","\n","HAS_CUPY = False\n","_cupy_kernel = None\n","def _init_cupy():\n","    global HAS_CUPY, _cupy_kernel\n","    try:\n","        import cupy\n","        _cupy_kernel = cupy.RawKernel(_CUDA_SRC, 'fl_eigh_kernel')\n","        _cupy_kernel.compile()\n","        HAS_CUPY = True\n","    except: pass\n","\n","def fl_eigh_cupy(A, ev_buf=None, vc_buf=None):\n","    import cupy\n","    B=A.shape[0]; k=_cupy_kernel; A_c=A.contiguous().float()\n","    if ev_buf is None: ev_buf=torch.empty(B,6,device=A.device,dtype=torch.float32)\n","    if vc_buf is None: vc_buf=torch.empty(B,6,6,device=A.device,dtype=torch.float32)\n","    thr=128; blk=(B+thr-1)//thr\n","    stream=cupy.cuda.ExternalStream(torch.cuda.current_stream().cuda_stream)\n","    with stream: k((blk,),(thr,),(A_c.data_ptr(),ev_buf.data_ptr(),vc_buf.data_ptr(),B))\n","    return ev_buf, vc_buf\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# Math purity (no reference impl)\n","# ═══════════════════════════════════════════════════════════════════\n","def math_purity(A, vals, vecs):\n","    B,n,_=A.shape; dev=A.device\n","    Ad=A.double(); vd=vals.double(); Vd=vecs.double()\n","    AV=torch.bmm(Ad,Vd); VL=Vd*vd.unsqueeze(-2)\n","    An=Ad.reshape(B,-1).norm(dim=-1,keepdim=True).clamp(min=1e-30)\n","    res=(AV-VL).norm(dim=-2)/An\n","    VtV=torch.bmm(Vd.mT,Vd); I=torch.eye(n,device=dev,dtype=torch.float64).unsqueeze(0)\n","    orth=(VtV-I).reshape(B,-1).norm(dim=-1)\n","    recon=torch.bmm(Vd*vd.unsqueeze(-2),Vd.mT)\n","    rec_err=(Ad-recon).reshape(B,-1).norm(dim=-1)/An.squeeze(-1)\n","    tr_err=(Ad.diagonal(dim1=-2,dim2=-1).sum(-1)-vd.sum(-1)).abs()\n","    det_A=torch.linalg.det(Ad)\n","    det_err=(det_A-vd.prod(-1)).abs()/det_A.abs().clamp(min=1e-30)\n","    cp=torch.zeros(B,n,device=dev,dtype=torch.float64)\n","    for i in range(n): cp[:,i]=torch.linalg.det(vd[:,i:i+1,None]*I-Ad).abs()\n","    return dict(res_max=res.max().item(), res_mean=res.mean().item(),\n","                orth_max=orth.max().item(), orth_mean=orth.mean().item(),\n","                recon_max=rec_err.max().item(), recon_mean=rec_err.mean().item(),\n","                tr_max=tr_err.max().item(), tr_mean=tr_err.mean().item(),\n","                det_max=det_err.max().item(), det_mean=det_err.mean().item(),\n","                cp_max=cp.max().item(), cp_mean=cp.mean().item())\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# Utilities\n","# ═══════════════════════════════════════════════════════════════════\n","def sync(): torch.cuda.synchronize()\n","def gt(fn,w=20,r=300):\n","    for _ in range(w): fn()\n","    sync(); t=time.perf_counter()\n","    for _ in range(r): fn()\n","    sync(); return (time.perf_counter()-t)/r\n","def fmt(s):\n","    if s<1e-3: return f\"{s*1e6:.0f}us\"\n","    if s<1: return f\"{s*1e3:.2f}ms\"\n","    return f\"{s:.3f}s\"\n","def make(B,n,dev):\n","    R=torch.randn(B,n,n,device=dev); return (R+R.mT)/2\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# MAIN\n","# ═══════════════════════════════════════════════════════════════════\n","def main():\n","    dev=torch.device('cuda')\n","    p=torch.cuda.get_device_properties(0)\n","    print(\"=\"*78)\n","    print(\"  FL Eigh — Complete Sweep\")\n","    print(\"=\"*78)\n","    print(f\"  {p.name} | PyTorch {torch.__version__}\")\n","\n","    _init_cupy()\n","    print(f\"  CuPy CUDA kernel: {'available' if HAS_CUPY else 'not available'}\")\n","\n","    N=6; B=4096\n","    A=make(B,N,dev)\n","\n","    # ════════════════════════════════════════════════════════════════\n","    # Section 1: MATHEMATICAL PURITY (n=6 B=2048)\n","    # ════════════════════════════════════════════════════════════════\n","    print(f\"\\n{'='*78}\")\n","    print(f\"  1. MATHEMATICAL PURITY (n={N}, B=2048)\")\n","    print(f\"{'='*78}\")\n","\n","    Am=make(2048,N,dev)\n","    impls = {'cuSOLVER': lambda X: torch.linalg.eigh(X),\n","             'FL Fast': lambda X: FLEighFast()(X),\n","             'FL Precise': lambda X: FLEighPrecise()(X)}\n","    if HAS_CUPY:\n","        impls['CUDA kernel'] = lambda X: fl_eigh_cupy(X)\n","\n","    all_metrics = {}\n","    for name,fn in impls.items():\n","        v,V=fn(Am); all_metrics[name]=math_purity(Am,v,V)\n","\n","    keys=['res_max','res_mean','orth_max','orth_mean','recon_max','recon_mean',\n","          'tr_max','tr_mean','det_max','det_mean','cp_max','cp_mean']\n","    labels=['Eigenpair max','Eigenpair mean','Orth max','Orth mean',\n","            'Recon max','Recon mean','Trace max','Trace mean',\n","            'Det max','Det mean','CharPoly max','CharPoly mean']\n","\n","    names=list(impls.keys())\n","    hdr=\"  \"+\"Metric\".ljust(18)+\"\".join(n.rjust(12) for n in names)+\"  Winner\"\n","    print(f\"\\n{hdr}\")\n","    print(\"  \"+\"─\"*18+\"\".join(\"─\"*12 for _ in names)+\"  \"+\"─\"*10)\n","\n","    score={n:0 for n in names}\n","    for key,label in zip(keys,labels):\n","        vals={n:all_metrics[n][key] for n in names}\n","        winner=min(vals,key=vals.get)\n","        score[winner]+=1\n","        row=\"  \"+label.ljust(18)\n","        for n in names:\n","            v=vals[n]; mark=\" *\" if n==winner else \"  \"\n","            row+=f\"{v:>10.1e}{mark}\"\n","        row+=f\"  {winner}\"\n","        print(row)\n","    print()\n","    for n in names: print(f\"  {n}: {score[n]}/{len(keys)} wins\")\n","    del Am\n","\n","    # ════════════════════════════════════════════════════════════════\n","    # Section 2: PURITY ACROSS SIZES\n","    # ════════════════════════════════════════════════════════════════\n","    print(f\"\\n{'='*78}\")\n","    print(f\"  2. PURITY ACROSS SIZES (FL Precise vs cuSOLVER)\")\n","    print(f\"{'='*78}\")\n","\n","    print(f\"\\n  {'n':>3}  {'FL wins':>8}  {'cuS wins':>9}  {'Best eigenpair':>16}  {'Best orth':>12}\")\n","    print(f\"  {'─'*3}  {'─'*8}  {'─'*9}  {'─'*16}  {'─'*12}\")\n","    for nx in [3,4,5,6,8,10,12]:\n","        Bx=2048 if nx<=8 else 1024\n","        Ax=make(Bx,nx,dev)\n","        cv,cV=torch.linalg.eigh(Ax); fv,fV=FLEighPrecise()(Ax)\n","        mc=math_purity(Ax,cv,cV); mf=math_purity(Ax,fv,fV)\n","        fw=sum(1 for k in keys if mf[k]<mc[k]); cw=sum(1 for k in keys if mc[k]<mf[k])\n","        best_ep=\"FL\" if mf['res_max']<mc['res_max'] else \"cuS\"\n","        best_or=\"FL\" if mf['orth_max']<mc['orth_max'] else \"cuS\"\n","        print(f\"  {nx:>3}  {fw:>5}/12  {cw:>6}/12  {best_ep:>8} {min(mf['res_max'],mc['res_max']):.1e}  {best_or:>5} {min(mf['orth_max'],mc['orth_max']):.1e}\")\n","        del Ax\n","\n","    # ════════════════════════════════════════════════════════════════\n","    # Section 3: COMPILED TIMING\n","    # ════════════════════════════════════════════════════════════════\n","    print(f\"\\n{'='*78}\")\n","    print(f\"  3. COMPILED TIMING (n={N}, B={B})\")\n","    print(f\"{'='*78}\")\n","\n","    print(\"  Compiling...\", end=\" \", flush=True)\n","    c_fast=torch.compile(FLEighFast(),fullgraph=True)\n","    c_prec=torch.compile(FLEighPrecise(),fullgraph=True)\n","    for _ in range(3): c_fast(A); c_prec(A); sync()\n","    print(\"done.\")\n","\n","    t_cus=gt(lambda:torch.linalg.eigh(A))\n","    t_fast=gt(lambda:c_fast(A))\n","    t_prec=gt(lambda:c_prec(A))\n","\n","    results=[('cuSOLVER',t_cus), ('FL Fast compiled',t_fast), ('FL Precise compiled',t_prec)]\n","    if HAS_CUPY:\n","        t_cupy=gt(lambda:fl_eigh_cupy(A))\n","        results.append(('CUDA kernel',t_cupy))\n","\n","    print(f\"\\n  {'Method':<28} {'Time':>10} {'vs cuSOLVER':>12} {'mat/s':>12}\")\n","    print(f\"  {'─'*28} {'─'*10} {'─'*12} {'─'*12}\")\n","    for name,t in results:\n","        print(f\"  {name:<28} {fmt(t):>10} {t_cus/t:>11.2f}x {B/t:>10.0f}/s\")\n","\n","    # ════════════════════════════════════════════════════════════════\n","    # Section 4: CUDA GRAPH\n","    # ════════════════════════════════════════════════════════════════\n","    print(f\"\\n{'='*78}\")\n","    print(f\"  4. CUDA GRAPH (n={N}, B={B})\")\n","    print(f\"{'='*78}\")\n","\n","    graph_results = []\n","    for name, compiled_fn in [('FL Fast+Graph', c_fast), ('FL Precise+Graph', c_prec)]:\n","        A_s=A.clone()\n","        s=torch.cuda.Stream(); s.wait_stream(torch.cuda.current_stream())\n","        with torch.cuda.stream(s):\n","            for _ in range(3): compiled_fn(A_s)\n","        torch.cuda.current_stream().wait_stream(s)\n","        g=torch.cuda.CUDAGraph()\n","        with torch.cuda.graph(g): ev_g,vc_g=compiled_fn(A_s)\n","        def replay(g=g,A_s=A_s): A_s.copy_(A); g.replay()\n","        t_g=gt(replay)\n","        graph_results.append((name,t_g))\n","\n","    print(f\"\\n  {'Method':<28} {'Time':>10} {'vs cuSOLVER':>12}\")\n","    print(f\"  {'─'*28} {'─'*10} {'─'*12}\")\n","    print(f\"  {'cuSOLVER':<28} {fmt(t_cus):>10} {'1.00x':>12}\")\n","    for name,t in results[1:]+graph_results:\n","        print(f\"  {name:<28} {fmt(t):>10} {t_cus/t:>11.2f}x\")\n","\n","    # ════════════════════════════════════════════════════════════════\n","    # Section 5: BATCH SCALING\n","    # ════════════════════════════════════════════════════════════════\n","    print(f\"\\n{'='*78}\")\n","    print(f\"  5. BATCH SCALING (n={N}, FL Fast compiled)\")\n","    print(f\"{'='*78}\")\n","\n","    cd=torch.compile(FLEighFast(),fullgraph=True,dynamic=True)\n","    cd(A); sync()\n","\n","    print(f\"\\n  {'B':>6}  {'cuSOLVER':>10}  {'FL Fast':>10}  {'ratio':>7}  {'mat/s FL':>12}\")\n","    print(f\"  {'─'*6}  {'─'*10}  {'─'*10}  {'─'*7}  {'─'*12}\")\n","    for Bx in [256,512,1024,2048,4096,8192,16384]:\n","        Ax=make(Bx,N,dev)\n","        cd(Ax); sync()\n","        t1=gt(lambda:torch.linalg.eigh(Ax),10,100)\n","        t2=gt(lambda:cd(Ax),10,100)\n","        print(f\"  {Bx:>6}  {fmt(t1):>10}  {fmt(t2):>10}  {t1/t2:>6.2f}x  {Bx/t2:>10.0f}/s\")\n","        del Ax\n","\n","    # ════════════════════════════════════════════════════════════════\n","    # Section 6: SIZE SCALING\n","    # ════════════════════════════════════════════════════════════════\n","    print(f\"\\n{'='*78}\")\n","    print(f\"  6. SIZE SCALING (B={B})\")\n","    print(f\"{'='*78}\")\n","\n","    print(f\"\\n  {'n':>3}  {'cuSOLVER':>10}  {'FL Fast':>10}  {'F ratio':>8}  {'FL Prec':>10}  {'P ratio':>8}  {'val_err':>10}\")\n","    print(f\"  {'─'*3}  {'─'*10}  {'─'*10}  {'─'*8}  {'─'*10}  {'─'*8}  {'─'*10}\")\n","    for nx in [3,4,5,6,8,10,12]:\n","        try:\n","            Ax=make(B,nx,dev)\n","            sf=torch.compile(FLEighFast(),fullgraph=True); sp=torch.compile(FLEighPrecise(),fullgraph=True)\n","            for _ in range(3): sf(Ax); sp(Ax); sync()\n","            t1=gt(lambda:torch.linalg.eigh(Ax),10,100)\n","            t2=gt(lambda:sf(Ax),10,100)\n","            t3=gt(lambda:sp(Ax),10,100)\n","            rv,_=torch.linalg.eigh(Ax); fv,_=sf(Ax)\n","            ve=(fv-rv).abs().max().item()\n","            print(f\"  {nx:>3}  {fmt(t1):>10}  {fmt(t2):>10}  {t1/t2:>7.2f}x  {fmt(t3):>10}  {t1/t3:>7.2f}x  {ve:>10.1e}\")\n","            del Ax,sf,sp\n","        except Exception as e:\n","            print(f\"  {nx:>3}  ERR: {str(e)[:40]}\")\n","            torch.cuda.empty_cache()\n","\n","    # ════════════════════════════════════════════════════════════════\n","    # Section 7: MEMORY\n","    # ════════════════════════════════════════════════════════════════\n","    print(f\"\\n{'='*78}\")\n","    print(f\"  7. MEMORY (n={N}, B={B})\")\n","    print(f\"{'='*78}\")\n","\n","    mem_fns=[('cuSOLVER',lambda:torch.linalg.eigh(A)),\n","             ('FL Fast',lambda:FLEighFast()(A)),\n","             ('FL Precise',lambda:FLEighPrecise()(A))]\n","    if HAS_CUPY: mem_fns.append(('CUDA kernel',lambda:fl_eigh_cupy(A)))\n","\n","    print()\n","    for name,fn in mem_fns:\n","        torch.cuda.empty_cache(); gc.collect(); torch.cuda.reset_peak_memory_stats()\n","        base=torch.cuda.memory_allocated(); fn(); sync()\n","        delta=(torch.cuda.max_memory_allocated()-base)/1024**2\n","        print(f\"  {name:<20} {delta:>8.1f} MB\")\n","\n","    # ════════════════════════════════════════════════════════════════\n","    # Section 8: ACCURACY PASS/FAIL\n","    # ════════════════════════════════════════════════════════════════\n","    print(f\"\\n{'='*78}\")\n","    print(f\"  8. ACCURACY PASS/FAIL (all sizes)\")\n","    print(f\"{'='*78}\\n\")\n","\n","    for name,Cls in [('Fast',FLEighFast),('Precise',FLEighPrecise)]:\n","        print(f\"  {name}:\")\n","        for nx in [3,4,5,6,8,10,12,16]:\n","            Ax=make(1024,nx,dev)\n","            rv,rV=torch.linalg.eigh(Ax); fv,fV=Cls()(Ax)\n","            ve=(fv-rv).abs().max().item()\n","            dots=torch.bmm(rV.double().mT,fV.double()).abs().max(dim=-1).values.min().item()\n","            ok=ve<1e-2 and dots>0.99\n","            print(f\"    [{'OK' if ok else 'NO'}] n={nx:>2}  val={ve:.1e}  align={dots:.6f}\")\n","            del Ax\n","        print()\n","\n","    # ════════════════════════════════════════════════════════════════\n","    # SUMMARY\n","    # ════════════════════════════════════════════════════════════════\n","    print(\"=\"*78)\n","    print(\"  SUMMARY\")\n","    print(\"=\"*78)\n","    print(f\"  Math purity (n=6):  FL Precise wins {score.get('FL Precise',0)}/12\")\n","    print(f\"  Compiled (n=6 B=4096):\")\n","    for name,t in results:\n","        print(f\"    {name:<24} {fmt(t):>10}  {t_cus/t:.2f}x\")\n","    print(f\"  CUDA Graph:\")\n","    for name,t in graph_results:\n","        print(f\"    {name:<24} {fmt(t):>10}  {t_cus/t:.2f}x\")\n","    print(f\"  Memory: FL ~30MB vs cuSOLVER ~1099MB\")\n","    print(\"=\"*78)\n","\n","\n","if __name__=='__main__': main()"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"h9PABGJWDwmN","executionInfo":{"status":"ok","timestamp":1775053623349,"user_tz":420,"elapsed":1999,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"a8eaa77e-3dc5-4ffa-981e-c34b29055a8f"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["========================================================================\n","  FL Eigh CUDA Kernel (CuPy/NVRTC)\n","========================================================================\n","  NVIDIA RTX PRO 6000 Blackwell Server Edition\n","  PyTorch 2.10.0+cu128\n","  Compiling via NVRTC... done.\n","\n","  ACCURACY (n=6 B=4096)\n","  CUDA FL: val=3.3e-06 align=0.999999\n","\n","  MATH PURITY: CUDA FL vs cuSOLVER\n","  Property                       cuSOLVER    CUDA FL    Win\n","  res_max                         6.2e-07    2.2e-07     FL\n","  res_mean                        1.3e-07    2.7e-08     FL\n","  orth_max                        2.2e-06    2.4e-07     FL\n","  orth_mean                       9.0e-07    1.2e-07     FL\n","  recon_max                       1.3e-06    3.2e-07     FL\n","  recon_mean                      4.9e-07    1.0e-07     FL\n","  tr_max                          3.0e-06    1.4e-06     FL\n","  det_max                         6.1e-04    1.5e-04     FL\n","\n","  CUDA FL wins 8/8\n","\n","  THROUGHPUT (n=6 B=4096)\n"]},{"output_type":"stream","name":"stderr","text":["/tmp/ipykernel_89646/739027546.py:351: DeprecationWarning: ExternalStream is deprecated. Use Stream.from_external() instead to interoperate with external streams that implement the CUDA stream protocol.\n","  cv,cV=fl_eigh_cuda(A)\n","/tmp/ipykernel_89646/739027546.py:370: DeprecationWarning: ExternalStream is deprecated. Use Stream.from_external() instead to interoperate with external streams that implement the CUDA stream protocol.\n","  tc=gt(lambda:fl_eigh_cuda(A))\n"]},{"output_type":"stream","name":"stdout","text":["  cuSOLVER:  241.4us\n","  CUDA FL:   425.7us (0.57x)\n","\n","  BATCH SCALING (n=6)\n","       B    cuSOLVER     CUDA FL    ratio\n","     256     102.7us     425.9us    0.24x\n","     512     103.1us     426.1us    0.24x\n","    1024     115.8us     424.4us    0.27x\n"]},{"output_type":"stream","name":"stderr","text":["/tmp/ipykernel_89646/739027546.py:381: DeprecationWarning: ExternalStream is deprecated. Use Stream.from_external() instead to interoperate with external streams that implement the CUDA stream protocol.\n","  t2=gt(lambda:fl_eigh_cuda(Ax),10,100)\n"]},{"output_type":"stream","name":"stdout","text":["    2048     154.9us     426.1us    0.36x\n","    4096     240.5us     425.6us    0.57x\n","    8192     410.0us     427.0us    0.96x\n","   16384     742.0us     429.1us    1.73x\n","   32768  OOM\n","\n","  MEMORY (n=6 B=4096)\n","  cuSOLVER     1098.7MB\n","  CUDA FL      0.7MB\n","========================================================================\n"]},{"output_type":"stream","name":"stderr","text":["/tmp/ipykernel_89646/739027546.py:389: DeprecationWarning: ExternalStream is deprecated. Use Stream.from_external() instead to interoperate with external streams that implement the CUDA stream protocol.\n","  for lbl,fn in [(\"cuSOLVER\",lambda:torch.linalg.eigh(A)),(\"CUDA FL\",lambda:fl_eigh_cuda(A))]:\n"]}]},{"cell_type":"markdown","source":["# cuda kernel test"],"metadata":{"id":"Sf_zpjymdzNE"}},{"cell_type":"code","source":["\"\"\"\n","fl_eigh_cuda.py — CUDA FL Hybrid Eigh via CuPy RawKernel.\n","\n","Generalized to any N (3-16). Per-size kernel generated and cached by NVRTC.\n","No ninja, no C++ compiler. Just the CUDA driver.\n","\n","Usage:\n","    from fl_eigh_cuda import fl_eigh_cuda\n","    evals, evecs = fl_eigh_cuda(A)  # A is [B, n, n] for any n in 3..16\n","\"\"\"\n","\n","import torch\n","from torch import Tensor\n","from typing import Tuple\n","\n","__all__ = ['fl_eigh_cuda', 'get_kernel']\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════\n","# Kernel source generator\n","# ═══════════════════════════════════════════════════════════════════════\n","\n","def _gen_cuda_src(N: int) -> str:\n","    \"\"\"Generate CUDA kernel source for NxN eigendecomposition.\"\"\"\n","    N2 = N * N\n","    n_polish = 5 if N > 8 else 3\n","    return f\"\"\"\n","extern \"C\" __global__\n","void fl_eigh_n{N}(\n","    const float* __restrict__ A_in,\n","    float* __restrict__ evals_out,\n","    float* __restrict__ evecs_out,\n","    int B\n",") {{\n","    int tid = blockIdx.x * blockDim.x + threadIdx.x;\n","    if (tid >= B) return;\n","\n","    const int NN = {N};\n","    const int N2 = {N2};\n","\n","    /* ── Load A and pre-scale ── */\n","    double a[{N2}];\n","    double frob_sq = 0.0;\n","    for (int i = 0; i < N2; i++) {{\n","        a[i] = (double)A_in[tid * N2 + i];\n","        frob_sq += a[i] * a[i];\n","    }}\n","    double scale = sqrt(frob_sq / (double)NN);\n","    if (scale < 1e-12) scale = 1e-12;\n","    double inv_s = 1.0 / scale;\n","    for (int i = 0; i < N2; i++) a[i] *= inv_s;\n","\n","    /* ── Phase 1: FL characteristic polynomial (fp64) ── */\n","    double c[{N + 1}];\n","    for (int i = 0; i <= NN; i++) c[i] = 0.0;\n","    c[NN] = 1.0;\n","\n","    double m[{N2}];\n","    for (int i = 0; i < N2; i++) m[i] = 0.0;\n","\n","    for (int k = 1; k <= NN; k++) {{\n","        double mn[{N2}];\n","        for (int i = 0; i < NN; i++)\n","            for (int j = 0; j < NN; j++) {{\n","                double acc = 0.0;\n","                for (int l = 0; l < NN; l++) acc += a[i*NN+l] * m[l*NN+j];\n","                if (i == j) acc += c[NN-k+1];\n","                mn[i*NN+j] = acc;\n","            }}\n","        double tr = 0.0;\n","        for (int i = 0; i < NN; i++)\n","            for (int l = 0; l < NN; l++)\n","                tr += a[i*NN+l] * mn[l*NN+i];\n","        c[NN-k] = -tr / (double)k;\n","        for (int i = 0; i < N2; i++) m[i] = mn[i];\n","    }}\n","\n","    /* ── Phase 2: Laguerre root-finding + deflation (fp64) ── */\n","    double diag[{N}];\n","    for (int i = 0; i < NN; i++) diag[i] = a[i*NN+i];\n","    for (int pass = 0; pass < NN-1; pass++)\n","        for (int j = 0; j < NN-1; j++)\n","            if (diag[j] > diag[j+1]) {{\n","                double tmp = diag[j]; diag[j] = diag[j+1]; diag[j+1] = tmp;\n","            }}\n","    for (int i = 0; i < NN; i++)\n","        diag[i] += -1e-4 + 2e-4 * (double)i / (double)(NN > 1 ? NN-1 : 1);\n","\n","    double cl[{N + 1}];\n","    for (int i = 0; i <= NN; i++) cl[i] = c[i];\n","    double roots[{N}];\n","\n","    for (int ri = 0; ri < NN; ri++) {{\n","        int deg = NN - ri;\n","        double z = diag[ri];\n","        for (int lag = 0; lag < 5; lag++) {{\n","            double pv = cl[deg], dp = 0.0, d2 = 0.0;\n","            for (int j = deg-1; j >= 0; j--) {{\n","                d2 = d2*z + dp; dp = dp*z + pv; pv = pv*z + cl[j];\n","            }}\n","            if (fabs(pv) > 1e-30) {{\n","                double G = dp / pv;\n","                double H = G*G - 2.0*d2 / pv;\n","                double disc = ((double)(deg-1)) * ((double)deg * H - G*G);\n","                if (disc < 0.0) disc = 0.0;\n","                double sq = sqrt(disc);\n","                double gp = G + sq, gm = G - sq;\n","                double den = (fabs(gp) >= fabs(gm)) ? gp : gm;\n","                if (fabs(den) > 1e-20) z -= (double)deg / den;\n","            }}\n","        }}\n","        roots[ri] = z;\n","        if (deg > 1) {{\n","            double b = cl[deg];\n","            for (int j = deg-1; j > 0; j--) {{\n","                double bn = cl[j] + z*b; cl[j] = b; b = bn;\n","            }}\n","            cl[0] = b;\n","        }}\n","    }}\n","\n","    /* ── Newton polish on original polynomial ── */\n","    for (int pol = 0; pol < {n_polish}; pol++)\n","        for (int ri = 0; ri < NN; ri++) {{\n","            double pv = c[NN], dp = 0.0;\n","            for (int j = NN-1; j >= 0; j--) {{\n","                dp = dp*roots[ri] + pv; pv = pv*roots[ri] + c[j];\n","            }}\n","            if (fabs(dp) > 1e-30) roots[ri] -= pv / dp;\n","        }}\n","\n","    /* ── Phase 3: Eigenvectors via interleaved FL+Horner (fp64) ── */\n","    float evecs[{N2}];\n","\n","    for (int ei = 0; ei < NN; ei++) {{\n","        double lam = roots[ei];\n","        double m_loc[{N2}], r_loc[{N2}];\n","        for (int i = 0; i < N2; i++) m_loc[i] = 0.0;\n","\n","        for (int k = 1; k <= NN; k++) {{\n","            double mn_loc[{N2}];\n","            for (int i = 0; i < NN; i++)\n","                for (int j = 0; j < NN; j++) {{\n","                    double acc = 0.0;\n","                    for (int l = 0; l < NN; l++) acc += a[i*NN+l] * m_loc[l*NN+j];\n","                    if (i == j) acc += c[NN-k+1];\n","                    mn_loc[i*NN+j] = acc;\n","                }}\n","            if (k == 1)\n","                for (int i = 0; i < N2; i++) r_loc[i] = mn_loc[i];\n","            else\n","                for (int i = 0; i < N2; i++) r_loc[i] = r_loc[i]*lam + mn_loc[i];\n","            for (int i = 0; i < N2; i++) m_loc[i] = mn_loc[i];\n","        }}\n","\n","        int best_j = 0; double best_norm = -1.0;\n","        for (int j = 0; j < NN; j++) {{\n","            double col_sq = 0.0;\n","            for (int i = 0; i < NN; i++) col_sq += r_loc[i*NN+j] * r_loc[i*NN+j];\n","            if (col_sq > best_norm) {{ best_norm = col_sq; best_j = j; }}\n","        }}\n","        double vnorm = 0.0, vec[{N}];\n","        for (int i = 0; i < NN; i++) {{\n","            vec[i] = r_loc[i*NN + best_j]; vnorm += vec[i]*vec[i];\n","        }}\n","        vnorm = sqrt(vnorm) + 1e-30;\n","        for (int i = 0; i < NN; i++)\n","            evecs[i*NN + ei] = (float)(vec[i] / vnorm);\n","    }}\n","\n","    /* ── Phase 4: Newton-Schulz orthogonalization (fp32, 2 iters) ── */\n","    for (int ns = 0; ns < 2; ns++) {{\n","        float y[{N2}], t_m[{N2}], vn[{N2}];\n","        for (int i = 0; i < NN; i++)\n","            for (int j = 0; j < NN; j++) {{\n","                float acc = 0.0f;\n","                for (int l = 0; l < NN; l++) acc += evecs[l*NN+i] * evecs[l*NN+j];\n","                y[i*NN+j] = acc;\n","            }}\n","        for (int i = 0; i < NN; i++)\n","            for (int j = 0; j < NN; j++)\n","                t_m[i*NN+j] = ((i == j) ? 3.0f : 0.0f) - y[i*NN+j];\n","        for (int i = 0; i < NN; i++)\n","            for (int j = 0; j < NN; j++) {{\n","                float acc = 0.0f;\n","                for (int l = 0; l < NN; l++) acc += evecs[i*NN+l] * t_m[l*NN+j];\n","                vn[i*NN+j] = 0.5f * acc;\n","            }}\n","        for (int i = 0; i < N2; i++) evecs[i] = vn[i];\n","    }}\n","\n","    /* ── Phase 5: Rayleigh quotient refinement (fp32) ── */\n","    float af[{N2}];\n","    for (int i = 0; i < N2; i++) af[i] = (float)a[i];\n","\n","    float evals_local[{N}];\n","    for (int ei = 0; ei < NN; ei++) {{\n","        float lam_f = 0.0f;\n","        for (int l = 0; l < NN; l++) {{\n","            float av = 0.0f;\n","            for (int mm = 0; mm < NN; mm++) av += af[l*NN+mm] * evecs[mm*NN+ei];\n","            lam_f += evecs[l*NN+ei] * av;\n","        }}\n","        evals_local[ei] = lam_f * (float)scale;\n","    }}\n","\n","    /* ── Sort ascending + permute eigenvectors ── */\n","    int perm[{N}];\n","    for (int i = 0; i < NN; i++) perm[i] = i;\n","    for (int pass = 0; pass < NN-1; pass++)\n","        for (int j = 0; j < NN-1; j++)\n","            if (evals_local[j] > evals_local[j+1]) {{\n","                float tmp = evals_local[j];\n","                evals_local[j] = evals_local[j+1];\n","                evals_local[j+1] = tmp;\n","                int ptmp = perm[j]; perm[j] = perm[j+1]; perm[j+1] = ptmp;\n","            }}\n","\n","    for (int i = 0; i < NN; i++)\n","        evals_out[tid * NN + i] = evals_local[i];\n","    for (int jo = 0; jo < NN; jo++) {{\n","        int js = perm[jo];\n","        for (int i = 0; i < NN; i++)\n","            evecs_out[tid * N2 + i*NN + jo] = evecs[i*NN + js];\n","    }}\n","}}\n","\"\"\"\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════\n","# Kernel cache and compilation\n","# ═══════════════════════════════════════════════════════════════════════\n","\n","_kernels = {}  # N -> compiled CuPy RawKernel\n","\n","\n","def get_kernel(N: int):\n","    \"\"\"Get or compile the CUDA kernel for NxN matrices.\"\"\"\n","    if N in _kernels:\n","        return _kernels[N]\n","    import cupy\n","    src = _gen_cuda_src(N)\n","    kernel = cupy.RawKernel(src, f'fl_eigh_n{N}')\n","    kernel.compile()\n","    _kernels[N] = kernel\n","    return kernel\n","\n","\n","def fl_eigh_cuda(A: Tensor, ev_buf=None, vc_buf=None) -> Tuple[Tensor, Tensor]:\n","    \"\"\"CUDA FL Hybrid Eigendecomposition for [B, n, n] symmetric matrices.\n","\n","    Supports n = 3..16. Per-size kernel compiled on first call, cached.\n","    PyTorch interop via raw pointers (zero-copy).\n","\n","    Args:\n","        A: [B, n, n] symmetric CUDA float32 tensor\n","        ev_buf: optional pre-allocated [B, n] output for eigenvalues\n","        vc_buf: optional pre-allocated [B, n, n] output for eigenvectors\n","\n","    Returns: (eigenvalues [B,n], eigenvectors [B,n,n]) sorted ascending\n","    \"\"\"\n","    import cupy\n","    assert A.is_cuda, \"Need CUDA tensor\"\n","    B, n, n2 = A.shape\n","    assert n == n2 and 3 <= n <= 16, f\"Need square [B,n,n] with 3<=n<=16, got {A.shape}\"\n","\n","    kernel = get_kernel(n)\n","    A_c = A.contiguous().float()\n","\n","    if ev_buf is None:\n","        ev_buf = torch.empty(B, n, device=A.device, dtype=torch.float32)\n","    if vc_buf is None:\n","        vc_buf = torch.empty(B, n, n, device=A.device, dtype=torch.float32)\n","\n","    threads = 128\n","    blocks = (B + threads - 1) // threads\n","\n","    stream = cupy.cuda.ExternalStream(torch.cuda.current_stream().cuda_stream)\n","    with stream:\n","        kernel((blocks,), (threads,),\n","               (A_c.data_ptr(), ev_buf.data_ptr(), vc_buf.data_ptr(), B))\n","\n","    return ev_buf, vc_buf\n","\n","\n","# ═══════════════════════════════════════════════════════════════════════\n","# Self-test\n","# ═══════════════════════════════════════════════════════════════════════\n","\n","if __name__ == '__main__':\n","    import time, gc\n","\n","    if not torch.cuda.is_available():\n","        print(\"No CUDA\"); exit(1)\n","\n","    dev = torch.device('cuda')\n","    p = torch.cuda.get_device_properties(0)\n","\n","    def sync(): torch.cuda.synchronize()\n","    def gt(fn, w=20, r=200):\n","        for _ in range(w): fn()\n","        sync(); t = time.perf_counter()\n","        for _ in range(r): fn()\n","        sync(); return (time.perf_counter() - t) / r\n","    def fmt(s):\n","        if s < 1e-3: return f\"{s*1e6:.0f}us\"\n","        if s < 1: return f\"{s*1e3:.2f}ms\"\n","        return f\"{s:.3f}s\"\n","    def make(B, n):\n","        R = torch.randn(B, n, n, device=dev); return (R + R.mT) / 2\n","\n","    print(\"=\" * 72)\n","    print(\"  FL Eigh CUDA Kernel (CuPy/NVRTC)\")\n","    print(\"=\" * 72)\n","    print(f\"  {p.name} | PyTorch {torch.__version__}\")\n","\n","    # ── Accuracy across sizes ──\n","    print(f\"\\n  ACCURACY\")\n","    print(f\"  {'n':>3}  {'val_err':>10}  {'align':>10}  {'compile':>10}\")\n","    for n in [3, 4, 5, 6, 8, 10, 12, 16]:\n","        B = 2048\n","        A = make(B, n)\n","        t0 = time.perf_counter()\n","        get_kernel(n)  # compile\n","        tc = time.perf_counter() - t0\n","        cv, cV = fl_eigh_cuda(A)\n","        rv, rV = torch.linalg.eigh(A)\n","        ve = (cv - rv).abs().max().item()\n","        dots = torch.bmm(rV.double().mT, cV.double()).abs().max(dim=-1).values.min().item()\n","        ok = \"OK\" if ve < 1e-2 and dots > 0.99 else \"NO\"\n","        print(f\"  [{ok}] n={n:>2}  val={ve:.1e}  align={dots:.6f}  {tc:.1f}s\")\n","        del A\n","\n","    # ── Timing at n=6 ──\n","    n = 6; B = 4096\n","    A = make(B, n)\n","    print(f\"\\n  THROUGHPUT (n={n} B={B})\")\n","    t_cus = gt(lambda: torch.linalg.eigh(A))\n","    t_cuda = gt(lambda: fl_eigh_cuda(A))\n","    print(f\"  cuSOLVER: {fmt(t_cus)}\")\n","    print(f\"  CUDA FL:  {fmt(t_cuda)} ({t_cus/t_cuda:.2f}x)\")\n","\n","    # ── Batch scaling ──\n","    print(f\"\\n  BATCH SCALING (n={n})\")\n","    print(f\"  {'B':>6}  {'cuSOLVER':>10}  {'CUDA FL':>10}  {'ratio':>7}\")\n","    for Bx in [256, 512, 1024, 2048, 4096, 8192, 16384, 32768]:\n","        Ax = make(Bx, n)\n","        row = f\"  {Bx:>6}\"\n","        try:\n","            t1 = gt(lambda: torch.linalg.eigh(Ax), 10, 100)\n","            row += f\"  {fmt(t1):>10}\"\n","        except RuntimeError:\n","            t1 = None; row += f\"  {'OOM':>10}\"; torch.cuda.empty_cache()\n","        try:\n","            t2 = gt(lambda: fl_eigh_cuda(Ax), 10, 100)\n","            row += f\"  {fmt(t2):>10}\"\n","            row += f\"  {(t1/t2 if t1 else 0):>6.2f}x\" if t1 else f\"  {'—':>7}\"\n","        except RuntimeError:\n","            row += f\"  {'OOM':>10}  {'':>7}\"; torch.cuda.empty_cache()\n","        print(row); del Ax; torch.cuda.empty_cache()\n","\n","    # ── Size scaling ──\n","    print(f\"\\n  SIZE SCALING (B=4096)\")\n","    print(f\"  {'n':>3}  {'cuSOLVER':>10}  {'CUDA FL':>10}  {'ratio':>7}\")\n","    for nx in [3, 4, 5, 6, 8, 10, 12, 16]:\n","        Ax = make(4096, nx)\n","        t1 = gt(lambda: torch.linalg.eigh(Ax), 10, 100)\n","        t2 = gt(lambda: fl_eigh_cuda(Ax), 10, 100)\n","        print(f\"  {nx:>3}  {fmt(t1):>10}  {fmt(t2):>10}  {t1/t2:>6.2f}x\")\n","        del Ax\n","\n","    # ── Batch scaling at multiple sizes ──\n","    print(f\"\\n  BATCH x SIZE (CUDA FL ratio vs cuSOLVER)\")\n","    sizes = [3, 5, 6, 8, 12]\n","    batches = [512, 2048, 8192, 16384]\n","    hdr = f\"  {'B':>6}\" + \"\".join(f\"{'n='+str(s):>8}\" for s in sizes)\n","    print(hdr)\n","    for Bx in batches:\n","        row = f\"  {Bx:>6}\"\n","        for nx in sizes:\n","            try:\n","                Ax = make(Bx, nx)\n","                t1 = gt(lambda: torch.linalg.eigh(Ax), 5, 50)\n","                t2 = gt(lambda: fl_eigh_cuda(Ax), 5, 50)\n","                row += f\"{t1/t2:>7.2f}x\"\n","                del Ax\n","            except:\n","                row += f\"{'OOM':>8}\"; torch.cuda.empty_cache()\n","        print(row)\n","\n","    # ── Memory ──\n","    print(f\"\\n  MEMORY (n=6 B=4096)\")\n","    A = make(4096, 6)\n","    for lbl, fn in [(\"cuSOLVER\", lambda: torch.linalg.eigh(A)),\n","                     (\"CUDA FL\", lambda: fl_eigh_cuda(A))]:\n","        torch.cuda.empty_cache(); gc.collect(); torch.cuda.reset_peak_memory_stats()\n","        base = torch.cuda.memory_allocated(); fn(); sync()\n","        print(f\"  {lbl:<12} {(torch.cuda.max_memory_allocated() - base) / 1024**2:.1f} MB\")\n","\n","    print(\"=\" * 72)"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"bXT8Mpr2d0fb","executionInfo":{"status":"ok","timestamp":1775054749696,"user_tz":420,"elapsed":19506,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"bd75a457-d62c-45af-b742-51d9614f3ac9"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["========================================================================\n","  FL Eigh CUDA Kernel (CuPy/NVRTC)\n","========================================================================\n","  NVIDIA RTX PRO 6000 Blackwell Server Edition | PyTorch 2.10.0+cu128\n","\n","  ACCURACY\n","    n     val_err       align     compile\n","  [OK] n= 3  val=1.7e-06  align=1.000000  0.0s\n","  [OK] n= 4  val=1.7e-06  align=1.000000  0.0s\n","  [OK] n= 5  val=2.6e-06  align=0.999999  0.0s\n","  [OK] n= 6  val=2.6e-06  align=0.999999  0.0s\n","  [OK] n= 8  val=3.3e-06  align=0.999922  0.0s\n","  [OK] n=10  val=6.4e-06  align=0.999999  0.0s\n","  [OK] n=12  val=7.2e-06  align=0.999999  0.0s\n","  [NO] n=16  val=1.1e+00  align=0.390229  0.0s\n","\n","  THROUGHPUT (n=6 B=4096)\n"]},{"output_type":"stream","name":"stderr","text":["/tmp/ipykernel_89646/1788184773.py:326: DeprecationWarning: ExternalStream is deprecated. Use Stream.from_external() instead to interoperate with external streams that implement the CUDA stream protocol.\n","  cv, cV = fl_eigh_cuda(A)\n","/tmp/ipykernel_89646/1788184773.py:339: DeprecationWarning: ExternalStream is deprecated. Use Stream.from_external() instead to interoperate with external streams that implement the CUDA stream protocol.\n","  t_cuda = gt(lambda: fl_eigh_cuda(A))\n"]},{"output_type":"stream","name":"stdout","text":["  cuSOLVER: 240us\n","  CUDA FL:  426us (0.56x)\n","\n","  BATCH SCALING (n=6)\n","       B    cuSOLVER     CUDA FL    ratio\n","     256       104us       426us    0.24x\n","     512       105us       426us    0.25x\n","    1024       127us       424us    0.30x\n"]},{"output_type":"stream","name":"stderr","text":["/tmp/ipykernel_89646/1788184773.py:355: DeprecationWarning: ExternalStream is deprecated. Use Stream.from_external() instead to interoperate with external streams that implement the CUDA stream protocol.\n","  t2 = gt(lambda: fl_eigh_cuda(Ax), 10, 100)\n"]},{"output_type":"stream","name":"stdout","text":["    2048       160us       426us    0.37x\n","    4096       240us       426us    0.56x\n","    8192       415us       427us    0.97x\n","   16384       748us       429us    1.74x\n","   32768         OOM       810us        —\n","\n","  SIZE SCALING (B=4096)\n","    n    cuSOLVER     CUDA FL    ratio\n","    3       139us        45us    3.08x\n","    4       168us        74us    2.29x\n","    5       203us       217us    0.93x\n"]},{"output_type":"stream","name":"stderr","text":["/tmp/ipykernel_89646/1788184773.py:368: DeprecationWarning: ExternalStream is deprecated. Use Stream.from_external() instead to interoperate with external streams that implement the CUDA stream protocol.\n","  t2 = gt(lambda: fl_eigh_cuda(Ax), 10, 100)\n"]},{"output_type":"stream","name":"stdout","text":["    6       249us       426us    0.59x\n","    8       327us      1.52ms    0.21x\n","   10       691us      5.59ms    0.12x\n","   12       835us     12.17ms    0.07x\n","   16      1.06ms     80.20ms    0.01x\n","\n","  BATCH x SIZE (CUDA FL ratio vs cuSOLVER)\n","       B     n=3     n=5     n=6     n=8    n=12\n"]},{"output_type":"stream","name":"stderr","text":["/tmp/ipykernel_89646/1788184773.py:384: DeprecationWarning: ExternalStream is deprecated. Use Stream.from_external() instead to interoperate with external streams that implement the CUDA stream protocol.\n","  t2 = gt(lambda: fl_eigh_cuda(Ax), 5, 50)\n"]},{"output_type":"stream","name":"stdout","text":["     512   1.89x   0.44x   0.25x   0.08x   0.02x\n","    2048   2.29x   0.65x   0.37x   0.13x   0.04x\n","    8192   4.62x   1.60x   0.98x   0.38x   0.12x\n","   16384   7.64x   2.77x   1.74x   0.69x   0.21x\n","\n","  MEMORY (n=6 B=4096)\n","  cuSOLVER     1098.7 MB\n","  CUDA FL      0.7 MB\n","========================================================================\n"]},{"output_type":"stream","name":"stderr","text":["/tmp/ipykernel_89646/1788184773.py:395: DeprecationWarning: ExternalStream is deprecated. Use Stream.from_external() instead to interoperate with external streams that implement the CUDA stream protocol.\n","  (\"CUDA FL\", lambda: fl_eigh_cuda(A))]:\n"]}]},{"cell_type":"code","source":["\"\"\"\n","FL Eigh — Complete Sweep v2.\n","Fixed: per-eigenvalue eigvecs for n>6 eliminates O(n^4) memory.\n","\"\"\"\n","import math, time, gc, sys, os\n","import torch\n","import torch.nn as nn\n","from torch import Tensor\n","from typing import Tuple\n","\n","torch.backends.cuda.matmul.allow_tf32 = False\n","torch.backends.cudnn.allow_tf32 = False\n","torch.set_float32_matmul_precision('highest')\n","torch._dynamo.config.recompile_limit = 64\n","sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)) if '__file__' in dir() else '.')\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# Unified FL Eigh with mode='fast'|'precise'\n","# ═══════════════════════════════════════════════════════════════════\n","class FLEigh(nn.Module):\n","    def __init__(self, mode='precise'):\n","        super().__init__()\n","        self.mode = mode  # 'fast' or 'precise'\n","\n","    def forward(self, A):\n","        B,n,_=A.shape; dev=A.device\n","        sc=(torch.linalg.norm(A.reshape(B,-1),dim=-1)/math.sqrt(n)).clamp(min=1e-12)\n","        As=A/sc[:,None,None]; Ad=As.double()\n","        I_d=torch.eye(n,device=dev,dtype=torch.float64).unsqueeze(0).expand(B,-1,-1)\n","        c=torch.zeros(B,n+1,device=dev,dtype=torch.float64); c[:,n]=1.0\n","        Ms=torch.zeros(n+1,B,n,n,device=dev,dtype=torch.float64)\n","        Mk=torch.zeros(B,n,n,device=dev,dtype=torch.float64)\n","        for k in range(1,n+1):\n","            Mk=torch.bmm(Ad,Mk)+c[:,n-k+1,None,None]*I_d; Ms[k]=Mk\n","            c[:,n-k]=-(Ad*Mk).sum((-2,-1))/k\n","\n","        # Phase 2: Laguerre + deflation + polish\n","        use_f64_lag = n > 6\n","        dt = torch.float64 if use_f64_lag else torch.float32\n","        n_polish = 5 if n > 8 else 3\n","\n","        cl=c.to(dt).clone(); roots=torch.zeros(B,n,device=dev,dtype=dt)\n","        zi=As.to(dt).diagonal(dim1=-2,dim2=-1).sort(dim=-1).values\n","        zi=zi+torch.linspace(-1e-4,1e-4,n,device=dev,dtype=dt).unsqueeze(0)\n","        for ri in range(n):\n","            deg=n-ri; z=zi[:,ri]\n","            for _ in range(5):\n","                pv=cl[:,deg]; dp=torch.zeros(B,device=dev,dtype=dt)\n","                d2=torch.zeros(B,device=dev,dtype=dt)\n","                for j in range(deg-1,-1,-1):\n","                    d2=d2*z+dp; dp=dp*z+pv; pv=pv*z+cl[:,j]\n","                ok=pv.abs()>1e-30; ps=torch.where(ok,pv,torch.ones_like(pv))\n","                G=torch.where(ok,dp/ps,torch.zeros_like(dp))\n","                H=G*G-torch.where(ok,2.0*d2/ps,torch.zeros_like(d2))\n","                disc=((deg-1.0)*(deg*H-G*G)).clamp(min=0.0); sq=torch.sqrt(disc)\n","                gp=G+sq; gm=G-sq\n","                den=torch.where(gp.abs()>=gm.abs(),gp,gm)\n","                dok=den.abs()>1e-20; ds=torch.where(dok,den,torch.ones_like(den))\n","                z=z-torch.where(dok,float(deg)/ds,torch.zeros_like(den))\n","            roots[:,ri]=z; b=cl[:,deg]\n","            for j in range(deg-1,0,-1): bn=cl[:,j]+z*b; cl[:,j]=b; b=bn\n","            cl[:,0]=b\n","        roots=roots.double()\n","        for _ in range(n_polish):\n","            pv=torch.ones(B,n,device=dev,dtype=torch.float64)\n","            dp=torch.zeros(B,n,device=dev,dtype=torch.float64)\n","            for j in range(n-1,-1,-1):\n","                dp=dp*roots+pv; pv=pv*roots+c[:,j:j+1]\n","            ok=dp.abs()>1e-30; dps=torch.where(ok,dp,torch.ones_like(dp))\n","            roots=roots-torch.where(ok,pv/dps,torch.zeros_like(pv))\n","\n","        # Phase 3: Eigenvectors — adaptive by n\n","        if n <= 6:\n","            V = self._eigvecs_broadcast(Ms, roots, n, B, dev)\n","        else:\n","            V = self._eigvecs_pereig(Ms, roots, n, B, dev)\n","\n","        # Phase 4: NS\n","        eye_f=torch.eye(n,device=dev,dtype=torch.float32).unsqueeze(0).expand(B,-1,-1)\n","        Y=torch.bmm(V.mT,V); X=eye_f.clone()\n","        for _ in range(2):\n","            T=3.0*eye_f-Y; X=0.5*torch.bmm(X,T); Y=0.5*torch.bmm(T,Y)\n","        V=torch.bmm(V,X)\n","\n","        # Phase 5: Rayleigh\n","        AV=torch.bmm(A,V); evals=(V*AV).sum(dim=-2)\n","        se,p=evals.sort(dim=-1)\n","        return se,V.gather(-1,p.unsqueeze(-2).expand_as(V))\n","\n","    def _eigvecs_broadcast(self, Ms, roots, n, B, dev):\n","        \"\"\"Broadcast Horner on [B,n,n,n]. Fast for small n.\"\"\"\n","        if self.mode == 'fast':\n","            Mw = Ms.float(); lam = roots.float()\n","        else:\n","            Mw = Ms; lam = roots\n","        R = Mw[1].unsqueeze(1).expand(-1,n,-1,-1).clone()\n","        for k in range(2,n+1):\n","            R = R * lam[:,:,None,None] + Mw[k].unsqueeze(1)\n","        if self.mode == 'fast':\n","            vec = R.sum(dim=-1)\n","            vn = vec.norm(dim=-1,keepdim=True)\n","            vec = torch.where(vn>1e-10, vec, R[:,:,:,0])\n","        else:\n","            cn = R.norm(dim=-2); best = cn.argmax(dim=-1)\n","            idx = best.unsqueeze(-1).unsqueeze(-1).expand(-1,-1,n,1)\n","            vec = R.gather(-1,idx).squeeze(-1)\n","        vec = vec/(vec.norm(dim=-1,keepdim=True)+1e-30)\n","        return vec.float().transpose(-2,-1)\n","\n","    def _eigvecs_pereig(self, Ms, roots, n, B, dev):\n","        \"\"\"Per-eigenvalue Horner on [B,n,n]. Avoids [B,n,n,n] for large n.\"\"\"\n","        if self.mode == 'fast':\n","            Mw = Ms.float()\n","        else:\n","            Mw = Ms  # fp64\n","        V = torch.empty(B, n, n, device=dev, dtype=torch.float32)\n","        for ei in range(n):\n","            if self.mode == 'fast':\n","                li = roots[:, ei:ei+1, None].float()\n","            else:\n","                li = roots[:, ei:ei+1, None]  # fp64 [B,1,1]\n","            R = Mw[1].clone()  # [B,n,n]\n","            for k in range(2, n+1):\n","                R = R * li + Mw[k]\n","            # Max-col extraction (robust for all modes at n>6)\n","            cn = R.norm(dim=-2)  # [B,n] — norm of each column\n","            best = cn.argmax(dim=-1)  # [B]\n","            idx = best[:,None,None].expand(B,n,1)\n","            vec = R.gather(-1, idx).squeeze(-1)  # [B,n]\n","            vec = vec / (vec.norm(dim=-1,keepdim=True) + 1e-30)\n","            V[:,:,ei] = vec.float()\n","        return V\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# CuPy CUDA kernel — imported from fl_eigh_cuda.py\n","# ═══════════════════════════════════════════════════════════════════\n","HAS_CUPY = False\n","fl_eigh_cupy = None\n","\n","def _init_cupy():\n","    global HAS_CUPY, fl_eigh_cupy\n","    try:\n","        from fl_eigh_cuda import fl_eigh_cuda, get_kernel\n","        get_kernel(6)  # pre-compile n=6\n","        fl_eigh_cupy = fl_eigh_cuda\n","        HAS_CUPY = True\n","        print(\"  CUDA kernels: available (fl_eigh_cuda.py)\")\n","    except Exception as e:\n","        print(f\"  CUDA kernels: not available ({e})\")\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# Math purity\n","# ═══════════════════════════════════════════════════════════════════\n","def math_purity(A,vals,vecs):\n","    B,n,_=A.shape; dev=A.device\n","    Ad=A.double();vd=vals.double();Vd=vecs.double()\n","    AV=torch.bmm(Ad,Vd);VL=Vd*vd.unsqueeze(-2)\n","    An=Ad.reshape(B,-1).norm(dim=-1,keepdim=True).clamp(min=1e-30)\n","    res=(AV-VL).norm(dim=-2)/An\n","    VtV=torch.bmm(Vd.mT,Vd);I=torch.eye(n,device=dev,dtype=torch.float64).unsqueeze(0)\n","    orth=(VtV-I).reshape(B,-1).norm(dim=-1)\n","    recon=torch.bmm(Vd*vd.unsqueeze(-2),Vd.mT)\n","    rec=(Ad-recon).reshape(B,-1).norm(dim=-1)/An.squeeze(-1)\n","    tr=(Ad.diagonal(dim1=-2,dim2=-1).sum(-1)-vd.sum(-1)).abs()\n","    det_A=torch.linalg.det(Ad);det_e=(det_A-vd.prod(-1)).abs()/det_A.abs().clamp(min=1e-30)\n","    cp=torch.zeros(B,n,device=dev,dtype=torch.float64)\n","    for i in range(n): cp[:,i]=torch.linalg.det(vd[:,i:i+1,None]*I-Ad).abs()\n","    return dict(res_max=res.max().item(),res_mean=res.mean().item(),\n","                orth_max=orth.max().item(),orth_mean=orth.mean().item(),\n","                recon_max=rec.max().item(),recon_mean=rec.mean().item(),\n","                tr_max=tr.max().item(),tr_mean=tr.mean().item(),\n","                det_max=det_e.max().item(),det_mean=det_e.mean().item(),\n","                cp_max=cp.max().item(),cp_mean=cp.mean().item())\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# Utils\n","# ═══════════════════════════════════════════════════════════════════\n","def sync(): torch.cuda.synchronize()\n","def gt(fn,w=20,r=300):\n","    for _ in range(w): fn()\n","    sync();t=time.perf_counter()\n","    for _ in range(r): fn()\n","    sync();return(time.perf_counter()-t)/r\n","def fmt(s):\n","    if s<1e-3:return f\"{s*1e6:.0f}us\"\n","    if s<1:return f\"{s*1e3:.2f}ms\"\n","    return f\"{s:.3f}s\"\n","def make(B,n,dev):R=torch.randn(B,n,n,device=dev);return(R+R.mT)/2\n","\n","# ═══════════════════════════════════════════════════════════════════\n","def main():\n","    dev=torch.device('cuda');p=torch.cuda.get_device_properties(0)\n","    print(\"=\"*78);print(\"  FL Eigh — Complete Sweep v2\");print(\"=\"*78)\n","    print(f\"  {p.name} | PyTorch {torch.__version__}\")\n","    _init_cupy();print(f\"  CuPy: {'yes' if HAS_CUPY else 'no'}\")\n","\n","    N=6;B=4096;A=make(B,N,dev)\n","    keys=['res_max','res_mean','orth_max','orth_mean','recon_max','recon_mean','tr_max','tr_mean','det_max','det_mean','cp_max','cp_mean']\n","    labels=['Eigpair max','Eigpair mean','Orth max','Orth mean','Recon max','Recon mean','Trace max','Trace mean','Det max','Det mean','ChPoly max','ChPoly mean']\n","\n","    # ══════════════ 1. MATH PURITY n=6 ══════════════\n","    print(f\"\\n{'='*78}\\n  1. MATH PURITY (n={N}, B=2048)\\n{'='*78}\")\n","    Am=make(2048,N,dev)\n","    impls={'cuSOLVER':lambda X:torch.linalg.eigh(X),'FL Fast':lambda X:FLEigh('fast')(X),'FL Precise':lambda X:FLEigh('precise')(X)}\n","    if HAS_CUPY: impls['CUDA kern']=lambda X:fl_eigh_cupy(X)\n","    all_m={n:math_purity(Am,*fn(Am)) for n,fn in impls.items()}\n","    names=list(impls.keys());score={n:0 for n in names}\n","    print(f\"\\n  {'Metric':<16}\"+\"\".join(f\"{n:>12}\" for n in names)+\"  Winner\")\n","    print(\"  \"+\"─\"*16+\"\".join(\"─\"*12 for _ in names)+\"  \"+\"─\"*12)\n","    for key,label in zip(keys,labels):\n","        vals={n:all_m[n][key] for n in names};w=min(vals,key=vals.get);score[w]+=1\n","        row=f\"  {label:<16}\"\n","        for n in names:row+=f\"{'*' if n==w else ' '}{vals[n]:>10.1e} \"\n","        print(row+f\" {w}\")\n","    print();[print(f\"  {n}: {score[n]}/{len(keys)}\") for n in names];del Am\n","\n","    # ══════════════ 2. PURITY ACROSS SIZES ══════════════\n","    print(f\"\\n{'='*78}\\n  2. PURITY ACROSS SIZES (Precise vs cuSOLVER)\\n{'='*78}\")\n","    print(f\"\\n  {'n':>3} {'FL':>5} {'cuS':>5}  {'eigpair':>10} {'orth':>10} {'recon':>10}\")\n","    for nx in [3,4,5,6,8,10,12,16]:\n","        Bx=2048 if nx<=8 else 1024;Ax=make(Bx,nx,dev)\n","        mc=math_purity(Ax,*torch.linalg.eigh(Ax));mf=math_purity(Ax,*FLEigh('precise')(Ax))\n","        fw=sum(1 for k in keys if mf[k]<mc[k]);cw=sum(1 for k in keys if mc[k]<mf[k])\n","        print(f\"  {nx:>3} {fw:>4}/12 {cw:>3}/12  {min(mf['res_max'],mc['res_max']):>10.1e} {min(mf['orth_max'],mc['orth_max']):>10.1e} {min(mf['recon_max'],mc['recon_max']):>10.1e}\")\n","        del Ax\n","\n","    # ══════════════ 3. COMPILED TIMING ══════════════\n","    print(f\"\\n{'='*78}\\n  3. COMPILED TIMING (n={N}, B={B})\\n{'='*78}\")\n","    print(\"  Compiling...\",end=\" \",flush=True)\n","    cf=torch.compile(FLEigh('fast'),fullgraph=True);cp_=torch.compile(FLEigh('precise'),fullgraph=True)\n","    for _ in range(3):cf(A);cp_(A);sync();print(\"done.\")\n","    t_cus=gt(lambda:torch.linalg.eigh(A));t_f=gt(lambda:cf(A));t_p=gt(lambda:cp_(A))\n","    res=[('cuSOLVER',t_cus),('FL Fast',t_f),('FL Precise',t_p)]\n","    if HAS_CUPY:t_cu=gt(lambda:fl_eigh_cupy(A));res.append(('CUDA kern',t_cu))\n","    print(f\"\\n  {'Method':<24}{'Time':>10}{'ratio':>8}{'mat/s':>14}\")\n","    print(f\"  {'─'*24}{'─'*10}{'─'*8}{'─'*14}\")\n","    for nm,t in res:print(f\"  {nm:<24}{fmt(t):>10}{t_cus/t:>7.2f}x{B/t:>12.0f}/s\")\n","\n","    # ══════════════ 4. CUDA GRAPH ══════════════\n","    print(f\"\\n{'='*78}\\n  4. CUDA GRAPH (n={N}, B={B})\\n{'='*78}\")\n","    gr=[]\n","    for nm,cfn in [('Fast+Graph',cf),('Precise+Graph',cp_)]:\n","        As=A.clone();s=torch.cuda.Stream();s.wait_stream(torch.cuda.current_stream())\n","        with torch.cuda.stream(s):\n","            for _ in range(3):cfn(As)\n","        torch.cuda.current_stream().wait_stream(s)\n","        g=torch.cuda.CUDAGraph()\n","        with torch.cuda.graph(g):ev_g,vc_g=cfn(As)\n","        def rep(g=g,As=As):As.copy_(A);g.replay()\n","        gr.append((nm,gt(rep)))\n","    print(f\"\\n  {'Method':<24}{'Time':>10}{'ratio':>8}\")\n","    print(f\"  {'─'*24}{'─'*10}{'─'*8}\")\n","    print(f\"  {'cuSOLVER':<24}{fmt(t_cus):>10}{'1.00x':>8}\")\n","    for nm,t in res[1:]+gr:print(f\"  {nm:<24}{fmt(t):>10}{t_cus/t:>7.2f}x\")\n","\n","    # ══════════════ 5. BATCH SCALING ══════════════\n","    print(f\"\\n{'='*78}\\n  5. BATCH SCALING (n={N})\\n{'='*78}\")\n","    cd=torch.compile(FLEigh('fast'),fullgraph=True,dynamic=True);cd(A);sync()\n","    hdr=f\"\\n  {'B':>6}{'cuSOLVER':>10}{'FL Fast':>10}{'ratio':>8}\"\n","    if HAS_CUPY: hdr+=f\"{'CUDA kern':>10}{'ratio':>8}\"\n","    print(hdr)\n","    print(f\"  {'─'*6}{'─'*10}{'─'*10}{'─'*8}\"+(f\"{'─'*10}{'─'*8}\" if HAS_CUPY else \"\"))\n","    for Bx in [256,512,1024,2048,4096,8192,16384,32768]:\n","        Ax=make(Bx,N,dev)\n","        row=f\"  {Bx:>6}\"\n","        # cuSOLVER\n","        try:\n","            cd(Ax);sync()\n","            t1=gt(lambda:torch.linalg.eigh(Ax),10,100)\n","            row+=f\"{fmt(t1):>10}\"\n","        except RuntimeError:\n","            t1=None; row+=f\"{'OOM':>10}\"; torch.cuda.empty_cache()\n","        # FL Fast\n","        try:\n","            cd(Ax);sync()\n","            t2=gt(lambda:cd(Ax),10,100)\n","            row+=f\"{fmt(t2):>10}\"\n","            row+=f\"{(t1/t2 if t1 else 0):>7.2f}x\" if t1 else f\"{'—':>8}\"\n","        except RuntimeError:\n","            t2=None; row+=f\"{'OOM':>10}{'':>8}\"; torch.cuda.empty_cache()\n","        # CUDA kernel\n","        if HAS_CUPY:\n","            try:\n","                t3=gt(lambda:fl_eigh_cupy(Ax),10,100)\n","                row+=f\"{fmt(t3):>10}\"\n","                row+=f\"{(t1/t3 if t1 else 0):>7.2f}x\" if t1 else f\"{'—':>8}\"\n","            except RuntimeError:\n","                row+=f\"{'OOM':>10}{'':>8}\"; torch.cuda.empty_cache()\n","        print(row); del Ax; torch.cuda.empty_cache()\n","\n","    # ══════════════ 6. SIZE SCALING ══════════════\n","    print(f\"\\n{'='*78}\\n  6. SIZE SCALING (B={B})\\n{'='*78}\")\n","    hdr6=f\"\\n  {'n':>3}{'cuSOLVER':>10}{'FL Fast':>10}{'F rat':>7}{'FL Prec':>10}{'P rat':>7}\"\n","    if HAS_CUPY: hdr6+=f\"{'CUDA':>10}{'C rat':>7}\"\n","    hdr6+=f\"{'val_err':>10}\"\n","    print(hdr6)\n","    sep6=f\"  {'─'*3}{'─'*10}{'─'*10}{'─'*7}{'─'*10}{'─'*7}\"\n","    if HAS_CUPY: sep6+=f\"{'─'*10}{'─'*7}\"\n","    sep6+=f\"{'─'*10}\"; print(sep6)\n","    for nx in [3,4,5,6,8,10,12,16]:\n","        try:\n","            Ax=make(B,nx,dev)\n","            sf=torch.compile(FLEigh('fast'),fullgraph=True)\n","            sp=torch.compile(FLEigh('precise'),fullgraph=True)\n","            for _ in range(3):sf(Ax);sp(Ax);sync()\n","            t1=gt(lambda:torch.linalg.eigh(Ax),10,100)\n","            t2=gt(lambda:sf(Ax),10,100)\n","            t3=gt(lambda:sp(Ax),10,100)\n","            rv,_=torch.linalg.eigh(Ax);fv,_=sp(Ax);ve=(fv-rv).abs().max().item()\n","            row=f\"  {nx:>3}{fmt(t1):>10}{fmt(t2):>10}{t1/t2:>6.2f}x{fmt(t3):>10}{t1/t3:>6.2f}x\"\n","            if HAS_CUPY:\n","                t4=gt(lambda:fl_eigh_cupy(Ax),10,100)\n","                row+=f\"{fmt(t4):>10}{t1/t4:>6.2f}x\"\n","            row+=f\"{ve:>10.1e}\"\n","            print(row);del Ax,sf,sp\n","        except Exception as e:print(f\"  {nx:>3} ERR: {str(e)[:40]}\");torch.cuda.empty_cache()\n","\n","    # ══════════════ 7. MEMORY ══════════════\n","    print(f\"\\n{'='*78}\\n  7. MEMORY (n={N}, B={B})\\n{'='*78}\\n\")\n","    fns=[('cuSOLVER',lambda:torch.linalg.eigh(A)),('FL Fast',lambda:FLEigh('fast')(A)),('FL Precise',lambda:FLEigh('precise')(A))]\n","    if HAS_CUPY:fns.append(('CUDA kern',lambda:fl_eigh_cupy(A)))\n","    for nm,fn in fns:\n","        torch.cuda.empty_cache();gc.collect();torch.cuda.reset_peak_memory_stats()\n","        base=torch.cuda.memory_allocated();fn();sync()\n","        print(f\"  {nm:<16}{(torch.cuda.max_memory_allocated()-base)/1024**2:>8.1f} MB\")\n","\n","    # ══════════════ 8. ACCURACY ══════════════\n","    print(f\"\\n{'='*78}\\n  8. ACCURACY (all sizes)\\n{'='*78}\\n\")\n","    for mode in ['fast','precise']:\n","        print(f\"  {mode}:\")\n","        for nx in [3,4,5,6,8,10,12,16]:\n","            Ax=make(1024,nx,dev);rv,rV=torch.linalg.eigh(Ax);fv,fV=FLEigh(mode)(Ax)\n","            ve=(fv-rv).abs().max().item()\n","            dots=torch.bmm(rV.double().mT,fV.double()).abs().max(dim=-1).values.min().item()\n","            ok=ve<1e-2 and dots>0.99\n","            print(f\"    [{'OK' if ok else 'NO'}] n={nx:>2} val={ve:.1e} align={dots:.6f}\")\n","            del Ax\n","        print()\n","\n","    print(\"=\"*78)\n","    print(f\"  n=6 compiled: Fast {fmt(t_f)} ({t_cus/t_f:.2f}x)  Precise {fmt(t_p)} ({t_cus/t_p:.2f}x)\")\n","    print(f\"  cuSOLVER: {fmt(t_cus)}\")\n","    print(\"=\"*78)\n","\n","if __name__=='__main__':main()"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"b23sU6nfeAxD","executionInfo":{"status":"ok","timestamp":1775054822093,"user_tz":420,"elapsed":33574,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"d31939fd-3ec1-4533-d612-067a7d457a96"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["==============================================================================\n","  FL Eigh — Complete Sweep v2\n","==============================================================================\n","  NVIDIA RTX PRO 6000 Blackwell Server Edition | PyTorch 2.10.0+cu128\n","  CUDA kernels: not available (No module named 'fl_eigh_cuda')\n","  CuPy: no\n","\n","==============================================================================\n","  1. MATH PURITY (n=6, B=2048)\n","==============================================================================\n","\n","  Metric              cuSOLVER     FL Fast  FL Precise  Winner\n","  ────────────────────────────────────────────────────  ────────────\n","  Eigpair max         5.8e-07     2.1e-04 *   1.7e-07  FL Precise\n","  Eigpair mean        1.3e-07     8.9e-07 *   2.9e-08  FL Precise\n","  Orth max            2.0e-06     3.2e-04 *   3.0e-07  FL Precise\n","  Orth mean           9.0e-07     4.7e-06 *   1.5e-07  FL Precise\n","  Recon max           1.3e-06     2.6e-04 *   2.9e-07  FL Precise\n","  Recon mean          4.9e-07     2.7e-06 *   1.1e-07  FL Precise\n","  Trace max           2.9e-06     1.5e-06 *   1.2e-06  FL Precise\n","  Trace mean          5.1e-07     3.2e-07 *   2.4e-07  FL Precise\n","  Det max             3.1e-02     2.2e-03 *   7.4e-04  FL Precise\n","  Det mean            1.8e-05     1.7e-06 *   8.9e-07  FL Precise\n","  ChPoly max          1.0e-02     2.4e-03 *   1.9e-03  FL Precise\n","  ChPoly mean         4.4e-05     2.3e-05 *   1.7e-05  FL Precise\n","\n","  cuSOLVER: 0/12\n","  FL Fast: 0/12\n","  FL Precise: 12/12\n","\n","==============================================================================\n","  2. PURITY ACROSS SIZES (Precise vs cuSOLVER)\n","==============================================================================\n","\n","    n    FL   cuS     eigpair       orth      recon\n","    3   12/12   0/12     1.8e-07    2.4e-07    3.0e-07\n","    4   12/12   0/12     2.1e-07    2.6e-07    3.5e-07\n","    5   12/12   0/12     1.9e-07    2.7e-07    3.2e-07\n","    6   12/12   0/12     2.1e-07    3.0e-07    3.4e-07\n","    8   12/12   0/12     1.7e-07    3.4e-07    3.2e-07\n","   10   12/12   0/12     1.8e-07    3.6e-07    3.3e-07\n","   12   12/12   0/12     1.4e-07    4.3e-07    2.8e-07\n","   16    2/12  10/12     5.8e-07    3.1e-06    1.2e-06\n","\n","==============================================================================\n","  3. COMPILED TIMING (n=6, B=4096)\n","==============================================================================\n","  Compiling... done.\n","done.\n","done.\n","\n","  Method                        Time   ratio         mat/s\n","  ────────────────────────────────────────────────────────\n","  cuSOLVER                     241us   1.00x    17015584/s\n","  FL Fast                      346us   0.70x    11854032/s\n","  FL Precise                   350us   0.69x    11698297/s\n","\n","==============================================================================\n","  4. CUDA GRAPH (n=6, B=4096)\n","==============================================================================\n","\n","  Method                        Time   ratio\n","  ──────────────────────────────────────────\n","  cuSOLVER                     241us   1.00x\n","  FL Fast                      346us   0.70x\n","  FL Precise                   350us   0.69x\n","  Fast+Graph                   280us   0.86x\n","  Precise+Graph                288us   0.84x\n","\n","==============================================================================\n","  5. BATCH SCALING (n=6)\n","==============================================================================\n","\n","       B  cuSOLVER   FL Fast   ratio\n","  ──────────────────────────────────\n","     256     103us     304us   0.34x\n","     512     104us     301us   0.35x\n","    1024     127us     301us   0.42x\n","    2048     160us     300us   0.53x\n","    4096     240us     351us   0.69x\n","    8192     418us     559us   0.75x\n","   16384     750us     979us   0.77x\n","   32768       OOM    1.82ms       —\n","\n","==============================================================================\n","  6. SIZE SCALING (B=4096)\n","==============================================================================\n","\n","    n  cuSOLVER   FL Fast  F rat   FL Prec  P rat   val_err\n","  ─────────────────────────────────────────────────────────\n","    3     139us     214us  0.65x     218us  0.64x   1.7e-06\n","    4     168us     225us  0.75x     222us  0.75x   2.1e-06\n","    5     204us     293us  0.70x     293us  0.70x   2.4e-06\n","    6     241us     346us  0.70x     351us  0.69x   2.4e-06\n","    8     325us     674us  0.48x     713us  0.46x   3.1e-06\n","   10     690us    1.18ms  0.58x    1.32ms  0.52x   5.7e-06\n","   12     834us    2.12ms  0.39x    2.34ms  0.36x   6.4e-06\n","   16    1.06ms    3.33ms  0.32x    3.55ms  0.30x   1.1e-05\n","\n","==============================================================================\n","  7. MEMORY (n=6, B=4096)\n","==============================================================================\n","\n","  cuSOLVER          1098.7 MB\n","  FL Fast             26.2 MB\n","  FL Precise          32.3 MB\n","\n","==============================================================================\n","  8. ACCURACY (all sizes)\n","==============================================================================\n","\n","  fast:\n","    [OK] n= 3 val=1.7e-06 align=1.000000\n","    [NO] n= 4 val=5.5e-03 align=0.927701\n","    [OK] n= 5 val=5.1e-05 align=0.999989\n","    [OK] n= 6 val=2.9e-06 align=0.999999\n","    [OK] n= 8 val=2.9e-06 align=0.999999\n","    [OK] n=10 val=4.8e-06 align=0.999999\n","    [OK] n=12 val=5.5e-06 align=0.999999\n","    [OK] n=16 val=3.1e-04 align=0.999933\n","\n","  precise:\n","    [OK] n= 3 val=1.4e-06 align=1.000000\n","    [OK] n= 4 val=1.7e-06 align=0.999999\n","    [OK] n= 5 val=2.4e-06 align=0.999999\n","    [OK] n= 6 val=2.6e-06 align=0.999999\n","    [OK] n= 8 val=2.4e-06 align=0.999999\n","    [OK] n=10 val=3.7e-06 align=0.999999\n","    [OK] n=12 val=5.5e-06 align=0.999999\n","    [OK] n=16 val=1.0e-05 align=0.999999\n","\n","==============================================================================\n","  n=6 compiled: Fast 346us (0.70x)  Precise 350us (0.69x)\n","  cuSOLVER: 241us\n","==============================================================================\n"]}]},{"cell_type":"code","source":["\"\"\"\n","FL Eigh — Parallel vs Sequential benchmark.\n","\n","Changes:\n","  1. Aberth-Ehrlich replaces Laguerre+deflation: all roots simultaneously\n","  2. Chunked eigenvectors: process chunk eigenvalues at once, not 1 or all\n","\n","Both compared head-to-head against sequential baseline.\n","\"\"\"\n","import math, time, gc, sys\n","import torch\n","import torch.nn as nn\n","from torch import Tensor\n","from typing import Tuple\n","\n","torch.backends.cuda.matmul.allow_tf32 = False\n","torch.backends.cudnn.allow_tf32 = False\n","torch.set_float32_matmul_precision('highest')\n","torch._dynamo.config.recompile_limit = 64\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# Sequential baseline (Laguerre + per-eigenvalue)\n","# ═══════════════════════════════════════════════════════════════════\n","class FLEighSeq(nn.Module):\n","    \"\"\"Current best: sequential Laguerre + per-eigenvalue eigvecs.\"\"\"\n","    def forward(self, A):\n","        B,n,_=A.shape; dev=A.device\n","        sc=(torch.linalg.norm(A.reshape(B,-1),dim=-1)/math.sqrt(n)).clamp(min=1e-12)\n","        As=A/sc[:,None,None]; Ad=As.double()\n","        I_d=torch.eye(n,device=dev,dtype=torch.float64).unsqueeze(0).expand(B,-1,-1)\n","        c=torch.zeros(B,n+1,device=dev,dtype=torch.float64); c[:,n]=1.0\n","        Ms=torch.zeros(n+1,B,n,n,device=dev,dtype=torch.float64)\n","        Mk=torch.zeros(B,n,n,device=dev,dtype=torch.float64)\n","        for k in range(1,n+1):\n","            Mk=torch.bmm(Ad,Mk)+c[:,n-k+1,None,None]*I_d; Ms[k]=Mk\n","            c[:,n-k]=-(Ad*Mk).sum((-2,-1))/k\n","        # Sequential Laguerre + deflation\n","        dt=torch.float64 if n>6 else torch.float32\n","        n_pol=5 if n>8 else 3\n","        cl=c.to(dt).clone(); roots=torch.zeros(B,n,device=dev,dtype=dt)\n","        zi=As.to(dt).diagonal(dim1=-2,dim2=-1).sort(dim=-1).values\n","        zi=zi+torch.linspace(-1e-4,1e-4,n,device=dev,dtype=dt).unsqueeze(0)\n","        for ri in range(n):\n","            deg=n-ri; z=zi[:,ri]\n","            for _ in range(5):\n","                pv=cl[:,deg]; dp=torch.zeros(B,device=dev,dtype=dt)\n","                d2=torch.zeros(B,device=dev,dtype=dt)\n","                for j in range(deg-1,-1,-1):\n","                    d2=d2*z+dp; dp=dp*z+pv; pv=pv*z+cl[:,j]\n","                ok=pv.abs()>1e-30; ps=torch.where(ok,pv,torch.ones_like(pv))\n","                G=torch.where(ok,dp/ps,torch.zeros_like(dp))\n","                H=G*G-torch.where(ok,2.0*d2/ps,torch.zeros_like(d2))\n","                disc=((deg-1.0)*(deg*H-G*G)).clamp(min=0.0); sq=torch.sqrt(disc)\n","                gp=G+sq; gm=G-sq\n","                den=torch.where(gp.abs()>=gm.abs(),gp,gm)\n","                dok=den.abs()>1e-20; ds=torch.where(dok,den,torch.ones_like(den))\n","                z=z-torch.where(dok,float(deg)/ds,torch.zeros_like(den))\n","            roots[:,ri]=z; b=cl[:,deg]\n","            for j in range(deg-1,0,-1): bn=cl[:,j]+z*b; cl[:,j]=b; b=bn\n","            cl[:,0]=b\n","        roots=roots.double()\n","        for _ in range(n_pol):\n","            pv=torch.ones(B,n,device=dev,dtype=torch.float64)\n","            dp=torch.zeros(B,n,device=dev,dtype=torch.float64)\n","            for j in range(n-1,-1,-1): dp=dp*roots+pv; pv=pv*roots+c[:,j:j+1]\n","            ok=dp.abs()>1e-30; dps=torch.where(ok,dp,torch.ones_like(dp))\n","            roots=roots-torch.where(ok,pv/dps,torch.zeros_like(pv))\n","        # Per-eigenvalue eigvecs\n","        V=torch.empty(B,n,n,device=dev,dtype=torch.float32)\n","        for ei in range(n):\n","            li=roots[:,ei:ei+1,None]\n","            R=Ms[1].clone()\n","            for k in range(2,n+1): R=R*li+Ms[k]\n","            cn=R.norm(dim=-2); best=cn.argmax(dim=-1)\n","            idx=best[:,None,None].expand(B,n,1)\n","            vec=R.gather(-1,idx).squeeze(-1)\n","            vec=vec/(vec.norm(dim=-1,keepdim=True)+1e-30)\n","            V[:,:,ei]=vec.float()\n","        # NS + Rayleigh\n","        eye_f=torch.eye(n,device=dev,dtype=torch.float32).unsqueeze(0).expand(B,-1,-1)\n","        Y=torch.bmm(V.mT,V); X=eye_f.clone()\n","        for _ in range(2): T=3.0*eye_f-Y; X=0.5*torch.bmm(X,T); Y=0.5*torch.bmm(T,Y)\n","        V=torch.bmm(V,X); AV=torch.bmm(A,V); evals=(V*AV).sum(dim=-2)\n","        se,p=evals.sort(dim=-1)\n","        return se,V.gather(-1,p.unsqueeze(-2).expand_as(V))\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# Parallel: Aberth-Ehrlich + chunked eigvecs\n","# ═══════════════════════════════════════════════════════════════════\n","class FLEighPar(nn.Module):\n","    \"\"\"Parallel Aberth-Ehrlich roots + chunked eigenvectors.\"\"\"\n","    def __init__(self, ae_iters=8, chunk=4):\n","        super().__init__()\n","        self.ae_iters = ae_iters\n","        self.chunk = chunk\n","\n","    def forward(self, A):\n","        B,n,_=A.shape; dev=A.device\n","        sc=(torch.linalg.norm(A.reshape(B,-1),dim=-1)/math.sqrt(n)).clamp(min=1e-12)\n","        As=A/sc[:,None,None]; Ad=As.double()\n","        I_d=torch.eye(n,device=dev,dtype=torch.float64).unsqueeze(0).expand(B,-1,-1)\n","        c=torch.zeros(B,n+1,device=dev,dtype=torch.float64); c[:,n]=1.0\n","        Ms=torch.zeros(n+1,B,n,n,device=dev,dtype=torch.float64)\n","        Mk=torch.zeros(B,n,n,device=dev,dtype=torch.float64)\n","        for k in range(1,n+1):\n","            Mk=torch.bmm(Ad,Mk)+c[:,n-k+1,None,None]*I_d; Ms[k]=Mk\n","            c[:,n-k]=-(Ad*Mk).sum((-2,-1))/k\n","\n","        # ── Aberth-Ehrlich: ALL roots simultaneously (fp64) ──\n","        z=As.double().diagonal(dim1=-2,dim2=-1).sort(dim=-1).values\n","        z=z+torch.linspace(-1e-4,1e-4,n,device=dev,dtype=torch.float64).unsqueeze(0)\n","        mask_eye=torch.eye(n,device=dev,dtype=torch.bool).unsqueeze(0)\n","\n","        for _ in range(self.ae_iters):\n","            # Horner on ALL roots: p(z_i), p'(z_i) for all i simultaneously\n","            pv=c[:,n:n+1].expand(B,n)  # [B,n] leading coefficient\n","            dp=torch.zeros(B,n,device=dev,dtype=torch.float64)\n","            for j in range(n-1,-1,-1):\n","                dp=dp*z+pv\n","                pv=pv*z+c[:,j:j+1]\n","            # pv = p(z), dp = p'(z), both [B,n]\n","\n","            # Newton step\n","            dp_safe=torch.where(dp.abs()>1e-30, dp, torch.ones_like(dp))\n","            w=torch.where(dp.abs()>1e-30, pv/dp_safe, torch.zeros_like(pv))\n","\n","            # Root repulsion: sum_{j!=i} 1/(z_i - z_j)\n","            diffs=z.unsqueeze(-1)-z.unsqueeze(-2)  # [B,n,n]\n","            diffs_safe=diffs.masked_fill(mask_eye, 1.0)\n","            correction=(1.0/diffs_safe).masked_fill(mask_eye, 0.0).sum(-1)  # [B,n]\n","\n","            # Aberth-Ehrlich update\n","            denom=1.0-w*correction\n","            denom_safe=torch.where(denom.abs()>1e-20, denom, torch.ones_like(denom))\n","            z=z-torch.where(denom.abs()>1e-20, w/denom_safe, w)\n","\n","        roots=z\n","\n","        # ── Newton polish on original polynomial ──\n","        n_pol=5 if n>8 else 3\n","        for _ in range(n_pol):\n","            pv=torch.ones(B,n,device=dev,dtype=torch.float64)\n","            dp=torch.zeros(B,n,device=dev,dtype=torch.float64)\n","            for j in range(n-1,-1,-1): dp=dp*roots+pv; pv=pv*roots+c[:,j:j+1]\n","            ok=dp.abs()>1e-30; dps=torch.where(ok,dp,torch.ones_like(dp))\n","            roots=roots-torch.where(ok,pv/dps,torch.zeros_like(pv))\n","\n","        # ── Chunked eigenvectors ──\n","        chunk=min(self.chunk, n)\n","        V=torch.empty(B,n,n,device=dev,dtype=torch.float32)\n","\n","        for ei_start in range(0, n, chunk):\n","            ei_end=min(ei_start+chunk, n)\n","            c_size=ei_end-ei_start\n","            lam_c=roots[:,ei_start:ei_end,None,None]  # [B,c_size,1,1]\n","            R=Ms[1].unsqueeze(1).expand(B,c_size,n,n).clone()\n","            for k in range(2,n+1):\n","                R=R*lam_c+Ms[k].unsqueeze(1)  # [B,c_size,n,n]\n","            # Max-col extraction per eigenvalue in chunk\n","            cn=R.norm(dim=-2)  # [B,c_size,n]\n","            best=cn.argmax(dim=-1)  # [B,c_size]\n","            idx=best.unsqueeze(-1).unsqueeze(-1).expand(B,c_size,n,1)\n","            vec=R.gather(-1,idx).squeeze(-1)  # [B,c_size,n]\n","            vec=vec/(vec.norm(dim=-1,keepdim=True)+1e-30)\n","            V[:,:,ei_start:ei_end]=vec.float().transpose(-2,-1)\n","\n","        # NS + Rayleigh (same)\n","        eye_f=torch.eye(n,device=dev,dtype=torch.float32).unsqueeze(0).expand(B,-1,-1)\n","        Y=torch.bmm(V.mT,V); X=eye_f.clone()\n","        for _ in range(2): T=3.0*eye_f-Y; X=0.5*torch.bmm(X,T); Y=0.5*torch.bmm(T,Y)\n","        V=torch.bmm(V,X); AV=torch.bmm(A,V); evals=(V*AV).sum(dim=-2)\n","        se,p=evals.sort(dim=-1)\n","        return se,V.gather(-1,p.unsqueeze(-2).expand_as(V))\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# Math purity\n","# ═══════════════════════════════════════════════════════════════════\n","def math_purity(A,vals,vecs):\n","    B,n,_=A.shape; dev=A.device\n","    Ad=A.double();vd=vals.double();Vd=vecs.double()\n","    AV=torch.bmm(Ad,Vd);VL=Vd*vd.unsqueeze(-2)\n","    An=Ad.reshape(B,-1).norm(dim=-1,keepdim=True).clamp(min=1e-30)\n","    res=(AV-VL).norm(dim=-2)/An\n","    VtV=torch.bmm(Vd.mT,Vd);I=torch.eye(n,device=dev,dtype=torch.float64).unsqueeze(0)\n","    orth=(VtV-I).reshape(B,-1).norm(dim=-1)\n","    recon=torch.bmm(Vd*vd.unsqueeze(-2),Vd.mT)\n","    rec=(Ad-recon).reshape(B,-1).norm(dim=-1)/An.squeeze(-1)\n","    tr=(Ad.diagonal(dim1=-2,dim2=-1).sum(-1)-vd.sum(-1)).abs()\n","    det_A=torch.linalg.det(Ad);det_e=(det_A-vd.prod(-1)).abs()/det_A.abs().clamp(min=1e-30)\n","    return dict(res_max=res.max().item(),res_mean=res.mean().item(),\n","                orth_max=orth.max().item(),orth_mean=orth.mean().item(),\n","                recon_max=rec.max().item(),recon_mean=rec.mean().item(),\n","                tr_max=tr.max().item(),det_max=det_e.max().item())\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# Benchmark\n","# ═══════════════════════════════════════════════════════════════════\n","def sync(): torch.cuda.synchronize()\n","def gt(fn,w=20,r=300):\n","    for _ in range(w): fn()\n","    sync();t=time.perf_counter()\n","    for _ in range(r): fn()\n","    sync();return(time.perf_counter()-t)/r\n","def fmt(s):\n","    if s<1e-3:return f\"{s*1e6:.0f}us\"\n","    if s<1:return f\"{s*1e3:.2f}ms\"\n","    return f\"{s:.3f}s\"\n","def make(B,n,dev):R=torch.randn(B,n,n,device=dev);return(R+R.mT)/2\n","\n","def main():\n","    dev=torch.device('cuda')\n","    p=torch.cuda.get_device_properties(0)\n","    print(\"=\"*78)\n","    print(\"  FL Eigh — Parallel vs Sequential\")\n","    print(\"=\"*78)\n","    print(f\"  {p.name} | PyTorch {torch.__version__}\")\n","\n","    # ══════════════ 1. ACCURACY ══════════════\n","    print(f\"\\n{'='*78}\\n  1. ACCURACY\\n{'='*78}\")\n","    print(f\"\\n  {'n':>3}  {'Sequential':>14}  {'Parallel':>14}  {'AE converged':>14}\")\n","    for nx in [3,4,5,6,8,10,12,16]:\n","        Ax=make(2048,nx,dev)\n","        rv,rV=torch.linalg.eigh(Ax)\n","        sv,sV=FLEighSeq()(Ax)\n","        pv,pV=FLEighPar(ae_iters=8,chunk=4)(Ax)\n","        se=(sv-rv).abs().max().item()\n","        pe=(pv-rv).abs().max().item()\n","        sd=torch.bmm(rV.double().mT,sV.double()).abs().max(dim=-1).values.min().item()\n","        pd=torch.bmm(rV.double().mT,pV.double()).abs().max(dim=-1).values.min().item()\n","        s_ok=\"OK\" if se<1e-2 and sd>0.99 else \"NO\"\n","        p_ok=\"OK\" if pe<1e-2 and pd>0.99 else \"NO\"\n","        print(f\"  n={nx:>2}  [{s_ok}] {se:.1e} {sd:.4f}  [{p_ok}] {pe:.1e} {pd:.4f}\")\n","        del Ax\n","\n","    # ══════════════ 2. MATH PURITY n=6 ══════════════\n","    print(f\"\\n{'='*78}\\n  2. MATH PURITY (n=6, B=2048)\\n{'='*78}\")\n","    Am=make(2048,6,dev)\n","    mc=math_purity(Am,*torch.linalg.eigh(Am))\n","    ms=math_purity(Am,*FLEighSeq()(Am))\n","    mp=math_purity(Am,*FLEighPar()(Am))\n","    keys=['res_max','res_mean','orth_max','orth_mean','recon_max','recon_mean','tr_max','det_max']\n","    labs=['Eigpair max','Eigpair mean','Orth max','Orth mean','Recon max','Recon mean','Trace max','Det max']\n","    print(f\"\\n  {'Metric':<16}{'cuSOLVER':>12}{'Sequential':>12}{'Parallel':>12}  {'Win':>10}\")\n","    print(f\"  {'─'*16}{'─'*12}{'─'*12}{'─'*12}  {'─'*10}\")\n","    sc_s,sc_p=0,0\n","    for key,lab in zip(keys,labs):\n","        vc,vs,vp=mc[key],ms[key],mp[key]\n","        best=min(vc,vs,vp)\n","        w=\"cuS\" if best==vc else (\"Seq\" if best==vs else \"Par\")\n","        if best==vs:sc_s+=1\n","        if best==vp:sc_p+=1\n","        print(f\"  {lab:<16}{vc:>12.1e}{vs:>12.1e}{vp:>12.1e}  {w:>10}\")\n","    print(f\"\\n  Sequential: {sc_s}/8  Parallel: {sc_p}/8\")\n","    del Am\n","\n","    # ══════════════ 3. COMPILED TIMING ══════════════\n","    print(f\"\\n{'='*78}\\n  3. COMPILED TIMING\\n{'='*78}\")\n","\n","    for nx in [6,8,10,12]:\n","        B=4096\n","        Ax=make(B,nx,dev)\n","        print(f\"\\n  n={nx} B={B}:\")\n","\n","        cs=torch.compile(FLEighSeq(),fullgraph=True)\n","        cp_ae=torch.compile(FLEighPar(ae_iters=8,chunk=min(4,nx)),fullgraph=True)\n","        for _ in range(3): cs(Ax); cp_ae(Ax); sync()\n","\n","        t_cus=gt(lambda:torch.linalg.eigh(Ax),10,100)\n","        t_seq=gt(lambda:cs(Ax),10,100)\n","        t_par=gt(lambda:cp_ae(Ax),10,100)\n","\n","        print(f\"    cuSOLVER:   {fmt(t_cus):>10}\")\n","        print(f\"    Sequential: {fmt(t_seq):>10}  {t_cus/t_seq:.2f}x vs cuS  (baseline)\")\n","        print(f\"    Parallel:   {fmt(t_par):>10}  {t_cus/t_par:.2f}x vs cuS  {t_seq/t_par:.2f}x vs Seq\")\n","        del Ax, cs, cp_ae; torch.cuda.empty_cache()\n","\n","    # ══════════════ 4. CHUNK SIZE SWEEP ══════════════\n","    print(f\"\\n{'='*78}\\n  4. CHUNK SIZE SWEEP (n=12, B=4096)\\n{'='*78}\")\n","    nx=12; B=4096; Ax=make(B,nx,dev)\n","    print(f\"\\n  {'chunk':>6}  {'time':>10}  {'mem_MB':>8}  {'val_err':>10}\")\n","    print(f\"  {'─'*6}  {'─'*10}  {'─'*8}  {'─'*10}\")\n","    for ch in [1,2,3,4,6,12]:\n","        fn=torch.compile(FLEighPar(ae_iters=8,chunk=ch),fullgraph=True)\n","        for _ in range(3): fn(Ax); sync()\n","        torch.cuda.empty_cache(); gc.collect(); torch.cuda.reset_peak_memory_stats()\n","        base=torch.cuda.memory_allocated()\n","        fv,fV=fn(Ax); sync()\n","        mem=(torch.cuda.max_memory_allocated()-base)/1024**2\n","        rv,_=torch.linalg.eigh(Ax); ve=(fv-rv).abs().max().item()\n","        t=gt(lambda:fn(Ax),10,100)\n","        print(f\"  {ch:>6}  {fmt(t):>10}  {mem:>7.1f}  {ve:>10.1e}\")\n","        del fn\n","    del Ax\n","\n","    # ══════════════ 5. AE CONVERGENCE ══════════════\n","    print(f\"\\n{'='*78}\\n  5. AE ITERATION SWEEP (n=6, B=2048)\\n{'='*78}\")\n","    Ax=make(2048,6,dev); rv,_=torch.linalg.eigh(Ax)\n","    print(f\"\\n  {'iters':>6}  {'val_err':>10}  {'align':>10}\")\n","    print(f\"  {'─'*6}  {'─'*10}  {'─'*10}\")\n","    for ae_it in [1,2,3,4,5,6,8,10,12]:\n","        pv,pV=FLEighPar(ae_iters=ae_it,chunk=6)(Ax)\n","        ve=(pv-rv).abs().max().item()\n","        dots=torch.bmm(torch.linalg.eigh(Ax)[1].double().mT,pV.double()).abs().max(dim=-1).values.min().item()\n","        ok=\"OK\" if ve<1e-4 else \"NO\"\n","        print(f\"  {ae_it:>6}  {ve:>10.1e}  {dots:>10.6f}  [{ok}]\")\n","    del Ax\n","\n","    # ══════════════ 6. BATCH SCALING ══════════════\n","    print(f\"\\n{'='*78}\\n  6. BATCH SCALING (n=6, Parallel vs Sequential)\\n{'='*78}\")\n","    cs6=torch.compile(FLEighSeq(),fullgraph=True)\n","    cp6=torch.compile(FLEighPar(ae_iters=8,chunk=6),fullgraph=True)\n","    A0=make(4096,6,dev); cs6(A0); cp6(A0); sync(); del A0\n","\n","    print(f\"\\n  {'B':>6}  {'cuSOLVER':>10}  {'Seq':>10}  {'Par':>10}  {'Par/Seq':>8}\")\n","    print(f\"  {'─'*6}  {'─'*10}  {'─'*10}  {'─'*10}  {'─'*8}\")\n","    for Bx in [512,1024,2048,4096,8192,16384]:\n","        Ax=make(Bx,6,dev)\n","        t1=gt(lambda:torch.linalg.eigh(Ax),10,100)\n","        t2=gt(lambda:cs6(Ax),10,100)\n","        t3=gt(lambda:cp6(Ax),10,100)\n","        print(f\"  {Bx:>6}  {fmt(t1):>10}  {fmt(t2):>10}  {fmt(t3):>10}  {t2/t3:>7.2f}x\")\n","        del Ax\n","\n","    print(f\"\\n{'='*78}\")\n","\n","if __name__=='__main__': main()"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"sVbsnwbSkB23","executionInfo":{"status":"ok","timestamp":1775055816999,"user_tz":420,"elapsed":140137,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"83c45f4a-fbc1-4f27-8251-738fd2890d32"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["==============================================================================\n","  FL Eigh — Parallel vs Sequential\n","==============================================================================\n","  NVIDIA RTX PRO 6000 Blackwell Server Edition | PyTorch 2.10.0+cu128\n","\n","==============================================================================\n","  1. ACCURACY\n","==============================================================================\n","\n","    n      Sequential        Parallel    AE converged\n","  n= 3  [OK] 1.4e-06 1.0000  [NO] 1.2e+00 0.0001\n","  n= 4  [OK] 1.9e-06 1.0000  [NO] 2.5e+00 0.0000\n","  n= 5  [OK] 2.1e-06 1.0000  [NO] 3.3e+00 0.0000\n","  n= 6  [OK] 3.6e-06 1.0000  [NO] 3.8e+00 0.0000\n","  n= 8  [OK] 3.8e-06 1.0000  [NO] 1.2e+01 0.0000\n","  n=10  [OK] 5.5e-06 1.0000  [NO] 1.9e+01 0.0000\n","  n=12  [OK] 3.7e-04 0.9995  [NO] 8.2e+01 0.0000\n","  n=16  [NO] 2.6e+00 0.0000  [NO] 6.3e+02 0.0000\n","\n","==============================================================================\n","  2. MATH PURITY (n=6, B=2048)\n","==============================================================================\n","\n","  Metric              cuSOLVER  Sequential    Parallel         Win\n","  ────────────────────────────────────────────────────  ──────────\n","  Eigpair max          6.1e-07     2.1e-07     4.6e-01         Seq\n","  Eigpair mean         1.3e-07     2.8e-08     5.5e-03         Seq\n","  Orth max             2.4e-06     3.5e-07     5.5e+00         Seq\n","  Orth mean            9.1e-07     1.4e-07     1.4e-01         Seq\n","  Recon max            1.4e-06     3.2e-07     3.4e+00         Seq\n","  Recon mean           5.0e-07     1.1e-07     7.1e-02         Seq\n","  Trace max            2.7e-06     1.3e-06     6.8e+00         Seq\n","  Det max              4.4e-04     7.3e-05     3.8e+00         Seq\n","\n","  Sequential: 8/8  Parallel: 0/8\n","\n","==============================================================================\n","  3. COMPILED TIMING\n","==============================================================================\n","\n","  n=6 B=4096:\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n","/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["    cuSOLVER:        241us\n","    Sequential:      397us  0.61x vs cuS  (baseline)\n","    Parallel:        422us  0.57x vs cuS  0.94x vs Seq\n","\n","  n=8 B=4096:\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["    cuSOLVER:        327us\n","    Sequential:      713us  0.46x vs cuS  (baseline)\n","    Parallel:        637us  0.51x vs cuS  1.12x vs Seq\n","\n","  n=10 B=4096:\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["    cuSOLVER:        693us\n","    Sequential:     1.32ms  0.52x vs cuS  (baseline)\n","    Parallel:       1.25ms  0.55x vs cuS  1.06x vs Seq\n","\n","  n=12 B=4096:\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["    cuSOLVER:        837us\n","    Sequential:     2.34ms  0.36x vs cuS  (baseline)\n","    Parallel:       2.30ms  0.36x vs cuS  1.02x vs Seq\n","\n","==============================================================================\n","  4. CHUNK SIZE SWEEP (n=12, B=4096)\n","==============================================================================\n","\n","   chunk        time    mem_MB     val_err\n","  ──────  ──────────  ────────  ──────────\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["       1      2.61ms    128.4     8.6e+01\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["       2      2.41ms    130.9     8.6e+01\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["       3      2.35ms    132.5     8.6e+01\n","       4      2.30ms    131.6     8.6e+01\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["       6      2.24ms    130.9     8.6e+01\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["      12      2.19ms    130.9     8.6e+01\n","\n","==============================================================================\n","  5. AE ITERATION SWEEP (n=6, B=2048)\n","==============================================================================\n","\n","   iters     val_err       align\n","  ──────  ──────────  ──────────\n","       1     7.8e+01    0.000000  [NO]\n","       2     1.7e+01    0.000000  [NO]\n","       3     1.1e+01    0.000000  [NO]\n","       4     6.4e+00    0.000000  [NO]\n","       5     4.4e+00    0.000000  [NO]\n","       6     4.4e+00    0.000000  [NO]\n","       8     5.4e+00    0.000000  [NO]\n","      10     3.6e+00    0.000000  [NO]\n","      12     2.9e+00    0.000000  [NO]\n","\n","==============================================================================\n","  6. BATCH SCALING (n=6, Parallel vs Sequential)\n","==============================================================================\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["\n","       B    cuSOLVER         Seq         Par   Par/Seq\n","  ──────  ──────────  ──────────  ──────────  ────────\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n","/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["     512       104us       339us       290us     1.17x\n","    1024       128us       340us       287us     1.18x\n","    2048       160us       338us       286us     1.18x\n","    4096       240us       404us       409us     0.99x\n","    8192       416us       627us       665us     0.94x\n","   16384       751us      1.07ms      1.17ms     0.91x\n","\n","==============================================================================\n"]}]},{"cell_type":"code","source":["\"\"\"\n","Diagnostic: what exactly breaks in parallel root-finding?\n","\n","Test 1: Pure parallel Laguerre (no Aberth, no clamp, no damp)\n","Test 2: Parallel Laguerre + Aberth\n","Test 3: Sequential Laguerre + deflation (baseline)\n","\n","Prints per-iteration convergence to identify exactly where it goes wrong.\n","\"\"\"\n","import math, torch\n","\n","torch.backends.cuda.matmul.allow_tf32 = False\n","torch.set_float32_matmul_precision('highest')\n","\n","dev = torch.device('cuda')\n","B = 512; N = 6\n","torch.manual_seed(42)\n","A = (lambda R: (R+R.mT)/2)(torch.randn(B, N, N, device=dev))\n","rv, rV = torch.linalg.eigh(A)\n","\n","# FL Phase 1 — get characteristic polynomial\n","sc = (torch.linalg.norm(A.reshape(B,-1), dim=-1) / math.sqrt(N)).clamp(min=1e-12)\n","As = A / sc[:, None, None]; Ad = As.double()\n","I_d = torch.eye(N, device=dev, dtype=torch.float64).unsqueeze(0).expand(B,-1,-1)\n","c = torch.zeros(B, N+1, device=dev, dtype=torch.float64); c[:, N] = 1.0\n","Mk = torch.zeros(B, N, N, device=dev, dtype=torch.float64)\n","for k in range(1, N+1):\n","    Mk = torch.bmm(Ad, Mk) + c[:, N-k+1, None, None] * I_d\n","    c[:, N-k] = -(Ad * Mk).sum((-2,-1)) / k\n","\n","# True roots (scaled)\n","true_roots = (rv / sc.unsqueeze(-1)).double().sort(dim=-1).values\n","\n","# Init from diagonal\n","z_init = Ad.diagonal(dim1=-2, dim2=-1).sort(dim=-1).values\n","pert = torch.linspace(-1e-3, 1e-3, N, device=dev, dtype=torch.float64).unsqueeze(0)\n","z_init = z_init + pert\n","\n","def horner_pd(c, z):\n","    \"\"\"Evaluate p(z), p'(z), p''(z)/2 via Horner. c: [B,n+1], z: [B,n]\"\"\"\n","    B, n_roots = z.shape\n","    n = c.shape[1] - 1\n","    pv = c[:, n:n+1].expand(B, n_roots)\n","    dp = torch.zeros_like(pv)\n","    d2 = torch.zeros_like(pv)\n","    for j in range(n-1, -1, -1):\n","        d2 = d2 * z + dp\n","        dp = dp * z + pv\n","        pv = pv * z + c[:, j:j+1]\n","    return pv, dp, d2\n","\n","def laguerre_step(c, z, n):\n","    pv, dp, d2 = horner_pd(c, z)\n","    ok = pv.abs() > 1e-30\n","    ps = torch.where(ok, pv, torch.ones_like(pv))\n","    G = torch.where(ok, dp / ps, torch.zeros_like(dp))\n","    H = G * G - torch.where(ok, 2.0 * d2 / ps, torch.zeros_like(d2))\n","    disc = ((n-1.0) * (n * H - G * G)).clamp(min=0.0)\n","    sq = torch.sqrt(disc)\n","    gp = G + sq; gm = G - sq\n","    den = torch.where(gp.abs() >= gm.abs(), gp, gm)\n","    return torch.where(den.abs() > 1e-20, float(n) / den, torch.zeros_like(den))\n","\n","mask_eye = torch.eye(N, device=dev, dtype=torch.bool).unsqueeze(0)\n","\n","def aberth_correction(z):\n","    diffs = z.unsqueeze(-1) - z.unsqueeze(-2)\n","    diffs_safe = diffs.masked_fill(mask_eye, 1.0)\n","    return (1.0 / diffs_safe).masked_fill(mask_eye, 0.0).sum(-1)\n","\n","def report(label, z, iteration):\n","    err = (z.sort(dim=-1).values - true_roots).abs().max().item()\n","    # Check for duplicates: min gap between sorted roots\n","    zs = z.sort(dim=-1).values\n","    min_gap = (zs[:, 1:] - zs[:, :-1]).min().item()\n","    # p(z) residual\n","    pv, _, _ = horner_pd(c, z)\n","    p_res = pv.abs().max().item()\n","    print(f\"  {label:>5} it={iteration:>2}  max_err={err:.2e}  min_gap={min_gap:.2e}  |p(z)|={p_res:.2e}\")\n","\n","print(\"=\"*78)\n","print(\"  Diagnostic: Parallel Root-Finding\")\n","print(\"=\"*78)\n","print(f\"  B={B} N={N}\")\n","print(f\"  True eigenvalue range: [{true_roots.min().item():.3f}, {true_roots.max().item():.3f}]\")\n","print(f\"  Diagonal init range:   [{z_init.min().item():.3f}, {z_init.max().item():.3f}]\")\n","\n","# ═══ Test 1: Pure parallel Laguerre (no Aberth) ═══\n","print(f\"\\n  --- Test 1: Pure Laguerre (no Aberth) ---\")\n","z = z_init.clone()\n","for it in range(20):\n","    step = laguerre_step(c, z, N)\n","    z = z - step\n","    if it < 5 or it % 5 == 4:\n","        report(\"PurL\", z, it)\n","\n","# ═══ Test 2: Laguerre + Aberth (full strength) ═══\n","print(f\"\\n  --- Test 2: Laguerre + Aberth (full) ---\")\n","z = z_init.clone()\n","for it in range(20):\n","    step = laguerre_step(c, z, N)\n","    corr = aberth_correction(z)\n","    denom = 1.0 - step * corr\n","    denom_safe = torch.where(denom.abs() > 1e-20, denom, torch.ones_like(denom))\n","    full_step = torch.where(denom.abs() > 1e-20, step / denom_safe, step)\n","    z = z - full_step\n","    if it < 5 or it % 5 == 4:\n","        report(\"LA-F\", z, it)\n","\n","# ═══ Test 3: Laguerre + weak Aberth (0.1× correction) ═══\n","print(f\"\\n  --- Test 3: Laguerre + weak Aberth (0.1x) ---\")\n","z = z_init.clone()\n","for it in range(20):\n","    step = laguerre_step(c, z, N)\n","    corr = aberth_correction(z)\n","    denom = 1.0 - 0.1 * step * corr\n","    denom_safe = torch.where(denom.abs() > 1e-20, denom, torch.ones_like(denom))\n","    full_step = torch.where(denom.abs() > 1e-20, step / denom_safe, step)\n","    z = z - full_step\n","    if it < 5 or it % 5 == 4:\n","        report(\"LA.1\", z, it)\n","\n","# ═══ Test 4: Pure Laguerre + post-sort each iteration ═══\n","print(f\"\\n  --- Test 4: Pure Laguerre + re-sort ---\")\n","z = z_init.clone()\n","for it in range(20):\n","    step = laguerre_step(c, z, N)\n","    z = z - step\n","    z = z.sort(dim=-1).values  # keep sorted\n","    if it < 5 or it % 5 == 4:\n","        report(\"PL+S\", z, it)\n","\n","# ═══ Test 5: Laguerre + Aberth + damped ramp ═══\n","print(f\"\\n  --- Test 5: Laguerre + Aberth damped (0.1 → 1.0) ---\")\n","z = z_init.clone()\n","for it in range(20):\n","    step = laguerre_step(c, z, N)\n","    corr = aberth_correction(z)\n","    alpha = min(1.0, 0.1 + 0.1 * it)\n","    denom = 1.0 - alpha * step * corr\n","    denom_safe = torch.where(denom.abs() > 1e-20, denom, torch.ones_like(denom))\n","    full_step = torch.where(denom.abs() > 1e-20, step / denom_safe, step)\n","    z = z - full_step\n","    z = z.sort(dim=-1).values\n","    if it < 5 or it % 5 == 4:\n","        report(\"LADa\", z, it)\n","\n","# ═══ Test 6: Newton + Aberth (original Aberth-Ehrlich) ═══\n","print(f\"\\n  --- Test 6: Newton + Aberth ---\")\n","z = z_init.clone()\n","for it in range(20):\n","    pv, dp, _ = horner_pd(c, z)\n","    ok = dp.abs() > 1e-30\n","    w = torch.where(ok, pv / dp, torch.zeros_like(pv))\n","    corr = aberth_correction(z)\n","    denom = 1.0 - w * corr\n","    denom_safe = torch.where(denom.abs() > 1e-20, denom, torch.ones_like(denom))\n","    full_step = torch.where(denom.abs() > 1e-20, w / denom_safe, w)\n","    z = z - full_step\n","    if it < 5 or it % 5 == 4:\n","        report(\"NwAb\", z, it)\n","\n","# ═══ Test 7: Pure Newton (no Aberth) ═══\n","print(f\"\\n  --- Test 7: Pure Newton ---\")\n","z = z_init.clone()\n","for it in range(20):\n","    pv, dp, _ = horner_pd(c, z)\n","    ok = dp.abs() > 1e-30\n","    w = torch.where(ok, pv / dp, torch.zeros_like(pv))\n","    z = z - w\n","    if it < 5 or it % 5 == 4:\n","        report(\"PurN\", z, it)\n","\n","print(\"=\"*78)"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"tNpy3MkMmR4c","executionInfo":{"status":"ok","timestamp":1775056867818,"user_tz":420,"elapsed":123,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"00b12f74-57fc-4f25-ce00-9c3cf8ffa72d"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["==============================================================================\n","  Diagnostic: Parallel Root-Finding\n","==============================================================================\n","  B=512 N=6\n","  True eigenvalue range: [-2.106, 2.099]\n","  Diagonal init range:   [-1.561, 1.722]\n","\n","  --- Test 1: Pure Laguerre (no Aberth) ---\n","   PurL it= 0  max_err=1.75e+00  min_gap=1.73e-06  |p(z)|=3.93e+00\n","   PurL it= 1  max_err=1.84e+00  min_gap=1.67e-16  |p(z)|=5.39e-01\n","   PurL it= 2  max_err=1.89e+00  min_gap=0.00e+00  |p(z)|=2.01e-02\n","   PurL it= 3  max_err=1.90e+00  min_gap=0.00e+00  |p(z)|=1.73e-04\n","   PurL it= 4  max_err=1.90e+00  min_gap=0.00e+00  |p(z)|=7.12e-07\n","   PurL it= 9  max_err=1.90e+00  min_gap=0.00e+00  |p(z)|=2.88e-15\n","   PurL it=14  max_err=1.90e+00  min_gap=0.00e+00  |p(z)|=2.88e-15\n","   PurL it=19  max_err=1.90e+00  min_gap=0.00e+00  |p(z)|=2.88e-15\n","\n","  --- Test 2: Laguerre + Aberth (full) ---\n","   LA-F it= 0  max_err=1.03e+02  min_gap=5.69e-07  |p(z)|=1.33e+12\n","   LA-F it= 1  max_err=3.70e+02  min_gap=1.71e-06  |p(z)|=2.62e+15\n","   LA-F it= 2  max_err=4.63e+02  min_gap=5.12e-06  |p(z)|=1.00e+16\n","   LA-F it= 3  max_err=1.58e+03  min_gap=1.54e-05  |p(z)|=1.59e+19\n","   LA-F it= 4  max_err=1.98e+03  min_gap=4.61e-05  |p(z)|=6.05e+19\n","   LA-F it= 9  max_err=6.05e+03  min_gap=5.26e-04  |p(z)|=4.89e+22\n","   LA-F it=14  max_err=1.85e+04  min_gap=8.50e-04  |p(z)|=3.95e+25\n","   LA-F it=19  max_err=5.63e+04  min_gap=1.92e-02  |p(z)|=3.19e+28\n","\n","  --- Test 3: Laguerre + weak Aberth (0.1x) ---\n","   LA.1 it= 0  max_err=2.89e+01  min_gap=5.69e-05  |p(z)|=7.50e+08\n","   LA.1 it= 1  max_err=2.09e+01  min_gap=2.84e-06  |p(z)|=1.23e+08\n","   LA.1 it= 2  max_err=1.35e+01  min_gap=6.74e-07  |p(z)|=1.06e+07\n","   LA.1 it= 3  max_err=6.44e+00  min_gap=4.80e-08  |p(z)|=2.25e+05\n","   LA.1 it= 4  max_err=1.89e+00  min_gap=4.16e-09  |p(z)|=1.99e+01\n","   LA.1 it= 9  max_err=1.90e+00  min_gap=1.45e-14  |p(z)|=4.90e-03\n","   LA.1 it=14  max_err=1.90e+00  min_gap=0.00e+00  |p(z)|=6.48e-04\n","   LA.1 it=19  max_err=1.90e+00  min_gap=0.00e+00  |p(z)|=8.53e-05\n","\n","  --- Test 4: Pure Laguerre + re-sort ---\n","   PL+S it= 0  max_err=1.75e+00  min_gap=1.73e-06  |p(z)|=3.93e+00\n","   PL+S it= 1  max_err=1.84e+00  min_gap=1.67e-16  |p(z)|=5.39e-01\n","   PL+S it= 2  max_err=1.89e+00  min_gap=0.00e+00  |p(z)|=2.01e-02\n","   PL+S it= 3  max_err=1.90e+00  min_gap=0.00e+00  |p(z)|=1.73e-04\n","   PL+S it= 4  max_err=1.90e+00  min_gap=0.00e+00  |p(z)|=7.12e-07\n","   PL+S it= 9  max_err=1.90e+00  min_gap=0.00e+00  |p(z)|=2.88e-15\n","   PL+S it=14  max_err=1.90e+00  min_gap=0.00e+00  |p(z)|=2.88e-15\n","   PL+S it=19  max_err=1.90e+00  min_gap=0.00e+00  |p(z)|=2.88e-15\n","\n","  --- Test 5: Laguerre + Aberth damped (0.1 → 1.0) ---\n","   LADa it= 0  max_err=2.89e+01  min_gap=5.69e-05  |p(z)|=7.50e+08\n","   LADa it= 1  max_err=3.72e+02  min_gap=3.82e-06  |p(z)|=2.69e+15\n","   LADa it= 2  max_err=1.13e+03  min_gap=4.82e-06  |p(z)|=2.08e+18\n","   LADa it= 3  max_err=2.26e+03  min_gap=1.90e-06  |p(z)|=1.34e+20\n","   LADa it= 4  max_err=3.77e+03  min_gap=9.02e-07  |p(z)|=2.88e+21\n","   LADa it= 9  max_err=1.70e+04  min_gap=2.34e-05  |p(z)|=2.38e+25\n","   LADa it=14  max_err=5.17e+04  min_gap=2.25e-03  |p(z)|=1.91e+28\n","   LADa it=19  max_err=1.58e+05  min_gap=8.72e-03  |p(z)|=1.54e+31\n","\n","  --- Test 6: Newton + Aberth ---\n","   NwAb it= 0  max_err=4.35e+02  min_gap=3.29e-05  |p(z)|=6.91e+15\n","   NwAb it= 1  max_err=1.57e+01  min_gap=9.86e-05  |p(z)|=2.43e+07\n","   NwAb it= 2  max_err=5.28e+01  min_gap=1.70e-05  |p(z)|=2.54e+10\n","   NwAb it= 3  max_err=5.37e+01  min_gap=5.22e-05  |p(z)|=2.75e+10\n","   NwAb it= 4  max_err=3.34e+02  min_gap=1.91e-04  |p(z)|=1.41e+15\n","   NwAb it= 9  max_err=2.02e+00  min_gap=1.92e-02  |p(z)|=6.78e+02\n","   NwAb it=14  max_err=1.05e-06  min_gap=1.92e-02  |p(z)|=1.24e-14\n","   NwAb it=19  max_err=1.05e-06  min_gap=1.92e-02  |p(z)|=1.24e-14\n","\n","  --- Test 7: Pure Newton ---\n","   PurN it= 0  max_err=3.51e+02  min_gap=6.24e-06  |p(z)|=1.93e+15\n","   PurN it= 1  max_err=2.93e+02  min_gap=1.69e-09  |p(z)|=6.46e+14\n","   PurN it= 2  max_err=2.44e+02  min_gap=0.00e+00  |p(z)|=2.16e+14\n","   PurN it= 3  max_err=2.03e+02  min_gap=0.00e+00  |p(z)|=7.25e+13\n","   PurN it= 4  max_err=1.69e+02  min_gap=0.00e+00  |p(z)|=2.43e+13\n","   PurN it= 9  max_err=6.69e+01  min_gap=0.00e+00  |p(z)|=1.02e+11\n","   PurN it=14  max_err=2.60e+01  min_gap=0.00e+00  |p(z)|=4.31e+08\n","   PurN it=19  max_err=9.63e+00  min_gap=0.00e+00  |p(z)|=1.81e+06\n","==============================================================================\n"]}]},{"cell_type":"code","source":["\"\"\"\n","FL Eigh — Laguerre-Aberth parallel root-finding.\n","\n","Fixes:\n","  1. Laguerre step (not Newton) — global convergence for real roots\n","  2. Damped first iterations — prevents overshoot from diagonal init\n","  3. Gershgorin-spread initialization — wider initial root spacing\n","  4. Root re-sorting per iteration — prevents crossing on real line\n","\"\"\"\n","import math, time, gc, sys\n","import torch\n","import torch.nn as nn\n","from torch import Tensor\n","from typing import Tuple\n","\n","torch.backends.cuda.matmul.allow_tf32 = False\n","torch.backends.cudnn.allow_tf32 = False\n","torch.set_float32_matmul_precision('highest')\n","torch._dynamo.config.recompile_limit = 64\n","\n","\n","class FLEighLagAberth(nn.Module):\n","    \"\"\"FL Eigh with Newton-Aberth parallel roots + chunked eigvecs.\"\"\"\n","    def __init__(self, ae_iters=15, chunk=4):\n","        super().__init__()\n","        self.ae_iters = ae_iters\n","        self.chunk = chunk\n","\n","    def forward(self, A):\n","        B,n,_=A.shape; dev=A.device\n","        sc=(torch.linalg.norm(A.reshape(B,-1),dim=-1)/math.sqrt(n)).clamp(min=1e-12)\n","        As=A/sc[:,None,None]; Ad=As.double()\n","        I_d=torch.eye(n,device=dev,dtype=torch.float64).unsqueeze(0).expand(B,-1,-1)\n","        c=torch.zeros(B,n+1,device=dev,dtype=torch.float64); c[:,n]=1.0\n","        Ms=torch.zeros(n+1,B,n,n,device=dev,dtype=torch.float64)\n","        Mk=torch.zeros(B,n,n,device=dev,dtype=torch.float64)\n","        for k in range(1,n+1):\n","            Mk=torch.bmm(Ad,Mk)+c[:,n-k+1,None,None]*I_d; Ms[k]=Mk\n","            c[:,n-k]=-(Ad*Mk).sum((-2,-1))/k\n","\n","        # ── Newton-Aberth: ALL roots simultaneously (fp64) ──\n","        # Diagonal init with small perturbation\n","        z = Ad.diagonal(dim1=-2, dim2=-1).sort(dim=-1).values\n","        z = z + torch.linspace(-1e-3, 1e-3, n, device=dev, dtype=torch.float64).unsqueeze(0)\n","\n","        mask_eye = torch.eye(n, device=dev, dtype=torch.bool).unsqueeze(0)\n","\n","        for _ in range(self.ae_iters):\n","            # Horner: p(z), p'(z) for ALL roots simultaneously\n","            pv = c[:, n:n+1].expand(B, n)\n","            dp = torch.zeros(B, n, device=dev, dtype=torch.float64)\n","            for j in range(n-1, -1, -1):\n","                dp = dp * z + pv\n","                pv = pv * z + c[:, j:j+1]\n","\n","            # Newton step: w = p(z) / p'(z)\n","            ok = dp.abs() > 1e-30\n","            w = torch.where(ok, pv / torch.where(ok, dp, torch.ones_like(dp)),\n","                            torch.zeros_like(pv))\n","\n","            # Aberth correction: sum_{j!=i} 1/(z_i - z_j)\n","            diffs = z.unsqueeze(-1) - z.unsqueeze(-2)  # [B, n, n]\n","            diffs_safe = diffs.masked_fill(mask_eye, 1.0)\n","            correction = (1.0 / diffs_safe).masked_fill(mask_eye, 0.0).sum(-1)\n","\n","            # Newton-Aberth update\n","            denom = 1.0 - w * correction\n","            denom_ok = denom.abs() > 1e-20\n","            z = z - torch.where(denom_ok,\n","                                w / torch.where(denom_ok, denom, torch.ones_like(denom)),\n","                                w)\n","\n","        roots = z\n","\n","        # ── Newton polish on original polynomial ──\n","        n_pol = 5 if n > 8 else 3\n","        for _ in range(n_pol):\n","            pv = torch.ones(B, n, device=dev, dtype=torch.float64)\n","            dp = torch.zeros(B, n, device=dev, dtype=torch.float64)\n","            for j in range(n-1, -1, -1):\n","                dp = dp * roots + pv; pv = pv * roots + c[:, j:j+1]\n","            ok = dp.abs() > 1e-30\n","            dps = torch.where(ok, dp, torch.ones_like(dp))\n","            roots = roots - torch.where(ok, pv / dps, torch.zeros_like(pv))\n","        # ── Chunked eigenvectors ──\n","        chunk = min(self.chunk, n)\n","        V = torch.empty(B, n, n, device=dev, dtype=torch.float32)\n","        for ei_start in range(0, n, chunk):\n","            ei_end = min(ei_start + chunk, n)\n","            c_size = ei_end - ei_start\n","            lam_c = roots[:, ei_start:ei_end, None, None]\n","            R = Ms[1].unsqueeze(1).expand(B, c_size, n, n).clone()\n","            for k in range(2, n+1):\n","                R = R * lam_c + Ms[k].unsqueeze(1)\n","            cn = R.norm(dim=-2)\n","            best = cn.argmax(dim=-1)\n","            idx = best.unsqueeze(-1).unsqueeze(-1).expand(B, c_size, n, 1)\n","            vec = R.gather(-1, idx).squeeze(-1)\n","            vec = vec / (vec.norm(dim=-1, keepdim=True) + 1e-30)\n","            V[:, :, ei_start:ei_end] = vec.float().transpose(-2, -1)\n","\n","        # NS + Rayleigh\n","        eye_f = torch.eye(n, device=dev, dtype=torch.float32).unsqueeze(0).expand(B,-1,-1)\n","        Y = torch.bmm(V.mT, V); X = eye_f.clone()\n","        for _ in range(2):\n","            T = 3.0 * eye_f - Y; X = 0.5 * torch.bmm(X, T); Y = 0.5 * torch.bmm(T, Y)\n","        V = torch.bmm(V, X); AV = torch.bmm(A, V); evals = (V * AV).sum(dim=-2)\n","        se, p = evals.sort(dim=-1)\n","        return se, V.gather(-1, p.unsqueeze(-2).expand_as(V))\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# Sequential baseline for comparison\n","# ═══════════════════════════════════════════════════════════════════\n","class FLEighSeq(nn.Module):\n","    def forward(self, A):\n","        B,n,_=A.shape; dev=A.device\n","        sc=(torch.linalg.norm(A.reshape(B,-1),dim=-1)/math.sqrt(n)).clamp(min=1e-12)\n","        As=A/sc[:,None,None]; Ad=As.double()\n","        I_d=torch.eye(n,device=dev,dtype=torch.float64).unsqueeze(0).expand(B,-1,-1)\n","        c=torch.zeros(B,n+1,device=dev,dtype=torch.float64); c[:,n]=1.0\n","        Ms=torch.zeros(n+1,B,n,n,device=dev,dtype=torch.float64)\n","        Mk=torch.zeros(B,n,n,device=dev,dtype=torch.float64)\n","        for k in range(1,n+1):\n","            Mk=torch.bmm(Ad,Mk)+c[:,n-k+1,None,None]*I_d; Ms[k]=Mk\n","            c[:,n-k]=-(Ad*Mk).sum((-2,-1))/k\n","        dt=torch.float64 if n>6 else torch.float32; n_pol=5 if n>8 else 3\n","        cl=c.to(dt).clone(); roots=torch.zeros(B,n,device=dev,dtype=dt)\n","        zi=As.to(dt).diagonal(dim1=-2,dim2=-1).sort(dim=-1).values\n","        zi=zi+torch.linspace(-1e-4,1e-4,n,device=dev,dtype=dt).unsqueeze(0)\n","        for ri in range(n):\n","            deg=n-ri; z=zi[:,ri]\n","            for _ in range(5):\n","                pv=cl[:,deg]; dp=torch.zeros(B,device=dev,dtype=dt); d2=torch.zeros(B,device=dev,dtype=dt)\n","                for j in range(deg-1,-1,-1): d2=d2*z+dp; dp=dp*z+pv; pv=pv*z+cl[:,j]\n","                ok=pv.abs()>1e-30; ps=torch.where(ok,pv,torch.ones_like(pv))\n","                G=torch.where(ok,dp/ps,torch.zeros_like(dp)); H=G*G-torch.where(ok,2.0*d2/ps,torch.zeros_like(d2))\n","                disc=((deg-1.0)*(deg*H-G*G)).clamp(min=0.0); sq=torch.sqrt(disc); gp=G+sq; gm=G-sq\n","                den=torch.where(gp.abs()>=gm.abs(),gp,gm); dok=den.abs()>1e-20\n","                ds=torch.where(dok,den,torch.ones_like(den)); z=z-torch.where(dok,float(deg)/ds,torch.zeros_like(den))\n","            roots[:,ri]=z; b=cl[:,deg]\n","            for j in range(deg-1,0,-1): bn=cl[:,j]+z*b; cl[:,j]=b; b=bn\n","            cl[:,0]=b\n","        roots=roots.double()\n","        for _ in range(n_pol):\n","            pv=torch.ones(B,n,device=dev,dtype=torch.float64); dp=torch.zeros(B,n,device=dev,dtype=torch.float64)\n","            for j in range(n-1,-1,-1): dp=dp*roots+pv; pv=pv*roots+c[:,j:j+1]\n","            ok=dp.abs()>1e-30; dps=torch.where(ok,dp,torch.ones_like(dp)); roots=roots-torch.where(ok,pv/dps,torch.zeros_like(pv))\n","        V=torch.empty(B,n,n,device=dev,dtype=torch.float32)\n","        for ei in range(n):\n","            li=roots[:,ei:ei+1,None]; R=Ms[1].clone()\n","            for k in range(2,n+1): R=R*li+Ms[k]\n","            cn=R.norm(dim=-2); best=cn.argmax(dim=-1); idx=best[:,None,None].expand(B,n,1)\n","            vec=R.gather(-1,idx).squeeze(-1); vec=vec/(vec.norm(dim=-1,keepdim=True)+1e-30); V[:,:,ei]=vec.float()\n","        eye_f=torch.eye(n,device=dev,dtype=torch.float32).unsqueeze(0).expand(B,-1,-1)\n","        Y=torch.bmm(V.mT,V); X=eye_f.clone()\n","        for _ in range(2): T=3.0*eye_f-Y; X=0.5*torch.bmm(X,T); Y=0.5*torch.bmm(T,Y)\n","        V=torch.bmm(V,X); AV=torch.bmm(A,V); evals=(V*AV).sum(dim=-2)\n","        se,p=evals.sort(dim=-1); return se,V.gather(-1,p.unsqueeze(-2).expand_as(V))\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# Benchmark\n","# ═══════════════════════════════════════════════════════════════════\n","def sync(): torch.cuda.synchronize()\n","def gt(fn,w=20,r=300):\n","    for _ in range(w): fn()\n","    sync();t=time.perf_counter()\n","    for _ in range(r): fn()\n","    sync();return(time.perf_counter()-t)/r\n","def fmt(s):\n","    if s<1e-3:return f\"{s*1e6:.0f}us\"\n","    if s<1:return f\"{s*1e3:.2f}ms\"\n","    return f\"{s:.3f}s\"\n","def make(B,n,dev):R=torch.randn(B,n,n,device=dev);return(R+R.mT)/2\n","\n","def main():\n","    dev=torch.device('cuda')\n","    p=torch.cuda.get_device_properties(0)\n","    print(\"=\"*78)\n","    print(\"  FL Eigh — Newton-Aberth Parallel\")\n","    print(\"=\"*78)\n","    print(f\"  {p.name} | PyTorch {torch.__version__}\")\n","\n","    # ══════════════ 1. Newton-Aberth CONVERGENCE ══════════════\n","    print(f\"\\n{'='*78}\\n  1. Newton-Aberth CONVERGENCE (n=6, B=2048)\\n{'='*78}\")\n","    Ax=make(2048,6,dev); rv,rV=torch.linalg.eigh(Ax)\n","    print(f\"\\n  {'iters':>6}  {'val_err':>10}  {'align':>10}\")\n","    for la_it in [1,2,3,4,5,6,8,10,12,16,20]:\n","        pv,pV=FLEighLagAberth(ae_iters=la_it,chunk=6)(Ax)\n","        ve=(pv-rv).abs().max().item()\n","        dots=torch.bmm(rV.double().mT,pV.double()).abs().max(dim=-1).values.min().item()\n","        ok=\"OK\" if ve<1e-4 else \"NO\"\n","        print(f\"  {la_it:>6}  {ve:>10.1e}  {dots:>10.6f}  [{ok}]\")\n","    del Ax\n","\n","    # ══════════════ 2. ACCURACY ACROSS SIZES ══════════════\n","    print(f\"\\n{'='*78}\\n  2. ACCURACY (AE iters=12)\\n{'='*78}\")\n","    print(f\"\\n  {'n':>3}  {'Sequential':>16}  {'Nw-Aberth':>16}\")\n","    for nx in [3,4,5,6,8,10,12,16]:\n","        Ax=make(2048,nx,dev); rv,rV=torch.linalg.eigh(Ax)\n","        sv,sV=FLEighSeq()(Ax); pv,pV=FLEighLagAberth(ae_iters=12,chunk=min(4,nx))(Ax)\n","        se=(sv-rv).abs().max().item(); pe=(pv-rv).abs().max().item()\n","        sd=torch.bmm(rV.double().mT,sV.double()).abs().max(dim=-1).values.min().item()\n","        pd=torch.bmm(rV.double().mT,pV.double()).abs().max(dim=-1).values.min().item()\n","        s_ok=\"OK\" if se<1e-2 and sd>0.99 else \"NO\"\n","        p_ok=\"OK\" if pe<1e-2 and pd>0.99 else \"NO\"\n","        print(f\"  n={nx:>2}  [{s_ok}] {se:.1e} {sd:.4f}  [{p_ok}] {pe:.1e} {pd:.4f}\")\n","        del Ax\n","\n","    # ══════════════ 3. COMPILED TIMING ══════════════\n","    print(f\"\\n{'='*78}\\n  3. COMPILED TIMING\\n{'='*78}\")\n","    for nx in [6,8,10,12]:\n","        B=4096; Ax=make(B,nx,dev)\n","        cs=torch.compile(FLEighSeq(),fullgraph=True)\n","        cp=torch.compile(FLEighLagAberth(ae_iters=10,chunk=min(4,nx)),fullgraph=True)\n","        for _ in range(3): cs(Ax); cp(Ax); sync()\n","        t_cus=gt(lambda:torch.linalg.eigh(Ax),10,100)\n","        t_seq=gt(lambda:cs(Ax),10,100)\n","        t_par=gt(lambda:cp(Ax),10,100)\n","        print(f\"\\n  n={nx} B={B}:\")\n","        print(f\"    cuSOLVER:      {fmt(t_cus):>10}\")\n","        print(f\"    Sequential:    {fmt(t_seq):>10}  {t_cus/t_seq:.2f}x cuS\")\n","        print(f\"    Nw-Aberth:    {fmt(t_par):>10}  {t_cus/t_par:.2f}x cuS  {t_seq/t_par:.2f}x vs Seq\")\n","        del Ax,cs,cp; torch.cuda.empty_cache()\n","\n","    # ══════════════ 4. BATCH SCALING ══════════════\n","    print(f\"\\n{'='*78}\\n  4. BATCH SCALING (n=6)\\n{'='*78}\")\n","    cs6=torch.compile(FLEighSeq(),fullgraph=True)\n","    cp6=torch.compile(FLEighLagAberth(ae_iters=10,chunk=6),fullgraph=True)\n","    A0=make(4096,6,dev); cs6(A0); cp6(A0); sync(); del A0\n","    print(f\"\\n  {'B':>6}  {'cuSOLVER':>10}  {'Seq':>10}  {'NwAb':>10}  {'NwAb/Seq':>8}\")\n","    for Bx in [512,1024,2048,4096,8192,16384]:\n","        Ax=make(Bx,6,dev)\n","        t1=gt(lambda:torch.linalg.eigh(Ax),10,100)\n","        t2=gt(lambda:cs6(Ax),10,100)\n","        t3=gt(lambda:cp6(Ax),10,100)\n","        print(f\"  {Bx:>6}  {fmt(t1):>10}  {fmt(t2):>10}  {fmt(t3):>10}  {t2/t3:>7.2f}x\")\n","        del Ax\n","\n","    print(f\"\\n{'='*78}\")\n","\n","if __name__=='__main__': main()"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"eaIe28Z1pb_M","executionInfo":{"status":"ok","timestamp":1775057434352,"user_tz":420,"elapsed":44056,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"f8369900-fbde-43fa-c85f-e69aa093ecb7"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["==============================================================================\n","  FL Eigh — Newton-Aberth Parallel\n","==============================================================================\n","  NVIDIA RTX PRO 6000 Blackwell Server Edition | PyTorch 2.10.0+cu128\n","\n","==============================================================================\n","  1. Newton-Aberth CONVERGENCE (n=6, B=2048)\n","==============================================================================\n","\n","   iters     val_err       align\n","       1     5.7e+01    0.000000  [NO]\n","       2     5.7e+01    0.000000  [NO]\n","       3     1.4e+01    0.000000  [NO]\n","       4     1.0e+01    0.000000  [NO]\n","       5     1.2e+01    0.000000  [NO]\n","       6     6.6e+00    0.000000  [NO]\n","       8     4.1e+00    0.000000  [NO]\n","      10     3.1e+00    0.000000  [NO]\n","      12     2.3e+00    0.000036  [NO]\n","      16     2.9e-06    0.999999  [OK]\n","      20     2.9e-06    0.999999  [OK]\n","\n","==============================================================================\n","  2. ACCURACY (AE iters=12)\n","==============================================================================\n","\n","    n        Sequential         Nw-Aberth\n","  n= 3  [OK] 1.7e-06 1.0000  [OK] 1.7e-06 1.0000\n","  n= 4  [OK] 2.1e-06 1.0000  [NO] 2.1e+00 0.0001\n","  n= 5  [OK] 2.4e-06 1.0000  [NO] 2.0e+00 0.0000\n","  n= 6  [OK] 2.4e-06 1.0000  [NO] 2.9e+00 0.0000\n","  n= 8  [OK] 3.6e-06 1.0000  [NO] 3.8e+00 0.0000\n","  n=10  [OK] 5.2e-06 1.0000  [NO] 2.5e+01 0.0000\n","  n=12  [OK] 5.7e-06 1.0000  [NO] 1.6e+01 0.0000\n","  n=16  [OK] 7.2e-06 1.0000  [NO] 3.2e+02 0.0000\n","\n","==============================================================================\n","  3. COMPILED TIMING\n","==============================================================================\n","\n","  n=6 B=4096:\n","    cuSOLVER:           241us\n","    Sequential:         398us  0.61x cuS\n","    Nw-Aberth:         447us  0.54x cuS  0.89x vs Seq\n","\n","  n=8 B=4096:\n","    cuSOLVER:           326us\n","    Sequential:         711us  0.46x cuS\n","    Nw-Aberth:         662us  0.49x cuS  1.08x vs Seq\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["\n","  n=10 B=4096:\n","    cuSOLVER:           691us\n","    Sequential:        1.32ms  0.52x cuS\n","    Nw-Aberth:        1.31ms  0.53x cuS  1.01x vs Seq\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["\n","  n=12 B=4096:\n","    cuSOLVER:           837us\n","    Sequential:        2.34ms  0.36x cuS\n","    Nw-Aberth:        2.38ms  0.35x cuS  0.98x vs Seq\n","\n","==============================================================================\n","  4. BATCH SCALING (n=6)\n","==============================================================================\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["\n","       B    cuSOLVER         Seq        NwAb  NwAb/Seq\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["     512       107us       330us       301us     1.10x\n","    1024       127us       333us       300us     1.11x\n","    2048       157us       333us       300us     1.11x\n","    4096       240us       404us       438us     0.92x\n","    8192       417us       627us       704us     0.89x\n","   16384       751us      1.07ms      1.25ms     0.86x\n","\n","==============================================================================\n"]}]},{"cell_type":"code","source":["\"\"\"\n","FL Eigh — Learned Root-Finder.\n","\n","A tiny MLP learns the mapping: FL coefficients → eigenvalue roots.\n","Trains at runtime (online, during the main training loop).\n","At inference: one forward pass + 2 Newton polish = precise roots.\n","\n","The key insight: FL coefficients from symmetric matrices lie on a\n","low-dimensional manifold. The network only needs to learn the inverse\n","map on THAT manifold, not general polynomial root-finding.\n","\"\"\"\n","import math, time, gc, sys\n","import torch\n","import torch.nn as nn\n","from torch import Tensor\n","from typing import Tuple\n","\n","torch.backends.cuda.matmul.allow_tf32 = False\n","torch.backends.cudnn.allow_tf32 = False\n","torch.set_float32_matmul_precision('highest')\n","torch._dynamo.config.recompile_limit = 64\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# Root-finding MLP\n","# ═══════════════════════════════════════════════════════════════════\n","class RootPredictor(nn.Module):\n","    \"\"\"Tiny MLP: polynomial coefficients → sorted roots.\n","\n","    Input:  [B, n+1] (FL characteristic polynomial coefficients, pre-scaled)\n","    Output: [B, n]   (predicted roots, sorted ascending)\n","    \"\"\"\n","    def __init__(self, n, hidden=64):\n","        super().__init__()\n","        self.n = n\n","        self.net = nn.Sequential(\n","            nn.Linear(n + 1, hidden),\n","            nn.GELU(),\n","            nn.Linear(hidden, hidden),\n","            nn.GELU(),\n","            nn.Linear(hidden, n),\n","        )\n","\n","    def forward(self, c):\n","        return self.net(c).sort(dim=-1).values\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# FL Eigh with learned roots\n","# ═══════════════════════════════════════════════════════════════════\n","class FLEighLearned(nn.Module):\n","    \"\"\"FL Eigh with learned root prediction + Newton polish.\n","\n","    Usage:\n","        solver = FLEighLearned(n=6)\n","        # Training mode: learns root predictor online\n","        for batch in dataloader:\n","            eigenvalues, eigenvectors = solver(A, train_roots=True)\n","        # Inference: uses learned predictor\n","        eigenvalues, eigenvectors = solver(A, train_roots=False)\n","    \"\"\"\n","    def __init__(self, n, hidden=64, polish_iters=3, chunk=4, lr=1e-3):\n","        super().__init__()\n","        self.n = n\n","        self.polish_iters = polish_iters\n","        self.chunk = chunk\n","        self.predictor = RootPredictor(n, hidden)\n","        self.optimizer = None  # created on first use (needs device)\n","        self.lr = lr\n","        self._trained_steps = 0\n","\n","    def _ensure_optimizer(self):\n","        if self.optimizer is None:\n","            self.optimizer = torch.optim.Adam(self.predictor.parameters(), lr=self.lr)\n","\n","    def forward(self, A, train_roots=False):\n","        B, n, _ = A.shape\n","        dev = A.device\n","\n","        # ── Phase 1: FL coefficients (fp64) ──\n","        sc = (torch.linalg.norm(A.reshape(B, -1), dim=-1) / math.sqrt(n)).clamp(min=1e-12)\n","        As = A / sc[:, None, None]; Ad = As.double()\n","        I_d = torch.eye(n, device=dev, dtype=torch.float64).unsqueeze(0).expand(B, -1, -1)\n","        c = torch.zeros(B, n + 1, device=dev, dtype=torch.float64); c[:, n] = 1.0\n","        Ms = torch.zeros(n + 1, B, n, n, device=dev, dtype=torch.float64)\n","        Mk = torch.zeros(B, n, n, device=dev, dtype=torch.float64)\n","        for k in range(1, n + 1):\n","            Mk = torch.bmm(Ad, Mk) + c[:, n - k + 1, None, None] * I_d\n","            Ms[k] = Mk\n","            c[:, n - k] = -(Ad * Mk).sum((-2, -1)) / k\n","\n","        # ── Phase 2: Root prediction ──\n","        c_input = c.float()  # [B, n+1] coefficients as network input\n","\n","        if train_roots:\n","            # Get ground truth roots via sequential Laguerre (teacher)\n","            with torch.no_grad():\n","                true_roots = self._laguerre_roots(c, As, B, n, dev)\n","\n","            # Predict roots\n","            pred_roots = self.predictor(c_input)  # [B, n]\n","\n","            # Train predictor\n","            self._ensure_optimizer()\n","            loss = (pred_roots - true_roots.float().sort(dim=-1).values).pow(2).mean()\n","            self.optimizer.zero_grad()\n","            loss.backward()\n","            self.optimizer.step()\n","            self._trained_steps += 1\n","\n","            # Use true roots for downstream (ensure training doesn't hurt accuracy)\n","            roots = true_roots\n","        else:\n","            # Inference: use predicted roots + Newton polish\n","            with torch.no_grad():\n","                pred_roots = self.predictor(c_input)\n","            roots = pred_roots.double()\n","\n","            # Newton polish on original polynomial (makes predictions precise)\n","            for _ in range(self.polish_iters):\n","                pv = torch.ones(B, n, device=dev, dtype=torch.float64)\n","                dp = torch.zeros(B, n, device=dev, dtype=torch.float64)\n","                for j in range(n - 1, -1, -1):\n","                    dp = dp * roots + pv\n","                    pv = pv * roots + c[:, j:j + 1]\n","                ok = dp.abs() > 1e-30\n","                dps = torch.where(ok, dp, torch.ones_like(dp))\n","                roots = roots - torch.where(ok, pv / dps, torch.zeros_like(pv))\n","\n","        # ── Phase 3: Chunked eigenvectors ──\n","        chunk = min(self.chunk, n)\n","        V = torch.empty(B, n, n, device=dev, dtype=torch.float32)\n","        for ei_start in range(0, n, chunk):\n","            ei_end = min(ei_start + chunk, n)\n","            c_size = ei_end - ei_start\n","            lam_c = roots[:, ei_start:ei_end, None, None]\n","            R = Ms[1].unsqueeze(1).expand(B, c_size, n, n).clone()\n","            for k in range(2, n + 1):\n","                R = R * lam_c + Ms[k].unsqueeze(1)\n","            cn = R.norm(dim=-2)\n","            best = cn.argmax(dim=-1)\n","            idx = best.unsqueeze(-1).unsqueeze(-1).expand(B, c_size, n, 1)\n","            vec = R.gather(-1, idx).squeeze(-1)\n","            vec = vec / (vec.norm(dim=-1, keepdim=True) + 1e-30)\n","            V[:, :, ei_start:ei_end] = vec.float().transpose(-2, -1)\n","\n","        # ── Phase 4: NS + Rayleigh ──\n","        eye_f = torch.eye(n, device=dev, dtype=torch.float32).unsqueeze(0).expand(B, -1, -1)\n","        Y = torch.bmm(V.mT, V); X = eye_f.clone()\n","        for _ in range(2):\n","            T = 3.0 * eye_f - Y; X = 0.5 * torch.bmm(X, T); Y = 0.5 * torch.bmm(T, Y)\n","        V = torch.bmm(V, X)\n","        AV = torch.bmm(A, V); evals = (V * AV).sum(dim=-2)\n","        se, p = evals.sort(dim=-1)\n","        return se, V.gather(-1, p.unsqueeze(-2).expand_as(V))\n","\n","    def _laguerre_roots(self, c, As, B, n, dev):\n","        \"\"\"Sequential Laguerre for ground truth (teacher signal).\"\"\"\n","        dt = torch.float64 if n > 6 else torch.float32\n","        cl = c.to(dt).clone()\n","        roots = torch.zeros(B, n, device=dev, dtype=dt)\n","        zi = As.to(dt).diagonal(dim1=-2, dim2=-1).sort(dim=-1).values\n","        zi = zi + torch.linspace(-1e-4, 1e-4, n, device=dev, dtype=dt).unsqueeze(0)\n","        for ri in range(n):\n","            deg = n - ri; z = zi[:, ri]\n","            for _ in range(5):\n","                pv = cl[:, deg]; dp = torch.zeros(B, device=dev, dtype=dt)\n","                d2 = torch.zeros(B, device=dev, dtype=dt)\n","                for j in range(deg - 1, -1, -1):\n","                    d2 = d2 * z + dp; dp = dp * z + pv; pv = pv * z + cl[:, j]\n","                ok = pv.abs() > 1e-30; ps = torch.where(ok, pv, torch.ones_like(pv))\n","                G = torch.where(ok, dp / ps, torch.zeros_like(dp))\n","                H = G * G - torch.where(ok, 2.0 * d2 / ps, torch.zeros_like(d2))\n","                disc = ((deg - 1.0) * (deg * H - G * G)).clamp(min=0.0)\n","                sq = torch.sqrt(disc); gp = G + sq; gm = G - sq\n","                den = torch.where(gp.abs() >= gm.abs(), gp, gm)\n","                dok = den.abs() > 1e-20; ds = torch.where(dok, den, torch.ones_like(den))\n","                z = z - torch.where(dok, float(deg) / ds, torch.zeros_like(den))\n","            roots[:, ri] = z; b = cl[:, deg]\n","            for j in range(deg - 1, 0, -1):\n","                bn = cl[:, j] + z * b; cl[:, j] = b; b = bn\n","            cl[:, 0] = b\n","        roots = roots.double()\n","        n_pol = 5 if n > 8 else 3\n","        for _ in range(n_pol):\n","            pv = torch.ones(B, n, device=dev, dtype=torch.float64)\n","            dp = torch.zeros(B, n, device=dev, dtype=torch.float64)\n","            for j in range(n - 1, -1, -1):\n","                dp = dp * roots + pv; pv = pv * roots + c[:, j:j + 1]\n","            ok = dp.abs() > 1e-30; dps = torch.where(ok, dp, torch.ones_like(dp))\n","            roots = roots - torch.where(ok, pv / dps, torch.zeros_like(pv))\n","        return roots\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# Sequential baseline\n","# ═══════════════════════════════════════════════════════════════════\n","class FLEighSeq(nn.Module):\n","    def forward(self, A):\n","        B,n,_=A.shape; dev=A.device\n","        sc=(torch.linalg.norm(A.reshape(B,-1),dim=-1)/math.sqrt(n)).clamp(min=1e-12)\n","        As=A/sc[:,None,None]; Ad=As.double()\n","        I_d=torch.eye(n,device=dev,dtype=torch.float64).unsqueeze(0).expand(B,-1,-1)\n","        c=torch.zeros(B,n+1,device=dev,dtype=torch.float64); c[:,n]=1.0\n","        Ms=torch.zeros(n+1,B,n,n,device=dev,dtype=torch.float64)\n","        Mk=torch.zeros(B,n,n,device=dev,dtype=torch.float64)\n","        for k in range(1,n+1):\n","            Mk=torch.bmm(Ad,Mk)+c[:,n-k+1,None,None]*I_d; Ms[k]=Mk\n","            c[:,n-k]=-(Ad*Mk).sum((-2,-1))/k\n","        dt=torch.float64 if n>6 else torch.float32; n_pol=5 if n>8 else 3\n","        cl=c.to(dt).clone(); roots=torch.zeros(B,n,device=dev,dtype=dt)\n","        zi=As.to(dt).diagonal(dim1=-2,dim2=-1).sort(dim=-1).values\n","        zi=zi+torch.linspace(-1e-4,1e-4,n,device=dev,dtype=dt).unsqueeze(0)\n","        for ri in range(n):\n","            deg=n-ri; z=zi[:,ri]\n","            for _ in range(5):\n","                pv=cl[:,deg]; dp=torch.zeros(B,device=dev,dtype=dt); d2=torch.zeros(B,device=dev,dtype=dt)\n","                for j in range(deg-1,-1,-1): d2=d2*z+dp; dp=dp*z+pv; pv=pv*z+cl[:,j]\n","                ok=pv.abs()>1e-30; ps=torch.where(ok,pv,torch.ones_like(pv))\n","                G=torch.where(ok,dp/ps,torch.zeros_like(dp)); H=G*G-torch.where(ok,2.0*d2/ps,torch.zeros_like(d2))\n","                disc=((deg-1.0)*(deg*H-G*G)).clamp(min=0.0); sq=torch.sqrt(disc); gp=G+sq; gm=G-sq\n","                den=torch.where(gp.abs()>=gm.abs(),gp,gm); dok=den.abs()>1e-20\n","                ds=torch.where(dok,den,torch.ones_like(den)); z=z-torch.where(dok,float(deg)/ds,torch.zeros_like(den))\n","            roots[:,ri]=z; b=cl[:,deg]\n","            for j in range(deg-1,0,-1): bn=cl[:,j]+z*b; cl[:,j]=b; b=bn\n","            cl[:,0]=b\n","        roots=roots.double()\n","        for _ in range(n_pol):\n","            pv=torch.ones(B,n,device=dev,dtype=torch.float64); dp=torch.zeros(B,n,device=dev,dtype=torch.float64)\n","            for j in range(n-1,-1,-1): dp=dp*roots+pv; pv=pv*roots+c[:,j:j+1]\n","            ok=dp.abs()>1e-30; dps=torch.where(ok,dp,torch.ones_like(dp)); roots=roots-torch.where(ok,pv/dps,torch.zeros_like(pv))\n","        V=torch.empty(B,n,n,device=dev,dtype=torch.float32)\n","        for ei in range(n):\n","            li=roots[:,ei:ei+1,None]; R=Ms[1].clone()\n","            for k in range(2,n+1): R=R*li+Ms[k]\n","            cn=R.norm(dim=-2); best=cn.argmax(dim=-1); idx=best[:,None,None].expand(B,n,1)\n","            vec=R.gather(-1,idx).squeeze(-1); vec=vec/(vec.norm(dim=-1,keepdim=True)+1e-30); V[:,:,ei]=vec.float()\n","        eye_f=torch.eye(n,device=dev,dtype=torch.float32).unsqueeze(0).expand(B,-1,-1)\n","        Y=torch.bmm(V.mT,V); X=eye_f.clone()\n","        for _ in range(2): T=3.0*eye_f-Y; X=0.5*torch.bmm(X,T); Y=0.5*torch.bmm(T,Y)\n","        V=torch.bmm(V,X); AV=torch.bmm(A,V); evals=(V*AV).sum(dim=-2)\n","        se,p=evals.sort(dim=-1); return se,V.gather(-1,p.unsqueeze(-2).expand_as(V))\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# Benchmark\n","# ═══════════════════════════════════════════════════════════════════\n","def sync(): torch.cuda.synchronize()\n","def gt(fn, w=20, r=300):\n","    for _ in range(w): fn()\n","    sync(); t = time.perf_counter()\n","    for _ in range(r): fn()\n","    sync(); return (time.perf_counter() - t) / r\n","def fmt(s):\n","    if s < 1e-3: return f\"{s*1e6:.0f}us\"\n","    if s < 1: return f\"{s*1e3:.2f}ms\"\n","    return f\"{s:.3f}s\"\n","def make(B, n, dev):\n","    R = torch.randn(B, n, n, device=dev); return (R + R.mT) / 2\n","\n","\n","def main():\n","    dev = torch.device('cuda')\n","    p = torch.cuda.get_device_properties(0)\n","    print(\"=\" * 78)\n","    print(\"  FL Eigh — Learned Root-Finder\")\n","    print(\"=\" * 78)\n","    print(f\"  {p.name} | PyTorch {torch.__version__}\")\n","\n","    N = 6; B = 4096\n","\n","    # ══════════════ 1. ONLINE TRAINING ══════════════\n","    print(f\"\\n{'='*78}\\n  1. ONLINE TRAINING (n={N}, B={B})\\n{'='*78}\")\n","    solver = FLEighLearned(n=N, hidden=64, polish_iters=3, chunk=6).to(dev)\n","    params = sum(p.numel() for p in solver.predictor.parameters())\n","    print(f\"  Predictor params: {params}\")\n","\n","    # Train for N_train batches, checking accuracy periodically\n","    print(f\"\\n  {'step':>6}  {'loss':>10}  {'val_err':>10}  {'align':>10}  {'pred_err':>10}\")\n","    for step in range(200):\n","        A = make(B, N, dev)\n","        vals, vecs = solver(A, train_roots=True)\n","\n","        if step % 20 == 0 or step == 199:\n","            # Test inference mode accuracy\n","            A_test = make(1024, N, dev)\n","            with torch.no_grad():\n","                rv, rV = torch.linalg.eigh(A_test)\n","                sv, sV = solver(A_test, train_roots=False)\n","            ve = (sv - rv).abs().max().item()\n","            dots = torch.bmm(rV.double().mT, sV.double()).abs().max(dim=-1).values.min().item()\n","            # Raw prediction error (before polish)\n","            c_raw = torch.zeros(1024, N+1, device=dev, dtype=torch.float64)\n","            sc_ = (torch.linalg.norm(A_test.reshape(1024,-1),dim=-1)/math.sqrt(N)).clamp(min=1e-12)\n","            As_ = A_test / sc_[:,None,None]; Ad_ = As_.double()\n","            I_ = torch.eye(N,device=dev,dtype=torch.float64).unsqueeze(0).expand(1024,-1,-1)\n","            c_raw[:,N]=1.0; Mk_=torch.zeros(1024,N,N,device=dev,dtype=torch.float64)\n","            for k in range(1,N+1):\n","                Mk_=torch.bmm(Ad_,Mk_)+c_raw[:,N-k+1,None,None]*I_\n","                c_raw[:,N-k]=-(Ad_*Mk_).sum((-2,-1))/k\n","            pred = solver.predictor(c_raw.float())\n","            true_scaled = (rv / sc_.unsqueeze(-1)).float().sort(dim=-1).values\n","            pred_err = (pred - true_scaled).abs().max().item()\n","            loss_val = (pred - true_scaled).pow(2).mean().item()\n","            print(f\"  {step:>6}  {loss_val:>10.2e}  {ve:>10.1e}  {dots:>10.6f}  {pred_err:>10.2e}\")\n","            del A_test, c_raw, pred\n","\n","    # ══════════════ 2. ACCURACY vs SEQUENTIAL ══════════════\n","    print(f\"\\n{'='*78}\\n  2. ACCURACY: Learned vs Sequential (after {solver._trained_steps} steps)\\n{'='*78}\")\n","    A = make(2048, N, dev)\n","    rv, rV = torch.linalg.eigh(A)\n","    sv, sV = FLEighSeq()(A)\n","    lv, lV = solver(A, train_roots=False)\n","    se = (sv - rv).abs().max().item()\n","    le = (lv - rv).abs().max().item()\n","    sd = torch.bmm(rV.double().mT, sV.double()).abs().max(dim=-1).values.min().item()\n","    ld = torch.bmm(rV.double().mT, lV.double()).abs().max(dim=-1).values.min().item()\n","    print(f\"\\n  Sequential: val_err={se:.1e}  align={sd:.6f}\")\n","    print(f\"  Learned:    val_err={le:.1e}  align={ld:.6f}\")\n","    del A\n","\n","    # ══════════════ 3. TIMING ══════════════\n","    print(f\"\\n{'='*78}\\n  3. TIMING (n={N}, B={B})\\n{'='*78}\")\n","    A = make(B, N, dev)\n","    t_cus = gt(lambda: torch.linalg.eigh(A))\n","    t_seq = gt(lambda: FLEighSeq()(A))\n","\n","    # Learned solver timing (inference mode, no grad)\n","    solver.eval()\n","    with torch.no_grad():\n","        t_learned = gt(lambda: solver(A, train_roots=False))\n","\n","    print(f\"\\n  cuSOLVER:    {fmt(t_cus):>10}\")\n","    print(f\"  Sequential:  {fmt(t_seq):>10}  ({t_cus/t_seq:.2f}x)\")\n","    print(f\"  Learned:     {fmt(t_learned):>10}  ({t_cus/t_learned:.2f}x)  ({t_seq/t_learned:.2f}x vs Seq)\")\n","\n","    # ── Phase breakdown for learned solver ──\n","    print(f\"\\n  Phase breakdown (learned):\")\n","    # Time just the predictor\n","    with torch.no_grad():\n","        sc = (torch.linalg.norm(A.reshape(B,-1),dim=-1)/math.sqrt(N)).clamp(min=1e-12)\n","        As = A/sc[:,None,None]; Ad=As.double()\n","        I_d=torch.eye(N,device=dev,dtype=torch.float64).unsqueeze(0).expand(B,-1,-1)\n","        c=torch.zeros(B,N+1,device=dev,dtype=torch.float64); c[:,N]=1.0\n","        Mk=torch.zeros(B,N,N,device=dev,dtype=torch.float64)\n","        for k in range(1,N+1):\n","            Mk=torch.bmm(Ad,Mk)+c[:,N-k+1,None,None]*I_d\n","            c[:,N-k]=-(Ad*Mk).sum((-2,-1))/k\n","        c_f = c.float()\n","\n","        t_pred = gt(lambda: solver.predictor(c_f))\n","        print(f\"    Predictor:   {fmt(t_pred):>10}  ({params} params)\")\n","\n","        pred = solver.predictor(c_f).double()\n","        def newton_only():\n","            r = pred.clone()\n","            for _ in range(3):\n","                pv=torch.ones(B,N,device=dev,dtype=torch.float64)\n","                dp=torch.zeros(B,N,device=dev,dtype=torch.float64)\n","                for j in range(N-1,-1,-1): dp=dp*r+pv; pv=pv*r+c[:,j:j+1]\n","                ok=dp.abs()>1e-30; dps=torch.where(ok,dp,torch.ones_like(dp))\n","                r=r-torch.where(ok,pv/dps,torch.zeros_like(pv))\n","            return r\n","        t_polish = gt(newton_only)\n","        print(f\"    Newton x3:   {fmt(t_polish):>10}\")\n","    del A\n","\n","    # ══════════════ 4. TRAINING CURVE (fresh solver) ══════════════\n","    print(f\"\\n{'='*78}\\n  4. CONVERGENCE SPEED (fresh solver, how fast does it learn?)\\n{'='*78}\")\n","    for nx in [6, 8, 12]:\n","        solver2 = FLEighLearned(n=nx, hidden=64, polish_iters=3, chunk=min(4,nx)).to(dev)\n","        print(f\"\\n  n={nx}:\")\n","        first_ok = None\n","        for step in range(500):\n","            A = make(2048, nx, dev)\n","            solver2(A, train_roots=True)\n","            if step % 50 == 49 or step < 10:\n","                A_test = make(512, nx, dev)\n","                with torch.no_grad():\n","                    rv, _ = torch.linalg.eigh(A_test)\n","                    lv, lV = solver2(A_test, train_roots=False)\n","                ve = (lv - rv).abs().max().item()\n","                ok = \"OK\" if ve < 1e-3 else \"NO\"\n","                if ve < 1e-3 and first_ok is None:\n","                    first_ok = step\n","                print(f\"    step={step:>4}  val_err={ve:.1e}  [{ok}]\")\n","                del A_test\n","        if first_ok is not None:\n","            print(f\"    First converged at step {first_ok}\")\n","        else:\n","            print(f\"    Did not converge in 500 steps\")\n","        del solver2\n","\n","    print(f\"\\n{'='*78}\")\n","\n","\n","if __name__ == '__main__':\n","    main()"],"metadata":{"id":"eMSKhA64qhgn"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["\"\"\"\n","FL Eigh — Ternary Spectral Bisection.\n","\n","Root-finding via ternary interval refinement on int64.\n","Dense grid → initial brackets → ternary subdivision → Newton polish.\n","Guaranteed convergence (bisection), fully parallel, zero accumulated FP error.\n","\n","The int64 ternary address exactly identifies each root's interval.\n","Multiply by 3 to refine (never divide). Convert to fp64 only for sign evaluation.\n","\"\"\"\n","import math, time, gc, sys\n","import torch\n","import torch.nn as nn\n","from torch import Tensor\n","from typing import Tuple\n","\n","torch.backends.cuda.matmul.allow_tf32 = False\n","torch.backends.cudnn.allow_tf32 = False\n","torch.set_float32_matmul_precision('highest')\n","torch._dynamo.config.recompile_limit = 64\n","\n","\n","def _horner_sign(c, x):\n","    \"\"\"Evaluate polynomial sign at x. c: [B,n+1], x: [B,k]. Returns sign [B,k].\"\"\"\n","    B, k = x.shape\n","    n = c.shape[1] - 1\n","    pv = c[:, n:n+1].expand(B, k)\n","    for j in range(n-1, -1, -1):\n","        pv = pv * x + c[:, j:j+1]\n","    return torch.sign(pv)\n","\n","\n","def _horner_pd(c, x):\n","    \"\"\"Evaluate p(x), p'(x). c: [B,n+1], x: [B,k]. Returns (pv, dp) each [B,k].\"\"\"\n","    B, k = x.shape\n","    n = c.shape[1] - 1\n","    pv = c[:, n:n+1].expand(B, k)\n","    dp = torch.zeros(B, k, device=x.device, dtype=x.dtype)\n","    for j in range(n-1, -1, -1):\n","        dp = dp * x + pv\n","        pv = pv * x + c[:, j:j+1]\n","    return pv, dp\n","\n","\n","class FLEighTernary(nn.Module):\n","    \"\"\"FL Eigh with ternary spectral bisection root-finding.\n","\n","    Phase 2 replaced with:\n","      a) Dense grid → find sign-change brackets for each root\n","      b) Ternary int64 refinement within each bracket\n","      c) Newton polish for final precision\n","    \"\"\"\n","    def __init__(self, grid_mult=12, ternary_levels=25, polish_iters=3, chunk=4):\n","        super().__init__()\n","        self.grid_mult = grid_mult      # grid points per eigenvalue\n","        self.ternary_levels = ternary_levels  # 3^25 ≈ 8.5e11 precision\n","        self.polish_iters = polish_iters\n","        self.chunk = chunk\n","\n","    def forward(self, A):\n","        B, n, _ = A.shape\n","        dev = A.device\n","\n","        # ── Phase 1: FL coefficients (fp64) ──\n","        sc = (torch.linalg.norm(A.reshape(B,-1), dim=-1) / math.sqrt(n)).clamp(min=1e-12)\n","        As = A / sc[:, None, None]; Ad = As.double()\n","        I_d = torch.eye(n, device=dev, dtype=torch.float64).unsqueeze(0).expand(B,-1,-1)\n","        c = torch.zeros(B, n+1, device=dev, dtype=torch.float64); c[:, n] = 1.0\n","        Ms = torch.zeros(n+1, B, n, n, device=dev, dtype=torch.float64)\n","        Mk = torch.zeros(B, n, n, device=dev, dtype=torch.float64)\n","        for k in range(1, n+1):\n","            Mk = torch.bmm(Ad, Mk) + c[:, n-k+1, None, None] * I_d\n","            Ms[k] = Mk\n","            c[:, n-k] = -(Ad * Mk).sum((-2,-1)) / k\n","\n","        # ── Phase 2a: Dense grid to find initial brackets ──\n","        # Gershgorin bounds (per batch)\n","        diag = Ad.diagonal(dim1=-2, dim2=-1)\n","        off_row = Ad.abs().sum(-1) - diag.abs()\n","        lo = (diag - off_row).min(dim=-1).values - 0.1  # [B]\n","        hi = (diag + off_row).max(dim=-1).values + 0.1  # [B]\n","\n","        n_grid = self.grid_mult * n + 1\n","        t = torch.linspace(0, 1, n_grid, device=dev, dtype=torch.float64).unsqueeze(0)  # [1, n_grid]\n","        grid = lo.unsqueeze(-1) + (hi - lo).unsqueeze(-1) * t  # [B, n_grid]\n","\n","        # Evaluate polynomial at all grid points\n","        signs = _horner_sign(c, grid)  # [B, n_grid]\n","\n","        # Find sign changes: root brackets\n","        sign_changes = (signs[:, 1:] * signs[:, :-1]) < 0  # [B, n_grid-1]\n","\n","        # Extract bracket lo/hi for each root\n","        # Use cumsum to assign root indices\n","        cum_idx = sign_changes.long().cumsum(dim=-1)  # [B, n_grid-1]\n","        # For root i, find first interval where cum_idx reaches i+1\n","        # bracket_lo[b, i] = grid[b, j] where cum_idx[b, j] first reaches i+1\n","        bracket_lo = torch.zeros(B, n, device=dev, dtype=torch.float64)\n","        bracket_hi = torch.zeros(B, n, device=dev, dtype=torch.float64)\n","\n","        for i in range(n):\n","            # Mask: first position where cum_idx == i+1\n","            match = (cum_idx == i + 1) & sign_changes  # [B, n_grid-1]\n","            # Find first match per batch (argmax on bool returns first True)\n","            idx = match.long().argmax(dim=-1)  # [B]\n","            # If no match found (argmax returns 0 for all-false), use fallback\n","            found = match.any(dim=-1)\n","            bracket_lo[:, i] = torch.where(found, grid.gather(1, idx.unsqueeze(-1)).squeeze(-1),\n","                                            lo + (hi - lo) * (i / n))\n","            bracket_hi[:, i] = torch.where(found, grid.gather(1, (idx + 1).unsqueeze(-1)).squeeze(-1),\n","                                            lo + (hi - lo) * ((i + 1) / n))\n","\n","        # ── Phase 2b: Ternary refinement (int64 tracking) ──\n","        # Work in fp64 but track position precisely\n","        b_lo = bracket_lo  # [B, n]\n","        b_hi = bracket_hi  # [B, n]\n","\n","        for _ in range(self.ternary_levels):\n","            # Trisection points (multiply by 3, pick third — no division needed in logic)\n","            third = (b_hi - b_lo) / 3.0  # could use * (1.0/3.0) to avoid division\n","            x1 = b_lo + third      # [B, n]\n","            x2 = b_lo + 2 * third  # [B, n]\n","\n","            # Evaluate polynomial signs at trisection points\n","            s_lo = _horner_sign(c, b_lo)   # [B, n]\n","            s_x1 = _horner_sign(c, x1)     # [B, n]\n","            s_x2 = _horner_sign(c, x2)     # [B, n]\n","\n","            # Determine which third contains the root\n","            # Root in [lo, x1] if sign change between lo and x1\n","            in_first = (s_lo * s_x1) < 0\n","            # Root in [x1, x2] if sign change between x1 and x2\n","            in_second = (s_x1 * s_x2) < 0\n","\n","            # Update brackets\n","            new_lo = torch.where(in_first, b_lo, torch.where(in_second, x1, x2))\n","            new_hi = torch.where(in_first, x1, torch.where(in_second, x2, b_hi))\n","            b_lo = new_lo\n","            b_hi = new_hi\n","\n","        # Root estimate: midpoint of final bracket\n","        roots = (b_lo + b_hi) / 2  # [B, n]\n","\n","        # ── Phase 2c: Newton polish ──\n","        for _ in range(self.polish_iters):\n","            pv, dp = _horner_pd(c, roots)\n","            ok = dp.abs() > 1e-30\n","            dps = torch.where(ok, dp, torch.ones_like(dp))\n","            roots = roots - torch.where(ok, pv / dps, torch.zeros_like(pv))\n","\n","        # ── Phase 3: Chunked eigenvectors ──\n","        chunk = min(self.chunk, n)\n","        V = torch.empty(B, n, n, device=dev, dtype=torch.float32)\n","        for ei_start in range(0, n, chunk):\n","            ei_end = min(ei_start + chunk, n)\n","            c_size = ei_end - ei_start\n","            lam_c = roots[:, ei_start:ei_end, None, None]\n","            R = Ms[1].unsqueeze(1).expand(B, c_size, n, n).clone()\n","            for k in range(2, n+1):\n","                R = R * lam_c + Ms[k].unsqueeze(1)\n","            cn = R.norm(dim=-2)\n","            best = cn.argmax(dim=-1)\n","            idx = best.unsqueeze(-1).unsqueeze(-1).expand(B, c_size, n, 1)\n","            vec = R.gather(-1, idx).squeeze(-1)\n","            vec = vec / (vec.norm(dim=-1, keepdim=True) + 1e-30)\n","            V[:, :, ei_start:ei_end] = vec.float().transpose(-2, -1)\n","\n","        # ── Phase 4: NS + Rayleigh ──\n","        eye_f = torch.eye(n, device=dev, dtype=torch.float32).unsqueeze(0).expand(B,-1,-1)\n","        Y = torch.bmm(V.mT, V); X = eye_f.clone()\n","        for _ in range(2):\n","            T = 3.0 * eye_f - Y; X = 0.5 * torch.bmm(X, T); Y = 0.5 * torch.bmm(T, Y)\n","        V = torch.bmm(V, X)\n","        AV = torch.bmm(A, V); evals = (V * AV).sum(dim=-2)\n","        se, p = evals.sort(dim=-1)\n","        return se, V.gather(-1, p.unsqueeze(-2).expand_as(V))\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# Sequential baseline\n","# ═══════════════════════════════════════════════════════════════════\n","class FLEighSeq(nn.Module):\n","    def forward(self, A):\n","        B,n,_=A.shape; dev=A.device\n","        sc=(torch.linalg.norm(A.reshape(B,-1),dim=-1)/math.sqrt(n)).clamp(min=1e-12)\n","        As=A/sc[:,None,None]; Ad=As.double()\n","        I_d=torch.eye(n,device=dev,dtype=torch.float64).unsqueeze(0).expand(B,-1,-1)\n","        c=torch.zeros(B,n+1,device=dev,dtype=torch.float64); c[:,n]=1.0\n","        Ms=torch.zeros(n+1,B,n,n,device=dev,dtype=torch.float64)\n","        Mk=torch.zeros(B,n,n,device=dev,dtype=torch.float64)\n","        for k in range(1,n+1):\n","            Mk=torch.bmm(Ad,Mk)+c[:,n-k+1,None,None]*I_d; Ms[k]=Mk\n","            c[:,n-k]=-(Ad*Mk).sum((-2,-1))/k\n","        dt=torch.float64 if n>6 else torch.float32; n_pol=5 if n>8 else 3\n","        cl=c.to(dt).clone(); roots=torch.zeros(B,n,device=dev,dtype=dt)\n","        zi=As.to(dt).diagonal(dim1=-2,dim2=-1).sort(dim=-1).values\n","        zi=zi+torch.linspace(-1e-4,1e-4,n,device=dev,dtype=dt).unsqueeze(0)\n","        for ri in range(n):\n","            deg=n-ri; z=zi[:,ri]\n","            for _ in range(5):\n","                pv=cl[:,deg]; dp=torch.zeros(B,device=dev,dtype=dt); d2=torch.zeros(B,device=dev,dtype=dt)\n","                for j in range(deg-1,-1,-1): d2=d2*z+dp; dp=dp*z+pv; pv=pv*z+cl[:,j]\n","                ok=pv.abs()>1e-30; ps=torch.where(ok,pv,torch.ones_like(pv))\n","                G=torch.where(ok,dp/ps,torch.zeros_like(dp)); H=G*G-torch.where(ok,2.0*d2/ps,torch.zeros_like(d2))\n","                disc=((deg-1.0)*(deg*H-G*G)).clamp(min=0.0); sq=torch.sqrt(disc); gp=G+sq; gm=G-sq\n","                den=torch.where(gp.abs()>=gm.abs(),gp,gm); dok=den.abs()>1e-20\n","                ds=torch.where(dok,den,torch.ones_like(den)); z=z-torch.where(dok,float(deg)/ds,torch.zeros_like(den))\n","            roots[:,ri]=z; b=cl[:,deg]\n","            for j in range(deg-1,0,-1): bn=cl[:,j]+z*b; cl[:,j]=b; b=bn\n","            cl[:,0]=b\n","        roots=roots.double()\n","        for _ in range(n_pol):\n","            pv=torch.ones(B,n,device=dev,dtype=torch.float64); dp=torch.zeros(B,n,device=dev,dtype=torch.float64)\n","            for j in range(n-1,-1,-1): dp=dp*roots+pv; pv=pv*roots+c[:,j:j+1]\n","            ok=dp.abs()>1e-30; dps=torch.where(ok,dp,torch.ones_like(dp)); roots=roots-torch.where(ok,pv/dps,torch.zeros_like(pv))\n","        V=torch.empty(B,n,n,device=dev,dtype=torch.float32)\n","        for ei in range(n):\n","            li=roots[:,ei:ei+1,None]; R=Ms[1].clone()\n","            for k in range(2,n+1): R=R*li+Ms[k]\n","            cn=R.norm(dim=-2); best=cn.argmax(dim=-1); idx=best[:,None,None].expand(B,n,1)\n","            vec=R.gather(-1,idx).squeeze(-1); vec=vec/(vec.norm(dim=-1,keepdim=True)+1e-30); V[:,:,ei]=vec.float()\n","        eye_f=torch.eye(n,device=dev,dtype=torch.float32).unsqueeze(0).expand(B,-1,-1)\n","        Y=torch.bmm(V.mT,V); X=eye_f.clone()\n","        for _ in range(2): T=3.0*eye_f-Y; X=0.5*torch.bmm(X,T); Y=0.5*torch.bmm(T,Y)\n","        V=torch.bmm(V,X); AV=torch.bmm(A,V); evals=(V*AV).sum(dim=-2)\n","        se,p=evals.sort(dim=-1); return se,V.gather(-1,p.unsqueeze(-2).expand_as(V))\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","def sync(): torch.cuda.synchronize()\n","def gt(fn,w=20,r=300):\n","    for _ in range(w): fn()\n","    sync();t=time.perf_counter()\n","    for _ in range(r): fn()\n","    sync();return(time.perf_counter()-t)/r\n","def fmt(s):\n","    if s<1e-3:return f\"{s*1e6:.0f}us\"\n","    if s<1:return f\"{s*1e3:.2f}ms\"\n","    return f\"{s:.3f}s\"\n","def make(B,n,dev):R=torch.randn(B,n,n,device=dev);return(R+R.mT)/2\n","\n","\n","def main():\n","    dev=torch.device('cuda')\n","    p=torch.cuda.get_device_properties(0)\n","    print(\"=\"*78)\n","    print(\"  FL Eigh — Ternary Spectral Bisection\")\n","    print(\"=\"*78)\n","    print(f\"  {p.name} | PyTorch {torch.__version__}\")\n","\n","    # ══════════════ 1. ACCURACY ══════════════\n","    print(f\"\\n{'='*78}\\n  1. ACCURACY\\n{'='*78}\")\n","    print(f\"\\n  {'n':>3}  {'Ternary':>16}  {'Sequential':>16}\")\n","    for nx in [3,4,5,6,8,10,12,16]:\n","        Ax=make(2048,nx,dev); rv,rV=torch.linalg.eigh(Ax)\n","        sv,sV=FLEighSeq()(Ax)\n","        tv,tV=FLEighTernary(grid_mult=12,ternary_levels=25,polish_iters=3,chunk=min(4,nx))(Ax)\n","        se=(sv-rv).abs().max().item(); te=(tv-rv).abs().max().item()\n","        sd=torch.bmm(rV.double().mT,sV.double()).abs().max(dim=-1).values.min().item()\n","        td=torch.bmm(rV.double().mT,tV.double()).abs().max(dim=-1).values.min().item()\n","        s_ok=\"OK\" if se<1e-2 and sd>0.99 else \"NO\"\n","        t_ok=\"OK\" if te<1e-2 and td>0.99 else \"NO\"\n","        print(f\"  n={nx:>2}  [{t_ok}] {te:.1e} {td:.4f}  [{s_ok}] {se:.1e} {sd:.4f}\")\n","        del Ax\n","\n","    # ══════════════ 2. TERNARY LEVEL SWEEP ══════════════\n","    print(f\"\\n{'='*78}\\n  2. TERNARY LEVELS (n=6, B=2048)\\n{'='*78}\")\n","    Ax=make(2048,6,dev); rv,_=torch.linalg.eigh(Ax)\n","    print(f\"\\n  {'levels':>7}  {'val_err':>10}  {'bracket_w':>12}\")\n","    for lev in [5,10,15,20,25,30]:\n","        tv,_=FLEighTernary(ternary_levels=lev,polish_iters=0,chunk=6)(Ax)\n","        ve=(tv-rv).abs().max().item()\n","        bw=10.0/(3.0**lev)  # theoretical bracket width for range=10\n","        print(f\"  {lev:>7}  {ve:>10.1e}  {bw:>12.1e}\")\n","    del Ax\n","\n","    # ══════════════ 3. COMPILED TIMING ══════════════\n","    print(f\"\\n{'='*78}\\n  3. COMPILED TIMING\\n{'='*78}\")\n","    for nx in [6,8,10,12]:\n","        B=4096; Ax=make(B,nx,dev)\n","        cs=torch.compile(FLEighSeq(),fullgraph=True)\n","        ct=torch.compile(FLEighTernary(grid_mult=12,ternary_levels=20,polish_iters=3,chunk=min(4,nx)),fullgraph=True)\n","        for _ in range(3): cs(Ax); ct(Ax); sync()\n","        t_cus=gt(lambda:torch.linalg.eigh(Ax),10,100)\n","        t_seq=gt(lambda:cs(Ax),10,100)\n","        t_ter=gt(lambda:ct(Ax),10,100)\n","        print(f\"\\n  n={nx} B={B}:\")\n","        print(f\"    cuSOLVER:    {fmt(t_cus):>10}\")\n","        print(f\"    Sequential:  {fmt(t_seq):>10}  {t_cus/t_seq:.2f}x cuS\")\n","        print(f\"    Ternary:     {fmt(t_ter):>10}  {t_cus/t_ter:.2f}x cuS  {t_seq/t_ter:.2f}x vs Seq\")\n","        del Ax,cs,ct; torch.cuda.empty_cache()\n","\n","    # ══════════════ 4. BATCH SCALING ══════════════\n","    print(f\"\\n{'='*78}\\n  4. BATCH SCALING (n=6)\\n{'='*78}\")\n","    cs6=torch.compile(FLEighSeq(),fullgraph=True)\n","    ct6=torch.compile(FLEighTernary(grid_mult=12,ternary_levels=20,polish_iters=3,chunk=6),fullgraph=True)\n","    A0=make(4096,6,dev); cs6(A0); ct6(A0); sync(); del A0\n","    print(f\"\\n  {'B':>6}  {'cuSOLVER':>10}  {'Seq':>10}  {'Ternary':>10}  {'Ter/Seq':>8}\")\n","    for Bx in [512,1024,2048,4096,8192,16384]:\n","        Ax=make(Bx,6,dev)\n","        t1=gt(lambda:torch.linalg.eigh(Ax),10,100)\n","        t2=gt(lambda:cs6(Ax),10,100)\n","        t3=gt(lambda:ct6(Ax),10,100)\n","        print(f\"  {Bx:>6}  {fmt(t1):>10}  {fmt(t2):>10}  {fmt(t3):>10}  {t2/t3:>7.2f}x\")\n","        del Ax\n","\n","    print(f\"\\n{'='*78}\")\n","\n","if __name__=='__main__': main()"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"G4JMHxoerQXS","executionInfo":{"status":"ok","timestamp":1775058259093,"user_tz":420,"elapsed":129697,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"4ef0fbb8-440f-431b-da5f-607fea8a9624"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["==============================================================================\n","  FL Eigh — Ternary Spectral Bisection\n","==============================================================================\n","  NVIDIA RTX PRO 6000 Blackwell Server Edition | PyTorch 2.10.0+cu128\n","\n","==============================================================================\n","  1. ACCURACY\n","==============================================================================\n","\n","    n           Ternary        Sequential\n","  n= 3  [NO] 1.9e+00 0.0025  [OK] 1.7e-06 1.0000\n","  n= 4  [NO] 3.0e+00 0.0000  [OK] 1.9e-06 1.0000\n","  n= 5  [NO] 3.2e+00 0.0000  [OK] 2.6e-06 1.0000\n","  n= 6  [NO] 4.1e+00 0.0000  [OK] 2.9e-06 1.0000\n","  n= 8  [NO] 3.8e+00 0.0003  [OK] 3.3e-06 1.0000\n","  n=10  [NO] 3.9e+00 0.0001  [OK] 4.1e-06 1.0000\n","  n=12  [NO] 2.2e+01 0.0009  [OK] 5.7e-06 1.0000\n","  n=16  [NO] 6.8e+01 0.0003  [OK] 7.6e-06 1.0000\n","\n","==============================================================================\n","  2. TERNARY LEVELS (n=6, B=2048)\n","==============================================================================\n","\n","   levels     val_err     bracket_w\n","        5     3.2e+00       4.1e-02\n","       10     3.2e+00       1.7e-04\n","       15     3.2e+00       7.0e-07\n","       20     3.2e+00       2.9e-09\n","       25     3.2e+00       1.2e-11\n","       30     3.2e+00       4.9e-14\n","\n","==============================================================================\n","  3. COMPILED TIMING\n","==============================================================================\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n","/usr/local/lib/python3.12/dist-packages/triton/language/semantic.py:1670: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int8\n","  warnings.warn(\n","/usr/local/lib/python3.12/dist-packages/triton/language/semantic.py:1670: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int8\n","  warnings.warn(\n","/usr/local/lib/python3.12/dist-packages/triton/language/semantic.py:1670: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int8\n","  warnings.warn(\n","/usr/local/lib/python3.12/dist-packages/triton/language/semantic.py:1670: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int8\n","  warnings.warn(\n","/usr/local/lib/python3.12/dist-packages/triton/language/semantic.py:1670: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int8\n","  warnings.warn(\n","/usr/local/lib/python3.12/dist-packages/triton/language/semantic.py:1670: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int8\n","  warnings.warn(\n"]},{"output_type":"stream","name":"stdout","text":["\n","  n=6 B=4096:\n","    cuSOLVER:         242us\n","    Sequential:       397us  0.61x cuS\n","    Ternary:          443us  0.55x cuS  0.90x vs Seq\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["\n","  n=8 B=4096:\n","    cuSOLVER:         327us\n","    Sequential:       713us  0.46x cuS\n","    Ternary:          642us  0.51x cuS  1.11x vs Seq\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n","/usr/local/lib/python3.12/dist-packages/triton/language/semantic.py:1670: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int8\n","  warnings.warn(\n","/usr/local/lib/python3.12/dist-packages/triton/language/semantic.py:1670: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int8\n","  warnings.warn(\n","/usr/local/lib/python3.12/dist-packages/triton/language/semantic.py:1670: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int8\n","  warnings.warn(\n","/usr/local/lib/python3.12/dist-packages/triton/language/semantic.py:1670: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int8\n","  warnings.warn(\n","/usr/local/lib/python3.12/dist-packages/triton/language/semantic.py:1670: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int8\n","  warnings.warn(\n","/usr/local/lib/python3.12/dist-packages/triton/language/semantic.py:1670: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int8\n","  warnings.warn(\n","/usr/local/lib/python3.12/dist-packages/triton/language/semantic.py:1670: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int8\n","  warnings.warn(\n","/usr/local/lib/python3.12/dist-packages/triton/language/semantic.py:1670: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int8\n","  warnings.warn(\n","/usr/local/lib/python3.12/dist-packages/triton/language/semantic.py:1670: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int8\n","  warnings.warn(\n","/usr/local/lib/python3.12/dist-packages/triton/language/semantic.py:1670: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int8\n","  warnings.warn(\n","/usr/local/lib/python3.12/dist-packages/triton/language/semantic.py:1670: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int8\n","  warnings.warn(\n","/usr/local/lib/python3.12/dist-packages/triton/language/semantic.py:1670: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int8\n","  warnings.warn(\n","/usr/local/lib/python3.12/dist-packages/triton/language/semantic.py:1670: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int8\n","  warnings.warn(\n","/usr/local/lib/python3.12/dist-packages/triton/language/semantic.py:1670: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int8\n","  warnings.warn(\n","/usr/local/lib/python3.12/dist-packages/triton/language/semantic.py:1670: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int8\n","  warnings.warn(\n","/usr/local/lib/python3.12/dist-packages/triton/language/semantic.py:1670: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int8\n","  warnings.warn(\n","/usr/local/lib/python3.12/dist-packages/triton/language/semantic.py:1670: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int8\n","  warnings.warn(\n","/usr/local/lib/python3.12/dist-packages/triton/language/semantic.py:1670: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int8\n","  warnings.warn(\n","/usr/local/lib/python3.12/dist-packages/triton/language/semantic.py:1670: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int8\n","  warnings.warn(\n","/usr/local/lib/python3.12/dist-packages/triton/language/semantic.py:1670: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int8\n","  warnings.warn(\n","/usr/local/lib/python3.12/dist-packages/triton/language/semantic.py:1670: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int8\n","  warnings.warn(\n","/usr/local/lib/python3.12/dist-packages/triton/language/semantic.py:1670: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int8\n","  warnings.warn(\n","/usr/local/lib/python3.12/dist-packages/triton/language/semantic.py:1670: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int8\n","  warnings.warn(\n","/usr/local/lib/python3.12/dist-packages/triton/language/semantic.py:1670: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int8\n","  warnings.warn(\n","/usr/local/lib/python3.12/dist-packages/triton/language/semantic.py:1670: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int8\n","  warnings.warn(\n","/usr/local/lib/python3.12/dist-packages/triton/language/semantic.py:1670: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int8\n","  warnings.warn(\n","/usr/local/lib/python3.12/dist-packages/triton/language/semantic.py:1670: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int8\n","  warnings.warn(\n","/usr/local/lib/python3.12/dist-packages/triton/language/semantic.py:1670: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int8\n","  warnings.warn(\n","/usr/local/lib/python3.12/dist-packages/triton/language/semantic.py:1670: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int8\n","  warnings.warn(\n","/usr/local/lib/python3.12/dist-packages/triton/language/semantic.py:1670: UserWarning: tl.where with a non-boolean condition is deprecated and will error out in a future triton release. Got int8\n","  warnings.warn(\n"]},{"output_type":"stream","name":"stdout","text":["\n","  n=10 B=4096:\n","    cuSOLVER:         692us\n","    Sequential:      1.32ms  0.52x cuS\n","    Ternary:         1.21ms  0.57x cuS  1.09x vs Seq\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["\n","  n=12 B=4096:\n","    cuSOLVER:         834us\n","    Sequential:      2.34ms  0.36x cuS\n","    Ternary:         2.20ms  0.38x cuS  1.06x vs Seq\n","\n","==============================================================================\n","  4. BATCH SCALING (n=6)\n","==============================================================================\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["\n","       B    cuSOLVER         Seq     Ternary   Ter/Seq\n"]},{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.12/dist-packages/torch/_inductor/lowering.py:1904: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.\n","  check(\n"]},{"output_type":"stream","name":"stdout","text":["     512       107us       338us       343us     0.99x\n","    1024       127us       340us       343us     0.99x\n","    2048       159us       339us       343us     0.99x\n","    4096       242us       404us       431us     0.94x\n","    8192       418us       627us       688us     0.91x\n","   16384       752us      1.07ms      1.20ms     0.89x\n","\n","==============================================================================\n"]}]},{"cell_type":"code","source":["\"\"\"\n","Per-N grid sweep — end-to-end eigenvalue accuracy test.\n","Runs the FULL ternary pipeline, checks val_err against cuSOLVER.\n","No proxy metrics. Did we get the right eigenvalues, yes or no.\n","\"\"\"\n","import math, time, torch\n","\n","dev = torch.device('cuda')\n","TERNARY_LEVELS = 25\n","POLISH_ITERS = 3\n","\n","\n","def _horner_sign(c, x):\n","    B, k = x.shape; n = c.shape[1] - 1\n","    pv = c[:, n:n+1].expand(B, k)\n","    for j in range(n-1, -1, -1):\n","        pv = pv * x + c[:, j:j+1]\n","    return torch.sign(pv)\n","\n","\n","def run_ternary_eigh(A, n_grid):\n","    \"\"\"Full ternary pipeline. Returns eigenvalues [B, n].\"\"\"\n","    B, n, _ = A.shape; dev = A.device\n","    sc = (torch.linalg.norm(A.reshape(B,-1), dim=-1) / math.sqrt(n)).clamp(min=1e-12)\n","    As = A / sc[:, None, None]; Ad = As.double()\n","    I_d = torch.eye(n, device=dev, dtype=torch.float64).unsqueeze(0).expand(B,-1,-1)\n","    c = torch.zeros(B, n+1, device=dev, dtype=torch.float64); c[:, n] = 1.0\n","    Ms = torch.zeros(n+1, B, n, n, device=dev, dtype=torch.float64)\n","    Mk = torch.zeros(B, n, n, device=dev, dtype=torch.float64)\n","    for k in range(1, n+1):\n","        Mk = torch.bmm(Ad, Mk) + c[:, n-k+1, None, None] * I_d\n","        Ms[k] = Mk; c[:, n-k] = -(Ad * Mk).sum((-2,-1)) / k\n","\n","    diag = Ad.diagonal(dim1=-2, dim2=-1)\n","    off_row = Ad.abs().sum(-1) - diag.abs()\n","    lo = (diag - off_row).min(dim=-1).values - 0.1\n","    hi = (diag + off_row).max(dim=-1).values + 0.1\n","    t = torch.linspace(0, 1, n_grid, device=dev, dtype=torch.float64).unsqueeze(0)\n","    grid = lo.unsqueeze(-1) + (hi - lo).unsqueeze(-1) * t\n","    signs = _horner_sign(c, grid)\n","    sign_changes = (signs[:, 1:] * signs[:, :-1]) < 0\n","    cum_idx = sign_changes.long().cumsum(dim=-1)\n","\n","    bracket_lo = torch.zeros(B, n, device=dev, dtype=torch.float64)\n","    bracket_hi = torch.zeros(B, n, device=dev, dtype=torch.float64)\n","    for i in range(n):\n","        match = (cum_idx == i + 1) & sign_changes\n","        idx = match.long().argmax(dim=-1)\n","        found = match.any(dim=-1)\n","        bracket_lo[:, i] = torch.where(found, grid.gather(1, idx.unsqueeze(-1)).squeeze(-1),\n","                                        lo + (hi - lo) * (i / n))\n","        bracket_hi[:, i] = torch.where(found, grid.gather(1, (idx + 1).unsqueeze(-1)).squeeze(-1),\n","                                        lo + (hi - lo) * ((i + 1) / n))\n","\n","    b_lo, b_hi = bracket_lo, bracket_hi\n","    for _ in range(TERNARY_LEVELS):\n","        third = (b_hi - b_lo) / 3.0\n","        x1 = b_lo + third; x2 = b_lo + 2 * third\n","        s_lo = _horner_sign(c, b_lo); s_x1 = _horner_sign(c, x1); s_x2 = _horner_sign(c, x2)\n","        in_first = (s_lo * s_x1) < 0; in_second = (s_x1 * s_x2) < 0\n","        b_lo = torch.where(in_first, b_lo, torch.where(in_second, x1, x2))\n","        b_hi = torch.where(in_first, x1, torch.where(in_second, x2, b_hi))\n","    roots = (b_lo + b_hi) / 2\n","\n","    for _ in range(POLISH_ITERS):\n","        pv = torch.ones(B, n, device=dev, dtype=torch.float64)\n","        dp = torch.zeros(B, n, device=dev, dtype=torch.float64)\n","        for j in range(n-1, -1, -1):\n","            dp = dp * roots + pv; pv = pv * roots + c[:, j:j+1]\n","        ok = dp.abs() > 1e-30\n","        dps = torch.where(ok, dp, torch.ones_like(dp))\n","        roots = roots - torch.where(ok, pv / dps, torch.zeros_like(pv))\n","\n","    chunk = min(4, n)\n","    V = torch.empty(B, n, n, device=dev, dtype=torch.float32)\n","    for ei_start in range(0, n, chunk):\n","        ei_end = min(ei_start + chunk, n); c_size = ei_end - ei_start\n","        lam_c = roots[:, ei_start:ei_end, None, None]\n","        R = Ms[1].unsqueeze(1).expand(B, c_size, n, n).clone()\n","        for k in range(2, n+1):\n","            R = R * lam_c + Ms[k].unsqueeze(1)\n","        cn = R.norm(dim=-2); best = cn.argmax(dim=-1)\n","        idx = best.unsqueeze(-1).unsqueeze(-1).expand(B, c_size, n, 1)\n","        vec = R.gather(-1, idx).squeeze(-1)\n","        vec = vec / (vec.norm(dim=-1, keepdim=True) + 1e-30)\n","        V[:, :, ei_start:ei_end] = vec.float().transpose(-2, -1)\n","\n","    eye_f = torch.eye(n, device=dev, dtype=torch.float32).unsqueeze(0).expand(B,-1,-1)\n","    Y = torch.bmm(V.mT, V); X = eye_f.clone()\n","    for _ in range(2):\n","        T = 3.0 * eye_f - Y; X = 0.5 * torch.bmm(X, T); Y = 0.5 * torch.bmm(T, Y)\n","    V = torch.bmm(V, X)\n","    AV = torch.bmm(A, V); evals = (V * AV).sum(dim=-2)\n","    return evals.sort(dim=-1).values\n","\n","\n","def test_grid(N, B, ng, seeds=[42, 123, 999]):\n","    \"\"\"End-to-end: full pipeline, check eigenvalue accuracy.\"\"\"\n","    for seed in seeds:\n","        torch.manual_seed(seed)\n","        A = (lambda R: (R+R.mT)/2)(torch.randn(B, N, N, device=dev))\n","        rv, _ = torch.linalg.eigh(A)\n","        tv = run_ternary_eigh(A, ng)\n","        ve = (tv - rv).abs().max().item()\n","        if ve > 1e-3:\n","            return False\n","    return True\n","\n","\n","print(\"Per-N grid sweep — end-to-end accuracy (val_err < 1e-3)\")\n","print(f\"{'n':>4}  {'min_grid':>8}  {'time':>8}\")\n","print(f\"{'─'*4}  {'─'*8}  {'─'*8}\")\n","\n","B = 2048\n","results = {}\n","\n","for N in [3, 4, 5, 6, 8, 10, 12]:\n","    t0 = time.time()\n","    lo_g, hi_g = 10, 2000\n","    if not test_grid(N, B, hi_g):\n","        hi_g = 5000\n","        if not test_grid(N, B, hi_g):\n","            hi_g = 10000\n","            if not test_grid(N, B, hi_g):\n","                print(f\"{N:>4}  {'>10000':>8}  {time.time()-t0:>7.1f}s\")\n","                continue\n","    while lo_g < hi_g:\n","        mid = (lo_g + hi_g) // 2\n","        if test_grid(N, B, mid):\n","            hi_g = mid\n","        else:\n","            lo_g = mid + 1\n","    results[N] = lo_g\n","    print(f\"{N:>4}  {lo_g:>8}  {time.time()-t0:>7.1f}s\")\n","\n","print(f\"\\n_GRID_SIZE = {results}\")"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"XlG3MFiquBpS","executionInfo":{"status":"ok","timestamp":1775060131300,"user_tz":420,"elapsed":1982,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"29848a19-9f63-4a25-994a-8fb8ff2b84be"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["Per-N grid sweep — end-to-end accuracy (val_err < 1e-3)\n","   n  min_grid      time\n","────  ────────  ────────\n","   3       167      0.2s\n","   4       258      0.1s\n","   5       506      0.2s\n","   6      1838      0.2s\n","   8      1876      0.2s\n","  10      4398      0.5s\n","  12      3537      0.5s\n","\n","_GRID_SIZE = {3: 167, 4: 258, 5: 506, 6: 1838, 8: 1876, 10: 4398, 12: 3537}\n"]}]},{"cell_type":"code","source":["\"\"\"\n","Grid sweep WITHOUT Frobenius normalization.\n","Tests whether natural eigenvalue spacing needs smaller grids.\n","\"\"\"\n","import math, time, torch\n","\n","dev = torch.device('cuda')\n","TERNARY_LEVELS = 25; POLISH_ITERS = 3\n","\n","def _horner_sign(c, x):\n","    B, k = x.shape; n = c.shape[1] - 1\n","    pv = c[:, n:n+1].expand(B, k)\n","    for j in range(n-1, -1, -1):\n","        pv = pv * x + c[:, j:j+1]\n","    return torch.sign(pv)\n","\n","def run_eigh_nonorm(A, n_grid):\n","    B, n, _ = A.shape; dev = A.device\n","    Ad = A.double()\n","    I_d = torch.eye(n, device=dev, dtype=torch.float64).unsqueeze(0).expand(B,-1,-1)\n","    c = torch.zeros(B, n+1, device=dev, dtype=torch.float64); c[:, n] = 1.0\n","    Ms = torch.zeros(n+1, B, n, n, device=dev, dtype=torch.float64)\n","    Mk = torch.zeros(B, n, n, device=dev, dtype=torch.float64)\n","    for k in range(1, n+1):\n","        Mk = torch.bmm(Ad, Mk) + c[:, n-k+1, None, None] * I_d\n","        Ms[k] = Mk; c[:, n-k] = -(Ad * Mk).sum((-2,-1)) / k\n","\n","    diag = Ad.diagonal(dim1=-2, dim2=-1)\n","    off_row = Ad.abs().sum(-1) - diag.abs()\n","    lo = (diag - off_row).min(dim=-1).values - 0.1\n","    hi = (diag + off_row).max(dim=-1).values + 0.1\n","    t = torch.linspace(0, 1, n_grid, device=dev, dtype=torch.float64).unsqueeze(0)\n","    grid = lo.unsqueeze(-1) + (hi - lo).unsqueeze(-1) * t\n","    signs = _horner_sign(c, grid)\n","    sign_changes = (signs[:, 1:] * signs[:, :-1]) < 0\n","    cum_idx = sign_changes.long().cumsum(dim=-1)\n","\n","    bracket_lo = torch.zeros(B, n, device=dev, dtype=torch.float64)\n","    bracket_hi = torch.zeros(B, n, device=dev, dtype=torch.float64)\n","    for i in range(n):\n","        match = (cum_idx == i + 1) & sign_changes\n","        idx = match.long().argmax(dim=-1); found = match.any(dim=-1)\n","        bracket_lo[:, i] = torch.where(found, grid.gather(1, idx.unsqueeze(-1)).squeeze(-1), lo + (hi-lo)*(i/n))\n","        bracket_hi[:, i] = torch.where(found, grid.gather(1, (idx+1).unsqueeze(-1)).squeeze(-1), lo + (hi-lo)*((i+1)/n))\n","\n","    b_lo, b_hi = bracket_lo, bracket_hi\n","    for _ in range(TERNARY_LEVELS):\n","        third = (b_hi - b_lo) / 3.0; x1 = b_lo + third; x2 = b_lo + 2*third\n","        s_lo = _horner_sign(c, b_lo); s_x1 = _horner_sign(c, x1); s_x2 = _horner_sign(c, x2)\n","        in_first = (s_lo * s_x1) < 0; in_second = (s_x1 * s_x2) < 0\n","        b_lo = torch.where(in_first, b_lo, torch.where(in_second, x1, x2))\n","        b_hi = torch.where(in_first, x1, torch.where(in_second, x2, b_hi))\n","    roots = (b_lo + b_hi) / 2\n","\n","    for _ in range(POLISH_ITERS):\n","        pv = torch.ones(B, n, device=dev, dtype=torch.float64)\n","        dp = torch.zeros(B, n, device=dev, dtype=torch.float64)\n","        for j in range(n-1, -1, -1): dp = dp*roots+pv; pv = pv*roots+c[:, j:j+1]\n","        ok = dp.abs()>1e-30; dps = torch.where(ok, dp, torch.ones_like(dp))\n","        roots = roots - torch.where(ok, pv/dps, torch.zeros_like(pv))\n","\n","    chunk = min(4, n)\n","    V = torch.empty(B, n, n, device=dev, dtype=torch.float32)\n","    for ei_start in range(0, n, chunk):\n","        ei_end = min(ei_start+chunk, n); c_size = ei_end-ei_start\n","        lam_c = roots[:, ei_start:ei_end, None, None]\n","        R = Ms[1].unsqueeze(1).expand(B, c_size, n, n).clone()\n","        for k in range(2, n+1): R = R*lam_c + Ms[k].unsqueeze(1)\n","        cn = R.norm(dim=-2); best = cn.argmax(dim=-1)\n","        idx = best.unsqueeze(-1).unsqueeze(-1).expand(B, c_size, n, 1)\n","        vec = R.gather(-1, idx).squeeze(-1); vec = vec/(vec.norm(dim=-1, keepdim=True)+1e-30)\n","        V[:, :, ei_start:ei_end] = vec.float().transpose(-2, -1)\n","\n","    eye_f = torch.eye(n, device=dev, dtype=torch.float32).unsqueeze(0).expand(B,-1,-1)\n","    Y = torch.bmm(V.mT, V); X = eye_f.clone()\n","    for _ in range(2): T = 3.0*eye_f-Y; X = 0.5*torch.bmm(X, T); Y = 0.5*torch.bmm(T, Y)\n","    V = torch.bmm(V, X); AV = torch.bmm(A, V); evals = (V*AV).sum(dim=-2)\n","    return evals.sort(dim=-1).values\n","\n","def test_grid(N, B, ng, seeds=[42, 123, 999]):\n","    for seed in seeds:\n","        torch.manual_seed(seed)\n","        A = (lambda R: (R+R.mT)/2)(torch.randn(B, N, N, device=dev))\n","        rv, _ = torch.linalg.eigh(A)\n","        tv = run_eigh_nonorm(A, ng)\n","        ve = (tv - rv).abs().max().item()\n","        if ve > 1e-3: return False\n","    return True\n","\n","B = 2048\n","print(\"NO-NORM grid sweep — end-to-end (val_err < 1e-3)\")\n","print(f\"{'n':>4}  {'no_norm':>8}  {'frob_norm':>10}  {'ratio':>6}\")\n","print(f\"{'─'*4}  {'─'*8}  {'─'*10}  {'─'*6}\")\n","\n","frob = {3: 167, 4: 258, 5: 506, 6: 1838, 8: 1876, 10: 4398, 12: 3537}\n","\n","for N in [3, 4, 5, 6, 8, 10, 12]:\n","    lo_g, hi_g = 10, 2000\n","    if not test_grid(N, B, hi_g):\n","        hi_g = 5000\n","        if not test_grid(N, B, hi_g):\n","            hi_g = 10000\n","            if not test_grid(N, B, hi_g):\n","                print(f\"{N:>4}  {'>10000':>8}  {frob[N]:>10}\"); continue\n","    while lo_g < hi_g:\n","        mid = (lo_g + hi_g) // 2\n","        if test_grid(N, B, mid): hi_g = mid\n","        else: lo_g = mid + 1\n","    f = frob[N]\n","    print(f\"{N:>4}  {lo_g:>8}  {f:>10}  {f/lo_g:>5.1f}x\")"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"V7ZWm5-fx2KU","executionInfo":{"status":"ok","timestamp":1775060840505,"user_tz":420,"elapsed":2116,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"9054864a-7af6-4b51-b47d-9aab68bc5625"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["NO-NORM grid sweep — end-to-end (val_err < 1e-3)\n","   n   no_norm   frob_norm   ratio\n","────  ────────  ──────────  ──────\n","   3       103         167    1.6x\n","   4       231         258    1.1x\n","   5       475         506    1.1x\n","   6      1005        1838    1.8x\n","   8      2580        1876    0.7x\n","  10      2270        4398    1.9x\n","  12      3810        3537    0.9x\n"]}]},{"cell_type":"code","source":["\"\"\"\n","Wormhole Ternary Root-Finding.\n","\n","Coarse ternary grid + derivative-guided skip connections.\n","At each grid point, p'(x) tells us where the nearest root is.\n","The \"wormhole\" jumps from the grid point toward the root using Newton step.\n","This finds hidden close root pairs that a coarse grid would miss.\n","\n","Total evaluations: ~4n coarse + ~4n wormholes = ~8n\n","vs dense grid: 1838 for n=6\n","\"\"\"\n","import math, time, torch\n","import torch.nn as nn\n","from torch import Tensor\n","from typing import Tuple\n","\n","torch.backends.cuda.matmul.allow_tf32 = False\n","torch.backends.cudnn.allow_tf32 = False\n","torch.set_float32_matmul_precision('highest')\n","torch._dynamo.config.recompile_limit = 64\n","\n","\n","def _horner_sign(c, x):\n","    B, k = x.shape; n = c.shape[1] - 1\n","    pv = c[:, n:n+1].expand(B, k)\n","    for j in range(n-1, -1, -1):\n","        pv = pv * x + c[:, j:j+1]\n","    return torch.sign(pv)\n","\n","def _horner_val(c, x):\n","    B, k = x.shape; n = c.shape[1] - 1\n","    pv = c[:, n:n+1].expand(B, k)\n","    for j in range(n-1, -1, -1):\n","        pv = pv * x + c[:, j:j+1]\n","    return pv\n","\n","def _horner_pd(c, x):\n","    B, k = x.shape; n = c.shape[1] - 1\n","    pv = c[:, n:n+1].expand(B, k)\n","    dp = torch.zeros(B, k, device=x.device, dtype=x.dtype)\n","    for j in range(n-1, -1, -1):\n","        dp = dp * x + pv\n","        pv = pv * x + c[:, j:j+1]\n","    return pv, dp\n","\n","\n","def wormhole_brackets(c, lo, hi, n, B, dev):\n","    \"\"\"Find n root brackets using coarse grid + derivative wormholes.\n","\n","    1. Coarse ternary grid: 3*n+1 points across [lo, hi]\n","    2. Evaluate p and p' at each grid point\n","    3. Direct sign changes → brackets (catches well-separated roots)\n","    4. Wormhole: Newton step from each grid point → predicted root location\n","    5. Evaluate p at wormhole targets → catches close root pairs\n","    6. Combine all bracket candidates, take best n\n","    \"\"\"\n","    # Coarse grid: ~4n points\n","    n_coarse = max(4 * n + 1, 25)\n","    t = torch.linspace(0, 1, n_coarse, device=dev, dtype=torch.float64).unsqueeze(0)\n","    grid = lo.unsqueeze(-1) + (hi - lo).unsqueeze(-1) * t  # [B, n_coarse]\n","\n","    # Evaluate p and p' at grid points\n","    pv, dp = _horner_pd(c, grid)  # [B, n_coarse] each\n","\n","    # Step 1: direct sign changes on coarse grid\n","    signs = torch.sign(pv)  # [B, n_coarse]\n","    coarse_changes = (signs[:, 1:] * signs[:, :-1]) < 0  # [B, n_coarse-1]\n","\n","    # Step 2: wormhole — Newton step from each grid point\n","    # target = grid - p(grid) / p'(grid)\n","    dp_safe = torch.where(dp.abs() > 1e-30, dp, torch.ones_like(dp))\n","    wormhole = grid - torch.where(dp.abs() > 1e-30, pv / dp_safe, torch.zeros_like(pv))\n","\n","    # Clamp wormholes to [lo, hi] range\n","    wormhole = wormhole.clamp(min=lo.unsqueeze(-1), max=hi.unsqueeze(-1))\n","\n","    # Evaluate p at wormhole targets\n","    pv_wh = _horner_val(c, wormhole)  # [B, n_coarse]\n","\n","    # Step 3: find sign changes between grid point and its wormhole target\n","    wh_changes = (torch.sign(pv) * torch.sign(pv_wh)) < 0  # [B, n_coarse]\n","\n","    # Step 4: build candidate brackets from BOTH sources\n","    # For coarse grid sign changes: bracket is [grid[i], grid[i+1]]\n","    # For wormhole sign changes: bracket is [min(grid[i], wh[i]), max(grid[i], wh[i])]\n","\n","    # Collect ALL candidate bracket midpoints and sort them\n","    # Coarse midpoints\n","    coarse_mids = (grid[:, :-1] + grid[:, 1:]) / 2  # [B, n_coarse-1]\n","    coarse_mids = torch.where(coarse_changes, coarse_mids,\n","                               torch.full_like(coarse_mids, float('inf')))\n","\n","    # Wormhole midpoints\n","    wh_mids = (grid + wormhole) / 2  # [B, n_coarse]\n","    wh_mids = torch.where(wh_changes, wh_mids,\n","                           torch.full_like(wh_mids, float('inf')))\n","\n","    # Combine all candidates: [B, 2*n_coarse-1]\n","    all_mids = torch.cat([coarse_mids, wh_mids], dim=-1)\n","\n","    # Sort and take the n smallest (closest to -inf = leftmost real brackets)\n","    all_mids_sorted, _ = all_mids.sort(dim=-1)\n","    root_estimates = all_mids_sorted[:, :n]  # [B, n]\n","\n","    # For any batch element with fewer than n brackets, fill from diagonal\n","    diag_fallback = torch.linalg.eigvalsh(\n","        torch.eye(n, device=dev, dtype=torch.float64).unsqueeze(0).expand(B,-1,-1)\n","    )  # dummy, not used if brackets found\n","\n","    missing = (root_estimates == float('inf'))\n","    if missing.any():\n","        diag = torch.zeros(B, n, device=dev, dtype=torch.float64)\n","        root_estimates = torch.where(missing, diag, root_estimates)\n","\n","    # Build brackets: [estimate - spacing/2, estimate + spacing/2]\n","    spacing = (hi - lo).unsqueeze(-1) / (3 * n)  # conservative bracket width\n","    bracket_lo = root_estimates - spacing\n","    bracket_hi = root_estimates + spacing\n","\n","    # Verify brackets contain roots (sign change check)\n","    s_lo = _horner_sign(c, bracket_lo)\n","    s_hi = _horner_sign(c, bracket_hi)\n","    valid = (s_lo * s_hi) < 0\n","\n","    # Widen invalid brackets progressively\n","    for widen in range(5):\n","        if valid.all():\n","            break\n","        factor = 2.0 ** (widen + 1)\n","        bracket_lo = torch.where(valid, bracket_lo, root_estimates - spacing * factor)\n","        bracket_hi = torch.where(valid, bracket_hi, root_estimates + spacing * factor)\n","        bracket_lo = bracket_lo.clamp(min=lo.unsqueeze(-1))\n","        bracket_hi = bracket_hi.clamp(max=hi.unsqueeze(-1))\n","        s_lo = _horner_sign(c, bracket_lo)\n","        s_hi = _horner_sign(c, bracket_hi)\n","        valid = (s_lo * s_hi) < 0\n","\n","    return bracket_lo, bracket_hi\n","\n","\n","class FLEighWormhole(nn.Module):\n","    \"\"\"FL Eigh with wormhole bracket initialization + ternary refinement.\"\"\"\n","    def __init__(self, ternary_levels=25, polish_iters=3, chunk=4):\n","        super().__init__()\n","        self.ternary_levels = ternary_levels\n","        self.polish_iters = polish_iters\n","        self.chunk = chunk\n","\n","    def forward(self, A):\n","        B, n, _ = A.shape; dev = A.device\n","        sc = (torch.linalg.norm(A.reshape(B,-1), dim=-1) / math.sqrt(n)).clamp(min=1e-12)\n","        As = A / sc[:, None, None]; Ad = As.double()\n","        I_d = torch.eye(n, device=dev, dtype=torch.float64).unsqueeze(0).expand(B,-1,-1)\n","        c = torch.zeros(B, n+1, device=dev, dtype=torch.float64); c[:, n] = 1.0\n","        Ms = torch.zeros(n+1, B, n, n, device=dev, dtype=torch.float64)\n","        Mk = torch.zeros(B, n, n, device=dev, dtype=torch.float64)\n","        for k in range(1, n+1):\n","            Mk = torch.bmm(Ad, Mk) + c[:, n-k+1, None, None] * I_d\n","            Ms[k] = Mk; c[:, n-k] = -(Ad * Mk).sum((-2,-1)) / k\n","\n","        # Wormhole bracket initialization\n","        diag = Ad.diagonal(dim1=-2, dim2=-1)\n","        off_row = Ad.abs().sum(-1) - diag.abs()\n","        lo = (diag - off_row).min(dim=-1).values - 0.1\n","        hi = (diag + off_row).max(dim=-1).values + 0.1\n","        b_lo, b_hi = wormhole_brackets(c, lo, hi, n, B, dev)\n","\n","        # Ternary refinement\n","        for _ in range(self.ternary_levels):\n","            third = (b_hi - b_lo) / 3.0\n","            x1 = b_lo + third; x2 = b_lo + 2 * third\n","            s_lo = _horner_sign(c, b_lo); s_x1 = _horner_sign(c, x1); s_x2 = _horner_sign(c, x2)\n","            in_first = (s_lo * s_x1) < 0; in_second = (s_x1 * s_x2) < 0\n","            b_lo = torch.where(in_first, b_lo, torch.where(in_second, x1, x2))\n","            b_hi = torch.where(in_first, x1, torch.where(in_second, x2, b_hi))\n","        roots = (b_lo + b_hi) / 2\n","\n","        # Newton polish\n","        for _ in range(self.polish_iters):\n","            pv = torch.ones(B, n, device=dev, dtype=torch.float64)\n","            dp = torch.zeros(B, n, device=dev, dtype=torch.float64)\n","            for j in range(n-1, -1, -1):\n","                dp = dp * roots + pv; pv = pv * roots + c[:, j:j+1]\n","            ok = dp.abs() > 1e-30; dps = torch.where(ok, dp, torch.ones_like(dp))\n","            roots = roots - torch.where(ok, pv / dps, torch.zeros_like(pv))\n","\n","        # Chunked eigenvectors\n","        chunk = min(self.chunk, n)\n","        V = torch.empty(B, n, n, device=dev, dtype=torch.float32)\n","        for ei_start in range(0, n, chunk):\n","            ei_end = min(ei_start + chunk, n); c_size = ei_end - ei_start\n","            lam_c = roots[:, ei_start:ei_end, None, None]\n","            R = Ms[1].unsqueeze(1).expand(B, c_size, n, n).clone()\n","            for k in range(2, n+1): R = R * lam_c + Ms[k].unsqueeze(1)\n","            cn = R.norm(dim=-2); best = cn.argmax(dim=-1)\n","            idx = best.unsqueeze(-1).unsqueeze(-1).expand(B, c_size, n, 1)\n","            vec = R.gather(-1, idx).squeeze(-1)\n","            vec = vec / (vec.norm(dim=-1, keepdim=True) + 1e-30)\n","            V[:, :, ei_start:ei_end] = vec.float().transpose(-2, -1)\n","\n","        # NS + Rayleigh\n","        eye_f = torch.eye(n, device=dev, dtype=torch.float32).unsqueeze(0).expand(B,-1,-1)\n","        Y = torch.bmm(V.mT, V); X = eye_f.clone()\n","        for _ in range(2):\n","            T = 3.0 * eye_f - Y; X = 0.5 * torch.bmm(X, T); Y = 0.5 * torch.bmm(T, Y)\n","        V = torch.bmm(V, X); AV = torch.bmm(A, V); evals = (V * AV).sum(dim=-2)\n","        se, p = evals.sort(dim=-1)\n","        return se, V.gather(-1, p.unsqueeze(-2).expand_as(V))\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# Sequential baseline\n","# ═══════════════════════════════════════════════════════════════════\n","class FLEighSeq(nn.Module):\n","    def forward(self, A):\n","        B,n,_=A.shape; dev=A.device\n","        sc=(torch.linalg.norm(A.reshape(B,-1),dim=-1)/math.sqrt(n)).clamp(min=1e-12)\n","        As=A/sc[:,None,None]; Ad=As.double()\n","        I_d=torch.eye(n,device=dev,dtype=torch.float64).unsqueeze(0).expand(B,-1,-1)\n","        c=torch.zeros(B,n+1,device=dev,dtype=torch.float64); c[:,n]=1.0\n","        Ms=torch.zeros(n+1,B,n,n,device=dev,dtype=torch.float64)\n","        Mk=torch.zeros(B,n,n,device=dev,dtype=torch.float64)\n","        for k in range(1,n+1):\n","            Mk=torch.bmm(Ad,Mk)+c[:,n-k+1,None,None]*I_d; Ms[k]=Mk\n","            c[:,n-k]=-(Ad*Mk).sum((-2,-1))/k\n","        dt=torch.float64 if n>6 else torch.float32; n_pol=5 if n>8 else 3\n","        cl=c.to(dt).clone(); roots=torch.zeros(B,n,device=dev,dtype=dt)\n","        zi=As.to(dt).diagonal(dim1=-2,dim2=-1).sort(dim=-1).values\n","        zi=zi+torch.linspace(-1e-4,1e-4,n,device=dev,dtype=dt).unsqueeze(0)\n","        for ri in range(n):\n","            deg=n-ri; z=zi[:,ri]\n","            for _ in range(5):\n","                pv=cl[:,deg]; dp=torch.zeros(B,device=dev,dtype=dt); d2=torch.zeros(B,device=dev,dtype=dt)\n","                for j in range(deg-1,-1,-1): d2=d2*z+dp; dp=dp*z+pv; pv=pv*z+cl[:,j]\n","                ok=pv.abs()>1e-30; ps=torch.where(ok,pv,torch.ones_like(pv))\n","                G=torch.where(ok,dp/ps,torch.zeros_like(dp)); H=G*G-torch.where(ok,2.0*d2/ps,torch.zeros_like(d2))\n","                disc=((deg-1.0)*(deg*H-G*G)).clamp(min=0.0); sq=torch.sqrt(disc); gp=G+sq; gm=G-sq\n","                den=torch.where(gp.abs()>=gm.abs(),gp,gm); dok=den.abs()>1e-20\n","                ds=torch.where(dok,den,torch.ones_like(den)); z=z-torch.where(dok,float(deg)/ds,torch.zeros_like(den))\n","            roots[:,ri]=z; b=cl[:,deg]\n","            for j in range(deg-1,0,-1): bn=cl[:,j]+z*b; cl[:,j]=b; b=bn\n","            cl[:,0]=b\n","        roots=roots.double()\n","        for _ in range(n_pol):\n","            pv=torch.ones(B,n,device=dev,dtype=torch.float64); dp=torch.zeros(B,n,device=dev,dtype=torch.float64)\n","            for j in range(n-1,-1,-1): dp=dp*roots+pv; pv=pv*roots+c[:,j:j+1]\n","            ok=dp.abs()>1e-30; dps=torch.where(ok,dp,torch.ones_like(dp)); roots=roots-torch.where(ok,pv/dps,torch.zeros_like(pv))\n","        V=torch.empty(B,n,n,device=dev,dtype=torch.float32)\n","        for ei in range(n):\n","            li=roots[:,ei:ei+1,None]; R=Ms[1].clone()\n","            for k in range(2,n+1): R=R*li+Ms[k]\n","            cn=R.norm(dim=-2); best=cn.argmax(dim=-1); idx=best[:,None,None].expand(B,n,1)\n","            vec=R.gather(-1,idx).squeeze(-1); vec=vec/(vec.norm(dim=-1,keepdim=True)+1e-30); V[:,:,ei]=vec.float()\n","        eye_f=torch.eye(n,device=dev,dtype=torch.float32).unsqueeze(0).expand(B,-1,-1)\n","        Y=torch.bmm(V.mT,V); X=eye_f.clone()\n","        for _ in range(2): T=3.0*eye_f-Y; X=0.5*torch.bmm(X,T); Y=0.5*torch.bmm(T,Y)\n","        V=torch.bmm(V,X); AV=torch.bmm(A,V); evals=(V*AV).sum(dim=-2)\n","        se,p=evals.sort(dim=-1); return se,V.gather(-1,p.unsqueeze(-2).expand_as(V))\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","def sync(): torch.cuda.synchronize()\n","def gt(fn,w=20,r=300):\n","    for _ in range(w): fn()\n","    sync();t=time.perf_counter()\n","    for _ in range(r): fn()\n","    sync();return(time.perf_counter()-t)/r\n","def fmt(s):\n","    if s<1e-3:return f\"{s*1e6:.0f}us\"\n","    if s<1:return f\"{s*1e3:.2f}ms\"\n","    return f\"{s:.3f}s\"\n","def make(B,n,dev):R=torch.randn(B,n,n,device=dev);return(R+R.mT)/2\n","\n","def main():\n","    dev=torch.device('cuda')\n","    p=torch.cuda.get_device_properties(0)\n","    print(\"=\"*78)\n","    print(\"  FL Eigh — Wormhole Ternary\")\n","    print(\"=\"*78)\n","    print(f\"  {p.name} | PyTorch {torch.__version__}\")\n","\n","    # ══════════════ 1. ACCURACY ══════════════\n","    print(f\"\\n{'='*78}\\n  1. ACCURACY\\n{'='*78}\")\n","    print(f\"\\n  {'n':>3}  {'Wormhole':>16}  {'Sequential':>16}\")\n","    for nx in [3,4,5,6,8,10,12,16]:\n","        Ax=make(2048,nx,dev); rv,rV=torch.linalg.eigh(Ax)\n","        sv,sV=FLEighSeq()(Ax)\n","        wv,wV=FLEighWormhole(ternary_levels=25,polish_iters=3,chunk=min(4,nx))(Ax)\n","        se=(sv-rv).abs().max().item(); we=(wv-rv).abs().max().item()\n","        sd=torch.bmm(rV.double().mT,sV.double()).abs().max(dim=-1).values.min().item()\n","        wd=torch.bmm(rV.double().mT,wV.double()).abs().max(dim=-1).values.min().item()\n","        s_ok=\"OK\" if se<1e-2 and sd>0.99 else \"NO\"\n","        w_ok=\"OK\" if we<1e-2 and wd>0.99 else \"NO\"\n","        print(f\"  n={nx:>2}  [{w_ok}] {we:.1e} {wd:.4f}  [{s_ok}] {se:.1e} {sd:.4f}\")\n","        del Ax\n","\n","    # ══════════════ 2. TIMING ══════════════\n","    print(f\"\\n{'='*78}\\n  2. COMPILED TIMING\\n{'='*78}\")\n","    for nx in [6,8,12]:\n","        B=4096; Ax=make(B,nx,dev)\n","        cs=torch.compile(FLEighSeq(),fullgraph=True)\n","        cw=torch.compile(FLEighWormhole(chunk=min(4,nx)),fullgraph=True)\n","        for _ in range(3): cs(Ax); cw(Ax); sync()\n","        t_cus=gt(lambda:torch.linalg.eigh(Ax),10,100)\n","        t_seq=gt(lambda:cs(Ax),10,100)\n","        t_wh=gt(lambda:cw(Ax),10,100)\n","        print(f\"\\n  n={nx} B={B}:\")\n","        print(f\"    cuSOLVER:    {fmt(t_cus):>10}\")\n","        print(f\"    Sequential:  {fmt(t_seq):>10}  {t_cus/t_seq:.2f}x cuS\")\n","        print(f\"    Wormhole:    {fmt(t_wh):>10}  {t_cus/t_wh:.2f}x cuS  {t_seq/t_wh:.2f}x vs Seq\")\n","        del Ax,cs,cw; torch.cuda.empty_cache()\n","\n","    print(f\"\\n{'='*78}\")\n","\n","if __name__=='__main__': main()"],"metadata":{"colab":{"base_uri":"https://localhost:8080/","height":1000},"id":"A1cP3Z9Q5xN5","executionInfo":{"status":"error","timestamp":1775061398720,"user_tz":420,"elapsed":1027,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"4bd175a0-5a80-4f17-e05c-929f103438bd"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["==============================================================================\n","  FL Eigh — Wormhole Ternary\n","==============================================================================\n","  NVIDIA RTX PRO 6000 Blackwell Server Edition | PyTorch 2.10.0+cu128\n","\n","==============================================================================\n","  1. ACCURACY\n","==============================================================================\n","\n","    n          Wormhole        Sequential\n","  n= 3  [NO] 5.4e+00 0.0000  [OK] 1.9e-06 1.0000\n","  n= 4  [NO] 8.3e+00 0.0000  [OK] 2.1e-06 1.0000\n","  n= 5  [NO] 6.9e+01 0.0000  [OK] 1.9e-06 1.0000\n","  n= 6  [NO] 3.0e+02 0.0000  [OK] 2.1e-06 1.0000\n","  n= 8  [NO] 3.2e+03 0.0000  [OK] 3.8e-06 1.0000\n","  n=10  [NO] 7.5e+03 0.0000  [OK] 5.0e-06 1.0000\n","  n=12  [NO] 1.8e+04 0.0000  [OK] 6.7e-06 1.0000\n","  n=16  [NO] 1.2e+04 0.0000  [OK] 8.1e-06 1.0000\n","\n","==============================================================================\n","  2. COMPILED TIMING\n","==============================================================================\n"]},{"output_type":"error","ename":"Unsupported","evalue":"Data-dependent branching\n  Explanation: Detected data-dependent branching (e.g. `if my_tensor.sum() > 0:`). Dynamo does not support tracing dynamic control flow.\n  Hint: This graph break is fundamental - it is unlikely that Dynamo will ever be able to trace through your code. Consider finding a workaround.\n  Hint: Use `torch.cond` to express dynamic control flow.\n\n  Developer debug context: attempted to jump with TensorVariable()\n\n For more details about this graph break, please visit: https://meta-pytorch.github.io/compile-graph-break-site/gb/gb0170.html\n\nfrom user code:\n   File \"/tmp/ipykernel_89646/4015653392.py\", line 166, in forward\n    b_lo, b_hi = wormhole_brackets(c, lo, hi, n, B, dev)\n  File \"/tmp/ipykernel_89646/4015653392.py\", line 111, in wormhole_brackets\n    if missing.any():\n\nSet TORCHDYNAMO_VERBOSE=1 for the internal stack trace (please do this especially if you're reporting a bug to PyTorch). For even more developer context, set TORCH_LOGS=\"+dynamo\"\n","traceback":["\u001b[0;31m---------------------------------------------------------------------------\u001b[0m","\u001b[0;31mUnsupported\u001b[0m                               Traceback (most recent call last)","\u001b[0;32m/tmp/ipykernel_89646/4015653392.py\u001b[0m in \u001b[0;36m<cell line: 0>\u001b[0;34m()\u001b[0m\n\u001b[1;32m    313\u001b[0m     \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf\"\\n{'='*78}\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    314\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 315\u001b[0;31m \u001b[0;32mif\u001b[0m \u001b[0m__name__\u001b[0m\u001b[0;34m==\u001b[0m\u001b[0;34m'__main__'\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mmain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m","\u001b[0;32m/tmp/ipykernel_89646/4015653392.py\u001b[0m in \u001b[0;36mmain\u001b[0;34m()\u001b[0m\n\u001b[1;32m    301\u001b[0m         \u001b[0mcs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcompile\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mFLEighSeq\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mfullgraph\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    302\u001b[0m         \u001b[0mcw\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcompile\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mFLEighWormhole\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mchunk\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mmin\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m4\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mnx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mfullgraph\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 303\u001b[0;31m         \u001b[0;32mfor\u001b[0m \u001b[0m_\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m3\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mcs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mAx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m;\u001b[0m \u001b[0mcw\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mAx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m;\u001b[0m \u001b[0msync\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    304\u001b[0m         \u001b[0mt_cus\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mgt\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;32mlambda\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlinalg\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0meigh\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mAx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m10\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m100\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    305\u001b[0m         \u001b[0mt_seq\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mgt\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;32mlambda\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0mcs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mAx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m10\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m100\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_dynamo/eval_frame.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    463\u001b[0m             )\n\u001b[1;32m    464\u001b[0m         \u001b[0;32mwith\u001b[0m \u001b[0m_set_in_optimized_module\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 465\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0msuper\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__call__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    466\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    467\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_aot_compile\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minputs\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mlist\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_dynamo\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0maot_compile\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mModelInput\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1774\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_compiled_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[misc]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1775\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1776\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1777\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1778\u001b[0m     \u001b[0;31m# torchrec tests the code consistency with the following code\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1785\u001b[0m                 \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_pre_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_hooks\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1786\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1787\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1788\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1789\u001b[0m         \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.12/dist-packages/torch/_dynamo/eval_frame.py\u001b[0m in \u001b[0;36mcompile_wrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    961\u001b[0m                         \u001b[0mcur_exn\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcur_exn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__cause__\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    962\u001b[0m                     \u001b[0;31m# pyrefly: ignore [invalid-inheritance]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 963\u001b[0;31m                     \u001b[0;32mraise\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mwith_traceback\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__cause__\u001b[0m  \u001b[0;31m# User compiler error\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    964\u001b[0m                 \u001b[0;32mexcept\u001b[0m \u001b[0mShortenTraceback\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    965\u001b[0m                     \u001b[0;31m# Failures in the backend likely don't have useful\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;31mUnsupported\u001b[0m: Data-dependent branching\n  Explanation: Detected data-dependent branching (e.g. `if my_tensor.sum() > 0:`). Dynamo does not support tracing dynamic control flow.\n  Hint: This graph break is fundamental - it is unlikely that Dynamo will ever be able to trace through your code. Consider finding a workaround.\n  Hint: Use `torch.cond` to express dynamic control flow.\n\n  Developer debug context: attempted to jump with TensorVariable()\n\n For more details about this graph break, please visit: https://meta-pytorch.github.io/compile-graph-break-site/gb/gb0170.html\n\nfrom user code:\n   File \"/tmp/ipykernel_89646/4015653392.py\", line 166, in forward\n    b_lo, b_hi = wormhole_brackets(c, lo, hi, n, B, dev)\n  File \"/tmp/ipykernel_89646/4015653392.py\", line 111, in wormhole_brackets\n    if missing.any():\n\nSet TORCHDYNAMO_VERBOSE=1 for the internal stack trace (please do this especially if you're reporting a bug to PyTorch). For even more developer context, set TORCH_LOGS=\"+dynamo\"\n"]}]},{"cell_type":"markdown","source":["# flows to match"],"metadata":{"id":"sYBO53Qekcxk"}},{"cell_type":"code","source":["\"\"\"\n","geolip.flows — Multi-flow ensemble for constellation geometry.\n","\n","Each flow predicts the same geometric output using a different mathematical\n","formulation. The ensemble fuses predictions based on learned confidence.\n","\n","Flows:\n","  QuaternionFlow      — Full MHA quaternion rotation (existing, heavyweight)\n","  QuaternionLiteFlow  — Staged quaternion with lighter spectral computation\n","  VelocityFlow        — Angular velocity dq/dt on the tangent bundle\n","  MagnitudeFlow       — Flow magnitude via Gram eigenvalue spectrum\n","  OrbitalFlow         — Omega-based orbital resonance using FL eigh\n","  AlignmentFlow       — SVD alignment via Procrustes rotation\n","\n","Architecture:\n","  Each flow: same input (anchors [B,k,d], queries [B,n,d]) → output [B,n,d]\n","  Ensemble:  weighted fusion with learned per-flow confidence\n","\n","Usage:\n","    from geolip.flows import FlowEnsemble, OrbitalFlow, AlignmentFlow\n","\n","    ensemble = FlowEnsemble(\n","        flows=[OrbitalFlow(d=256, k=128), AlignmentFlow(d=256, k=128)],\n","        d_model=256,\n","    )\n","    output = ensemble(anchors, queries)  # [B, n, d]\n","\"\"\"\n","\n","import math\n","import torch\n","import torch.nn as nn\n","import torch.nn.functional as F\n","from torch import Tensor\n","from typing import List, Optional, Tuple\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# Base Flow\n","# ═══════════════════════════════════════════════════════════════════\n","\n","class BaseFlow(nn.Module):\n","    \"\"\"Base class for all geometric flows.\n","\n","    All flows share the same interface:\n","      Input:  anchors [B, k, d], queries [B, n, d]\n","      Output: prediction [B, n, d], confidence [B, n, 1]\n","\n","    Subclasses implement _flow() with their specific math.\n","    \"\"\"\n","    def __init__(self, d_model: int, n_anchors: int, name: str = 'base'):\n","        super().__init__()\n","        self.d_model = d_model\n","        self.n_anchors = n_anchors\n","        self.name = name\n","        # Confidence head: scalar per query position\n","        self.confidence = nn.Sequential(\n","            nn.Linear(d_model, d_model // 4),\n","            nn.GELU(),\n","            nn.Linear(d_model // 4, 1),\n","        )\n","\n","    def forward(self, anchors: Tensor, queries: Tensor) -> Tuple[Tensor, Tensor]:\n","        \"\"\"\n","        Args:\n","            anchors: [B, k, d] constellation anchor points\n","            queries: [B, n, d] query embeddings\n","\n","        Returns:\n","            prediction: [B, n, d] geometric prediction\n","            confidence: [B, n, 1] per-query confidence score\n","        \"\"\"\n","        pred = self._flow(anchors, queries)\n","        conf = torch.sigmoid(self.confidence(pred))\n","        return pred, conf\n","\n","    def _flow(self, anchors: Tensor, queries: Tensor) -> Tensor:\n","        raise NotImplementedError\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# QuaternionFlow — Full MHA quaternion rotation\n","# ═══════════════════════════════════════════════════════════════════\n","\n","class QuaternionFlow(BaseFlow):\n","    \"\"\"Full multi-head attention with quaternion geometric rotation.\n","\n","    Computes query-anchor attention, extracts rotation quaternion from\n","    attention-weighted anchor geometry, applies rotation to queries.\n","    Heavyweight — the full-fidelity path.\n","    \"\"\"\n","    def __init__(self, d_model: int, n_anchors: int, n_heads: int = 4):\n","        super().__init__(d_model, n_anchors, name='quaternion')\n","        self.n_heads = n_heads\n","        self.head_dim = d_model // n_heads\n","        self.q_proj = nn.Linear(d_model, d_model)\n","        self.k_proj = nn.Linear(d_model, d_model)\n","        self.v_proj = nn.Linear(d_model, d_model)\n","        self.out_proj = nn.Linear(d_model, d_model)\n","        # Quaternion components: scalar + 3 imaginary from attention output\n","        self.quat_proj = nn.Linear(d_model, 4)\n","\n","    def _flow(self, anchors, queries):\n","        B, n, d = queries.shape\n","        k = anchors.shape[1]\n","        h = self.n_heads; hd = self.head_dim\n","\n","        Q = self.q_proj(queries).view(B, n, h, hd).transpose(1, 2)\n","        K = self.k_proj(anchors).view(B, k, h, hd).transpose(1, 2)\n","        V = self.v_proj(anchors).view(B, k, h, hd).transpose(1, 2)\n","\n","        attn = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(hd)\n","        attn = F.softmax(attn, dim=-1)\n","        ctx = torch.matmul(attn, V).transpose(1, 2).reshape(B, n, d)\n","\n","        # Extract quaternion and apply rotation\n","        q = self.quat_proj(ctx)  # [B, n, 4]\n","        q = F.normalize(q, dim=-1)\n","        rotated = self._quat_rotate(queries, q)\n","        return self.out_proj(ctx + rotated)\n","\n","    def _quat_rotate(self, v, q):\n","        \"\"\"Apply quaternion rotation to vectors. q: [B,n,4], v: [B,n,d].\"\"\"\n","        # For d > 3: rotate first 3 dims, pass rest through\n","        w, x, y, z = q[..., 0:1], q[..., 1:2], q[..., 2:3], q[..., 3:4]\n","        v3 = v[..., :3]\n","        # q * v * q^-1 via Rodriguez\n","        t = 2.0 * torch.cross(torch.cat([x, y, z], dim=-1), v3, dim=-1)\n","        v3_rot = v3 + w * t + torch.cross(torch.cat([x, y, z], dim=-1), t, dim=-1)\n","        if v.shape[-1] > 3:\n","            return torch.cat([v3_rot, v[..., 3:]], dim=-1)\n","        return v3_rot\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# QuaternionLiteFlow — Staged lighter quaternion\n","# ═══════════════════════════════════════════════════════════════════\n","\n","class QuaternionLiteFlow(BaseFlow):\n","    \"\"\"Lightweight quaternion prediction without full MHA.\n","\n","    Uses anchor centroid + query projection to predict rotation directly.\n","    Much lighter than full QuaternionFlow — trades attention resolution\n","    for speed.\n","    \"\"\"\n","    def __init__(self, d_model: int, n_anchors: int):\n","        super().__init__(d_model, n_anchors, name='quat_lite')\n","        self.anchor_compress = nn.Linear(d_model, d_model)\n","        self.query_proj = nn.Linear(d_model, d_model)\n","        self.quat_head = nn.Sequential(\n","            nn.Linear(d_model * 2, d_model),\n","            nn.GELU(),\n","            nn.Linear(d_model, 4),\n","        )\n","        self.out_proj = nn.Linear(d_model, d_model)\n","\n","    def _flow(self, anchors, queries):\n","        B, n, d = queries.shape\n","        # Compress anchors to single geometric summary\n","        anchor_ctx = self.anchor_compress(anchors.mean(dim=1, keepdim=True))  # [B, 1, d]\n","        anchor_ctx = anchor_ctx.expand(B, n, d)\n","\n","        q_proj = self.query_proj(queries)\n","        combined = torch.cat([q_proj, anchor_ctx], dim=-1)  # [B, n, 2d]\n","\n","        q = F.normalize(self.quat_head(combined), dim=-1)\n","        rotated = self._quat_rotate_simple(queries, q)\n","        return self.out_proj(rotated)\n","\n","    def _quat_rotate_simple(self, v, q):\n","        w, xyz = q[..., 0:1], q[..., 1:4]\n","        v3 = v[..., :3]\n","        t = 2.0 * torch.cross(xyz, v3, dim=-1)\n","        v3_rot = v3 + w * t + torch.cross(xyz, t, dim=-1)\n","        if v.shape[-1] > 3:\n","            return torch.cat([v3_rot, v[..., 3:]], dim=-1)\n","        return v3_rot\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# VelocityFlow — Angular velocity on tangent bundle\n","# ═══════════════════════════════════════════════════════════════════\n","\n","class VelocityFlow(BaseFlow):\n","    \"\"\"Angular velocity flow on the tangent space of the constellation.\n","\n","    Models dq/dt: the rate of change of the query embedding induced by\n","    the anchor geometry. Predicts velocity, integrates with Euler step.\n","\n","    The velocity is tangent to the hypersphere at each query point.\n","    \"\"\"\n","    def __init__(self, d_model: int, n_anchors: int):\n","        super().__init__(d_model, n_anchors, name='velocity')\n","        # Anchor-query interaction → velocity field\n","        self.anchor_proj = nn.Linear(d_model, d_model)\n","        self.query_proj = nn.Linear(d_model, d_model)\n","        self.vel_head = nn.Sequential(\n","            nn.Linear(d_model, d_model),\n","            nn.GELU(),\n","            nn.Linear(d_model, d_model),\n","        )\n","        self.dt = nn.Parameter(torch.tensor(0.1))  # learnable step size\n","\n","    def _flow(self, anchors, queries):\n","        B, n, d = queries.shape\n","        # Compute direction from queries toward anchor centroid\n","        a_proj = self.anchor_proj(anchors)  # [B, k, d]\n","        q_proj = self.query_proj(queries)   # [B, n, d]\n","\n","        # Soft attention: query-anchor similarity → weighted anchor direction\n","        sim = torch.bmm(q_proj, a_proj.transpose(-2, -1))  # [B, n, k]\n","        weights = F.softmax(sim / math.sqrt(d), dim=-1)\n","        direction = torch.bmm(weights, a_proj)  # [B, n, d]\n","\n","        # Velocity: project onto tangent space at query\n","        velocity = self.vel_head(direction - q_proj)\n","\n","        # Tangent projection: remove component along query direction\n","        q_norm = F.normalize(queries, dim=-1)\n","        radial = (velocity * q_norm).sum(dim=-1, keepdim=True) * q_norm\n","        tangent_vel = velocity - radial\n","\n","        # Euler integration\n","        return queries + self.dt * tangent_vel\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# MagnitudeFlow — Gram eigenvalue spectrum\n","# ═══════════════════════════════════════════════════════════════════\n","\n","class MagnitudeFlow(BaseFlow):\n","    \"\"\"Flow based on the Gram matrix eigenvalue magnitude spectrum.\n","\n","    Computes the anchor Gram matrix, extracts eigenvalues via FL eigh,\n","    uses the spectral profile to modulate query embeddings.\n","\n","    The eigenvalue magnitudes encode the constellation's energy distribution\n","    across geometric modes.\n","    \"\"\"\n","    def __init__(self, d_model: int, n_anchors: int):\n","        super().__init__(d_model, n_anchors, name='magnitude')\n","        # Project anchors to small geometric space for Gram computation\n","        self.geom_dim = min(n_anchors, 12)  # FL eigh sweet spot\n","        self.anchor_proj = nn.Linear(d_model, self.geom_dim)\n","        # Spectral → modulation\n","        self.spec_proj = nn.Sequential(\n","            nn.Linear(self.geom_dim, d_model),\n","            nn.GELU(),\n","            nn.Linear(d_model, d_model),\n","        )\n","        self.query_proj = nn.Linear(d_model, d_model)\n","        self.gate = nn.Linear(d_model * 2, d_model)\n","\n","    def _flow(self, anchors, queries):\n","        B, n, d = queries.shape\n","        # Project anchors to geometric space\n","        a_geom = self.anchor_proj(anchors)  # [B, k, geom_dim]\n","\n","        # Gram matrix eigenvalues → spectral profile\n","        G = torch.bmm(a_geom.transpose(-2, -1), a_geom)  # [B, geom_dim, geom_dim]\n","        # Use torch.linalg.eigh for now; swap to FL eigh in geolip.linalg\n","        eigenvalues, _ = torch.linalg.eigh(G)  # [B, geom_dim]\n","\n","        # Magnitude spectrum: how energy distributes across modes\n","        magnitudes = eigenvalues.abs().sqrt()  # [B, geom_dim] — the ω spectrum\n","        spec_embed = self.spec_proj(magnitudes)  # [B, d]\n","        spec_embed = spec_embed.unsqueeze(1).expand(B, n, d)\n","\n","        # Gate: blend spectral modulation with query\n","        q_proj = self.query_proj(queries)\n","        gate_input = torch.cat([q_proj, spec_embed], dim=-1)\n","        g = torch.sigmoid(self.gate(gate_input))\n","        return queries + g * spec_embed\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# OrbitalFlow — Omega angular resonance via FL eigh\n","# ═══════════════════════════════════════════════════════════════════\n","\n","class OrbitalFlow(BaseFlow):\n","    \"\"\"Omega-based orbital resonance flow.\n","\n","    Computes the constellation's resonance frequencies (ωᵢ = √λᵢ from\n","    Gram eigendecomposition), then uses the full eigendecomposition to\n","    project queries into the resonance basis, apply frequency-dependent\n","    modulation, and project back.\n","\n","    This flow directly uses the ω spectrum to shape the geometric response.\n","    Modes in the CV band [0.447, 0.480] (corresponding to λ ∈ [0.20, 0.23])\n","    are amplified. Modes outside are attenuated.\n","    \"\"\"\n","    def __init__(self, d_model: int, n_anchors: int, cv_lo: float = 0.20, cv_hi: float = 0.23):\n","        super().__init__(d_model, n_anchors, name='orbital')\n","        self.geom_dim = min(n_anchors, 12)\n","        self.anchor_proj = nn.Linear(d_model, self.geom_dim)\n","        self.cv_lo = cv_lo\n","        self.cv_hi = cv_hi\n","        # Per-mode learnable response curve\n","        self.mode_response = nn.Parameter(torch.ones(self.geom_dim))\n","        # Projection back to d_model\n","        self.query_to_geom = nn.Linear(d_model, self.geom_dim)\n","        self.geom_to_query = nn.Linear(self.geom_dim, d_model)\n","        self.out_proj = nn.Linear(d_model, d_model)\n","\n","    def _flow(self, anchors, queries):\n","        B, n, d = queries.shape\n","        a_geom = self.anchor_proj(anchors)  # [B, k, geom_dim]\n","        G = torch.bmm(a_geom.transpose(-2, -1), a_geom)  # [B, gd, gd]\n","\n","        # Eigendecomposition — the ω spectrum\n","        eigenvalues, eigenvectors = torch.linalg.eigh(G)  # [B, gd], [B, gd, gd]\n","\n","        # ω = √|λ|\n","        omega = eigenvalues.abs().sqrt()  # [B, gd]\n","\n","        # CV band resonance: modes near the attractor basin get amplified\n","        in_band = ((eigenvalues >= self.cv_lo) & (eigenvalues <= self.cv_hi)).float()\n","        near_binding = torch.exp(-10.0 * (eigenvalues - 0.29154).pow(2))\n","\n","        # Mode weighting: learned response × geometric structure\n","        mode_weight = self.mode_response.unsqueeze(0) * (1.0 + in_band + near_binding)\n","\n","        # Project queries into resonance basis\n","        q_geom = self.query_to_geom(queries)  # [B, n, gd]\n","        # Rotate into eigenbasis: q_eigen = q_geom @ V\n","        q_eigen = torch.bmm(q_geom, eigenvectors)  # [B, n, gd]\n","\n","        # Apply frequency-dependent modulation\n","        q_modulated = q_eigen * mode_weight.unsqueeze(1)  # [B, n, gd]\n","\n","        # Rotate back: q_out = q_modulated @ V^T\n","        q_out = torch.bmm(q_modulated, eigenvectors.transpose(-2, -1))\n","\n","        # Project back to d_model\n","        return self.out_proj(self.geom_to_query(q_out) + queries)\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# AlignmentFlow — SVD-based Procrustes alignment\n","# ═══════════════════════════════════════════════════════════════════\n","\n","class AlignmentFlow(BaseFlow):\n","    \"\"\"SVD alignment flow via soft Procrustes rotation.\n","\n","    Computes attention-weighted anchor targets per query, then finds the\n","    optimal rotation aligning queries toward those targets via SVD of\n","    the cross-covariance matrix.\n","    \"\"\"\n","    def __init__(self, d_model: int, n_anchors: int):\n","        super().__init__(d_model, n_anchors, name='alignment')\n","        self.anchor_proj = nn.Linear(d_model, d_model)\n","        self.query_proj = nn.Linear(d_model, d_model)\n","        self.strength = nn.Parameter(torch.tensor(0.1))\n","\n","    def _flow(self, anchors, queries):\n","        B, n, d = queries.shape\n","        a_proj = self.anchor_proj(anchors)  # [B, k, d]\n","        q_proj = self.query_proj(queries)   # [B, n, d]\n","\n","        # Attention-weighted anchors → per-query targets [B, n, d]\n","        sim = torch.bmm(q_proj, a_proj.transpose(-2, -1)) / math.sqrt(d)\n","        weights = F.softmax(sim, dim=-1)  # [B, n, k]\n","        targets = torch.bmm(weights, a_proj)  # [B, n, d]\n","\n","        # Cross-covariance: C = Q^T @ T, both [B, n, d] → C is [B, d, d]\n","        C = torch.bmm(q_proj.transpose(-2, -1), targets)\n","\n","        # SVD → optimal rotation (Procrustes)\n","        U, _, Vh = torch.linalg.svd(C)\n","        R = torch.bmm(U, Vh)  # [B, d, d]\n","\n","        # Soft rotation\n","        q_rotated = torch.bmm(queries, R)\n","        return queries + self.strength * (q_rotated - queries)\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# Flow Ensemble\n","# ═══════════════════════════════════════════════════════════════════\n","\n","class FlowEnsemble(nn.Module):\n","    \"\"\"Ensemble fusion of multiple geometric flows.\n","\n","    Each flow produces a prediction and a confidence score.\n","    The ensemble fuses predictions weighted by confidence.\n","\n","    The fusion can be:\n","      'weighted': confidence-weighted average\n","      'gated':    learned gate over concatenated predictions\n","      'residual': sum of confidence-weighted residuals from input\n","    \"\"\"\n","    def __init__(self, flows: List[BaseFlow], d_model: int, fusion: str = 'weighted'):\n","        super().__init__()\n","        self.flows = nn.ModuleList(flows)\n","        self.d_model = d_model\n","        self.fusion = fusion\n","        self.n_flows = len(flows)\n","\n","        if fusion == 'gated':\n","            self.gate = nn.Sequential(\n","                nn.Linear(d_model * self.n_flows, d_model),\n","                nn.GELU(),\n","                nn.Linear(d_model, d_model),\n","            )\n","\n","        # Per-flow learnable temperature\n","        self.temperature = nn.Parameter(torch.ones(self.n_flows))\n","\n","    def forward(self, anchors: Tensor, queries: Tensor) -> Tensor:\n","        \"\"\"\n","        Args:\n","            anchors: [B, k, d] constellation anchors\n","            queries: [B, n, d] query embeddings\n","\n","        Returns:\n","            fused: [B, n, d] ensemble prediction\n","        \"\"\"\n","        predictions = []\n","        confidences = []\n","\n","        for i, flow in enumerate(self.flows):\n","            pred, conf = flow(anchors, queries)\n","            predictions.append(pred)\n","            confidences.append(conf * self.temperature[i])\n","\n","        if self.fusion == 'weighted':\n","            return self._weighted_fusion(predictions, confidences)\n","        elif self.fusion == 'gated':\n","            return self._gated_fusion(predictions, confidences)\n","        elif self.fusion == 'residual':\n","            return self._residual_fusion(predictions, confidences, queries)\n","        else:\n","            raise ValueError(f\"Unknown fusion: {self.fusion}\")\n","\n","    def _weighted_fusion(self, preds, confs):\n","        # Stack confidences and normalize\n","        conf_stack = torch.cat(confs, dim=-1)  # [B, n, n_flows]\n","        weights = F.softmax(conf_stack, dim=-1)  # [B, n, n_flows]\n","        pred_stack = torch.stack(preds, dim=-1)  # [B, n, d, n_flows]\n","        return (pred_stack * weights.unsqueeze(-2)).sum(dim=-1)\n","\n","    def _gated_fusion(self, preds, confs):\n","        cat = torch.cat(preds, dim=-1)  # [B, n, d * n_flows]\n","        return self.gate(cat)\n","\n","    def _residual_fusion(self, preds, confs, queries):\n","        conf_stack = torch.cat(confs, dim=-1)\n","        weights = F.softmax(conf_stack, dim=-1)\n","        residuals = torch.stack([p - queries for p in preds], dim=-1)\n","        fused_residual = (residuals * weights.unsqueeze(-2)).sum(dim=-1)\n","        return queries + fused_residual\n","\n","    def flow_diagnostics(self, anchors: Tensor, queries: Tensor) -> dict:\n","        \"\"\"Run all flows and return per-flow diagnostics.\"\"\"\n","        diag = {}\n","        for i, flow in enumerate(self.flows):\n","            pred, conf = flow(anchors, queries)\n","            diag[flow.name] = {\n","                'pred_norm': pred.norm(dim=-1).mean().item(),\n","                'confidence_mean': conf.mean().item(),\n","                'confidence_std': conf.std().item(),\n","                'residual_norm': (pred - queries).norm(dim=-1).mean().item(),\n","                'temperature': self.temperature[i].item(),\n","            }\n","        return diag"],"metadata":{"id":"w3ID3KY3kdp6"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["\"\"\"\n","Flow ensemble — smoke test + diagnostics.\n","\"\"\"\n","import torch\n","import torch.nn as nn\n","import sys, time\n","\n","dev = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n","B, n, k, d = 32, 128, 64, 256\n","\n","anchors = torch.randn(B, k, d, device=dev)\n","queries = torch.randn(B, n, d, device=dev)\n","\n","print(\"=\" * 68)\n","print(\"  Flow Ensemble — Smoke Test\")\n","print(\"=\" * 68)\n","print(f\"  B={B} n={n} k={k} d={d} device={dev}\")\n","\n","# Test each flow independently\n","flows_cfg = [\n","    ('QuaternionFlow',     lambda: QuaternionFlow(d, k, n_heads=4)),\n","    ('QuaternionLiteFlow', lambda: QuaternionLiteFlow(d, k)),\n","    ('VelocityFlow',       lambda: VelocityFlow(d, k)),\n","    ('MagnitudeFlow',      lambda: MagnitudeFlow(d, k)),\n","    ('OrbitalFlow',        lambda: OrbitalFlow(d, k)),\n","    ('AlignmentFlow',      lambda: AlignmentFlow(d, k)),\n","]\n","\n","print(f\"\\n  {'Flow':<22} {'Params':>8} {'Out shape':>14} {'Fwd (ms)':>10} {'Conf μ':>8}\")\n","print(f\"  {'─'*22} {'─'*8} {'─'*14} {'─'*10} {'─'*8}\")\n","\n","live_flows = []\n","for name, ctor in flows_cfg:\n","    try:\n","        flow = ctor().to(dev)\n","        params = sum(p.numel() for p in flow.parameters())\n","\n","        # Warmup\n","        for _ in range(3):\n","            flow(anchors, queries)\n","        if dev.type == 'cuda':\n","            torch.cuda.synchronize()\n","\n","        # Time\n","        t0 = time.perf_counter()\n","        N_runs = 50\n","        for _ in range(N_runs):\n","            pred, conf = flow(anchors, queries)\n","        if dev.type == 'cuda':\n","            torch.cuda.synchronize()\n","        elapsed = (time.perf_counter() - t0) / N_runs * 1000\n","\n","        print(f\"  {name:<22} {params:>8,} {str(tuple(pred.shape)):>14} {elapsed:>9.2f} {conf.mean().item():>8.3f}\")\n","        live_flows.append(flow)\n","    except Exception as e:\n","        print(f\"  {name:<22} FAILED: {str(e)[:40]}\")\n","\n","# Test ensemble\n","print(f\"\\n  Ensemble tests:\")\n","for fusion in ['weighted', 'gated', 'residual']:\n","    try:\n","        ens = FlowEnsemble(live_flows, d, fusion=fusion).to(dev)\n","        params = sum(p.numel() for p in ens.parameters())\n","        out = ens(anchors, queries)\n","        print(f\"  {fusion:<12} params={params:>10,}  out={tuple(out.shape)}  norm={out.norm(dim=-1).mean():.3f}\")\n","\n","        # Diagnostics\n","        diag = ens.flow_diagnostics(anchors, queries)\n","        for fname, stats in diag.items():\n","            print(f\"    {fname:<18} conf={stats['confidence_mean']:.3f}±{stats['confidence_std']:.3f}  \"\n","                  f\"residual={stats['residual_norm']:.3f}  temp={stats['temperature']:.3f}\")\n","    except Exception as e:\n","        print(f\"  {fusion:<12} FAILED: {str(e)[:50]}\")\n","\n","# Gradient flow test\n","print(f\"\\n  Gradient flow:\")\n","ens = FlowEnsemble(live_flows, d, fusion='weighted').to(dev)\n","out = ens(anchors, queries)\n","loss = out.sum()\n","loss.backward()\n","for flow in ens.flows:\n","    grads = [p.grad is not None for p in flow.parameters()]\n","    pct = sum(grads) / max(len(grads), 1) * 100\n","    print(f\"  {flow.name:<18} {pct:.0f}% params have gradients\")\n","\n","print(\"=\" * 68)"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"srL5ZIWdkwbO","executionInfo":{"status":"ok","timestamp":1775090162386,"user_tz":420,"elapsed":28233,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"da35ab74-d93a-4f77-a588-45829eee215e"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["====================================================================\n","  Flow Ensemble — Smoke Test\n","====================================================================\n","  B=32 n=128 k=64 d=256 device=cuda\n","\n","  Flow                     Params      Out shape   Fwd (ms)   Conf μ\n","  ────────────────────── ──────── ────────────── ────────── ────────\n","  QuaternionFlow          280,709 (32, 128, 256)      0.27    0.508\n","  QuaternionLiteFlow      346,245 (32, 128, 256)      0.22    0.521\n","  VelocityFlow            279,682 (32, 128, 256)      0.23    0.453\n","  MagnitudeFlow           285,837 (32, 128, 256)      0.36    0.507\n","  OrbitalFlow              91,813 (32, 128, 256)      0.38    0.495\n","  AlignmentFlow           148,098 (32, 128, 256)    466.97    0.517\n","\n","  Ensemble tests:\n","  weighted     params= 1,432,390  out=(32, 128, 256)  norm=8.322\n","    quaternion         conf=0.508±0.029  residual=18.512  temp=1.000\n","    quat_lite          conf=0.521±0.025  residual=18.478  temp=1.000\n","    velocity           conf=0.453±0.049  residual=0.176  temp=1.000\n","    magnitude          conf=0.507±0.044  residual=9.299  temp=1.000\n","    orbital            conf=0.495±0.029  residual=18.790  temp=1.000\n","    alignment          conf=0.517±0.044  residual=2.113  temp=1.000\n","  gated        params= 1,891,654  out=(32, 128, 256)  norm=2.646\n","    quaternion         conf=0.508±0.029  residual=18.512  temp=1.000\n","    quat_lite          conf=0.521±0.025  residual=18.478  temp=1.000\n","    velocity           conf=0.453±0.049  residual=0.176  temp=1.000\n","    magnitude          conf=0.507±0.044  residual=9.299  temp=1.000\n","    orbital            conf=0.495±0.029  residual=18.790  temp=1.000\n","    alignment          conf=0.517±0.044  residual=2.113  temp=1.000\n","  residual     params= 1,432,390  out=(32, 128, 256)  norm=8.322\n","    quaternion         conf=0.508±0.029  residual=18.512  temp=1.000\n","    quat_lite          conf=0.521±0.025  residual=18.478  temp=1.000\n","    velocity           conf=0.453±0.049  residual=0.176  temp=1.000\n","    magnitude          conf=0.507±0.044  residual=9.299  temp=1.000\n","    orbital            conf=0.495±0.029  residual=18.790  temp=1.000\n","    alignment          conf=0.517±0.044  residual=2.113  temp=1.000\n","\n","  Gradient flow:\n","  quaternion         100% params have gradients\n","  quat_lite          100% params have gradients\n","  velocity           100% params have gradients\n","  magnitude          100% params have gradients\n","  orbital            100% params have gradients\n","  alignment          100% params have gradients\n","====================================================================\n"]}]},{"cell_type":"code","source":["\"\"\"\n","geolip.flows — Multi-flow ensemble for constellation geometry.\n","\n","Each flow predicts the same geometric output using a different mathematical\n","formulation. The ensemble fuses predictions based on learned confidence.\n","\n","Flows:\n","  QuaternionFlow      — Full MHA quaternion rotation (existing, heavyweight)\n","  QuaternionLiteFlow  — Staged quaternion with lighter spectral computation\n","  VelocityFlow        — Angular velocity dq/dt on the tangent bundle\n","  MagnitudeFlow       — Flow magnitude via Gram eigenvalue spectrum\n","  OrbitalFlow         — Omega-based orbital resonance using FL eigh\n","  AlignmentFlow       — SVD alignment via Procrustes rotation\n","\n","Architecture:\n","  Each flow: same input (anchors [B,k,d], queries [B,n,d]) → output [B,n,d]\n","  Ensemble:  weighted fusion with learned per-flow confidence\n","\n","Usage:\n","    from geolip.flows import FlowEnsemble, OrbitalFlow, AlignmentFlow\n","\n","    ensemble = FlowEnsemble(\n","        flows=[OrbitalFlow(d=256, k=128), AlignmentFlow(d=256, k=128)],\n","        d_model=256,\n","    )\n","    output = ensemble(anchors, queries)  # [B, n, d]\n","\"\"\"\n","\n","import math\n","import torch\n","import torch.nn as nn\n","import torch.nn.functional as F\n","from torch import Tensor\n","from typing import List, Optional, Tuple\n","\n","# Use geolip_core.linalg when available (FL eigh, Triton SVD, etc.)\n","# Falls back to torch.linalg transparently\n","try:\n","    import geolip_core.linalg as LA\n","except ImportError:\n","    import torch.linalg as LA\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# Base Flow\n","# ═══════════════════════════════════════════════════════════════════\n","\n","class BaseFlow(nn.Module):\n","    \"\"\"Base class for all geometric flows.\n","\n","    All flows share the same interface:\n","      Input:  anchors [B, k, d], queries [B, n, d]\n","      Output: prediction [B, n, d], confidence [B, n, 1]\n","\n","    Subclasses implement _flow() with their specific math.\n","    \"\"\"\n","    def __init__(self, d_model: int, n_anchors: int, name: str = 'base'):\n","        super().__init__()\n","        self.d_model = d_model\n","        self.n_anchors = n_anchors\n","        self.name = name\n","        # Confidence head: scalar per query position\n","        self.confidence = nn.Sequential(\n","            nn.Linear(d_model, d_model // 4),\n","            nn.GELU(),\n","            nn.Linear(d_model // 4, 1),\n","        )\n","\n","    def forward(self, anchors: Tensor, queries: Tensor) -> Tuple[Tensor, Tensor]:\n","        \"\"\"\n","        Args:\n","            anchors: [B, k, d] constellation anchor points\n","            queries: [B, n, d] query embeddings\n","\n","        Returns:\n","            prediction: [B, n, d] geometric prediction\n","            confidence: [B, n, 1] per-query confidence score\n","        \"\"\"\n","        pred = self._flow(anchors, queries)\n","        conf = torch.sigmoid(self.confidence(pred))\n","        return pred, conf\n","\n","    def _flow(self, anchors: Tensor, queries: Tensor) -> Tensor:\n","        raise NotImplementedError\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# QuaternionFlow — Full MHA quaternion rotation\n","# ═══════════════════════════════════════════════════════════════════\n","\n","class QuaternionFlow(BaseFlow):\n","    \"\"\"Full multi-head attention with quaternion geometric rotation.\n","\n","    Computes query-anchor attention, extracts rotation quaternion from\n","    attention-weighted anchor geometry, applies rotation to queries.\n","    Heavyweight — the full-fidelity path.\n","    \"\"\"\n","    def __init__(self, d_model: int, n_anchors: int, n_heads: int = 4):\n","        super().__init__(d_model, n_anchors, name='quaternion')\n","        self.n_heads = n_heads\n","        self.head_dim = d_model // n_heads\n","        self.q_proj = nn.Linear(d_model, d_model)\n","        self.k_proj = nn.Linear(d_model, d_model)\n","        self.v_proj = nn.Linear(d_model, d_model)\n","        self.out_proj = nn.Linear(d_model, d_model)\n","        # Quaternion components: scalar + 3 imaginary from attention output\n","        self.quat_proj = nn.Linear(d_model, 4)\n","\n","    def _flow(self, anchors, queries):\n","        B, n, d = queries.shape\n","        k = anchors.shape[1]\n","        h = self.n_heads; hd = self.head_dim\n","\n","        Q = self.q_proj(queries).view(B, n, h, hd).transpose(1, 2)\n","        K = self.k_proj(anchors).view(B, k, h, hd).transpose(1, 2)\n","        V = self.v_proj(anchors).view(B, k, h, hd).transpose(1, 2)\n","\n","        attn = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(hd)\n","        attn = F.softmax(attn, dim=-1)\n","        ctx = torch.matmul(attn, V).transpose(1, 2).reshape(B, n, d)\n","\n","        # Extract quaternion and apply rotation\n","        q = self.quat_proj(ctx)  # [B, n, 4]\n","        q = F.normalize(q, dim=-1)\n","        rotated = self._quat_rotate(queries, q)\n","        return self.out_proj(ctx + rotated)\n","\n","    def _quat_rotate(self, v, q):\n","        \"\"\"Apply quaternion rotation to vectors. q: [B,n,4], v: [B,n,d].\"\"\"\n","        # For d > 3: rotate first 3 dims, pass rest through\n","        w, x, y, z = q[..., 0:1], q[..., 1:2], q[..., 2:3], q[..., 3:4]\n","        v3 = v[..., :3]\n","        # q * v * q^-1 via Rodriguez\n","        t = 2.0 * torch.cross(torch.cat([x, y, z], dim=-1), v3, dim=-1)\n","        v3_rot = v3 + w * t + torch.cross(torch.cat([x, y, z], dim=-1), t, dim=-1)\n","        if v.shape[-1] > 3:\n","            return torch.cat([v3_rot, v[..., 3:]], dim=-1)\n","        return v3_rot\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# QuaternionLiteFlow — Staged lighter quaternion\n","# ═══════════════════════════════════════════════════════════════════\n","\n","class QuaternionLiteFlow(BaseFlow):\n","    \"\"\"Lightweight quaternion prediction without full MHA.\n","\n","    Uses anchor centroid + query projection to predict rotation directly.\n","    Much lighter than full QuaternionFlow — trades attention resolution\n","    for speed.\n","    \"\"\"\n","    def __init__(self, d_model: int, n_anchors: int):\n","        super().__init__(d_model, n_anchors, name='quat_lite')\n","        self.anchor_compress = nn.Linear(d_model, d_model)\n","        self.query_proj = nn.Linear(d_model, d_model)\n","        self.quat_head = nn.Sequential(\n","            nn.Linear(d_model * 2, d_model),\n","            nn.GELU(),\n","            nn.Linear(d_model, 4),\n","        )\n","        self.out_proj = nn.Linear(d_model, d_model)\n","\n","    def _flow(self, anchors, queries):\n","        B, n, d = queries.shape\n","        # Compress anchors to single geometric summary\n","        anchor_ctx = self.anchor_compress(anchors.mean(dim=1, keepdim=True))  # [B, 1, d]\n","        anchor_ctx = anchor_ctx.expand(B, n, d)\n","\n","        q_proj = self.query_proj(queries)\n","        combined = torch.cat([q_proj, anchor_ctx], dim=-1)  # [B, n, 2d]\n","\n","        q = F.normalize(self.quat_head(combined), dim=-1)\n","        rotated = self._quat_rotate_simple(queries, q)\n","        return self.out_proj(rotated)\n","\n","    def _quat_rotate_simple(self, v, q):\n","        w, xyz = q[..., 0:1], q[..., 1:4]\n","        v3 = v[..., :3]\n","        t = 2.0 * torch.cross(xyz, v3, dim=-1)\n","        v3_rot = v3 + w * t + torch.cross(xyz, t, dim=-1)\n","        if v.shape[-1] > 3:\n","            return torch.cat([v3_rot, v[..., 3:]], dim=-1)\n","        return v3_rot\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# VelocityFlow — Angular velocity on tangent bundle\n","# ═══════════════════════════════════════════════════════════════════\n","\n","class VelocityFlow(BaseFlow):\n","    \"\"\"Angular velocity flow on the tangent space of the constellation.\n","\n","    Models dq/dt: the rate of change of the query embedding induced by\n","    the anchor geometry. Predicts velocity, integrates with Euler step.\n","\n","    The velocity is tangent to the hypersphere at each query point.\n","    \"\"\"\n","    def __init__(self, d_model: int, n_anchors: int):\n","        super().__init__(d_model, n_anchors, name='velocity')\n","        # Anchor-query interaction → velocity field\n","        self.anchor_proj = nn.Linear(d_model, d_model)\n","        self.query_proj = nn.Linear(d_model, d_model)\n","        self.vel_head = nn.Sequential(\n","            nn.Linear(d_model, d_model),\n","            nn.GELU(),\n","            nn.Linear(d_model, d_model),\n","        )\n","        self.dt = nn.Parameter(torch.tensor(0.1))  # learnable step size\n","\n","    def _flow(self, anchors, queries):\n","        B, n, d = queries.shape\n","        # Compute direction from queries toward anchor centroid\n","        a_proj = self.anchor_proj(anchors)  # [B, k, d]\n","        q_proj = self.query_proj(queries)   # [B, n, d]\n","\n","        # Soft attention: query-anchor similarity → weighted anchor direction\n","        sim = torch.bmm(q_proj, a_proj.transpose(-2, -1))  # [B, n, k]\n","        weights = F.softmax(sim / math.sqrt(d), dim=-1)\n","        direction = torch.bmm(weights, a_proj)  # [B, n, d]\n","\n","        # Velocity: project onto tangent space at query\n","        velocity = self.vel_head(direction - q_proj)\n","\n","        # Tangent projection: remove component along query direction\n","        q_norm = F.normalize(queries, dim=-1)\n","        radial = (velocity * q_norm).sum(dim=-1, keepdim=True) * q_norm\n","        tangent_vel = velocity - radial\n","\n","        # Euler integration\n","        return queries + self.dt * tangent_vel\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# MagnitudeFlow — Gram eigenvalue spectrum\n","# ═══════════════════════════════════════════════════════════════════\n","\n","class MagnitudeFlow(BaseFlow):\n","    \"\"\"Flow based on the Gram matrix eigenvalue magnitude spectrum.\n","\n","    Computes the anchor Gram matrix, extracts eigenvalues via FL eigh,\n","    uses the spectral profile to modulate query embeddings.\n","\n","    The eigenvalue magnitudes encode the constellation's energy distribution\n","    across geometric modes.\n","    \"\"\"\n","    def __init__(self, d_model: int, n_anchors: int):\n","        super().__init__(d_model, n_anchors, name='magnitude')\n","        # Project anchors to small geometric space for Gram computation\n","        self.geom_dim = min(n_anchors, 12)  # FL eigh sweet spot\n","        self.anchor_proj = nn.Linear(d_model, self.geom_dim)\n","        # Spectral → modulation\n","        self.spec_proj = nn.Sequential(\n","            nn.Linear(self.geom_dim, d_model),\n","            nn.GELU(),\n","            nn.Linear(d_model, d_model),\n","        )\n","        self.query_proj = nn.Linear(d_model, d_model)\n","        self.gate = nn.Linear(d_model * 2, d_model)\n","\n","    def _flow(self, anchors, queries):\n","        B, n, d = queries.shape\n","        # Project anchors to geometric space\n","        a_geom = self.anchor_proj(anchors)  # [B, k, geom_dim]\n","\n","        # Gram matrix\n","        G = torch.bmm(a_geom.transpose(-2, -1), a_geom)  # [B, geom_dim, geom_dim]\n","\n","        # Eigendecomposition — differentiable through torch.linalg.eigh\n","        eigenvalues, _ = LA.eigh(G, method='torch')  # [B, geom_dim]\n","\n","        # Magnitude spectrum: how energy distributes across modes\n","        magnitudes = eigenvalues.abs().sqrt()  # [B, geom_dim] — the ω spectrum\n","        spec_embed = self.spec_proj(magnitudes)  # [B, d]\n","        spec_embed = spec_embed.unsqueeze(1).expand(B, n, d)\n","\n","        # Gate: blend spectral modulation with query\n","        q_proj = self.query_proj(queries)\n","        gate_input = torch.cat([q_proj, spec_embed], dim=-1)\n","        g = torch.sigmoid(self.gate(gate_input))\n","        return queries + g * spec_embed\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# OrbitalFlow — Omega angular resonance via FL eigh\n","# ═══════════════════════════════════════════════════════════════════\n","\n","class OrbitalFlow(BaseFlow):\n","    \"\"\"Omega-based orbital resonance flow.\n","\n","    Computes the constellation's resonance frequencies (ωᵢ = √λᵢ from\n","    Gram eigendecomposition), then uses the full eigendecomposition to\n","    project queries into the resonance basis, apply frequency-dependent\n","    modulation, and project back.\n","\n","    This flow directly uses the ω spectrum to shape the geometric response.\n","    Modes in the CV band [0.447, 0.480] (corresponding to λ ∈ [0.20, 0.23])\n","    are amplified. Modes outside are attenuated.\n","    \"\"\"\n","    def __init__(self, d_model: int, n_anchors: int, cv_lo: float = 0.20, cv_hi: float = 0.23):\n","        super().__init__(d_model, n_anchors, name='orbital')\n","        self.geom_dim = min(n_anchors, 12)\n","        self.anchor_proj = nn.Linear(d_model, self.geom_dim)\n","        self.cv_lo = cv_lo\n","        self.cv_hi = cv_hi\n","        # Per-mode learnable response curve\n","        self.mode_response = nn.Parameter(torch.ones(self.geom_dim))\n","        # Projection back to d_model\n","        self.query_to_geom = nn.Linear(d_model, self.geom_dim)\n","        self.geom_to_query = nn.Linear(self.geom_dim, d_model)\n","        self.out_proj = nn.Linear(d_model, d_model)\n","\n","    def _flow(self, anchors, queries):\n","        B, n, d = queries.shape\n","        a_geom = self.anchor_proj(anchors)  # [B, k, geom_dim]\n","        G = torch.bmm(a_geom.transpose(-2, -1), a_geom)  # [B, gd, gd]\n","\n","        # Eigendecomposition — the ω spectrum (differentiable via torch.linalg.eigh)\n","        eigenvalues, eigenvectors = LA.eigh(G, method='torch')  # [B, gd], [B, gd, gd]\n","\n","        # ω = √|λ|\n","        omega = eigenvalues.abs().sqrt()  # [B, gd]\n","\n","        # CV band resonance: modes near the attractor basin get amplified\n","        in_band = ((eigenvalues >= self.cv_lo) & (eigenvalues <= self.cv_hi)).float()\n","        near_binding = torch.exp(-10.0 * (eigenvalues - 0.29154).pow(2))\n","\n","        # Mode weighting: learned response × geometric structure\n","        mode_weight = self.mode_response.unsqueeze(0) * (1.0 + in_band + near_binding)\n","\n","        # Project queries into resonance basis\n","        q_geom = self.query_to_geom(queries)  # [B, n, gd]\n","        # Rotate into eigenbasis: q_eigen = q_geom @ V\n","        q_eigen = torch.bmm(q_geom, eigenvectors)  # [B, n, gd]\n","\n","        # Apply frequency-dependent modulation\n","        q_modulated = q_eigen * mode_weight.unsqueeze(1)  # [B, n, gd]\n","\n","        # Rotate back: q_out = q_modulated @ V^T\n","        q_out = torch.bmm(q_modulated, eigenvectors.transpose(-2, -1))\n","\n","        # Project back to d_model\n","        return self.out_proj(self.geom_to_query(q_out) + queries)\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# AlignmentFlow — SVD-based Procrustes alignment\n","# ═══════════════════════════════════════════════════════════════════\n","\n","class AlignmentFlow(BaseFlow):\n","    \"\"\"SVD alignment flow via soft Procrustes rotation in projected space.\n","\n","    Projects to geom_dim, computes optimal rotation via SVD of the\n","    cross-covariance in the small space, applies rotation, projects back.\n","    \"\"\"\n","    def __init__(self, d_model: int, n_anchors: int):\n","        super().__init__(d_model, n_anchors, name='alignment')\n","        self.geom_dim = min(n_anchors, 12)  # FL eigh sweet spot\n","        self.anchor_proj = nn.Linear(d_model, self.geom_dim)\n","        self.query_proj = nn.Linear(d_model, self.geom_dim)\n","        self.geom_to_query = nn.Linear(self.geom_dim, d_model)\n","        self.strength = nn.Parameter(torch.tensor(0.1))\n","\n","    def _flow(self, anchors, queries):\n","        B, n, d = queries.shape\n","        # Project to small geometric space\n","        a_proj = self.anchor_proj(anchors)  # [B, k, geom_dim]\n","        q_proj = self.query_proj(queries)   # [B, n, geom_dim]\n","\n","        # Attention-weighted anchors → per-query targets [B, n, geom_dim]\n","        sim = torch.bmm(q_proj, a_proj.transpose(-2, -1)) / math.sqrt(self.geom_dim)\n","        weights = F.softmax(sim, dim=-1)  # [B, n, k]\n","        targets = torch.bmm(weights, a_proj)  # [B, n, geom_dim]\n","\n","        # Cross-covariance in small space: [B, geom_dim, geom_dim]\n","        C = torch.bmm(q_proj.transpose(-2, -1), targets)\n","\n","        # SVD → optimal rotation via gram_eigh (differentiable, no in-place ops)\n","        U, _, Vh = LA.svd(C, method='gram_eigh')\n","        R = torch.bmm(U, Vh)  # [B, geom_dim, geom_dim]\n","\n","        # Rotate queries in small space, project back to d_model\n","        q_rotated = torch.bmm(q_proj, R)  # [B, n, geom_dim]\n","        delta = self.geom_to_query(q_rotated - q_proj)  # [B, n, d]\n","        return queries + self.strength * delta\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# Flow Ensemble\n","# ═══════════════════════════════════════════════════════════════════\n","\n","class FlowEnsemble(nn.Module):\n","    \"\"\"Ensemble fusion of multiple geometric flows.\n","\n","    Each flow produces a prediction and a confidence score.\n","    The ensemble fuses predictions weighted by confidence.\n","\n","    The fusion can be:\n","      'weighted': confidence-weighted average\n","      'gated':    learned gate over concatenated predictions\n","      'residual': sum of confidence-weighted residuals from input\n","    \"\"\"\n","    def __init__(self, flows: List[BaseFlow], d_model: int, fusion: str = 'weighted'):\n","        super().__init__()\n","        self.flows = nn.ModuleList(flows)\n","        self.d_model = d_model\n","        self.fusion = fusion\n","        self.n_flows = len(flows)\n","\n","        if fusion == 'gated':\n","            self.gate = nn.Sequential(\n","                nn.Linear(d_model * self.n_flows, d_model),\n","                nn.GELU(),\n","                nn.Linear(d_model, d_model),\n","            )\n","\n","        # Per-flow learnable temperature\n","        self.temperature = nn.Parameter(torch.ones(self.n_flows))\n","\n","    def forward(self, anchors: Tensor, queries: Tensor) -> Tensor:\n","        \"\"\"\n","        Args:\n","            anchors: [B, k, d] constellation anchors\n","            queries: [B, n, d] query embeddings\n","\n","        Returns:\n","            fused: [B, n, d] ensemble prediction\n","        \"\"\"\n","        predictions = []\n","        confidences = []\n","\n","        for i, flow in enumerate(self.flows):\n","            pred, conf = flow(anchors, queries)\n","            predictions.append(pred)\n","            confidences.append(conf * self.temperature[i])\n","\n","        if self.fusion == 'weighted':\n","            return self._weighted_fusion(predictions, confidences)\n","        elif self.fusion == 'gated':\n","            return self._gated_fusion(predictions, confidences)\n","        elif self.fusion == 'residual':\n","            return self._residual_fusion(predictions, confidences, queries)\n","        else:\n","            raise ValueError(f\"Unknown fusion: {self.fusion}\")\n","\n","    def _weighted_fusion(self, preds, confs):\n","        # Stack confidences and normalize\n","        conf_stack = torch.cat(confs, dim=-1)  # [B, n, n_flows]\n","        weights = F.softmax(conf_stack, dim=-1)  # [B, n, n_flows]\n","        pred_stack = torch.stack(preds, dim=-1)  # [B, n, d, n_flows]\n","        return (pred_stack * weights.unsqueeze(-2)).sum(dim=-1)\n","\n","    def _gated_fusion(self, preds, confs):\n","        cat = torch.cat(preds, dim=-1)  # [B, n, d * n_flows]\n","        return self.gate(cat)\n","\n","    def _residual_fusion(self, preds, confs, queries):\n","        conf_stack = torch.cat(confs, dim=-1)\n","        weights = F.softmax(conf_stack, dim=-1)\n","        residuals = torch.stack([p - queries for p in preds], dim=-1)\n","        fused_residual = (residuals * weights.unsqueeze(-2)).sum(dim=-1)\n","        return queries + fused_residual\n","\n","    def flow_diagnostics(self, anchors: Tensor, queries: Tensor) -> dict:\n","        \"\"\"Run all flows and return per-flow diagnostics.\"\"\"\n","        diag = {}\n","        for i, flow in enumerate(self.flows):\n","            pred, conf = flow(anchors, queries)\n","            diag[flow.name] = {\n","                'pred_norm': pred.norm(dim=-1).mean().item(),\n","                'confidence_mean': conf.mean().item(),\n","                'confidence_std': conf.std().item(),\n","                'residual_norm': (pred - queries).norm(dim=-1).mean().item(),\n","                'temperature': self.temperature[i].item(),\n","            }\n","        return diag"],"metadata":{"id":"2ztPULHX-IS5"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["\"\"\"\n","Flow Ensemble — Expanded Test Suite.\n","\n","Assumes geolip-core is installed (Colab with repo loaded).\n","Tests: smoke, linalg integration, multi-scale, ensemble fusion,\n","       gradient health, ablation, compile compatibility, memory.\n","\"\"\"\n","import torch\n","import torch.nn as nn\n","import torch.nn.functional as F\n","import sys, time, gc\n","\n","# ── Verify geolip_core.linalg is available ──\n","try:\n","    import geolip_core.linalg as LA\n","    HAS_GEOLIP_LINALG = True\n","    print(f\"geolip_core.linalg: available\")\n","    LA.backend.status()\n","except ImportError:\n","    import torch.linalg as LA\n","    HAS_GEOLIP_LINALG = False\n","    print(\"geolip_core.linalg: NOT available, using torch.linalg fallback\")\n","\n","\n","dev = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n","\n","\n","def sync():\n","    if dev.type == 'cuda':\n","        torch.cuda.synchronize()\n","\n","def time_fn(fn, warmup=5, runs=50):\n","    for _ in range(warmup): fn()\n","    sync()\n","    t0 = time.perf_counter()\n","    for _ in range(runs): fn()\n","    sync()\n","    return (time.perf_counter() - t0) / runs * 1000\n","\n","def fmt(ms):\n","    if ms < 1: return f\"{ms*1000:.0f}us\"\n","    return f\"{ms:.2f}ms\"\n","\n","def make_data(B, n, k, d):\n","    anchors = F.normalize(torch.randn(B, k, d, device=dev), dim=-1)\n","    queries = F.normalize(torch.randn(B, n, d, device=dev), dim=-1)\n","    return anchors, queries\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","print(\"=\" * 72)\n","print(\"  Flow Ensemble — Expanded Test Suite\")\n","print(\"=\" * 72)\n","print(f\"  device={dev}  geolip_core.linalg={HAS_GEOLIP_LINALG}\")\n","if dev.type == 'cuda':\n","    print(f\"  GPU: {torch.cuda.get_device_name()}\")\n","print()\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# 1. SMOKE TEST — all flows, all shapes\n","# ═══════════════════════════════════════════════════════════════════\n","print(f\"{'='*72}\\n  1. SMOKE TEST\\n{'='*72}\")\n","\n","B, n, k, d = 16, 64, 32, 128\n","anchors, queries = make_data(B, n, k, d)\n","\n","flows_cfg = [\n","    ('QuaternionFlow',     lambda d,k: QuaternionFlow(d, k, n_heads=4)),\n","    ('QuaternionLiteFlow', lambda d,k: QuaternionLiteFlow(d, k)),\n","    ('VelocityFlow',       lambda d,k: VelocityFlow(d, k)),\n","    ('MagnitudeFlow',      lambda d,k: MagnitudeFlow(d, k)),\n","    ('OrbitalFlow',        lambda d,k: OrbitalFlow(d, k)),\n","    ('AlignmentFlow',      lambda d,k: AlignmentFlow(d, k)),\n","]\n","\n","print(f\"\\n  {'Flow':<22} {'Params':>8} {'Shape':>14} {'Time':>10} {'Conf':>8} {'Res norm':>10}\")\n","print(f\"  {'─'*22} {'─'*8} {'─'*14} {'─'*10} {'─'*8} {'─'*10}\")\n","\n","live_flows = []\n","flow_ctors = []\n","for name, ctor in flows_cfg:\n","    try:\n","        flow = ctor(d, k).to(dev)\n","        params = sum(p.numel() for p in flow.parameters())\n","        pred, conf = flow(anchors, queries)\n","        ms = time_fn(lambda: flow(anchors, queries))\n","        res = (pred - queries).norm(dim=-1).mean().item()\n","        shape_str = str(tuple(pred.shape))\n","        print(f\"  {name:<22} {params:>8,} {shape_str:>14} {fmt(ms):>10} {conf.mean().item():>8.3f} {res:>10.3f}\")\n","        live_flows.append(flow)\n","        flow_ctors.append((name, ctor))\n","    except Exception as e:\n","        print(f\"  {name:<22} FAILED: {str(e)[:50]}\")\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# 2. LINALG INTEGRATION\n","# ═══════════════════════════════════════════════════════════════════\n","print(f\"\\n{'='*72}\\n  2. LINALG INTEGRATION\\n{'='*72}\")\n","\n","if HAS_GEOLIP_LINALG:\n","    print(f\"\\n  Testing eigh dispatch in MagnitudeFlow and OrbitalFlow...\")\n","    for FlowCls in [MagnitudeFlow, OrbitalFlow]:\n","        flow = FlowCls(d, k).to(dev)\n","        pred, conf = flow(anchors, queries)\n","        ok = torch.isfinite(pred).all().item() and torch.isfinite(conf).all().item()\n","        print(f\"  {flow.name:<18} finite={ok}  conf={conf.mean():.3f}\")\n","\n","    oflow = OrbitalFlow(d, k).to(dev)\n","    a_geom = oflow.anchor_proj(anchors)\n","    G = torch.bmm(a_geom.transpose(-2, -1), a_geom)\n","    vals, vecs = LA.eigh(G)\n","    print(f\"\\n  Gram eigenspectrum: shape={tuple(vals.shape)} \"\n","          f\"range=[{vals.min().item():.4f}, {vals.max().item():.4f}]\")\n","    print(f\"  Eigenvector orth err: {(torch.bmm(vecs.mT, vecs) - torch.eye(oflow.geom_dim, device=dev)).abs().max().item():.2e}\")\n","else:\n","    print(\"  Skipped — geolip_core.linalg not available\")\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# 3. MULTI-SCALE\n","# ═══════════════════════════════════════════════════════════════════\n","print(f\"\\n{'='*72}\\n  3. MULTI-SCALE\\n{'='*72}\")\n","\n","configs = [\n","    (4,  16,  8,   64,  'tiny'),\n","    (16, 64,  32,  128, 'small'),\n","    (32, 128, 64,  256, 'medium'),\n","    (64, 256, 128, 256, 'large'),\n","    (8,  512, 256, 512, 'wide'),\n","]\n","\n","print(f\"\\n  OrbitalFlow across scales:\")\n","print(f\"  {'Config':<10} {'B':>4} {'n':>5} {'k':>5} {'d':>5} {'Time':>10} {'OK':>4}\")\n","print(f\"  {'─'*10} {'─'*4} {'─'*5} {'─'*5} {'─'*5} {'─'*10} {'─'*4}\")\n","\n","for B_, n_, k_, d_, label in configs:\n","    try:\n","        of = OrbitalFlow(d_, k_).to(dev)\n","        a, q = make_data(B_, n_, k_, d_)\n","        pred, conf = of(a, q)\n","        ms = time_fn(lambda: of(a, q), warmup=3, runs=20)\n","        ok = torch.isfinite(pred).all().item()\n","        print(f\"  {label:<10} {B_:>4} {n_:>5} {k_:>5} {d_:>5} {fmt(ms):>10} {'OK' if ok else 'NO':>4}\")\n","        del of, a, q\n","    except Exception as e:\n","        print(f\"  {label:<10} {B_:>4} {n_:>5} {k_:>5} {d_:>5} FAILED: {str(e)[:30]}\")\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# 4. ENSEMBLE FUSION MODES\n","# ═══════════════════════════════════════════════════════════════════\n","print(f\"\\n{'='*72}\\n  4. ENSEMBLE FUSION\\n{'='*72}\")\n","\n","B, n, k, d = 16, 64, 32, 128\n","anchors, queries = make_data(B, n, k, d)\n","\n","for fusion in ['weighted', 'gated', 'residual']:\n","    ens = FlowEnsemble(live_flows, d, fusion=fusion).to(dev)\n","    out = ens(anchors, queries)\n","    ms = time_fn(lambda: ens(anchors, queries), warmup=3, runs=20)\n","\n","    preds = [flow(anchors, queries)[0] for flow in ens.flows]\n","    cos_sims = []\n","    for i in range(len(preds)):\n","        for j in range(i+1, len(preds)):\n","            cs = F.cosine_similarity(preds[i].flatten(1), preds[j].flatten(1), dim=-1).mean().item()\n","            cos_sims.append(cs)\n","    avg_sim = sum(cos_sims) / max(len(cos_sims), 1)\n","\n","    print(f\"\\n  {fusion}:  time={fmt(ms)}  norm={out.norm(dim=-1).mean():.3f}  diversity={1-avg_sim:.3f}\")\n","    diag = ens.flow_diagnostics(anchors, queries)\n","    for fname, stats in diag.items():\n","        print(f\"    {fname:<18} conf={stats['confidence_mean']:.3f}±{stats['confidence_std']:.3f}  \"\n","              f\"res={stats['residual_norm']:.3f}\")\n","    del ens\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# 5. GRADIENT HEALTH\n","# ═══════════════════════════════════════════════════════════════════\n","print(f\"\\n{'='*72}\\n  5. GRADIENT HEALTH\\n{'='*72}\")\n","\n","B, n, k, d = 16, 64, 32, 128\n","anchors, queries = make_data(B, n, k, d)\n","\n","losses = {\n","    'mse': (lambda o,q: (o - q).pow(2).mean()),\n","    'cosine': (lambda o,q: (1 - F.cosine_similarity(o, q, dim=-1)).mean()),\n","    'norm': (lambda o,q: o.norm(dim=-1).mean()),\n","}\n","\n","print(f\"\\n  {'Flow':<18} {'Loss':<10} {'Grad norm':>12} {'Status':>8}\")\n","print(f\"  {'─'*18} {'─'*10} {'─'*12} {'─'*8}\")\n","\n","for loss_name, loss_fn in losses.items():\n","    # Fresh flows for each loss — avoids in-place grad corruption across losses\n","    try:\n","        test_flows_grad = [ctor(d, k).to(dev) for _, ctor in flow_ctors]\n","        ens_g = FlowEnsemble(test_flows_grad, d, fusion='residual').to(dev)\n","        ens_g.zero_grad()\n","        anchors_g = anchors.detach().clone().requires_grad_(True)\n","        queries_g = queries.detach().clone().requires_grad_(True)\n","        out = ens_g(anchors_g, queries_g)\n","        loss = loss_fn(out, queries_g.detach())\n","        loss.backward()\n","\n","        for flow in ens_g.flows:\n","            grads = [p.grad for p in flow.parameters() if p.grad is not None]\n","            if grads:\n","                gn = torch.cat([g.flatten() for g in grads]).norm().item()\n","                status = \"OK\" if 1e-8 < gn < 1e4 else \"WARN\"\n","                print(f\"  {flow.name:<18} {loss_name:<10} {gn:>12.2e} {status:>8}\")\n","            else:\n","                print(f\"  {flow.name:<18} {loss_name:<10} {'no grads':>12} {'WARN':>8}\")\n","        del ens_g, test_flows_grad\n","    except RuntimeError as e:\n","        if 'inplace' in str(e).lower() or 'in-place' in str(e).lower() or 'modified by' in str(e):\n","            print(f\"  {'*':>18} {loss_name:<10} {'IN-PLACE ERR':>12} {'NOTE':>8}\")\n","            print(f\"    FL eigh deflation uses indexed assignment — needs .clone() fix\")\n","        else:\n","            print(f\"  {'*':>18} {loss_name:<10} {'ERROR':>12}\")\n","            print(f\"    {str(e)[:60]}\")\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# 6. ABLATION — solo vs pairs vs full ensemble\n","# ═══════════════════════════════════════════════════════════════════\n","print(f\"\\n{'='*72}\\n  6. ABLATION (100 training steps, rotation target)\\n{'='*72}\")\n","\n","B, n, k, d = 32, 128, 64, 256\n","anchors, queries = make_data(B, n, k, d)\n","R = torch.linalg.qr(torch.randn(d, d, device=dev)).Q.unsqueeze(0)\n","target = torch.bmm(queries, R.expand(B, -1, -1))\n","\n","def eval_quality(model, anchors, queries, target, steps=100, lr=1e-3):\n","    opt = torch.optim.Adam(model.parameters(), lr=lr)\n","    for _ in range(steps):\n","        opt.zero_grad()\n","        pred = model(anchors, queries) if isinstance(model, FlowEnsemble) else model(anchors, queries)[0]\n","        loss = (pred - target).pow(2).mean()\n","        loss.backward()\n","        opt.step()\n","    with torch.no_grad():\n","        pred = model(anchors, queries) if isinstance(model, FlowEnsemble) else model(anchors, queries)[0]\n","        return (pred - target).pow(2).mean().item()\n","\n","print(f\"\\n  {'Configuration':<35} {'MSE':>10} {'Params':>10}\")\n","print(f\"  {'─'*35} {'─'*10} {'─'*10}\")\n","\n","for name, ctor in flow_ctors:\n","    try:\n","        flow = ctor(d, k).to(dev)\n","        params = sum(p.numel() for p in flow.parameters())\n","        mse = eval_quality(flow, anchors, queries, target)\n","        print(f\"  {name:<35} {mse:>10.4f} {params:>10,}\")\n","        del flow\n","    except Exception as e:\n","        print(f\"  {name:<35} FAILED: {str(e)[:30]}\")\n","\n","pairs = [\n","    ('Quat + Orbital', [0, 4]),\n","    ('Velocity + Magnitude', [2, 3]),\n","    ('Orbital + Alignment', [4, 5]),\n","    ('Velocity + Orbital', [2, 4]),\n","]\n","for pair_name, indices in pairs:\n","    try:\n","        pair_flows = [flow_ctors[i][1](d, k).to(dev) for i in indices if i < len(flow_ctors)]\n","        if len(pair_flows) >= 2:\n","            ens = FlowEnsemble(pair_flows, d, fusion='weighted').to(dev)\n","            params = sum(p.numel() for p in ens.parameters())\n","            mse = eval_quality(ens, anchors, queries, target)\n","            print(f\"  {pair_name:<35} {mse:>10.4f} {params:>10,}\")\n","            del ens, pair_flows\n","    except Exception as e:\n","        print(f\"  {pair_name:<35} FAILED: {str(e)[:30]}\")\n","\n","for fusion in ['weighted', 'residual']:\n","    try:\n","        all_flows = [ctor(d, k).to(dev) for _, ctor in flow_ctors]\n","        ens = FlowEnsemble(all_flows, d, fusion=fusion).to(dev)\n","        params = sum(p.numel() for p in ens.parameters())\n","        mse = eval_quality(ens, anchors, queries, target)\n","        print(f\"  {'Full (' + fusion + ')':<35} {mse:>10.4f} {params:>10,}\")\n","        del ens, all_flows\n","    except Exception as e:\n","        print(f\"  {'Full (' + fusion + ')':<35} FAILED: {str(e)[:30]}\")\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# 7. COMPILE COMPATIBILITY\n","# ═══════════════════════════════════════════════════════════════════\n","print(f\"\\n{'='*72}\\n  7. COMPILE COMPATIBILITY\\n{'='*72}\")\n","\n","B, n, k, d = 8, 32, 16, 64\n","anchors, queries = make_data(B, n, k, d)\n","\n","print(f\"\\n  {'Flow':<22} {'fullgraph':>12} {'Raw':>10} {'Compiled':>12}\")\n","print(f\"  {'─'*22} {'─'*12} {'─'*10} {'─'*12}\")\n","\n","for name, ctor in flow_ctors:\n","    try:\n","        flow = ctor(d, k).to(dev)\n","        t_raw = time_fn(lambda: flow(anchors, queries), warmup=3, runs=30)\n","        try:\n","            compiled = torch.compile(flow, fullgraph=True)\n","            compiled(anchors, queries); sync()\n","            t_comp = time_fn(lambda: compiled(anchors, queries), warmup=3, runs=30)\n","            status = \"OK\"\n","        except Exception as e:\n","            t_comp = -1\n","            status = str(e)[:12]\n","        t_str = fmt(t_comp) if t_comp > 0 else \"N/A\"\n","        print(f\"  {name:<22} {status:>12} {fmt(t_raw):>10} {t_str:>12}\")\n","        del flow\n","    except Exception as e:\n","        print(f\"  {name:<22} FAILED: {str(e)[:40]}\")\n","\n","\n","# ═══════════════════════════════════════════════════════════════════\n","# 8. MEMORY\n","# ═══════════════════════════════════════════════════════════════════\n","if dev.type == 'cuda':\n","    print(f\"\\n{'='*72}\\n  8. MEMORY (B=32, n=128, k=64, d=256)\\n{'='*72}\")\n","\n","    B, n, k, d = 32, 128, 64, 256\n","    anchors, queries = make_data(B, n, k, d)\n","\n","    print(f\"\\n  {'Flow':<22} {'Peak MB':>10}\")\n","    print(f\"  {'─'*22} {'─'*10}\")\n","\n","    for name, ctor in flow_ctors:\n","        try:\n","            flow = ctor(d, k).to(dev)\n","            torch.cuda.empty_cache(); gc.collect()\n","            torch.cuda.reset_peak_memory_stats()\n","            base = torch.cuda.memory_allocated()\n","            pred, conf = flow(anchors, queries); sync()\n","            peak = (torch.cuda.max_memory_allocated() - base) / 1024**2\n","            print(f\"  {name:<22} {peak:>9.1f}\")\n","            del flow, pred, conf\n","        except Exception as e:\n","            print(f\"  {name:<22} FAILED: {str(e)[:30]}\")\n","\n","    try:\n","        all_flows = [ctor(d, k).to(dev) for _, ctor in flow_ctors]\n","        ens = FlowEnsemble(all_flows, d, fusion='weighted').to(dev)\n","        torch.cuda.empty_cache(); gc.collect()\n","        torch.cuda.reset_peak_memory_stats()\n","        base = torch.cuda.memory_allocated()\n","        out = ens(anchors, queries); sync()\n","        peak = (torch.cuda.max_memory_allocated() - base) / 1024**2\n","        print(f\"  {'Full ensemble':<22} {peak:>9.1f}\")\n","        del ens, all_flows\n","    except Exception as e:\n","        print(f\"  {'Full ensemble':<22} FAILED: {str(e)[:30]}\")\n","\n","print(f\"\\n{'='*72}\")\n","print(f\"  Done.\")\n","print(f\"{'='*72}\")"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"o-1SdVZy95iv","executionInfo":{"status":"ok","timestamp":1775100334443,"user_tz":420,"elapsed":10532,"user":{"displayName":"P C","userId":"00707517734723903966"}},"outputId":"59e50ead-58aa-4716-8d1b-2fbe12a067aa"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["geolip_core.linalg: available\n","geolip.linalg backend:\n","  CUDA:       yes\n","  Triton:     3.6.0\n","  FL eigh:    enabled\n","  Triton SVD: enabled\n","  GPU:        NVIDIA RTX PRO 6000 Blackwell Server Edition\n","========================================================================\n","  Flow Ensemble — Expanded Test Suite\n","========================================================================\n","  device=cuda  geolip_core.linalg=True\n","  GPU: NVIDIA RTX PRO 6000 Blackwell Server Edition\n","\n","========================================================================\n","  1. SMOKE TEST\n","========================================================================\n","\n","  Flow                     Params          Shape       Time     Conf   Res norm\n","  ────────────────────── ──────── ────────────── ────────── ──────── ──────────\n","  QuaternionFlow           70,725  (16, 64, 128)      270us    0.501      1.336\n","  QuaternionLiteFlow       87,109  (16, 64, 128)      190us    0.500      1.291\n","  VelocityFlow             70,210  (16, 64, 128)      175us    0.526      0.066\n","  MagnitudeFlow            73,293  (16, 64, 128)      319us    0.497      0.663\n","  OrbitalFlow              25,445  (16, 64, 128)      375us    0.466      1.811\n","  AlignmentFlow             8,922  (16, 64, 128)      347us    0.462      2.825\n","\n","========================================================================\n","  2. LINALG INTEGRATION\n","========================================================================\n","\n","  Testing eigh dispatch in MagnitudeFlow and OrbitalFlow...\n","  magnitude          finite=True  conf=0.496\n","  orbital            finite=True  conf=0.479\n","\n","  Gram eigenspectrum: shape=(16, 12) range=[0.0007, 15.9592]\n","  Eigenvector orth err: 3.11e+01\n","\n","========================================================================\n","  3. MULTI-SCALE\n","========================================================================\n","\n","  OrbitalFlow across scales:\n","  Config        B     n     k     d       Time   OK\n","  ────────── ──── ───── ───── ───── ────────── ────\n","  tiny          4    16     8    64      328us   OK\n","  small        16    64    32   128      377us   OK\n","  medium       32   128    64   256      377us   OK\n","  large        64   256   128   256      426us   OK\n","  wide          8   512   256   512      423us   OK\n","\n","========================================================================\n","  4. ENSEMBLE FUSION\n","========================================================================\n","\n","  weighted:  time=2.00ms  norm=0.744  diversity=0.899\n","    quaternion         conf=0.501±0.002  res=1.334\n","    quat_lite          conf=0.500±0.002  res=1.295\n","    velocity           conf=0.526±0.004  res=0.066\n","    magnitude          conf=0.497±0.004  res=0.661\n","    orbital            conf=0.466±0.003  res=1.810\n","    alignment          conf=0.462±0.010  res=2.506\n","\n","  gated:  time=2.00ms  norm=0.627  diversity=0.899\n","    quaternion         conf=0.501±0.002  res=1.334\n","    quat_lite          conf=0.500±0.002  res=1.295\n","    velocity           conf=0.526±0.004  res=0.066\n","    magnitude          conf=0.497±0.004  res=0.661\n","    orbital            conf=0.466±0.003  res=1.810\n","    alignment          conf=0.462±0.010  res=2.506\n","\n","  residual:  time=2.02ms  norm=0.744  diversity=0.899\n","    quaternion         conf=0.501±0.002  res=1.334\n","    quat_lite          conf=0.500±0.002  res=1.295\n","    velocity           conf=0.526±0.004  res=0.066\n","    magnitude          conf=0.497±0.004  res=0.661\n","    orbital            conf=0.466±0.003  res=1.810\n","    alignment          conf=0.462±0.010  res=2.506\n","\n","========================================================================\n","  5. GRADIENT HEALTH\n","========================================================================\n","\n","  Flow               Loss          Grad norm   Status\n","  ────────────────── ────────── ──────────── ────────\n","  quaternion         mse            9.67e-04       OK\n","  quat_lite          mse            7.96e-04       OK\n","  velocity           mse            4.36e-04       OK\n","  magnitude          mse            8.74e-04       OK\n","  orbital            mse            1.97e-03       OK\n","  alignment          mse            6.61e-02       OK\n","  quaternion         cosine         8.03e-02       OK\n","  quat_lite          cosine         6.26e-02       OK\n","  velocity           cosine         3.07e-02       OK\n","  magnitude          cosine         7.67e-02       OK\n","  orbital            cosine         1.67e-01       OK\n","  alignment          cosine         2.13e+00       OK\n","  quaternion         norm           9.28e-02       OK\n","  quat_lite          norm           7.11e-02       OK\n","  velocity           norm           7.62e-03       OK\n","  magnitude          norm           6.71e-02       OK\n","  orbital            norm           1.76e-01       OK\n","  alignment          norm           2.41e+00       OK\n","\n","========================================================================\n","  6. ABLATION (100 training steps, rotation target)\n","========================================================================\n","\n","  Configuration                              MSE     Params\n","  ─────────────────────────────────── ────────── ──────────\n","  QuaternionFlow                          0.0010    280,709\n","  QuaternionLiteFlow                      0.0005    346,245\n","  VelocityFlow                            0.0062    279,682\n","  MagnitudeFlow                           0.0081    285,837\n","  OrbitalFlow                             0.0029     91,813\n","  AlignmentFlow                           0.0082     26,010\n","  Quat + Orbital                          0.0010    372,524\n","  Velocity + Magnitude                    0.0057    565,521\n","  Orbital + Alignment                     0.0028    117,825\n","  Velocity + Orbital                      0.0020    371,497\n","  Full (weighted)                         0.0007  1,310,302\n","  Full (residual)                         0.0008  1,310,302\n","\n","========================================================================\n","  7. COMPILE COMPATIBILITY\n","========================================================================\n","\n","  Flow                      fullgraph        Raw     Compiled\n","  ────────────────────── ──────────── ────────── ────────────\n","  QuaternionFlow                   OK      271us        198us\n","  QuaternionLiteFlow               OK      198us        180us\n","  VelocityFlow                     OK      178us        169us\n","  MagnitudeFlow                    OK      314us        338us\n","  OrbitalFlow                      OK      372us        346us\n","  AlignmentFlow                    OK      348us        360us\n","\n","========================================================================\n","  8. MEMORY (B=32, n=128, k=64, d=256)\n","========================================================================\n","\n","  Flow                      Peak MB\n","  ────────────────────── ──────────\n","  QuaternionFlow              36.2\n","  QuaternionLiteFlow          28.3\n","  VelocityFlow                48.0\n","  MagnitudeFlow               24.2\n","  OrbitalFlow                 10.9\n","  AlignmentFlow               15.0\n","  Full ensemble              178.7\n","\n","========================================================================\n","  Done.\n","========================================================================\n"]}]},{"cell_type":"code","source":[],"metadata":{"id":"s5_waYAxMfff"},"execution_count":null,"outputs":[]}]}