mnist_transformer.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyPcx4JQWfJ6niDX4Gxp0EcF",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "5dfcb01e5ef7433fb6037567378cc8e6": {
          "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_434f45e109534c8cbf2c9a8062f25e47",
              "IPY_MODEL_6b336b23c5e44931b440315e48034e0c",
              "IPY_MODEL_172d7c6e42294286ba4ab33404f54d7e"
            ],
            "layout": "IPY_MODEL_f206e5296e834188a6723be58326f3d5"
          }
        },
        "434f45e109534c8cbf2c9a8062f25e47": {
          "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_2daa53b1f3d7434a953f0a3f86c23fa4",
            "placeholder": "​",
            "style": "IPY_MODEL_266265a026ce42598ddfc71e772b1de4",
            "value": "100%"
          }
        },
        "6b336b23c5e44931b440315e48034e0c": {
          "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_db26bc79362f45a89b2e3bca161bff54",
            "max": 938,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_a7f69acfd06b4c4195d4d9e890a11e81",
            "value": 938
          }
        },
        "172d7c6e42294286ba4ab33404f54d7e": {
          "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_5b6221a97b9b4f4e9dcfabce07d46a05",
            "placeholder": "​",
            "style": "IPY_MODEL_3badc02581f841dc8a6cc7d8b34ffb19",
            "value": " 938/938 [01:30&lt;00:00, 11.42it/s]"
          }
        },
        "f206e5296e834188a6723be58326f3d5": {
          "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
          }
        },
        "2daa53b1f3d7434a953f0a3f86c23fa4": {
          "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
          }
        },
        "266265a026ce42598ddfc71e772b1de4": {
          "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": ""
          }
        },
        "db26bc79362f45a89b2e3bca161bff54": {
          "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
          }
        },
        "a7f69acfd06b4c4195d4d9e890a11e81": {
          "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": ""
          }
        },
        "5b6221a97b9b4f4e9dcfabce07d46a05": {
          "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
          }
        },
        "3badc02581f841dc8a6cc7d8b34ffb19": {
          "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": ""
          }
        },
        "09aa98006d4c4ecfb0496432c539b3ed": {
          "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_8f9596be50a54cdb897ab9e4cc06d499",
              "IPY_MODEL_0ca0569a28bf44bcae3fe8f9d89aae79",
              "IPY_MODEL_a98ab8ad0232453a880ba06ac3b0130e"
            ],
            "layout": "IPY_MODEL_a13d7d92ec1e435d82b655084963981b"
          }
        },
        "8f9596be50a54cdb897ab9e4cc06d499": {
          "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_b3c71306e2e940459691114724728e11",
            "placeholder": "​",
            "style": "IPY_MODEL_b9c5b294365e4a5e8927657acf82904d",
            "value": "100%"
          }
        },
        "0ca0569a28bf44bcae3fe8f9d89aae79": {
          "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_b1a1dacbb4d041e3abbe1e36485a25ef",
            "max": 938,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_8929f30d46444027bc5cbf0fcb4d5b58",
            "value": 938
          }
        },
        "a98ab8ad0232453a880ba06ac3b0130e": {
          "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_e37a94be493f46279fa16c116bccf25b",
            "placeholder": "​",
            "style": "IPY_MODEL_350cb32a69854442904a8976d623fa3a",
            "value": " 938/938 [01:30&lt;00:00,  7.36it/s]"
          }
        },
        "a13d7d92ec1e435d82b655084963981b": {
          "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
          }
        },
        "b3c71306e2e940459691114724728e11": {
          "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
          }
        },
        "b9c5b294365e4a5e8927657acf82904d": {
          "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": ""
          }
        },
        "b1a1dacbb4d041e3abbe1e36485a25ef": {
          "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
          }
        },
        "8929f30d46444027bc5cbf0fcb4d5b58": {
          "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": ""
          }
        },
        "e37a94be493f46279fa16c116bccf25b": {
          "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
          }
        },
        "350cb32a69854442904a8976d623fa3a": {
          "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": ""
          }
        },
        "087f185326224638bd27eb9cea145c2e": {
          "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_0c827f76efac4255a32d175e4e2bb393",
              "IPY_MODEL_4f336eb55c4e4070a5e620b63d8ca8c9",
              "IPY_MODEL_a7fa34590d664828bfeafc4ad66bf540"
            ],
            "layout": "IPY_MODEL_bbeca2ec4fa94acfa810ad88034ff621"
          }
        },
        "0c827f76efac4255a32d175e4e2bb393": {
          "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_f0feaa17a0c4466e9fde788a3c2feed2",
            "placeholder": "​",
            "style": "IPY_MODEL_bcfda19f9e1e41b890935ff0b2aee14d",
            "value": "100%"
          }
        },
        "4f336eb55c4e4070a5e620b63d8ca8c9": {
          "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_5ae711dccee743b69063f739097d0d73",
            "max": 938,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_119b67f0e5a149108acdff2fd565f0e5",
            "value": 938
          }
        },
        "a7fa34590d664828bfeafc4ad66bf540": {
          "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_d39a3f047f784e50a3758e515509d46b",
            "placeholder": "​",
            "style": "IPY_MODEL_15887dad12f74ccaaa738efb6a0761a7",
            "value": " 938/938 [01:29&lt;00:00, 11.70it/s]"
          }
        },
        "bbeca2ec4fa94acfa810ad88034ff621": {
          "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
          }
        },
        "f0feaa17a0c4466e9fde788a3c2feed2": {
          "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
          }
        },
        "bcfda19f9e1e41b890935ff0b2aee14d": {
          "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": ""
          }
        },
        "5ae711dccee743b69063f739097d0d73": {
          "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
          }
        },
        "119b67f0e5a149108acdff2fd565f0e5": {
          "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": ""
          }
        },
        "d39a3f047f784e50a3758e515509d46b": {
          "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
          }
        },
        "15887dad12f74ccaaa738efb6a0761a7": {
          "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": ""
          }
        },
        "3e1d208b1f884a5db8bb989c5a4bae97": {
          "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_192092b8d6ca4c58806dbc877ceaf48b",
              "IPY_MODEL_881ab5bc3c83467eaa52a7cfb312ffd9",
              "IPY_MODEL_a3c02eb69d0c45faa66859a95d69fdc4"
            ],
            "layout": "IPY_MODEL_069b08f4e02f42acbdea3cacf386a2a7"
          }
        },
        "192092b8d6ca4c58806dbc877ceaf48b": {
          "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_ebabfd7c4a8c410abe2e247e9835223e",
            "placeholder": "​",
            "style": "IPY_MODEL_4bc2fa1a42864c638e7ca2ea6bce77ed",
            "value": "100%"
          }
        },
        "881ab5bc3c83467eaa52a7cfb312ffd9": {
          "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_85ac58b159b748d2b372917de3341fd5",
            "max": 938,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_67a4e7e15c2742d09d2fb13b8d4376d6",
            "value": 938
          }
        },
        "a3c02eb69d0c45faa66859a95d69fdc4": {
          "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_d9c6df1498164f4fabeee93693c5bf52",
            "placeholder": "​",
            "style": "IPY_MODEL_0113ccea44ab46ffab501656a9bb8a1f",
            "value": " 938/938 [01:29&lt;00:00, 11.70it/s]"
          }
        },
        "069b08f4e02f42acbdea3cacf386a2a7": {
          "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
          }
        },
        "ebabfd7c4a8c410abe2e247e9835223e": {
          "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
          }
        },
        "4bc2fa1a42864c638e7ca2ea6bce77ed": {
          "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": ""
          }
        },
        "85ac58b159b748d2b372917de3341fd5": {
          "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
          }
        },
        "67a4e7e15c2742d09d2fb13b8d4376d6": {
          "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": ""
          }
        },
        "d9c6df1498164f4fabeee93693c5bf52": {
          "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
          }
        },
        "0113ccea44ab46ffab501656a9bb8a1f": {
          "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": ""
          }
        },
        "f8dc96c2caac43bd80965e38690ac8eb": {
          "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_3d1fab038ec447ebbc7b0c9ea8a98d0e",
              "IPY_MODEL_12654616c7f141008fe693c81d62a65d",
              "IPY_MODEL_b5c5a0eb5aa344b0a17e44c897046241"
            ],
            "layout": "IPY_MODEL_6b2caaf4c4d94022bda20d8a7dabd65e"
          }
        },
        "3d1fab038ec447ebbc7b0c9ea8a98d0e": {
          "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_2f5fe6be68064b8e95694ca4429f25c5",
            "placeholder": "​",
            "style": "IPY_MODEL_e5429605d6bc4898a450b1afdd426127",
            "value": "100%"
          }
        },
        "12654616c7f141008fe693c81d62a65d": {
          "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_e37f81454ffd473e88cc63ee21e9f201",
            "max": 938,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_6558548b51344df5861c5b8bc9f7e01a",
            "value": 938
          }
        },
        "b5c5a0eb5aa344b0a17e44c897046241": {
          "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_4f611ee6ed504b62839957b4fd370ce4",
            "placeholder": "​",
            "style": "IPY_MODEL_037846bd7b66465f950e16fa0a26ff68",
            "value": " 938/938 [01:29&lt;00:00, 11.18it/s]"
          }
        },
        "6b2caaf4c4d94022bda20d8a7dabd65e": {
          "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
          }
        },
        "2f5fe6be68064b8e95694ca4429f25c5": {
          "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
          }
        },
        "e5429605d6bc4898a450b1afdd426127": {
          "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": ""
          }
        },
        "e37f81454ffd473e88cc63ee21e9f201": {
          "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
          }
        },
        "6558548b51344df5861c5b8bc9f7e01a": {
          "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": ""
          }
        },
        "4f611ee6ed504b62839957b4fd370ce4": {
          "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
          }
        },
        "037846bd7b66465f950e16fa0a26ff68": {
          "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": ""
          }
        },
        "c7e6f6623bd94d7ab1d5c546793661a2": {
          "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_05a816c78c344913944e000e49b515ac",
              "IPY_MODEL_d4366bb03a3c4f849194a0ace42ef744",
              "IPY_MODEL_9c25506af10a4b2683dbf192dbfcc3d2"
            ],
            "layout": "IPY_MODEL_7bbcaa2c26f343688bd7580ca09b939d"
          }
        },
        "05a816c78c344913944e000e49b515ac": {
          "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_f22e88b1e899460fbd47b43a06edec62",
            "placeholder": "​",
            "style": "IPY_MODEL_7def5a415fc840679011fc1603ede78e",
            "value": "100%"
          }
        },
        "d4366bb03a3c4f849194a0ace42ef744": {
          "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_f53d6e6c898b40f0ae1daaa9954efc0e",
            "max": 157,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_e33c5252ff0c4066976f48be3ec246e0",
            "value": 157
          }
        },
        "9c25506af10a4b2683dbf192dbfcc3d2": {
          "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_b6b2eed876204a6e8e4223a32902c8fc",
            "placeholder": "​",
            "style": "IPY_MODEL_7d7c7c10697a43edb853d7896b07f350",
            "value": " 157/157 [00:08&lt;00:00, 14.47it/s]"
          }
        },
        "7bbcaa2c26f343688bd7580ca09b939d": {
          "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
          }
        },
        "f22e88b1e899460fbd47b43a06edec62": {
          "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
          }
        },
        "7def5a415fc840679011fc1603ede78e": {
          "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": ""
          }
        },
        "f53d6e6c898b40f0ae1daaa9954efc0e": {
          "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
          }
        },
        "e33c5252ff0c4066976f48be3ec246e0": {
          "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": ""
          }
        },
        "b6b2eed876204a6e8e4223a32902c8fc": {
          "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
          }
        },
        "7d7c7c10697a43edb853d7896b07f350": {
          "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": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/foxtrotmike/CS909/blob/master/mnist_transformer.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Transformers for MNIST: A foundational tutorial\n",
        "By [Fayyaz Minhas](https://sites.google.com/view/fayyaz/home)\n",
        "\n",
        "Transformers have revolutionized the field of natural language processing (NLP) and beyond, offering significant improvements in tasks like text translation, summarization, and generation. At the heart of this architecture are several innovative components, each contributing to the transformer's ability to handle sequential data effectively, without the limitations of recurrence or convolution. This tutorial will introduce you to these key components, aiming to provide a solid understanding of how transformers work.\n",
        "\n",
        "## Multiple Multi-headed Attention Blocks\n",
        "The transformer model employs multiple blocks of multi-headed attention mechanisms. These are designed to enable the model to focus on different parts of the input sequence simultaneously, allowing it to capture a wide range of dependencies between words or features. Each \"head\" in these attention blocks can be thought of as an independent feature extraction mechanism, looking at the input from a different perspective. By having multiple heads, the model can aggregate diverse insights about the relationships within the data.\n",
        "\n",
        "## Layer Normalization\n",
        "Layer normalization is a technique used within transformers to stabilize the neural network's learning process. It is applied to the outputs of each sub-layer (such as attention or feed-forward layers) within the transformer blocks. By normalizing the activations across the features for each data sample in a mini-batch, layer normalization helps in speeding up training and improving the overall performance of the model.\n",
        "\n",
        "## Standardization Across Features of the Same Input\n",
        "This process, often implemented as part of layer normalization, involves adjusting the features of the input data to have a mean of zero and a standard deviation of one. Standardization ensures that all features contribute equally to the model's learning, preventing any one feature with larger numerical values from dominating the learning process.\n",
        "\n",
        "## Skip Connections\n",
        "Also known as residual connections, skip connections are a critical component of the transformer architecture. They allow the output of one layer to be added to the output of a later layer, helping to mitigate the vanishing gradient problem in deep neural networks. This promotes the flow of gradients through the network, making it easier to train deeper models.\n",
        "\n",
        "## Various Types of Positional Encodings\n",
        "Transformers do not inherently process sequence order, unlike recurrent neural networks. To address this, positional encodings are added to the input embeddings to provide the model with information about the order of the sequence elements. These encodings can be fixed (based on sine and cosine functions of different frequencies) or learnable parameters. This feature enables transformers to effectively capture the sequential nature of the input data.\n",
        "\n",
        "## “Class” Tokens\n",
        "In tasks such as sequence classification, a special “class” token is often prepended to the input sequence. The final state corresponding to this token is used as the aggregate representation of the entire sequence for classification purposes. This approach allows the model to use global sequence information effectively for making predictions.\n",
        "\n",
        "\n",
        "## Masking Strategies\n",
        "Masking is a crucial strategy in transformer models, particularly for tasks such as sentence completion or machine translation. It involves selectively hiding parts of the input or future tokens from the model during training. This forces the model to predict missing words or the next word in a sequence without relying on the actual next word, fostering a more robust understanding of context and relationships within the data. However, for vision transformer implementation below, we will be ignoring masking.\n",
        "\n",
        "By understanding these components and their interplay within the transformer architecture, you'll be equipped to dive deeper into the model's applications and nuances. Each element plays a pivotal role in the transformer's ability to learn from sequential data, making it a powerful tool in the machine learning practitioner's toolkit.\n",
        "\n",
        "# Coding\n",
        "\n",
        " we will explore the implementation and workings of a Vision Transformer (ViT) for image classification tasks. The Vision Transformer applies principles of the transformer architecture, originally designed for natural language processing, to the domain of computer vision. Let's break down the code into its components and understand how it transforms an image into a format suitable for the transformer, processes it through multiple attention blocks, and finally uses it for classification.\n",
        "\n",
        "## Image Patching and Positional Encoding\n",
        "The first step in the Vision Transformer is to convert an image into a sequence of flattened patches. This is akin to tokenizing a sentence in NLP tasks. The ViT class starts by ensuring the image dimensions are divisible by the patch size. Each image is then rearranged into non-overlapping patches using the rearrange function from einops. These patches are flattened and projected to a higher-dimensional space (the embedding dimension of the transformer) using a linear layer (self.patch_to_embedding). This process effectively treats each patch as a \"word\" in the context of NLP.\n",
        "\n",
        "Next, positional encodings are added to the patch embeddings to retain information about the original location of each patch within the image. The Vision Transformer also introduces a learnable \"class token\" at the beginning of the sequence. This class token is a unique embedding that is updated through the transformer layers and used for classification. The combination of patch embeddings, positional embeddings, and the class token forms the input to the transformer.\n",
        "\n",
        "## Multi-headed Attention and Transformer Blocks\n",
        "Within the Transformer class, the input sequence passes through multiple layers, each consisting of a multi-head attention block followed by a feed-forward neural network. Both of these components are encapsulated in residual connections (Residual) and preceded by layer normalization (PreNorm), following the standard transformer architecture.\n",
        "\n",
        "### Attention Mechanism:\n",
        "The multi-headed attention mechanism (Attention class) allows the model to focus on different parts of the image simultaneously. It does this by computing scaled dot-product attention across different \"heads,\" enabling the model to capture various aspects of the input.\n",
        "\n",
        "### Feed-Forward Networks:\n",
        "After the attention block, the data passes through a feed-forward network (FeedForward class) in each transformer layer. This network applies a simple sequence of linear transformations and non-linear activations to further process the information.\n",
        "\n",
        "## Classification\n",
        "After the input sequence has been processed by the transformer layers, the updated class token is extracted and passed through a final classifier (the mlp_head in the ViT class). This classifier consists of linear layers and GELU activations, culminating in a prediction over the defined number of classes.\n",
        "\n",
        "## Summary of the Approach\n",
        "* Image Patching:\n",
        " The input image is divided into patches, which are then flattened and projected to a higher dimension.\n",
        "* Positional Encoding:\n",
        "Positional embeddings are added to the patch embeddings to give the model information about the original position of each patch in the image. A learnable class token is also prepended to the sequence.\n",
        "* Transformer Processing:\n",
        "The sequence of embeddings is processed through several layers of multi-headed attention and feed-forward networks, all within residual connections and preceded by normalization.\n",
        "* Classification:\n",
        "The class token, now enriched with information from across the image, is used to make a final classification.\n",
        "\n",
        "This coding walkthrough illustrates the innovative way the Vision Transformer adapts the transformer architecture from text to images, treating image patches as analogous to words in a sentence. It highlights the transformer's versatility and its powerful capability for both NLP and computer vision tasks."
      ],
      "metadata": {
        "id": "GuML2b-yuOU5"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install einops"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "txmEqOQyyRjC",
        "outputId": "54be8677-e16a-4c38-cf90-0c14fef92d85"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting einops\n",
            "  Downloading einops-0.7.0-py3-none-any.whl (44 kB)\n",
            "\u001b[?25l     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/44.6 kB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m44.6/44.6 kB\u001b[0m \u001b[31m1.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hInstalling collected packages: einops\n",
            "Successfully installed einops-0.7.0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "from torchvision import transforms\n",
        "from tqdm.notebook import tqdm\n",
        "import matplotlib.pyplot as plt\n",
        "import torch.nn.functional as F\n",
        "from einops import rearrange\n",
        "\n",
        "class Residual(nn.Module):\n",
        "    # Implements a residual connection followed by a function fn.\n",
        "    def __init__(self, fn):\n",
        "        super().__init__()\n",
        "        self.fn = fn  # The function to apply before adding the input back.\n",
        "\n",
        "    def forward(self, x, **kwargs):\n",
        "        # Applies the function and adds the input x back to the result.\n",
        "        return self.fn(x, **kwargs) + x\n",
        "\n",
        "class PreNorm(nn.Module):\n",
        "    # Applies layer normalization before a function fn.\n",
        "    def __init__(self, dim, fn):\n",
        "        super().__init__()\n",
        "        self.norm = nn.LayerNorm(dim)  # Normalizes the input across the specified dimension.\n",
        "        self.fn = fn  # The function to apply after normalization.\n",
        "\n",
        "    def forward(self, x, **kwargs):\n",
        "        # Normalizes the input and then applies the function.\n",
        "        return self.fn(self.norm(x), **kwargs)\n",
        "\n",
        "class FeedForward(nn.Module):\n",
        "    # Implements a feedforward neural network with one hidden layer.\n",
        "    def __init__(self, dim, hidden_dim):\n",
        "        super().__init__()\n",
        "        self.net = nn.Sequential(\n",
        "            nn.Linear(dim, hidden_dim),  # Linear transformation to the hidden dimension.\n",
        "            nn.GELU(),  # GELU non-linearity.\n",
        "            nn.Linear(hidden_dim, dim)  # Linear transformation back to the original dimension.\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        # Passes the input through the feedforward network.\n",
        "        return self.net(x)\n",
        "\n",
        "class Attention(nn.Module):\n",
        "    # Implements the multi-head self-attention mechanism.\n",
        "    def __init__(self, dim, heads=8):\n",
        "        super().__init__()\n",
        "        self.heads = heads  # Number of attention heads.\n",
        "        self.scale = dim ** -0.5  # Scaling factor for the dot products.\n",
        "\n",
        "        # Linear layer to project the input into queries, keys, and values.\n",
        "        self.to_qkv = nn.Linear(dim, dim * 3, bias=False)\n",
        "        # Linear layer to project the concatenated outputs back to the original dimension.\n",
        "        self.to_out = nn.Linear(dim, dim)\n",
        "\n",
        "    def forward(self, x, mask=None):\n",
        "        b, n, _, h = *x.shape, self.heads\n",
        "        # Projects input to queries, keys, and values.\n",
        "        qkv = self.to_qkv(x)\n",
        "        # Rearranges the projections for multi-head attention processing.\n",
        "        q, k, v = rearrange(qkv, 'b n (qkv h d) -> qkv b h n d', qkv=3, h=h)\n",
        "\n",
        "        # Calculates the dot products of queries and keys, scales the result.\n",
        "        dots = torch.einsum('bhid,bhjd->bhij', q, k) * self.scale\n",
        "\n",
        "        if mask is not None:\n",
        "            # Applies masking if provided, for handling variable sequence lengths.\n",
        "            mask = F.pad(mask.flatten(1), (1, 0), value=True)\n",
        "            assert mask.shape[-1] == dots.shape[-1], 'mask has incorrect dimensions'\n",
        "            mask = mask[:, None, :] * mask[:, :, None]\n",
        "            dots.masked_fill_(~mask, float('-inf'))\n",
        "            del mask\n",
        "\n",
        "        # Applies softmax to get the attention weights.\n",
        "        attn = dots.softmax(dim=-1)\n",
        "\n",
        "        # Uses einsum to compute weighted sum of values.\n",
        "        out = torch.einsum('bhij,bhjd->bhid', attn, v)\n",
        "        # Rearranges the output back to the original input shape.\n",
        "        out = rearrange(out, 'b h n d -> b n (h d)')\n",
        "        # Projects the output back to the original dimension.\n",
        "        out = self.to_out(out)\n",
        "        return out\n",
        "\n",
        "class Transformer(nn.Module):\n",
        "    # Implements a sequence of transformer blocks.\n",
        "    def __init__(self, dim, depth, heads, mlp_dim):\n",
        "        super().__init__()\n",
        "        self.layers = nn.ModuleList([])\n",
        "        for _ in range(depth):\n",
        "            # Each block consists of a multi-head attention layer and a feedforward layer, both wrapped in residual and normalization layers.\n",
        "            self.layers.append(nn.ModuleList([\n",
        "                Residual(PreNorm(dim, Attention(dim, heads=heads))),\n",
        "                Residual(PreNorm(dim, FeedForward(dim, mlp_dim)))\n",
        "            ]))\n",
        "\n",
        "    def forward(self, x, mask=None):\n",
        "        # Processes the input through each block in sequence.\n",
        "        for attn, ff in self.layers:\n",
        "            x = attn(x, mask=mask)  # Applies attention.\n",
        "            x = ff(x)  # Applies feedforward network.\n",
        "        return x\n",
        "\n",
        "class ViT(nn.Module):\n",
        "    # Vision Transformer (ViT) for image classification.\n",
        "    def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, channels=3):\n",
        "        super().__init__()\n",
        "        assert image_size % patch_size == 0, 'image dimensions must be divisible by the patch size'\n",
        "        num_patches = (image_size // patch_size) ** 2\n",
        "        patch_dim = channels * patch_size ** 2  # Dimensionality of the flattened patch.\n",
        "\n",
        "        self.patch_size = patch_size\n",
        "\n",
        "        # Learnable position embeddings.\n",
        "        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))\n",
        "        # Linear layer to project flattened patches to the dimensionality of the transformer.\n",
        "        self.patch_to_embedding = nn.Linear(patch_dim, dim)\n",
        "        # Learnable class token.\n",
        "        self.cls_token = nn.Parameter(torch.randn(1, 1, dim))\n",
        "        # The transformer core.\n",
        "        self.transformer = Transformer(dim, depth, heads, mlp_dim)\n",
        "\n",
        "        self.to_cls_token = nn.Identity()\n",
        "\n",
        "        # The \"head\" of the ViT to make final class predictions.\n",
        "        self.mlp_head = nn.Sequential(\n",
        "            nn.Linear(dim, mlp_dim),\n",
        "            nn.GELU(),\n",
        "            nn.Linear(mlp_dim, num_classes)\n",
        "        )\n",
        "\n",
        "    def forward(self, img, mask=None):\n",
        "        p = self.patch_size\n",
        "\n",
        "        # Rearranges the input image into non-overlapping patches.\n",
        "        x = rearrange(img, 'b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1=p, p2=p)\n",
        "\n",
        "        # Projects patches to embeddings.\n",
        "        x = self.patch_to_embedding(x)\n",
        "\n",
        "        # Prepends the class token to the sequence of embedded patches.\n",
        "        cls_tokens = self.cls_token.expand(img.shape[0], -1, -1)\n",
        "        x = torch.cat((cls_tokens, x), dim=1)\n",
        "        # Adds position embeddings to the patch embeddings.\n",
        "        x += self.pos_embedding\n",
        "        # Processes the sequence through the transformer.\n",
        "        x = self.transformer(x, mask)\n",
        "\n",
        "        # Extracts the class token and applies the classification head.\n",
        "        x = self.to_cls_token(x[:, 0])\n",
        "        return self.mlp_head(x)\n"
      ],
      "metadata": {
        "id": "vSbV5BXS50lR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from torchvision.datasets import MNIST\n",
        "from torch.utils.data import DataLoader\n",
        "model = ViT(image_size=28, patch_size=7, num_classes=10, channels=1,\n",
        "            dim=64, depth=6, heads=8, mlp_dim=128)\n",
        "# Define a transform to normalize the data\n",
        "transform = transforms.Compose([\n",
        "    transforms.Resize((28, 28)),  # Ensure the image size is 28x28\n",
        "    transforms.ToTensor(),  # Convert images to tensors\n",
        "    transforms.Normalize((0.5,), (0.5,))  # Normalize the images\n",
        "])\n",
        "\n",
        "# Download and load the training data\n",
        "trainset = MNIST('', download=True, train=True, transform=transform)\n",
        "trainloader = DataLoader(trainset, batch_size=64, shuffle=True)\n",
        "\n",
        "# Download and load the test data\n",
        "testset = MNIST('', download=True, train=False, transform=transform)\n",
        "testloader = DataLoader(testset, batch_size=64, shuffle=False)\n",
        "import torch.optim as optim\n",
        "\n",
        "# Define the loss function and optimizer\n",
        "loss_function = nn.CrossEntropyLoss()\n",
        "optimizer = optim.Adam(model.parameters(), lr=0.003)\n",
        "\n",
        "epochs = 5\n",
        "for epoch in range(epochs):\n",
        "    running_loss = 0\n",
        "    for images, labels in tqdm(trainloader):\n",
        "        # Reset the gradients to zero\n",
        "        optimizer.zero_grad()\n",
        "\n",
        "        # Forward pass\n",
        "        output = model(images)\n",
        "        loss = loss_function(output, labels)\n",
        "\n",
        "        # Backward pass and optimize\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "        running_loss += loss.item()\n",
        "    else:\n",
        "        print(f\"Training loss: {running_loss/len(trainloader)}\")\n",
        "\n",
        "correct = 0\n",
        "total = 0\n",
        "with torch.no_grad():\n",
        "    model.eval()\n",
        "    for images, labels in tqdm(testloader):\n",
        "        outputs = model(images)\n",
        "        _, predicted = torch.max(outputs.data, 1)\n",
        "        total += labels.size(0)\n",
        "        correct += (predicted == labels).sum().item()\n",
        "\n",
        "print(f'Accuracy of the network on the 10000 test images: {100 * correct / total} %')\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 317,
          "referenced_widgets": [
            "5dfcb01e5ef7433fb6037567378cc8e6",
            "434f45e109534c8cbf2c9a8062f25e47",
            "6b336b23c5e44931b440315e48034e0c",
            "172d7c6e42294286ba4ab33404f54d7e",
            "f206e5296e834188a6723be58326f3d5",
            "2daa53b1f3d7434a953f0a3f86c23fa4",
            "266265a026ce42598ddfc71e772b1de4",
            "db26bc79362f45a89b2e3bca161bff54",
            "a7f69acfd06b4c4195d4d9e890a11e81",
            "5b6221a97b9b4f4e9dcfabce07d46a05",
            "3badc02581f841dc8a6cc7d8b34ffb19",
            "09aa98006d4c4ecfb0496432c539b3ed",
            "8f9596be50a54cdb897ab9e4cc06d499",
            "0ca0569a28bf44bcae3fe8f9d89aae79",
            "a98ab8ad0232453a880ba06ac3b0130e",
            "a13d7d92ec1e435d82b655084963981b",
            "b3c71306e2e940459691114724728e11",
            "b9c5b294365e4a5e8927657acf82904d",
            "b1a1dacbb4d041e3abbe1e36485a25ef",
            "8929f30d46444027bc5cbf0fcb4d5b58",
            "e37a94be493f46279fa16c116bccf25b",
            "350cb32a69854442904a8976d623fa3a",
            "087f185326224638bd27eb9cea145c2e",
            "0c827f76efac4255a32d175e4e2bb393",
            "4f336eb55c4e4070a5e620b63d8ca8c9",
            "a7fa34590d664828bfeafc4ad66bf540",
            "bbeca2ec4fa94acfa810ad88034ff621",
            "f0feaa17a0c4466e9fde788a3c2feed2",
            "bcfda19f9e1e41b890935ff0b2aee14d",
            "5ae711dccee743b69063f739097d0d73",
            "119b67f0e5a149108acdff2fd565f0e5",
            "d39a3f047f784e50a3758e515509d46b",
            "15887dad12f74ccaaa738efb6a0761a7",
            "3e1d208b1f884a5db8bb989c5a4bae97",
            "192092b8d6ca4c58806dbc877ceaf48b",
            "881ab5bc3c83467eaa52a7cfb312ffd9",
            "a3c02eb69d0c45faa66859a95d69fdc4",
            "069b08f4e02f42acbdea3cacf386a2a7",
            "ebabfd7c4a8c410abe2e247e9835223e",
            "4bc2fa1a42864c638e7ca2ea6bce77ed",
            "85ac58b159b748d2b372917de3341fd5",
            "67a4e7e15c2742d09d2fb13b8d4376d6",
            "d9c6df1498164f4fabeee93693c5bf52",
            "0113ccea44ab46ffab501656a9bb8a1f",
            "f8dc96c2caac43bd80965e38690ac8eb",
            "3d1fab038ec447ebbc7b0c9ea8a98d0e",
            "12654616c7f141008fe693c81d62a65d",
            "b5c5a0eb5aa344b0a17e44c897046241",
            "6b2caaf4c4d94022bda20d8a7dabd65e",
            "2f5fe6be68064b8e95694ca4429f25c5",
            "e5429605d6bc4898a450b1afdd426127",
            "e37f81454ffd473e88cc63ee21e9f201",
            "6558548b51344df5861c5b8bc9f7e01a",
            "4f611ee6ed504b62839957b4fd370ce4",
            "037846bd7b66465f950e16fa0a26ff68",
            "c7e6f6623bd94d7ab1d5c546793661a2",
            "05a816c78c344913944e000e49b515ac",
            "d4366bb03a3c4f849194a0ace42ef744",
            "9c25506af10a4b2683dbf192dbfcc3d2",
            "7bbcaa2c26f343688bd7580ca09b939d",
            "f22e88b1e899460fbd47b43a06edec62",
            "7def5a415fc840679011fc1603ede78e",
            "f53d6e6c898b40f0ae1daaa9954efc0e",
            "e33c5252ff0c4066976f48be3ec246e0",
            "b6b2eed876204a6e8e4223a32902c8fc",
            "7d7c7c10697a43edb853d7896b07f350"
          ]
        },
        "id": "xTqOq6rZyLvX",
        "outputId": "9b47e2f9-e841-4b78-8187-e494f7022039"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "  0%|          | 0/938 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "5dfcb01e5ef7433fb6037567378cc8e6"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Training loss: 0.34086699017694894\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "  0%|          | 0/938 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "09aa98006d4c4ecfb0496432c539b3ed"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Training loss: 0.13769662428869686\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "  0%|          | 0/938 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "087f185326224638bd27eb9cea145c2e"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Training loss: 0.11251812642276375\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "  0%|          | 0/938 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "3e1d208b1f884a5db8bb989c5a4bae97"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Training loss: 0.0998129521939419\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "  0%|          | 0/938 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "f8dc96c2caac43bd80965e38690ac8eb"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Training loss: 0.08331751398621068\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "  0%|          | 0/157 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "c7e6f6623bd94d7ab1d5c546793661a2"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Accuracy of the network on the 10000 test images: 97.64 %\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "qKKyMAseyfrY"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}