From 9b9aefca9fd5356e0a91cf2bd4d60f727e3576da Mon Sep 17 00:00:00 2001 From: torchxlabot2 Date: Thu, 2 May 2024 00:22:50 +0000 Subject: [PATCH] Update doc from commit 93ce054132ead2fd7720175d53f455f41575522d --- master/_modules/index.html | 3 +- master/_modules/torch_xla/core/functions.html | 3 +- master/_modules/torch_xla/core/xla_model.html | 3 +- .../distributed/parallel_loader.html | 3 +- .../distributed/xla_multiprocessing.html | 3 +- .../torch_xla/utils/serialization.html | 3 +- master/_modules/torch_xla/utils/utils.html | 3 +- master/_static/basic.css | 35 +---- master/_static/doctools.js | 130 ++-------------- master/_static/documentation_options.js | 2 +- master/_static/searchtools.js | 91 +++++++---- master/_static/sphinx_highlight.js | 144 ++++++++++++++++++ master/genindex.html | 3 +- master/index.html | 109 ++++++++++++- master/notes/source_of_recompilation.html | 3 +- master/objects.inv | Bin 724 -> 724 bytes master/py-modindex.html | 3 +- master/search.html | 3 +- master/searchindex.js | 2 +- 19 files changed, 346 insertions(+), 200 deletions(-) create mode 100644 master/_static/sphinx_highlight.js diff --git a/master/_modules/index.html b/master/_modules/index.html index 97c897a6553..628c68d4701 100644 --- a/master/_modules/index.html +++ b/master/_modules/index.html @@ -227,7 +227,7 @@
- master (2.4.0+git0a54b2b ) + master (2.4.0+git93ce054 )
@@ -387,6 +387,7 @@

All modules for which code is available

+ diff --git a/master/_modules/torch_xla/core/functions.html b/master/_modules/torch_xla/core/functions.html index 3e708d9d0f8..365cfe6ebc2 100644 --- a/master/_modules/torch_xla/core/functions.html +++ b/master/_modules/torch_xla/core/functions.html @@ -227,7 +227,7 @@
- master (2.4.0+git0a54b2b ) + master (2.4.0+git93ce054 )
@@ -578,6 +578,7 @@

Source code for torch_xla.core.functions

          
          
          
+         
      
 
   
diff --git a/master/_modules/torch_xla/core/xla_model.html b/master/_modules/torch_xla/core/xla_model.html
index b526746c6bb..b557c360d20 100644
--- a/master/_modules/torch_xla/core/xla_model.html
+++ b/master/_modules/torch_xla/core/xla_model.html
@@ -227,7 +227,7 @@
               
               
                 

-                  master (2.4.0+git0a54b2b )
+                  master (2.4.0+git93ce054 )
                 

               
             
@@ -1851,6 +1851,7 @@ 
Source code for torch_xla.core.xla_model
          
          
          
+         
      
 
   
diff --git a/master/_modules/torch_xla/distributed/parallel_loader.html b/master/_modules/torch_xla/distributed/parallel_loader.html
index 44170837f0e..8d8cd52b5d9 100644
--- a/master/_modules/torch_xla/distributed/parallel_loader.html
+++ b/master/_modules/torch_xla/distributed/parallel_loader.html
@@ -227,7 +227,7 @@
               
               
                 

-                  master (2.4.0+git0a54b2b )
+                  master (2.4.0+git93ce054 )
                 

               
             
@@ -592,6 +592,7 @@ 
Source code for torch_xla.distributed.parallel_loader

          
          
+         
      
 
   
diff --git a/master/_modules/torch_xla/distributed/xla_multiprocessing.html b/master/_modules/torch_xla/distributed/xla_multiprocessing.html
index 7e5aec1d7f1..59ebb2cd7e9 100644
--- a/master/_modules/torch_xla/distributed/xla_multiprocessing.html
+++ b/master/_modules/torch_xla/distributed/xla_multiprocessing.html
@@ -227,7 +227,7 @@
               
               
                 

-                  master (2.4.0+git0a54b2b )
+                  master (2.4.0+git93ce054 )
                 

               
             
@@ -503,6 +503,7 @@ 
Source code for torch_xla.distributed.xla_multiprocessing

          
          
+         
      
 
   
diff --git a/master/_modules/torch_xla/utils/serialization.html b/master/_modules/torch_xla/utils/serialization.html
index 2a319fe9357..31489964526 100644
--- a/master/_modules/torch_xla/utils/serialization.html
+++ b/master/_modules/torch_xla/utils/serialization.html
@@ -227,7 +227,7 @@
               
               
                 

-                  master (2.4.0+git0a54b2b )
+                  master (2.4.0+git93ce054 )
                 

               
             
@@ -483,6 +483,7 @@ 
Source code for torch_xla.utils.serialization

          
          
+         
      
 
   
diff --git a/master/_modules/torch_xla/utils/utils.html b/master/_modules/torch_xla/utils/utils.html
index d5e2d4b61e2..84937b64b27 100644
--- a/master/_modules/torch_xla/utils/utils.html
+++ b/master/_modules/torch_xla/utils/utils.html
@@ -227,7 +227,7 @@
               
               
                 

-                  master (2.4.0+git0a54b2b )
+                  master (2.4.0+git93ce054 )
                 

               
             
@@ -700,6 +700,7 @@ 
Source code for torch_xla.utils.utils
          
          
          
+         
      
 
   
diff --git a/master/_static/basic.css b/master/_static/basic.css
index 7243282d260..eeb0519a69b 100644
--- a/master/_static/basic.css
+++ b/master/_static/basic.css
@@ -236,7 +236,6 @@ div.body p, div.body dd, div.body li, div.body blockquote {
 a.headerlink {
     visibility: hidden;
 }
-
 a.brackets:before,
 span.brackets > a:before{
     content: "[";
@@ -247,6 +246,7 @@ span.brackets > a:after {
     content: "]";
 }
 
+
 h1:hover > a.headerlink,
 h2:hover > a.headerlink,
 h3:hover > a.headerlink,
@@ -334,13 +334,11 @@ aside.sidebar {
 p.sidebar-title {
     font-weight: bold;
 }
-
 div.admonition, div.topic, blockquote {
     clear: left;
 }
 
 /* -- topics ---------------------------------------------------------------- */
-
 div.topic {
     border: 1px solid #ccc;
     padding: 7px;
@@ -610,8 +608,6 @@ ol.simple p,
 ul.simple p {
     margin-bottom: 0;
 }
-
-/* Docutils 0.17 and older (footnotes & citations) */
 dl.footnote > dt,
 dl.citation > dt {
     float: left;
@@ -629,33 +625,6 @@ dl.citation > dd:after {
     clear: both;
 }
 
-/* Docutils 0.18+ (footnotes & citations) */
-aside.footnote > span,
-div.citation > span {
-    float: left;
-}
-aside.footnote > span:last-of-type,
-div.citation > span:last-of-type {
-  padding-right: 0.5em;
-}
-aside.footnote > p {
-  margin-left: 2em;
-}
-div.citation > p {
-  margin-left: 4em;
-}
-aside.footnote > p:last-of-type,
-div.citation > p:last-of-type {
-    margin-bottom: 0em;
-}
-aside.footnote > p:last-of-type:after,
-div.citation > p:last-of-type:after {
-    content: "";
-    clear: both;
-}
-
-/* Footnotes & citations ends */
-
 dl.field-list {
     display: grid;
     grid-template-columns: fit-content(30%) auto;
@@ -667,11 +636,11 @@ dl.field-list > dt {
     padding-left: 0.5em;
     padding-right: 5px;
 }
-
 dl.field-list > dt:after {
     content: ":";
 }
 
+
 dl.field-list > dd {
     padding-left: 0.5em;
     margin-top: 0em;
diff --git a/master/_static/doctools.js b/master/_static/doctools.js
index c3db08d1c38..527b876ca63 100644
--- a/master/_static/doctools.js
+++ b/master/_static/doctools.js
@@ -10,6 +10,13 @@
  */
 "use strict";
 
+const BLACKLISTED_KEY_CONTROL_ELEMENTS = new Set([
+  "TEXTAREA",
+  "INPUT",
+  "SELECT",
+  "BUTTON",
+]);
+
 const _ready = (callback) => {
   if (document.readyState !== "loading") {
     callback();
@@ -18,73 +25,11 @@ const _ready = (callback) => {
   }
 };
 
-/**
- * highlight a given string on a node by wrapping it in
- * span elements with the given class name.
- */
-const _highlight = (node, addItems, text, className) => {
-  if (node.nodeType === Node.TEXT_NODE) {
-    const val = node.nodeValue;
-    const parent = node.parentNode;
-    const pos = val.toLowerCase().indexOf(text);
-    if (
-      pos >= 0 &&
-      !parent.classList.contains(className) &&
-      !parent.classList.contains("nohighlight")
-    ) {
-      let span;
-
-      const closestNode = parent.closest("body, svg, foreignObject");
-      const isInSVG = closestNode && closestNode.matches("svg");
-      if (isInSVG) {
-        span = document.createElementNS("http://www.w3.org/2000/svg", "tspan");
-      } else {
-        span = document.createElement("span");
-        span.classList.add(className);
-      }
-
-      span.appendChild(document.createTextNode(val.substr(pos, text.length)));
-      parent.insertBefore(
-        span,
-        parent.insertBefore(
-          document.createTextNode(val.substr(pos + text.length)),
-          node.nextSibling
-        )
-      );
-      node.nodeValue = val.substr(0, pos);
-
-      if (isInSVG) {
-        const rect = document.createElementNS(
-          "http://www.w3.org/2000/svg",
-          "rect"
-        );
-        const bbox = parent.getBBox();
-        rect.x.baseVal.value = bbox.x;
-        rect.y.baseVal.value = bbox.y;
-        rect.width.baseVal.value = bbox.width;
-        rect.height.baseVal.value = bbox.height;
-        rect.setAttribute("class", className);
-        addItems.push({ parent: parent, target: rect });
-      }
-    }
-  } else if (node.matches && !node.matches("button, select, textarea")) {
-    node.childNodes.forEach((el) => _highlight(el, addItems, text, className));
-  }
-};
-const _highlightText = (thisNode, text, className) => {
-  let addItems = [];
-  _highlight(thisNode, addItems, text, className);
-  addItems.forEach((obj) =>
-    obj.parent.insertAdjacentElement("beforebegin", obj.target)
-  );
-};
-
 /**
  * Small JavaScript module for the documentation.
  */
 const Documentation = {
   init: () => {
-    Documentation.highlightSearchWords();
     Documentation.initDomainIndexTable();
     Documentation.initOnKeyListeners();
   },
@@ -126,51 +71,6 @@ const Documentation = {
     Documentation.LOCALE = catalog.locale;
   },
 
-  /**
-   * highlight the search words provided in the url in the text
-   */
-  highlightSearchWords: () => {
-    const highlight =
-      new URLSearchParams(window.location.search).get("highlight") || "";
-    const terms = highlight.toLowerCase().split(/\s+/).filter(x => x);
-    if (terms.length === 0) return; // nothing to do
-
-    // There should never be more than one element matching "div.body"
-    const divBody = document.querySelectorAll("div.body");
-    const body = divBody.length ? divBody[0] : document.querySelector("body");
-    window.setTimeout(() => {
-      terms.forEach((term) => _highlightText(body, term, "highlighted"));
-    }, 10);
-
-    const searchBox = document.getElementById("searchbox");
-    if (searchBox === null) return;
-    searchBox.appendChild(
-      document
-        .createRange()
-        .createContextualFragment(
-          '
' +
-            '' +
-            Documentation.gettext("Hide Search Matches") +
-            "
"
-        )
-    );
-  },
-
-  /**
-   * helper function to hide the search marks again
-   */
-  hideSearchWords: () => {
-    document
-      .querySelectorAll("#searchbox .highlight-link")
-      .forEach((el) => el.remove());
-    document
-      .querySelectorAll("span.highlighted")
-      .forEach((el) => el.classList.remove("highlighted"));
-    const url = new URL(window.location);
-    url.searchParams.delete("highlight");
-    window.history.replaceState({}, "", url);
-  },
-
   /**
    * helper function to focus on search bar
    */
@@ -210,15 +110,11 @@ const Documentation = {
     )
       return;
 
-    const blacklistedElements = new Set([
-      "TEXTAREA",
-      "INPUT",
-      "SELECT",
-      "BUTTON",
-    ]);
     document.addEventListener("keydown", (event) => {
-      if (blacklistedElements.has(document.activeElement.tagName)) return; // bail for input elements
-      if (event.altKey || event.ctrlKey || event.metaKey) return; // bail with special keys
+      // bail for input elements
+      if (BLACKLISTED_KEY_CONTROL_ELEMENTS.has(document.activeElement.tagName)) return;
+      // bail with special keys
+      if (event.altKey || event.ctrlKey || event.metaKey) return;
 
       if (!event.shiftKey) {
         switch (event.key) {
@@ -240,10 +136,6 @@ const Documentation = {
               event.preventDefault();
             }
             break;
-          case "Escape":
-            if (!DOCUMENTATION_OPTIONS.ENABLE_SEARCH_SHORTCUTS) break;
-            Documentation.hideSearchWords();
-            event.preventDefault();
         }
       }
 
diff --git a/master/_static/documentation_options.js b/master/_static/documentation_options.js
index df431cc682c..ac1f83182f2 100644
--- a/master/_static/documentation_options.js
+++ b/master/_static/documentation_options.js
@@ -10,5 +10,5 @@ var DOCUMENTATION_OPTIONS = {
     SOURCELINK_SUFFIX: '.txt',
     NAVIGATION_WITH_KEYS: false,
     SHOW_SEARCH_SUMMARY: true,
-    ENABLE_SEARCH_SHORTCUTS: false,
+    ENABLE_SEARCH_SHORTCUTS: true,
 };
\ No newline at end of file
diff --git a/master/_static/searchtools.js b/master/_static/searchtools.js
index ac4d5861f95..e89e34d4e77 100644
--- a/master/_static/searchtools.js
+++ b/master/_static/searchtools.js
@@ -57,14 +57,14 @@ const _removeChildren = (element) => {
 const _escapeRegExp = (string) =>
   string.replace(/[.*+\-?^${}()|[\]\\]/g, "\\$&"); // $& means the whole matched string
 
-const _displayItem = (item, highlightTerms, searchTerms) => {
+const _displayItem = (item, searchTerms) => {
   const docBuilder = DOCUMENTATION_OPTIONS.BUILDER;
   const docUrlRoot = DOCUMENTATION_OPTIONS.URL_ROOT;
   const docFileSuffix = DOCUMENTATION_OPTIONS.FILE_SUFFIX;
   const docLinkSuffix = DOCUMENTATION_OPTIONS.LINK_SUFFIX;
   const showSearchSummary = DOCUMENTATION_OPTIONS.SHOW_SEARCH_SUMMARY;
 
-  const [docName, title, anchor, descr] = item;
+  const [docName, title, anchor, descr, score, _filename] = item;
 
   let listItem = document.createElement("li");
   let requestUrl;
@@ -82,13 +82,12 @@ const _displayItem = (item, highlightTerms, searchTerms) => {
     requestUrl = docUrlRoot + docName + docFileSuffix;
     linkUrl = docName + docLinkSuffix;
   }
-  const params = new URLSearchParams();
-  params.set("highlight", [...highlightTerms].join(" "));
   let linkEl = listItem.appendChild(document.createElement("a"));
-  linkEl.href = linkUrl + "?" + params.toString() + anchor;
+  linkEl.href = linkUrl + anchor;
+  linkEl.dataset.score = score;
   linkEl.innerHTML = title;
   if (descr)
-    listItem.appendChild(document.createElement("span")).innerText =
+    listItem.appendChild(document.createElement("span")).innerHTML =
       " (" + descr + ")";
   else if (showSearchSummary)
     fetch(requestUrl)
@@ -96,7 +95,7 @@ const _displayItem = (item, highlightTerms, searchTerms) => {
       .then((data) => {
         if (data)
           listItem.appendChild(
-            Search.makeSearchSummary(data, searchTerms, highlightTerms)
+            Search.makeSearchSummary(data, searchTerms)
           );
       });
   Search.output.appendChild(listItem);
@@ -116,15 +115,14 @@ const _finishSearch = (resultCount) => {
 const _displayNextItem = (
   results,
   resultCount,
-  highlightTerms,
   searchTerms
 ) => {
   // results left, load the summary and display it
   // this is intended to be dynamic (don't sub resultsCount)
   if (results.length) {
-    _displayItem(results.pop(), highlightTerms, searchTerms);
+    _displayItem(results.pop(), searchTerms);
     setTimeout(
-      () => _displayNextItem(results, resultCount, highlightTerms, searchTerms),
+      () => _displayNextItem(results, resultCount, searchTerms),
       5
     );
   }
@@ -155,10 +153,8 @@ const Search = {
   _pulse_status: -1,
 
   htmlToText: (htmlString) => {
-    const htmlElement = document
-      .createRange()
-      .createContextualFragment(htmlString);
-    _removeChildren(htmlElement.querySelectorAll(".headerlink"));
+    const htmlElement = new DOMParser().parseFromString(htmlString, 'text/html');
+    htmlElement.querySelectorAll(".headerlink").forEach((el) => { el.remove() });
     const docContent = htmlElement.querySelector('[role="main"]');
     if (docContent !== undefined) return docContent.textContent;
     console.warn(
@@ -239,6 +235,12 @@ const Search = {
    * execute search (requires search index to be loaded)
    */
   query: (query) => {
+    const filenames = Search._index.filenames;
+    const docNames = Search._index.docnames;
+    const titles = Search._index.titles;
+    const allTitles = Search._index.alltitles;
+    const indexEntries = Search._index.indexentries;
+
     // stem the search terms and add them to the correct list
     const stemmer = new Stemmer();
     const searchTerms = new Set();
@@ -266,6 +268,10 @@ const Search = {
       }
     });
 
+    if (SPHINX_HIGHLIGHT_ENABLED) {  // set in sphinx_highlight.js
+      localStorage.setItem("sphinx_highlight_terms", [...highlightTerms].join(" "))
+    }
+
     // console.debug("SEARCH: searching for:");
     // console.info("required: ", [...searchTerms]);
     // console.info("excluded: ", [...excludedTerms]);
@@ -274,6 +280,40 @@ const Search = {
     let results = [];
     _removeChildren(document.getElementById("search-progress"));
 
+    const queryLower = query.toLowerCase();
+    for (const [title, foundTitles] of Object.entries(allTitles)) {
+      if (title.toLowerCase().includes(queryLower) && (queryLower.length >= title.length/2)) {
+        for (const [file, id] of foundTitles) {
+          let score = Math.round(100 * queryLower.length / title.length)
+          results.push([
+            docNames[file],
+            titles[file] !== title ? `${titles[file]} > ${title}` : title,
+            id !== null ? "#" + id : "",
+            null,
+            score,
+            filenames[file],
+          ]);
+        }
+      }
+    }
+
+    // search for explicit entries in index directives
+    for (const [entry, foundEntries] of Object.entries(indexEntries)) {
+      if (entry.includes(queryLower) && (queryLower.length >= entry.length/2)) {
+        for (const [file, id] of foundEntries) {
+          let score = Math.round(100 * queryLower.length / entry.length)
+          results.push([
+            docNames[file],
+            titles[file],
+            id ? "#" + id : "",
+            null,
+            score,
+            filenames[file],
+          ]);
+        }
+      }
+    }
+
     // lookup as object
     objectTerms.forEach((term) =>
       results.push(...Search.performObjectSearch(term, objectTerms))
@@ -320,7 +360,7 @@ const Search = {
     // console.info("search results:", Search.lastresults);
 
     // print the results
-    _displayNextItem(results, results.length, highlightTerms, searchTerms);
+    _displayNextItem(results, results.length, searchTerms);
   },
 
   /**
@@ -401,8 +441,8 @@ const Search = {
     // prepare search
     const terms = Search._index.terms;
     const titleTerms = Search._index.titleterms;
-    const docNames = Search._index.docnames;
     const filenames = Search._index.filenames;
+    const docNames = Search._index.docnames;
     const titles = Search._index.titles;
 
     const scoreMap = new Map();
@@ -499,16 +539,15 @@ const Search = {
   /**
    * helper function to return a node containing the
    * search summary for a given text. keywords is a list
-   * of stemmed words, highlightWords is the list of normal, unstemmed
-   * words. the first one is used to find the occurrence, the
-   * latter for highlighting it.
+   * of stemmed words.
    */
-  makeSearchSummary: (htmlText, keywords, highlightWords) => {
-    const text = Search.htmlToText(htmlText).toLowerCase();
+  makeSearchSummary: (htmlText, keywords) => {
+    const text = Search.htmlToText(htmlText);
     if (text === "") return null;
 
+    const textLower = text.toLowerCase();
     const actualStartPosition = [...keywords]
-      .map((k) => text.indexOf(k.toLowerCase()))
+      .map((k) => textLower.indexOf(k.toLowerCase()))
       .filter((i) => i > -1)
       .slice(-1)[0];
     const startWithContext = Math.max(actualStartPosition - 120, 0);
@@ -516,13 +555,9 @@ const Search = {
     const top = startWithContext === 0 ? "" : "...";
     const tail = startWithContext + 240 < text.length ? "..." : "";
 
-    let summary = document.createElement("div");
+    let summary = document.createElement("p");
     summary.classList.add("context");
-    summary.innerText = top + text.substr(startWithContext, 240).trim() + tail;
-
-    highlightWords.forEach((highlightWord) =>
-      _highlightText(summary, highlightWord, "highlighted")
-    );
+    summary.textContent = top + text.substr(startWithContext, 240).trim() + tail;
 
     return summary;
   },
diff --git a/master/_static/sphinx_highlight.js b/master/_static/sphinx_highlight.js
new file mode 100644
index 00000000000..aae669d7ea6
--- /dev/null
+++ b/master/_static/sphinx_highlight.js
@@ -0,0 +1,144 @@
+/* Highlighting utilities for Sphinx HTML documentation. */
+"use strict";
+
+const SPHINX_HIGHLIGHT_ENABLED = true
+
+/**
+ * highlight a given string on a node by wrapping it in
+ * span elements with the given class name.
+ */
+const _highlight = (node, addItems, text, className) => {
+  if (node.nodeType === Node.TEXT_NODE) {
+    const val = node.nodeValue;
+    const parent = node.parentNode;
+    const pos = val.toLowerCase().indexOf(text);
+    if (
+      pos >= 0 &&
+      !parent.classList.contains(className) &&
+      !parent.classList.contains("nohighlight")
+    ) {
+      let span;
+
+      const closestNode = parent.closest("body, svg, foreignObject");
+      const isInSVG = closestNode && closestNode.matches("svg");
+      if (isInSVG) {
+        span = document.createElementNS("http://www.w3.org/2000/svg", "tspan");
+      } else {
+        span = document.createElement("span");
+        span.classList.add(className);
+      }
+
+      span.appendChild(document.createTextNode(val.substr(pos, text.length)));
+      parent.insertBefore(
+        span,
+        parent.insertBefore(
+          document.createTextNode(val.substr(pos + text.length)),
+          node.nextSibling
+        )
+      );
+      node.nodeValue = val.substr(0, pos);
+
+      if (isInSVG) {
+        const rect = document.createElementNS(
+          "http://www.w3.org/2000/svg",
+          "rect"
+        );
+        const bbox = parent.getBBox();
+        rect.x.baseVal.value = bbox.x;
+        rect.y.baseVal.value = bbox.y;
+        rect.width.baseVal.value = bbox.width;
+        rect.height.baseVal.value = bbox.height;
+        rect.setAttribute("class", className);
+        addItems.push({ parent: parent, target: rect });
+      }
+    }
+  } else if (node.matches && !node.matches("button, select, textarea")) {
+    node.childNodes.forEach((el) => _highlight(el, addItems, text, className));
+  }
+};
+const _highlightText = (thisNode, text, className) => {
+  let addItems = [];
+  _highlight(thisNode, addItems, text, className);
+  addItems.forEach((obj) =>
+    obj.parent.insertAdjacentElement("beforebegin", obj.target)
+  );
+};
+
+/**
+ * Small JavaScript module for the documentation.
+ */
+const SphinxHighlight = {
+
+  /**
+   * highlight the search words provided in localstorage in the text
+   */
+  highlightSearchWords: () => {
+    if (!SPHINX_HIGHLIGHT_ENABLED) return;  // bail if no highlight
+
+    // get and clear terms from localstorage
+    const url = new URL(window.location);
+    const highlight =
+        localStorage.getItem("sphinx_highlight_terms")
+        || url.searchParams.get("highlight")
+        || "";
+    localStorage.removeItem("sphinx_highlight_terms")
+    url.searchParams.delete("highlight");
+    window.history.replaceState({}, "", url);
+
+    // get individual terms from highlight string
+    const terms = highlight.toLowerCase().split(/\s+/).filter(x => x);
+    if (terms.length === 0) return; // nothing to do
+
+    // There should never be more than one element matching "div.body"
+    const divBody = document.querySelectorAll("div.body");
+    const body = divBody.length ? divBody[0] : document.querySelector("body");
+    window.setTimeout(() => {
+      terms.forEach((term) => _highlightText(body, term, "highlighted"));
+    }, 10);
+
+    const searchBox = document.getElementById("searchbox");
+    if (searchBox === null) return;
+    searchBox.appendChild(
+      document
+        .createRange()
+        .createContextualFragment(
+          '
' +
+            '' +
+            _("Hide Search Matches") +
+            "
"
+        )
+    );
+  },
+
+  /**
+   * helper function to hide the search marks again
+   */
+  hideSearchWords: () => {
+    document
+      .querySelectorAll("#searchbox .highlight-link")
+      .forEach((el) => el.remove());
+    document
+      .querySelectorAll("span.highlighted")
+      .forEach((el) => el.classList.remove("highlighted"));
+    localStorage.removeItem("sphinx_highlight_terms")
+  },
+
+  initEscapeListener: () => {
+    // only install a listener if it is really needed
+    if (!DOCUMENTATION_OPTIONS.ENABLE_SEARCH_SHORTCUTS) return;
+
+    document.addEventListener("keydown", (event) => {
+      // bail for input elements
+      if (BLACKLISTED_KEY_CONTROL_ELEMENTS.has(document.activeElement.tagName)) return;
+      // bail with special keys
+      if (event.shiftKey || event.altKey || event.ctrlKey || event.metaKey) return;
+      if (DOCUMENTATION_OPTIONS.ENABLE_SEARCH_SHORTCUTS && (event.key === "Escape")) {
+        SphinxHighlight.hideSearchWords();
+        event.preventDefault();
+      }
+    });
+  },
+};
+
+_ready(SphinxHighlight.highlightSearchWords);
+_ready(SphinxHighlight.initEscapeListener);
diff --git a/master/genindex.html b/master/genindex.html
index c767cb7e329..748bae3980d 100644
--- a/master/genindex.html
+++ b/master/genindex.html
@@ -227,7 +227,7 @@
               
               
                 

-                  master (2.4.0+git0a54b2b )
+                  master (2.4.0+git93ce054 )
                 

               
             
@@ -632,6 +632,7 @@ 
X

          
          
          
+         
      
 
   
diff --git a/master/index.html b/master/index.html
index 1566b38fa32..2e46666dba2 100644
--- a/master/index.html
+++ b/master/index.html
@@ -227,7 +227,7 @@
               
               
                 

-                  master (2.4.0+git0a54b2b )
+                  master (2.4.0+git93ce054 )
                 

               
             
@@ -279,9 +279,56 @@
 
 
 
  • PyTorch/XLA API
 
    • 
@@ -3850,9 +3897,56 @@ 
    HuggingFace Llama 2 ExamplePyTorch/XLA API
 
@@ -4026,6 +4120,7 @@ 
    HuggingFace Llama 2 Example
          
          
+         
      
 
   
diff --git a/master/notes/source_of_recompilation.html b/master/notes/source_of_recompilation.html
index 743c6d50fee..0d8fc0efa03 100644
--- a/master/notes/source_of_recompilation.html
+++ b/master/notes/source_of_recompilation.html
@@ -227,7 +227,7 @@
               
               
                 
    
-                  master (2.4.0+git0a54b2b )
+                  master (2.4.0+git93ce054 )
                 
    
               
             
@@ -560,6 +560,7 @@ 
    Appendix:
          
          
+         
      
 
   
diff --git a/master/objects.inv b/master/objects.inv
index 13cce30cfd5dfed1e43aa850e4993d827cc520eb..7cde2c5cae603ab4b2212921a316dfcb48cd61a6 100644
GIT binary patch
delta 18
Zcmcb@dWCgD0K27ea;kx;$;R+nCICE%1?~U<

delta 18
Zcmcb@dWCgD0J}k=sY#Mi(#G&wCICGC1`z-N

diff --git a/master/py-modindex.html b/master/py-modindex.html
index ea0baee4ebd..e4b4c9c9003 100644
--- a/master/py-modindex.html
+++ b/master/py-modindex.html
@@ -230,7 +230,7 @@
               
               
                 
    
-                  master (2.4.0+git0a54b2b )
+                  master (2.4.0+git93ce054 )
                 
    
               
             
@@ -431,6 +431,7 @@ 
    Python Module Index
    
          
          
          
+         
      
 
   
diff --git a/master/search.html b/master/search.html
index 3f29a0c8f18..6f363bee24e 100644
--- a/master/search.html
+++ b/master/search.html
@@ -227,7 +227,7 @@
               
               
                 
    
-                  master (2.4.0+git0a54b2b )
+                  master (2.4.0+git93ce054 )
                 
    
               
             
@@ -394,6 +394,7 @@
          
          
          
+         
          
          
      
diff --git a/master/searchindex.js b/master/searchindex.js
index 088962ebf04..b8d529ad819 100644
--- a/master/searchindex.js
+++ b/master/searchindex.js
@@ -1 +1 @@
-Search.setIndex({"docnames": ["index", "notes/source_of_recompilation"], "filenames": ["index.rst", "notes/source_of_recompilation.md"], "titles": ["PyTorch on XLA Devices", "Source of recompilations in torch_xla"], "terms": {"like": [0, 1], "torch_xla": 0, "packag": 0, "thi": [0, 1], "document": 0, "describ": 0, "your": [0, 1], "add": [0, 1], "type": 0, "work": [0, 1], "just": [0, 1], "other": [0, 1], "For": [0, 1], "print": [0, 1], "import": 0, "core": 0, "t": [0, 1], "randn": 0, "xla_devic": 0, "look": 0, "familiar": 0, "same": [0, 1], "interfac": 0, "regular": 0, "few": [0, 1], "addit": 0, "initi": 0, "return": [0, 1], "mai": [0, 1], "depend": [0, 1], "oper": 0, "cuda": [0, 1], "ad": [0, 1], "togeth": 0, "t0": 0, "t1": 0, "Or": [0, 1], "matrix": 0, "multipli": 0, "mm": 0, "neural": 0, "network": 0, "modul": 0, "l_in": 0, "linear": 0, "nn": 0, "20": 0, "l_out": 0, "onli": [0, 1], "so": [0, 1], "input": 0, "floattensor": 0, "throw": 0, "error": 0, "sinc": [0, 1], "exist": 0, "one": [0, 1], "requir": [0, 1], "line": [0, 1], "specif": 0, "follow": [0, 1], "snippet": 0, "highlight": 0, "when": 0, "show": 0, "loss_fn": 0, "nllloss": 0, "sgd": 0, "lr": 0, "momentum": 0, "target": [0, 1], "train_load": 0, "zero_grad": 0, "loss": 0, "backward": 0, "step": [0, 1], "mark_step": 0, "easi": [0, 1], "switch": [0, 1], "definit": [0, 1], "dataload": 0, "loop": [0, 1], "ani": [0, 1], "coupl": 0, "acquir": 0, "mark": 0, "call": [0, 1], "end": 0, "each": [0, 1], "iter": 0, "caus": [0, 1], "its": 0, "graph": [0, 1], "updat": [0, 1], "see": [0, 1], "make": [0, 1], "acceler": 0, "parallel_load": 0, "pl": 0, "xla_multiprocess": 0, "xmp": 0, "def": 0, "_mp_fn": 0, "index": 0, "mp_device_load": 0, "mpdeviceload": 0, "optimizer_step": 0, "__name__": 0, "__main__": 0, "spawn": 0, "arg": 0, "There": [0, 1], "three": 0, "between": [0, 1], "previou": [0, 1], "let": 0, "go": 0, "over": 0, "abl": [0, 1], "access": [0, 1], "assign": 0, "v4": 0, "4": [0, 1], "being": 0, "up": [0, 1], "own": 0, "note": [0, 1], "you": 0, "0": [0, 1], "all": [0, 1], "becaus": 0, "doe": [0, 1], "mean": [0, 1], "function": 0, "have": [0, 1], "thread": 0, "doc": [0, 1], "onto": 0, "wrap": 0, "It": [0, 1], "preload": 0, "overlap": 0, "improv": 0, "also": [0, 1], "everi": [0, 1], "batches_per_execut": 0, "default": 0, "batch": 0, "yield": 0, "consolid": 0, "issu": 0, "comput": [0, 1], "pretti": [0, 1], "much": [0, 1], "all_reduce_gradi": 0, "reduc": 0, "remain": [0, 1], "start": 0, "retriev": [0, 1], "within": 0, "which": [0, 1], "ha": [0, 1], "parent": 0, "stack": [0, 1], "full": 0, "multiprocess": 0, "veri": [0, 1], "talk": 0, "about": [0, 1], "independ": 0, "bit": 0, "avail": [0, 1], "13": 0, "releas": 0, "befor": [0, 1], "pleas": 0, "our": [0, 1], "explain": [0, 1], "googl": 0, "cloud": 0, "basi": 0, "gcloud": 0, "command": 0, "project": 0, "howto": 0, "focu": [0, 1], "perspect": 0, "assum": [0, 1], "abov": [0, 1], "section": 0, "train_mnist_xla": 0, "py": 0, "If": [0, 1], "would": [0, 1], "ssh": 0, "tpuvm": 0, "pjrt_devic": 0, "python3": 0, "now": [0, 1], "order": 0, "16": 0, "need": [0, 1], "sure": 0, "usual": 0, "done": [0, 1], "scp": 0, "copi": 0, "time": [0, 1], "alpha": 0, "vm": 0, "zone": 0, "worker": 0, "outsid": 0, "underli": 0, "infrastructur": 0, "global": 0, "topolog": 0, "local": 0, "ordin": 0, "cross": 0, "commun": 0, "happen": [0, 1], "across": 0, "instead": [0, 1], "regard": 0, "inform": [0, 1], "complet": 0, "fakedata": 0, "But": [0, 1], "even": [0, 1], "though": 0, "act": 0, "lot": [0, 1], "intern": [0, 1], "uniqu": [0, 1], "launch": 0, "immedi": 0, "eagerli": [0, 1], "hand": 0, "thei": [0, 1], "record": 0, "until": 0, "result": 0, "defer": 0, "separ": 0, "might": [0, 1], "fuse": 0, "gener": [0, 1], "invis": 0, "caller": 0, "construct": 0, "send": 0, "them": [0, 1], "synchron": 0, "insert": 0, "barrier": 0, "explicitli": [0, 1], "paper": 0, "datatyp": 0, "In": [0, 1], "fact": 0, "handl": [0, 1], "float": [0, 1], "doubl": [0, 1], "behavior": 0, "control": 0, "xla_use_bf16": 0, "xla_downcast_bf16": 0, "By": [0, 1], "both": [0, 1], "set": [0, 1], "float32": 0, "directli": [0, 1], "map": 0, "bf16": 0, "primit": 0, "alwai": [0, 1], "regardless": 0, "actual": [0, 1], "re": [0, 1], "convers": [0, 1], "opaqu": 0, "back": 0, "trigger": [0, 1], "represent": 0, "expos": 0, "storag": 0, "appear": 0, "contigu": 0, "unlik": 0, "allow": 0, "adjust": 0, "better": [0, 1], "view": 0, "relationship": 0, "preserv": 0, "put": 0, "anoth": [0, 1], "wai": [0, 1], "onc": [0, 1], "again": 0, "appreci": 0, "accommod": 0, "transit": 0, "pt": 0, "were": [0, 1], "per": 0, "care": [0, 1], "must": 0, "taken": [0, 1], "recommend": 0, "recreat": 0, "after": [0, 1], "been": [0, 1], "destin": 0, "provid": [0, 1], "previous": 0, "state_dict": 0, "path": [0, 1], "case": 0, "master": 0, "where": [0, 1], "limit": 0, "compar": 0, "size": [0, 1], "footprint": 0, "serial": 0, "xser": 0, "stream": 0, "amount": [0, 1], "match": 0, "restor": 0, "load_state_dict": 0, "possibl": 0, "unavail": 0, "fail": 0, "under": 0, "activ": 0, "futur": [0, 1], "trace": [0, 1], "hlo": 0, "consum": [0, 1], "doesn": [0, 1], "persist": 0, "disk": 0, "reus": 0, "significantli": 0, "recompil": 0, "still": [0, 1], "occur": 0, "experiment": 0, "opt": 0, "through": [0, 1], "initialize_cach": 0, "xr": 0, "your_cache_path": 0, "readonli": 0, "fals": 0, "specifi": 0, "whether": 0, "write": 0, "share": 0, "mount": 0, "workload": 0, "n": 0, "none": 0, "devkind": 0, "given": [0, 1], "python": [0, 1], "int": [0, 1], "option": 0, "otherwis": [0, 1], "first": 0, "string": 0, "custom": [0, 1], "deprec": 0, "request": [0, 1], "get_xla_supported_devic": 0, "max_devic": 0, "list": 0, "kind": 0, "name": [0, 1], "maximum": 0, "xla_device_hw": 0, "hardwar": 0, "get_ordin": 0, "defval": 0, "replic": 0, "rang": 0, "xrt_world_siz": 0, "minu": 0, "valu": [0, 1], "ignor": 0, "get_local_ordin": 0, "is_master_ordin": 0, "true": [0, 1], "bool": 0, "while": [0, 1], "num_host": 0, "boolean": 0, "indic": [0, 1], "part": 0, "all_reduc": 0, "reduce_typ": 0, "scale": 0, "pin_layout": 0, "inplac": 0, "One": 0, "reduce_sum": 0, "reduce_mul": 0, "reduce_and": 0, "reduce_or": 0, "reduce_min": 0, "reduce_max": 0, "either": [0, 1], "op": [0, 1], "appli": 0, "repres": [0, 1], "replica": 0, "3": 0, "5": [0, 1], "6": [0, 1], "7": 0, "defin": 0, "two": [0, 1], "pin": 0, "pine": 0, "prevent": 0, "potenti": 0, "corrupt": 0, "particip": 0, "slightli": 0, "program": [0, 1], "unpin": 0, "messag": 0, "hlomodul": 0, "constrain": 0, "pass": 0, "hold": 0, "tupl": [0, 1], "itself": 0, "all_gath": 0, "dim": 0, "gather": 0, "along": 0, "dimens": 0, "all_to_al": 0, "split_dimens": 0, "concat_dimens": 0, "split_count": 0, "alltoal": 0, "http": 0, "www": 0, "tensorflow": [0, 1], "org": 0, "operation_semant": 0, "upon": 0, "split": 0, "concat": 0, "count": 0, "add_step_closur": 0, "closur": 0, "run_async": 0, "ones": [0, 1], "mani": [0, 1], "dure": 0, "consol": 0, "post": 0, "tensorboard": 0, "etc": [0, 1], "content": 0, "intermediari": 0, "inspect": 0, "point": [0, 1], "typic": 0, "ensur": [0, 1], "live": [0, 1], "alreadi": [0, 1], "materi": [0, 1], "includ": [0, 1], "captur": 0, "argument": 0, "queu": 0, "sequenti": 0, "advis": 0, "throttl": 0, "event": 0, "callabl": 0, "asynchron": 0, "wait_device_op": 0, "wait": 0, "async": 0, "whose": 0, "empti": 0, "optimizer_arg": 0, "parallelload": 0, "dataparallel": 0, "necessari": 0, "loader": 0, "next": [0, 1], "dict": 0, "dictionari": 0, "file_or_path": 0, "master_onli": 0, "global_mast": 0, "file": 0, "transfer": 0, "nest": 0, "object": 0, "overrid": 0, "locat": 0, "flag": 0, "sync": 0, "main": 0, "hang": 0, "tag": 0, "payload": 0, "b": [0, 1], "client": 0, "reach": 0, "server": 0, "effect": 0, "alia": 0, "xla_rendezv": 0, "join": 0, "byte": 0, "sent": 0, "exchang": 0, "i": [0, 1], "posit": 0, "do_on_ordin": 0, "contain": [0, 1], "interpret": [0, 1], "ran": 0, "mesh_reduc": 0, "reduce_fn": 0, "out": 0, "reduct": 0, "receiv": 0, "come": [0, 1], "set_rng_stat": 0, "seed": 0, "random": 0, "state": 0, "integ": 0, "rng": 0, "miss": 0, "get_rng_stat": 0, "get_memory_info": 0, "kb_free": 0, "free": 0, "kb": 0, "kb_total": 0, "total": [0, 1], "kei": 0, "get_stablehlo": 0, "str": 0, "stablehlo": 0, "format": 0, "dump": 0, "whole": [0, 1], "todo": 0, "lsy323": 0, "intermedi": 0, "investig": 0, "straightforward": 0, "identifi": 0, "To": [0, 1], "enabl": 0, "info": [0, 1], "env": 0, "var": 0, "xla_hlo_debug": 0, "root": [0, 1], "get_stablehlo_bytecod": 0, "bytecod": 0, "autograd": 0, "differenti": 0, "select": 0, "class": 0, "batchdim": 0, "loader_prefetch_s": 0, "device_prefetch_s": 0, "host_to_device_transfer_thread": 0, "input_shard": 0, "upload": 0, "th": 0, "sampl": 0, "len": 0, "max": [0, 1], "capac": 0, "queue": 0, "deposit": 0, "shardingspec": 0, "compat": 0, "per_device_load": 0, "structur": 0, "resid": 0, "fn": 0, "nproc": 0, "daemon": 0, "start_method": 0, "base": [0, 1], "At": 0, "moment": 0, "block": 0, "creation": 0, "mpmodelwrapp": 0, "minim": 0, "usag": [0, 1], "fork": 0, "henc": 0, "scope": 0, "insid": 0, "wrapped_model": 0, "mynetwork": 0, "advantag": 0, "page": 0, "origin": 0, "weight": 0, "second": 0, "lower": [0, 1], "system": 0, "mpserialexecutor": 0, "fashion": 0, "among": 0, "serial_exec": 0, "load_dataset": 0, "maybe_download_and_load": 0, "avoid": 0, "download": 0, "executor": 0, "dataset": 0, "lambda": 0, "tmp": 0, "wrt": 0, "samplegener": 0, "sample_count": 0, "place": 0, "synthet": 0, "datawrapp": 0, "illustr": 0, "e": [0, 1], "g": [0, 1], "solut": [0, 1], "howev": 0, "serv": 0, "brief": 0, "help": [0, 1], "reader": 0, "modif": 0, "supplement": 0, "fetch": 0, "ir": [0, 1], "encount": 0, "discuss": 0, "below": [0, 1], "lazytensor": 0, "machin": 0, "readabl": 0, "opcod": 0, "effici": 0, "fed": 0, "later": 0, "except": 0, "four": 0, "slice": 0, "attach": 0, "shape": 0, "evalu": [0, 1], "sometim": [0, 1], "lead": 0, "signific": 0, "log": 0, "callback": 0, "slower": 0, "xla_tensor_z": 0, "respons": 0, "These": 0, "cut": [0, 1], "transferfromdevic": 0, "special": 0, "instruct": 0, "tell": [0, 1], "break": [0, 1], "properti": [0, 1], "suppos": 0, "we": [0, 1], "tensors_on_devic": 0, "z": [0, 1], "without": 0, "subgraph": [0, 1], "correspond": 0, "introduc": 0, "smaller": [0, 1], "signal": 0, "far": 0, "submit": 0, "seen": 0, "small": [0, 1], "bottleneck": 0, "dynam": 0, "suitabl": 0, "pad": [0, 1], "fix": 0, "spent": 0, "trade": [0, 1], "off": 0, "consid": 0, "larg": [0, 1], "spend": 0, "fusion": 0, "long": [0, 1], "faster": [0, 1], "due": 0, "worth": [0, 1], "As": [0, 1], "could": [0, 1], "latter": 0, "job": 0, "imag": [0, 1], "nightli": 0, "runtime_vers": 0, "export": 0, "central2": 0, "project_id": 0, "id": 0, "accelerator_typ": 0, "32": 0, "tpu_nam": 0, "your_tpu_nam": 0, "subnetwork": 0, "tpusubnet": 0, "similar": 0, "pip3": 0, "googleapi": 0, "com": 0, "cp38": 0, "linux_x86_64": 0, "whl": 0, "librari": 0, "sudo": 0, "apt": 0, "libopenbla": 0, "dev": 0, "libgl1": 0, "guidelin": 0, "modifi": 0, "replac": 0, "remov": 0, "progress": 0, "bar": 0, "rememb": 0, "someth": 0, "txt2img": 0, "prompt": 0, "photograph": 0, "astronaut": 0, "ride": 0, "hors": 0, "diff": 0, "found": 0, "relat": 0, "written": 0, "precision_scop": 0, "addition": 0, "particular": 0, "configur": 0, "frozenclipembedd": 0, "therefor": 0, "well": [0, 1], "simplic": [0, 1], "tutori": 0, "self": 0, "ddim": 0, "schedul": 0, "top": 0, "attr": 0, "statement": [0, 1], "disabl": 0, "stop": 0, "fall": 0, "especi": 0, "slow": 0, "tri": 0, "huge": [0, 1], "difficult": 0, "degrad": 0, "easier": [0, 1], "readi": 0, "reason": 0, "cover": 0, "sd": 0, "xl": 0, "git": 0, "github": 0, "text_to_imag": 0, "inference_tpu_single_devic": 0, "speed": [0, 1], "inference_tpu_multidevic": 0, "warn": 0, "watch": 0, "made": [0, 1], "lora": 0, "model_id": 0, "stabilityai": 0, "9": [0, 1], "pipelin": 0, "dpmsolvermultistepschedul": 0, "pip": 0, "r": 0, "txt": 0, "invisible_watermark": 0, "transform": 0, "safetensor": 0, "agre": 0, "licens": 0, "card": 0, "account": 0, "token": 0, "cli": 0, "login": 0, "_your_copied_token__": 0, "readm": 0, "pipe": 0, "hour": 0, "wherea": 0, "likewis": 0, "gpt": 0, "15": 0, "min": 0, "epoch": 0, "becom": [0, 1], "subsequ": 0, "mayb": 0, "try": [0, 1], "notic": 0, "piec": 0, "__call__": 0, "commit": 0, "nonzero": [0, 1], "item": 0, "desir": 0, "down": 0, "rule": 0, "thumb": 0, "fit": 0, "durat": 0, "enough": [0, 1], "least": 0, "good": [0, 1], "constantli": 0, "idl": 0, "inference_tpu_": 0, "capture_profil": 0, "give": 0, "xp": 0, "measur": 0, "u": 0, "net": 0, "analyz": 0, "cannot": [0, 1], "portion": 0, "busi": 0, "middl": [0, 1], "scroll": 0, "find": 0, "occupi": 0, "demonstr": 0, "displai": 0, "largest": 0, "appropri": 0, "zoom": 0, "timelin": 0, "period": 0, "examin": 0, "thing": 0, "did": 0, "pipe_watermark": 0, "closer": 0, "preced": 0, "finish": 0, "proceed": 0, "watermark": 0, "numpi": 0, "arrai": 0, "cv2": 0, "pywt": 0, "leav": 0, "broken": 0, "rewrit": [0, 1], "those": 0, "rerun": 0, "wa": [0, 1], "scale_model_input": 0, "overal": 0, "my_funct": 0, "Then": 0, "your_script": 0, "magic": [0, 1], "treat": 0, "constant": 0, "xla_no_special_scalar": 0, "subject": 0, "softwar": 0, "peculiar": 0, "implement": [0, 1], "depth": 0, "want": [0, 1], "detial": 0, "__version__": 0, "cu121": 0, "100": 0, "t2": 0, "200": 0, "300": 0, "test_train_mp_imagenet": 0, "fake_data": 0, "branch": [0, 1], "rx": 0, "r2": 0, "conclud": 0, "correctli": 0, "diagnos": 0, "counter": 0, "extrem": 0, "bug": 0, "pt_xla_debug": 0, "xla_dynamo_debug": 0, "summari": 0, "compiletim": 0, "too": [0, 1], "frequent": 0, "21": 0, "11": [0, 1], "transferfromdevicetim": 0, "aten": [0, 1], "_ctc_loss": 0, "_ctc_loss_backward": 0, "open": 0, "23": 0, "12": 0, "hash": 0, "537d4b0264b029688281412214d252e9": 0, "588": 0, "320": 0, "frame": 0, "workspac": 0, "dk2": 0, "840": 0, "broadcast_master_param": 0, "1230": 0, "train_imagenet": 0, "261": 0, "365": 0, "_intern": 0, "176": 0, "_thread_fn": 0, "70": 0, "usr": 0, "lib": 0, "concurr": 0, "57": 0, "_worker": 0, "80": 0, "manual": 0, "exit": 0, "steptrac": 0, "region": 0, "decid": [0, 1], "often": [0, 1], "expect": [0, 1], "frequenc": 0, "pair": 0, "stabil": 0, "keep": [0, 1], "met": 0, "short": [0, 1], "short_metrics_report": 0, "metrics_report": 0, "destroi": 0, "term": [0, 1], "percentil": 0, "totalsampl": 0, "202": 0, "06m09s401ms746": 0, "001u": 0, "valuer": 0, "778ms572": 0, "062u": 0, "rate": 0, "425201": 0, "001ms32": 0, "778u": 0, "001ms61": 0, "283u": 0, "001ms79": 0, "236u": 0, "001ms110": 0, "973u": 0, "001ms228": 0, "773u": 0, "001ms339": 0, "183u": 0, "90": 0, "001ms434": 0, "305u": 0, "95": 0, "002ms921": 0, "063u": 0, "99": 0, "21s102ms853": 0, "173u": 0, "track": 0, "statu": 0, "cachedsynctensor": 0, "395": 0, "context": [0, 1], "area": 0, "rout": 0, "engin": 0, "qualifi": 0, "c": [0, 1], "namespac": 0, "33": 0, "than": [0, 1], "_local_scalar_dens": 0, "feel": 0, "clear_al": 0, "resourc": 0, "offici": 0, "colab": 0, "notebook": 0, "behav": 0, "semant": [0, 1], "constraint": 0, "suggest": 0, "certain": [0, 1], "pattern": [0, 1], "bad": 0, "mind": 0, "expens": [0, 1], "speedup": 0, "rest": [0, 1], "direct": 0, "indirect": 0, "mask": [0, 1], "thu": 0, "low": 0, "variat": 0, "don": [0, 1], "nativ": 0, "translat": 0, "slowdown": 0, "ask": [0, 1], "unless": [0, 1], "most": 0, "checkout": 0, "scalar": [0, 1], "substitut": 0, "flow": 0, "applic": 0, "clip_grad": 0, "norm": 0, "problemat": 0, "impact": [0, 1], "patch": 0, "clip_grad_norm_": 0, "dramat": 0, "els": [0, 1], "total_norm": 0, "zero": 0, "p": [0, 1], "param_norm": 0, "grad": 0, "norm_typ": 0, "add_": 0, "clip_coef": 0, "max_norm": 0, "1e": 0, "mul_": 0, "data_parallel": 0, "drop": 0, "last": 0, "stride": 0, "reconstruct": 0, "shallow": 0, "ty": 0, "_xlac": [0, 1], "_get_xla_tensors_text": [0, 1], "_get_xla_tensors_hlo": 0, "prior": 0, "degre": 0, "xla_ir_debug": 0, "_xla_ir": 0, "propag": 0, "metadata": 0, "xla_save_tensors_fil": 0, "realli": [0, 1], "big": [0, 1], "left": 0, "append": 0, "clean": 0, "sheet": 0, "xla_save_tensors_fmt": 0, "store": 0, "_xla_save_tensor": 0, "text": 0, "dot": 0, "graphviz": 0, "xla_flag": 0, "xla_dump_to": 0, "dir_nam": 0, "unoptim": 0, "optimz": 0, "xla_metrics_fil": 0, "xla_save_hlo_fil": 0, "offend": 0, "xla_sync_wait": 0, "forc": [0, 1], "xla_use_eager_debug_mod": 0, "bypass": 0, "higher": 0, "optimizaiton": 0, "skip": 0, "bifloat16": 0, "arithmet": 0, "accur": 0, "accumul": 0, "4096": 0, "dtype": [0, 1], "4097": 0, "averag": 0, "stai": 0, "fp32": 0, "xla_use_f16": 0, "float16": 0, "half": 0, "tf_cpp_log_thread_id": 0, "tf": [0, 1], "tf_cpp_vmodul": 0, "vlog": 0, "form": [0, 1], "tf_cpp_min_log_level": 0, "turn": 0, "tf_vlog": 0, "xla_dump_hlo_graph": 0, "rais": 0, "xla_util": 0, "cc": 0, "save1": 0, "xla_graph_executor": 0, "pjrt_computation_cli": 0, "pr": 0, "dir": 0, "pytorch_test_with_slow": 0, "test_torch": 0, "k": 0, "test_put_xla_uint8": 0, "torch_test_devic": 0, "pytorch_test_bas": 0, "migrat": 0, "jax": 0, "public": 0, "init": 0, "renam": 0, "regist": 0, "xla_backend": 0, "torchrun": 0, "init_method": 0, "plugin": 0, "xpu": 0, "neuron": 0, "continu": 0, "xrt_tpu_config": 0, "libtpu": 0, "30": 0, "thousand": 0, "preview": 0, "On": 0, "safe": 0, "broadcast": 0, "manual_se": 0, "pjrt_backend": 0, "os": 0, "ddp": 0, "dist": 0, "init_process_group": 0, "rank": 0, "world_siz": 0, "42": 0, "128": 0, "gradient_as_bucket_view": 0, "mseloss": 0, "001": 0, "confirm": 0, "localservic": 0, "localhost": 0, "51011": 0, "master_addr": 0, "master_port": 0, "12355": 0, "overhead": 0, "grpc": 0, "torchbench": 0, "observ": 0, "35": 0, "2048": 0, "chip": 0, "learn": 0, "test_train_mp_mnist": 0, "batch_siz": 0, "256": 0, "num_epoch": 0, "tpu_process_bound": 0, "tpu_visible_chip": 0, "r1": 0, "preinstal": 0, "docker_imag": 0, "gcr": 0, "io": 0, "authent": 0, "privat": 0, "gcp": 0, "repositori": 0, "auth": 0, "rm": 0, "privileg": 0, "simpli": 0, "gpu_num_devic": 0, "nnode": 0, "num_gpu_devic": 0, "pjrt_distribut": 0, "physic": 0, "number_gpu_vm": 0, "node_rank": 0, "current_node_rank": 0, "nproc_per_nod": 0, "number_local_gpu_devic": 0, "rdzv_endpoint": 0, "internal_ip_address": 0, "port": 0, "multinode_train": 0, "endpoint": 0, "ip": 0, "address": 0, "omit": 0, "machine_0": 0, "machine_1": 0, "machine_0_internal_ip_address": 0, "ident": 0, "although": [0, 1], "mostli": 0, "interchang": 0, "subtl": 0, "importantli": 0, "around": [0, 1], "architectur": 0, "latenc": 0, "deseri": 0, "gain": 0, "interact": 0, "normal": [0, 1], "plan": 0, "simpler": 0, "xla_dist": 0, "sdk": 0, "reimplement": 0, "collect": 0, "enhanc": 0, "substanti": 0, "vari": [0, 1], "consist": 0, "servic": 0, "connect": 0, "practic": [0, 1], "unreli": 0, "inbound": 0, "entir": 0, "restart": 0, "impos": 0, "unwant": 0, "permit": 0, "subset": 0, "old": 0, "alter": 0, "html": 0, "_": 0, "all_gather_object": 0, "gloo": 0, "backend": [0, 1], "new_group": 0, "subgroup": 0, "experi": 0, "altern": [0, 1], "less": [0, 1], "reliabl": 0, "blob": 0, "md": 0, "strongli": 0, "world": [0, 1], "queri": 0, "_all_gath": 0, "No": [0, 1], "int32": 0, "zeros_lik": 0, "get_world_s": 0, "task": 0, "175": 0, "chart": 0, "breakdown": 0, "tfrt": 0, "legaci": 0, "streamexecutor": 0, "tpu_legaci": 0, "shown": [0, 1], "vs": [0, 1], "comparison": 0, "jit": 0, "unmodifi": 0, "hook": 0, "biggest": 0, "right": [0, 1], "bridg": 0, "torchfx": 0, "recogn": [0, 1], "technolog": 0, "fx": 0, "openxla": 0, "a_xla": 0, "b_xla": 0, "compiled_cod": 0, "resnet18": 0, "torchvis": 0, "eval_model": 0, "xla_resnet18": 0, "eval": 0, "dynamo_resnet18": 0, "no_grad": 0, "resent18": 0, "binari": 0, "invok": 0, "bench": 0, "59": 0, "64": 0, "resnext50_32x4d": 0, "91": 0, "alexnet": 0, "28": 0, "mobilenet_v2": 0, "18": 0, "62": 0, "mnasnet1_0": 0, "68": 0, "vgg16": 0, "bert_pytorch": 0, "49": 0, "squeezenet1_1": 0, "29": 0, "timm_vision_transform": 0, "52": 0, "geomean": 0, "04": 0, "team": 0, "train_model": 0, "crossentropyloss": 0, "pred": 0, "train_model_main": 0, "dynamo_train_model": 0, "xla_optim": 0, "weight_decai": 0, "extract": 0, "07": 0, "43": 0, "19": 0, "81": 0, "87": 0, "41": 0, "fwd": 0, "bwd": 0, "e2": 0, "easili": [0, 1], "hide": 0, "cost": 0, "scenario": 0, "larger": 0, "forward": 0, "best": 0, "ideal": [0, 1], "promis": 0, "complex": 0, "tradit": 0, "hard": [0, 1], "expand": 0, "mention": [0, 1], "super": 0, "excit": 0, "upcom": 0, "heavili": 0, "invest": 0, "upstream": 0, "matur": 0, "stori": 0, "xlafullyshardeddataparallel": 0, "my_modul": 0, "adam": 0, "0001": 0, "sum": 0, "individu": 0, "layer": 0, "outer": 0, "wrapper": 0, "leftov": 0, "arxiv": 0, "ab": 0, "1910": 0, "02054": 0, "reshard_after_forward": 0, "test_train_mp_mnist_fsdp_with_ckpt": 0, "test_train_mp_imagenet_fsdp": 0, "interleav": 0, "submodul": 0, "inner": 0, "fsdpvitmodel": 0, "ronghanghu": 0, "vit_10b_fsdp_exampl": 0, "run_vit_train": 0, "checkpoint_modul": 0, "pull": 0, "3524": 0, "auto_wrap_polici": 0, "size_based_auto_wrap_polici": 0, "polici": 0, "100m": 0, "transformer_auto_wrap_polici": 0, "conv2d": 0, "partial": [0, 1], "transformer_layer_cl": 0, "auto_wrapper_cal": 0, "remateri": 0, "m": [0, 1], "kwarg": 0, "resum": 0, "get_shard_metadata": 0, "consolidate_sharded_model_checkpoint": 0, "stitch": 0, "ckpt": 0, "shard_metadata": 0, "ckpt_path": 0, "f": 0, "pth": 0, "bash": 0, "consolidate_sharded_ckpt": 0, "ckpt_prefix": 0, "your_sharded_checkpoint_fil": 0, "ckpt_suffix": 0, "_rank": 0, "inspir": 0, "fairscal": 0, "fullyshardeddataparallel": 0, "readthedoc": 0, "en": 0, "explicit": 0, "resort": 0, "approach": [0, 1], "newer": 0, "recurs": 0, "98": 0, "accuraci": 0, "drop_last": 0, "use_nested_fsdp": 0, "use_gradient_checkpoint": 0, "final_ckpt": 0, "75": 0, "1k": 0, "datadir": 0, "test_set_batch_s": 0, "eval_interv": 0, "num_warmup_epoch": 0, "lr_scheduler_divide_every_n_epoch": 0, "lr_scheduler_divisor": 0, "residu": 0, "algorithm": 0, "vision": 0, "vit": 0, "against": 0, "abil": [0, 1], "And": [0, 1], "who": 0, "eager": [0, 1], "mode": [0, 1], "know": [0, 1], "nccl": 0, "new_rank": 0, "ddp_model": 0, "final": 0, "launcher": 0, "demo_fn": 0, "everyth": [0, 1], "touch": 0, "plu": 0, "five": 0, "sy": 0, "tempfil": 0, "cleanup": 0, "destroy_process_group": 0, "toymodel": 0, "__init__": 0, "net1": 0, "1000000": 0, "relu": 0, "net2": 0, "demo_bas": 0, "assert": 0, "graident_as_bucket_view": 0, "label": 0, "run_demo": 0, "tot": 0, "statist": 0, "produc": [0, 1], "median": 0, "90th": 0, "std": 0, "cv": 0, "418": 0, "54": 0, "419": 0, "22": 0, "430": 0, "40": 0, "76": 0, "02": 0, "97": 0, "407": 0, "60": 0, "39": 0, "seem": 0, "extra": 0, "17864": 0, "20108": 0, "96": 0, "24351": 0, "74": 0, "5866": 0, "83": 0, "10701": 0, "11770": 0, "00": 0, "14313": 0, "78": 0, "3102": 0, "92": 0, "round": 0, "sens": 0, "amort": 0, "logdir": 0, "converg": 0, "achiev": 0, "48": 0, "caution": 0, "interest": 0, "enforc": [0, 1], "crash": 0, "nvidia": 0, "driver": 0, "publish": 0, "prebuilt": 0, "cuda11": 0, "config": 0, "central1": 0, "pkg": 0, "nightly_3": 0, "8_cuda_12": 0, "toolkit": 0, "datacent": 0, "latest": 0, "curl": 0, "fssl": 0, "libnvidia": 0, "gpgkei": 0, "gpg": 0, "dearmor": 0, "o": 0, "keyr": 0, "l": 0, "deb": 0, "sed": 0, "sign": 0, "tee": 0, "d": [0, 1], "ctk": 0, "systemctl": 0, "shm": 0, "16g": 0, "bin": 0, "exec": 0, "ps": 0, "awk": 0, "nr": 0, "visibl": [0, 1], "smi": 0, "verifi": 0, "20ab2c7a2d06": 0, "dec": 0, "06": 0, "24": 0, "2022": 0, "510": 0, "47": 0, "03": 0, "bu": 0, "disp": 0, "volatil": 0, "uncorr": 0, "ecc": 0, "fan": 0, "temp": 0, "perf": [0, 1], "pwr": 0, "cap": 0, "mig": 0, "tesla": 0, "v100": 0, "sxm2": 0, "00000000": 0, "36c": 0, "p0": 0, "38w": 0, "300w": 0, "0mib": 0, "16384mib": 0, "gi": 0, "pid": 0, "ld_library_path": 0, "echo": 0, "link": 0, "bashrc": 0, "lib64": 0, "cp310": 0, "manylinux_2_28_x86_64": 0, "prepar": 0, "begin": 0, "38": 0, "89059": 0, "82": 0, "globalr": 0, "79297": 0, "117": 0, "45": 0, "84": 0, "36": 0, "43628": 0, "281": 0, "83108": 0, "346": 0, "88": 0, "108": 0, "99023": 0, "373": 0, "132": 0, "56": 0, "92699": 0, "384": 0, "152": 0, "14": 0, "120": 0, "68816": 0, "388": 0, "169": 0, "09": 0, "use_cuda": 0, "xla_cuda": 0, "successfulli": 0, "gspmd": 0, "annot": 0, "ml": 0, "proper": 0, "hint": 0, "figur": 0, "strategi": 0, "non": [0, 1], "suit": 0, "framework": [0, 1], "popular": 0, "shift": 0, "unlock": 0, "advanc": 0, "use_spmd": 0, "b8b484515a97f74e013dcf38125c44d53a41f011": 0, "l214": 0, "is_spmd": 0, "mechan": 0, "mark_shard": 0, "src": 0, "union": 0, "partition_spec": 0, "logic": [0, 1], "xlatensor": 0, "rfc": 0, "np": 0, "xs": 0, "mesh_shap": 0, "num_devic": 0, "global_runtime_device_count": 0, "device_id": 0, "m1_shard": 0, "isinst": 0, "techniqu": 0, "comment": [0, 1], "simplelinear": 0, "fc1": 0, "enumer": 0, "sharidng": 0, "spatial": 0, "ouput": 0, "cluster": 0, "interconnect": 0, "deriv": 0, "sub": 0, "abstract": 0, "axis_nam": 0, "nad": 0, "get_logical_mesh": 0, "ordereddict": 0, "tile": 0, "intend": 0, "manipul": 0, "nice": 0, "arrang": 0, "involv": [0, 1], "center": 0, "dcn": 0, "hybridmesh": 0, "box": 0, "multislic": 0, "accept": 0, "ici_mesh_shap": 0, "dcn_mesh_shap": 0, "denot": 0, "device_mesh": 0, "axi": 0, "input_tensor": 0, "evenli": 0, "equival": 0, "3d": 0, "2d": [0, 1], "two_d_partially_repl": 0, "two_d_fully_shard": 0, "4d": 0, "four_d": 0, "delai": 0, "carri": [0, 1], "lazili": [0, 1], "satisfi": 0, "subclass": 0, "__torch_dispatch__": 0, "spmdpartition": 0, "global_tensor": 0, "strictli": 0, "local_shard": 0, "xlashard": 0, "4e8e5511555073ce8b6d1a436bf808c9333dcac6": 0, "xla_sharded_tensor": 0, "l12": 0, "ongo": 0, "effort": [0, 1], "distributedtensor": 0, "prototyp": 0, "proof": 0, "concept": 0, "distribute_tensor": 0, "devicemesh": 0, "big_tensor": 0, "100000": 0, "my_dtensor": 0, "tune": 0, "_input_sharding_": 0, "field": 0, "input_mesh": 0, "dedic": 0, "planner": 0, "spmdsaveplann": 0, "spmdloadplann": 0, "dist_cp": 0, "distributed_checkpoint": 0, "xc": 0, "storage_writ": 0, "filesystemwrit": 0, "checkpoint_dir": 0, "storage_read": 0, "filesystemread": 0, "manag": 0, "all_step": 0, "save_async": 0, "unblock": 0, "dispatch": 0, "preemption": 0, "detect": 0, "termin": 0, "provis": 0, "queuedresourc": 0, "autocheckpoint": 0, "chkpt_on_preempt": 0, "fsspec": 0, "filesystem": 0, "gc": 0, "chkpt_mgr": 0, "gs": 0, "my": [0, 1], "bucket": [0, 1], "tracked_step": 0, "choos": 0, "highest": 0, "best_step": 0, "discov": 0, "_after": 0, "_the": 0, "unnecessari": 0, "forth": 0, "decis": 0, "hardcod": 0, "leverag": 0, "num_gpu_machin": 0, "rank_of_current_machin": 0, "machine_0_ip_address": 0, "training_or_inference_script_using_spmd": 0, "xla_use_spmd": 0, "test_train_spmd_imagenet": 0, "0th": 0, "w": 0, "h": 0, "highli": 0, "param": 0, "named_paramet": 0, "quick": 0, "plai": 0, "512": 0, "throughput": 0, "am": 0, "roughli": 0, "4x": 0, "placement": 0, "visualize_shard": 0, "rich": 0, "generated_t": 0, "use_color": 0, "alt_text": 0, "asset": 0, "spmd_debug_1": 0, "png": 0, "style": 0, "partial_repl": 0, "envvar": 0, "xla_spmd_auto": 0, "_tensor": 0, "auto_polici": 0, "device_count": 0, "mymodul": 0, "sharded_model": 0, "behvaior": 0, "xla_auto_use_group_shard": 0, "reshard": 0, "xla_auto_spmd_mesh": 0, "unset": 0, "fsdpv2": 0, "epxress": 0, "famou": 0, "aim": [0, 1], "offer": 0, "enjoi": 0, "bring": 0, "tabl": 0, "review": 0, "spmd_fully_sharded_data_parallel": 0, "spmdfullyshardeddataparallel": 0, "autowrap": 0, "shard_output": 0, "categori": 0, "element": [0, 1], "logit": 0, "children": 0, "infinit": 0, "compil": 1, "xla": 1, "ar": 1, "static": 1, "word": 1, "chang": 1, "hurt": 1, "understand": 1, "debug": 1, "user": 1, "pov": 1, "sai": 1, "support": 1, "assur": 1, "gone": 1, "coverag": 1, "detail": 1, "explan": 1, "common": 1, "do": 1, "get": 1, "rid": 1, "mainli": 1, "problem": 1, "beginn": 1, "propos": 1, "here": 1, "reli": 1, "impract": 1, "assumpt": 1, "ye": 1, "exampl": 1, "differ": 1, "sentenc": 1, "length": 1, "ll": 1, "new": 1, "more": 1, "kinda": 1, "anti": 1, "pytorch": 1, "frontend": 1, "lazi": 1, "matter": 1, "cpu": 1, "workaround": 1, "okai": 1, "teach": 1, "data": 1, "torch": 1, "theoret": 1, "memori": 1, "sort": 1, "devic": 1, "node": 1, "obviou": 1, "s64": 1, "8": 1, "num_output": 1, "10": 1, "mul": 1, "inde": 1, "_get_xla_tensor_dimension_s": 1, "commonli": 1, "basic": 1, "correct": 1, "wrong": 1, "wors": 1, "take": 1, "probabl": 1, "how": 1, "upper": 1, "nit": 1, "api": 1, "rand": 1, "solv": 1, "should": 1, "kept": 1, "earli": 1, "accessor": 1, "implicitli": 1, "current": 1, "overload": 1, "IS": 1, "A": 1, "explod": 1, "cheap": 1, "ve": 1, "variant": 1, "hoc": 1, "think": 1, "verison": 1, "x": 1, "bla": 1, "blabla": 1, "execut": 1, "proce": 1, "choic": 1, "multipl": 1, "combin": 1, "wide": 1, "adopt": 1, "code": 1, "uglier": 1, "win": 1, "pars": 1, "automat": 1, "torchscript": 1, "somehow": 1, "merg": 1, "properli": 1, "haven": 1, "thought": 1, "trivial": 1, "side": 1, "That": 1, "why": 1, "hit": 1, "bandwidth": 1, "automag": 1, "gold": 1, "smart": 1, "trick": 1, "tbh": 1, "longer": 1, "y": 1, "The": 1, "unawar": 1, "hope": 1, "smash": 1, "an": 1, "blocker": 1, "ahead": 1, "nnc": 1, "symbol": 1, "optim": 1, "concret": 1, "kernel": 1, "exactli": 1, "transpos": 1, "With": 1, "brian": 1, "hirsh": 1, "bdhirsh": 1, "question": 1, "move": 1, "Is": 1, "stick": 1, "torch_warn": 1, "yea": 1, "hei": 1, "won": 1, "run": 1, "blaze": 1, "fast": 1, "isn": 1, "model": 1, "devirtu": 1, "sound": 1, "great": 1, "truth": 1, "irvalu": 1, "discrep": 1, "followup": 1, "1000": 1, "singl": 1, "presenc": 1, "get_dimention_s": 1, "didn": 1, "condit": 1, "exponenti": 1, "blowup": 1, "fewer": 1, "opportun": 1, "feasibl": 1, "annoi": 1, "variabl": 1, "wasn": 1, "materiz": 1, "combo": 1}, "objects": {"torch_xla.core": [[0, 0, 0, "-", "functions"], [0, 0, 0, "-", "xla_model"]], "torch_xla.core.functions": [[0, 1, 1, "", "all_gather"], [0, 1, 1, "", "all_reduce"]], "torch_xla.core.xla_model": [[0, 1, 1, "", "add_step_closure"], [0, 1, 1, "", "all_gather"], [0, 1, 1, "", "all_reduce"], [0, 1, 1, "", "all_to_all"], [0, 1, 1, "", "do_on_ordinals"], [0, 1, 1, "", "get_local_ordinal"], [0, 1, 1, "", "get_memory_info"], [0, 1, 1, "", "get_ordinal"], [0, 1, 1, "", "get_rng_state"], [0, 1, 1, "", "get_stablehlo"], [0, 1, 1, "", "get_stablehlo_bytecode"], [0, 1, 1, "", "get_xla_supported_devices"], [0, 1, 1, "", "is_master_ordinal"], [0, 1, 1, "", "mesh_reduce"], [0, 1, 1, "", "optimizer_step"], [0, 1, 1, "", "rendezvous"], [0, 1, 1, "", "save"], [0, 1, 1, "", "set_rng_state"], [0, 1, 1, "", "wait_device_ops"], [0, 1, 1, "", "xla_device"], [0, 1, 1, "", "xla_device_hw"], [0, 1, 1, "", "xrt_world_size"]], "torch_xla.distributed": [[0, 0, 0, "-", "parallel_loader"], [0, 0, 0, "-", "xla_multiprocessing"]], "torch_xla.distributed.parallel_loader": [[0, 2, 1, "", "ParallelLoader"]], "torch_xla.distributed.parallel_loader.ParallelLoader": [[0, 3, 1, "", "per_device_loader"]], "torch_xla.distributed.xla_multiprocessing": [[0, 2, 1, "", "MpModelWrapper"], [0, 2, 1, "", "MpSerialExecutor"], [0, 1, 1, "", "spawn"]], "torch_xla.distributed.xla_multiprocessing.MpModelWrapper": [[0, 3, 1, "", "to"]], "torch_xla.distributed.xla_multiprocessing.MpSerialExecutor": [[0, 3, 1, "", "run"]], "torch_xla.utils": [[0, 0, 0, "-", "serialization"], [0, 0, 0, "-", "utils"]], "torch_xla.utils.serialization": [[0, 1, 1, "", "load"], [0, 1, 1, "", "save"]], "torch_xla.utils.utils": [[0, 2, 1, "", "DataWrapper"], [0, 2, 1, "", "SampleGenerator"]]}, "objtypes": {"0": "py:module", "1": "py:function", "2": "py:class", "3": "py:method"}, "objnames": {"0": ["py", "module", "Python module"], "1": ["py", "function", "Python function"], "2": ["py", "class", "Python class"], "3": ["py", "method", "Python method"]}, "titleterms": {"pytorch": 0, "xla": 0, "devic": 0, "creat": 0, "an": 0, "tensor": [0, 1], "ar": 0, "run": 0, "model": 0, "singl": 0, "multipl": 0, "multi": 0, "process": 0, "tpu": 0, "pod": 0, "deep": 0, "dive": 0, "lazi": 0, "bfloat16": 0, "memori": 0, "layout": 0, "move": 0, "from": [0, 1], "cpu": 0, "save": 0, "load": 0, "compil": 0, "cach": 0, "further": 0, "read": 0, "api": 0, "xla_model": 0, "distribut": 0, "util": 0, "test": 0, "beginn": 0, "s": [0, 1], "guid": 0, "basic": 0, "high": 0, "level": 0, "understand": 0, "some": [0, 1], "detail": 0, "setup": 0, "convert": 0, "code": 0, "exampl": 0, "1": [0, 1], "stabl": 0, "diffus": 0, "infer": 0, "lightn": 0, "2": [0, 1], "hf": 0, "profil": 0, "perform": 0, "analysi": 0, "troubleshoot": 0, "saniti": 0, "check": 0, "version": 0, "A": 0, "simpl": 0, "calcul": 0, "resnet": 0, "With": 0, "fake": 0, "data": 0, "debug": 0, "tool": 0, "dynamo": 0, "auto": 0, "metric": 0, "execut": 0, "get": 0, "report": 0, "The": 0, "clear": 0, "known": 0, "caveat": 0, "quirk": 0, "more": 0, "environ": 0, "variabl": 0, "common": 0, "combin": 0, "reproduc": 0, "ci": 0, "cd": 0, "unit": 0, "failur": 0, "pjrt": 0, "runtim": 0, "tl": 0, "dr": 0, "benefit": 0, "quickstart": 0, "docker": 0, "gpu": 0, "node": 0, "train": 0, "differ": 0, "xrt": 0, "multithread": 0, "v2": 0, "v3": 0, "chang": 0, "xm": 0, "rendezv": 0, "torch": 0, "new": 0, "torchdynamo": 0, "integr": 0, "featur": 0, "gap": 0, "take": 0, "awai": 0, "fulli": 0, "shard": 0, "parallel": 0, "fsdp": 0, "script": 0, "mnist": 0, "imagenet": 0, "instal": 0, "clone": 0, "repo": 0, "8": 0, "50": 0, "10": 0, "billion": 0, "paramet": 0, "how": 0, "do": 0, "distributeddataparallel": 0, "background": 0, "motiv": 0, "us": [0, 1], "benchmark": 0, "resnet50": 0, "real": [0, 1], "disclaim": 0, "instanc": 0, "wheel": 0, "amp": 0, "automat": 0, "mix": 0, "precis": 0, "develop": 0, "build": 0, "sourc": [0, 1], "support": 0, "spmd": 0, "user": 0, "what": [0, 1], "design": 0, "overview": 0, "aannot": 0, "mesh": 0, "hybrid": 0, "partit": 0, "spec": 0, "xlashardedtensor": 0, "dtensor": 0, "awar": 0, "host": 0, "checkpoint": 0, "checkpointmanag": 0, "group": 0, "virtual": 0, "optim": 0, "number": 0, "refer": 0, "express": 0, "here": 0, "2x2": 0, "ax": 0, "x": 0, "y": 0, "can": [0, 1], "visual": 0, "visualize_tensor_shard": 0, "method": 0, "current": 0, "should": 0, "via": 0, "distribute_modul": 0, "output": [0, 1], "gradient": 0, "huggingfac": 0, "llama": 0, "recompil": 1, "torch_xla": 1, "let": 1, "first": 1, "start": 1, "fact": 1, "constraint": 1, "input": 1, "dataset": 1, "oper": 1, "bound": 1, "dynam": 1, "shape": 1, "fix": 1, "case": 1, "when": 1, "you": 1, "without": 1, "queri": 1, "its": 1, "dimens": 1, "3": 1, "control": 1, "flow": 1, "conclus": 1, "appendix": 1}, "envversion": {"sphinx.domains.c": 2, "sphinx.domains.changeset": 1, "sphinx.domains.citation": 1, "sphinx.domains.cpp": 6, "sphinx.domains.index": 1, "sphinx.domains.javascript": 2, "sphinx.domains.math": 2, "sphinx.domains.python": 3, "sphinx.domains.rst": 2, "sphinx.domains.std": 2, "sphinx.ext.intersphinx": 1, "sphinx.ext.todo": 2, "sphinx.ext.viewcode": 1, "sphinx": 56}})
\ No newline at end of file
+Search.setIndex({"docnames": ["index", "notes/source_of_recompilation"], "filenames": ["index.rst", "notes/source_of_recompilation.md"], "titles": ["PyTorch on XLA Devices", "Source of recompilations in torch_xla"], "terms": {"like": [0, 1], "torch_xla": 0, "packag": 0, "thi": [0, 1], "document": 0, "describ": 0, "your": [0, 1], "add": [0, 1], "type": 0, "work": [0, 1], "just": [0, 1], "other": [0, 1], "For": [0, 1], "print": [0, 1], "import": 0, "core": 0, "t": [0, 1], "randn": 0, "xla_devic": 0, "look": 0, "familiar": 0, "same": [0, 1], "interfac": 0, "regular": 0, "few": [0, 1], "addit": 0, "initi": 0, "return": [0, 1], "mai": [0, 1], "depend": [0, 1], "oper": 0, "cuda": [0, 1], "ad": [0, 1], "togeth": 0, "t0": 0, "t1": 0, "Or": [0, 1], "matrix": 0, "multipli": 0, "mm": 0, "neural": 0, "network": 0, "modul": 0, "l_in": 0, "linear": 0, "nn": 0, "20": 0, "l_out": 0, "onli": [0, 1], "so": [0, 1], "input": 0, "floattensor": 0, "throw": 0, "error": 0, "sinc": [0, 1], "exist": 0, "one": [0, 1], "requir": [0, 1], "line": [0, 1], "specif": 0, "follow": [0, 1], "snippet": 0, "highlight": 0, "when": 0, "show": 0, "loss_fn": 0, "nllloss": 0, "sgd": 0, "lr": 0, "momentum": 0, "target": [0, 1], "train_load": 0, "zero_grad": 0, "loss": 0, "backward": 0, "step": [0, 1], "mark_step": 0, "easi": [0, 1], "switch": [0, 1], "definit": [0, 1], "dataload": 0, "loop": [0, 1], "ani": [0, 1], "coupl": 0, "acquir": 0, "mark": 0, "call": [0, 1], "end": 0, "each": [0, 1], "iter": 0, "caus": [0, 1], "its": 0, "graph": [0, 1], "updat": [0, 1], "see": [0, 1], "make": [0, 1], "acceler": 0, "parallel_load": 0, "pl": 0, "xla_multiprocess": 0, "xmp": 0, "def": 0, "_mp_fn": 0, "index": 0, "mp_device_load": 0, "mpdeviceload": 0, "optimizer_step": 0, "__name__": 0, "__main__": 0, "spawn": 0, "arg": 0, "There": [0, 1], "three": 0, "between": [0, 1], "previou": [0, 1], "let": 0, "go": 0, "over": 0, "abl": [0, 1], "access": [0, 1], "assign": 0, "v4": 0, "4": [0, 1], "being": 0, "up": [0, 1], "own": 0, "note": [0, 1], "you": 0, "0": [0, 1], "all": [0, 1], "becaus": 0, "doe": [0, 1], "mean": [0, 1], "function": 0, "have": [0, 1], "thread": 0, "doc": [0, 1], "onto": 0, "wrap": 0, "It": [0, 1], "preload": 0, "overlap": 0, "improv": 0, "also": [0, 1], "everi": [0, 1], "batches_per_execut": 0, "default": 0, "batch": 0, "yield": 0, "consolid": 0, "issu": 0, "comput": [0, 1], "pretti": [0, 1], "much": [0, 1], "all_reduce_gradi": 0, "reduc": 0, "remain": [0, 1], "start": 0, "retriev": [0, 1], "within": 0, "which": [0, 1], "ha": [0, 1], "parent": 0, "stack": [0, 1], "full": 0, "multiprocess": 0, "veri": [0, 1], "talk": 0, "about": [0, 1], "independ": 0, "bit": 0, "avail": [0, 1], "13": 0, "releas": 0, "befor": [0, 1], "pleas": 0, "our": [0, 1], "explain": [0, 1], "googl": 0, "cloud": 0, "basi": 0, "gcloud": 0, "command": 0, "project": 0, "howto": 0, "focu": [0, 1], "perspect": 0, "assum": [0, 1], "abov": [0, 1], "section": 0, "train_mnist_xla": 0, "py": 0, "If": [0, 1], "would": [0, 1], "ssh": 0, "tpuvm": 0, "pjrt_devic": 0, "python3": 0, "now": [0, 1], "order": 0, "16": 0, "need": [0, 1], "sure": 0, "usual": 0, "done": [0, 1], "scp": 0, "copi": 0, "time": [0, 1], "alpha": 0, "vm": 0, "zone": 0, "worker": 0, "outsid": 0, "underli": 0, "infrastructur": 0, "global": 0, "topologi": 0, "local": 0, "ordin": 0, "cross": 0, "commun": 0, "happen": [0, 1], "across": 0, "instead": [0, 1], "regard": 0, "inform": [0, 1], "complet": 0, "fakedata": 0, "But": [0, 1], "even": [0, 1], "though": 0, "act": 0, "lot": [0, 1], "intern": [0, 1], "uniqu": [0, 1], "launch": 0, "immedi": 0, "eagerli": [0, 1], "hand": 0, "thei": [0, 1], "record": 0, "until": 0, "result": 0, "defer": 0, "separ": 0, "might": [0, 1], "fuse": 0, "gener": [0, 1], "invis": 0, "caller": 0, "construct": 0, "send": 0, "them": [0, 1], "synchron": 0, "insert": 0, "barrier": 0, "explicitli": [0, 1], "paper": 0, "datatyp": 0, "In": [0, 1], "fact": 0, "handl": [0, 1], "float": [0, 1], "doubl": [0, 1], "behavior": 0, "control": 0, "xla_use_bf16": 0, "xla_downcast_bf16": 0, "By": [0, 1], "both": [0, 1], "set": [0, 1], "float32": 0, "directli": [0, 1], "map": 0, "bf16": 0, "primit": 0, "alwai": [0, 1], "regardless": 0, "actual": [0, 1], "re": [0, 1], "convers": [0, 1], "opaqu": 0, "back": 0, "trigger": [0, 1], "represent": 0, "expos": 0, "storag": 0, "appear": 0, "contigu": 0, "unlik": 0, "allow": 0, "adjust": 0, "better": [0, 1], "view": 0, "relationship": 0, "preserv": 0, "put": 0, "anoth": [0, 1], "wai": [0, 1], "onc": [0, 1], "again": 0, "appreci": 0, "accommod": 0, "transit": 0, "pt": 0, "were": [0, 1], "per": 0, "care": [0, 1], "must": 0, "taken": [0, 1], "recommend": 0, "recreat": 0, "after": [0, 1], "been": [0, 1], "destin": 0, "provid": [0, 1], "previous": 0, "state_dict": 0, "path": [0, 1], "case": 0, "master": 0, "where": [0, 1], "limit": 0, "compar": 0, "size": [0, 1], "footprint": 0, "serial": 0, "xser": 0, "stream": 0, "amount": [0, 1], "match": 0, "restor": 0, "load_state_dict": 0, "possibl": 0, "unavail": 0, "fail": 0, "under": 0, "activ": 0, "futur": [0, 1], "trace": [0, 1], "hlo": 0, "consum": [0, 1], "doesn": [0, 1], "persist": 0, "disk": 0, "reus": 0, "significantli": 0, "recompil": 0, "still": [0, 1], "occur": 0, "experiment": 0, "opt": 0, "through": [0, 1], "initialize_cach": 0, "xr": 0, "your_cache_path": 0, "readonli": 0, "fals": 0, "specifi": 0, "whether": 0, "write": 0, "share": 0, "mount": 0, "workload": 0, "n": 0, "none": 0, "devkind": 0, "given": [0, 1], "python": [0, 1], "int": [0, 1], "option": 0, "otherwis": [0, 1], "first": 0, "string": 0, "custom": [0, 1], "deprec": 0, "request": [0, 1], "get_xla_supported_devic": 0, "max_devic": 0, "list": 0, "kind": 0, "name": [0, 1], "maximum": 0, "xla_device_hw": 0, "hardwar": 0, "get_ordin": 0, "defval": 0, "replic": 0, "rang": 0, "xrt_world_siz": 0, "minu": 0, "valu": [0, 1], "ignor": 0, "get_local_ordin": 0, "is_master_ordin": 0, "true": [0, 1], "bool": 0, "while": [0, 1], "num_host": 0, "boolean": 0, "indic": [0, 1], "part": 0, "all_reduc": 0, "reduce_typ": 0, "scale": 0, "pin_layout": 0, "inplac": 0, "One": 0, "reduce_sum": 0, "reduce_mul": 0, "reduce_and": 0, "reduce_or": 0, "reduce_min": 0, "reduce_max": 0, "either": [0, 1], "op": [0, 1], "appli": 0, "repres": [0, 1], "replica": 0, "3": 0, "5": [0, 1], "6": [0, 1], "7": 0, "defin": 0, "two": [0, 1], "pin": 0, "pine": 0, "prevent": 0, "potenti": 0, "corrupt": 0, "particip": 0, "slightli": 0, "program": [0, 1], "unpin": 0, "messag": 0, "hlomodul": 0, "constrain": 0, "pass": 0, "hold": 0, "tupl": [0, 1], "itself": 0, "all_gath": 0, "dim": 0, "gather": 0, "along": 0, "dimens": 0, "all_to_al": 0, "split_dimens": 0, "concat_dimens": 0, "split_count": 0, "alltoal": 0, "http": 0, "www": 0, "tensorflow": [0, 1], "org": 0, "operation_semant": 0, "upon": 0, "split": 0, "concat": 0, "count": 0, "add_step_closur": 0, "closur": 0, "run_async": 0, "ones": [0, 1], "mani": [0, 1], "dure": 0, "consol": 0, "post": 0, "tensorboard": 0, "etc": [0, 1], "content": 0, "intermediari": 0, "inspect": 0, "point": [0, 1], "typic": 0, "ensur": [0, 1], "live": [0, 1], "alreadi": [0, 1], "materi": [0, 1], "includ": [0, 1], "captur": 0, "argument": 0, "queu": 0, "sequenti": 0, "advis": 0, "throttl": 0, "event": 0, "callabl": 0, "asynchron": 0, "wait_device_op": 0, "wait": 0, "async": 0, "whose": 0, "empti": 0, "optimizer_arg": 0, "parallelload": 0, "dataparallel": 0, "necessari": 0, "loader": 0, "next": [0, 1], "dict": 0, "dictionari": 0, "file_or_path": 0, "master_onli": 0, "global_mast": 0, "file": 0, "transfer": 0, "nest": 0, "object": 0, "overrid": 0, "locat": 0, "flag": 0, "sync": 0, "main": 0, "hang": 0, "tag": 0, "payload": 0, "b": [0, 1], "client": 0, "reach": 0, "server": 0, "effect": 0, "alia": 0, "xla_rendezv": 0, "join": 0, "byte": 0, "sent": 0, "exchang": 0, "posit": 0, "do_on_ordin": 0, "contain": [0, 1], "interpret": [0, 1], "ran": 0, "mesh_reduc": 0, "reduce_fn": 0, "out": 0, "reduct": 0, "receiv": 0, "come": [0, 1], "set_rng_stat": 0, "seed": 0, "random": 0, "state": 0, "integ": 0, "rng": 0, "miss": 0, "get_rng_stat": 0, "get_memory_info": 0, "kb_free": 0, "free": 0, "kb": 0, "kb_total": 0, "total": [0, 1], "kei": 0, "get_stablehlo": 0, "str": 0, "stablehlo": 0, "format": 0, "dump": 0, "whole": [0, 1], "todo": 0, "lsy323": 0, "intermedi": 0, "investig": 0, "straightforward": 0, "identifi": 0, "To": [0, 1], "enabl": 0, "info": [0, 1], "env": 0, "var": 0, "xla_hlo_debug": 0, "root": [0, 1], "get_stablehlo_bytecod": 0, "bytecod": 0, "autograd": 0, "differenti": 0, "select": 0, "class": 0, "batchdim": 0, "loader_prefetch_s": 0, "device_prefetch_s": 0, "host_to_device_transfer_thread": 0, "input_shard": 0, "upload": 0, "th": 0, "sampl": 0, "len": 0, "max": [0, 1], "capac": 0, "queue": 0, "deposit": 0, "shardingspec": 0, "compat": 0, "per_device_load": 0, "structur": 0, "resid": 0, "fn": 0, "nproc": 0, "daemon": 0, "start_method": 0, "base": [0, 1], "At": 0, "moment": 0, "block": 0, "creation": 0, "mpmodelwrapp": 0, "minim": 0, "usag": [0, 1], "fork": 0, "henc": 0, "scope": 0, "insid": 0, "wrapped_model": 0, "mynetwork": 0, "advantag": 0, "page": 0, "origin": 0, "weight": 0, "second": 0, "lower": [0, 1], "system": 0, "mpserialexecutor": 0, "fashion": 0, "among": 0, "serial_exec": 0, "load_dataset": 0, "maybe_download_and_load": 0, "avoid": 0, "download": 0, "executor": 0, "dataset": 0, "lambda": 0, "tmp": 0, "wrt": 0, "samplegener": 0, "sample_count": 0, "place": 0, "synthet": 0, "datawrapp": 0, "illustr": 0, "e": [0, 1], "g": [0, 1], "solut": [0, 1], "howev": 0, "serv": 0, "brief": 0, "help": [0, 1], "reader": 0, "modif": 0, "supplement": 0, "fetch": 0, "ir": [0, 1], "encount": 0, "discuss": 0, "below": [0, 1], "lazytensor": 0, "machin": 0, "readabl": 0, "opcod": 0, "effici": 0, "fed": 0, "later": 0, "except": 0, "four": 0, "slice": 0, "attach": 0, "shape": 0, "evalu": [0, 1], "sometim": [0, 1], "lead": 0, "signific": 0, "log": 0, "callback": 0, "slower": 0, "xla_tensor_z": 0, "respons": 0, "These": 0, "cut": [0, 1], "transferfromdevic": 0, "special": 0, "instruct": 0, "tell": [0, 1], "break": [0, 1], "properti": [0, 1], "suppos": 0, "we": [0, 1], "tensors_on_devic": 0, "z": [0, 1], "without": 0, "subgraph": [0, 1], "correspond": 0, "introduc": 0, "smaller": [0, 1], "signal": 0, "far": 0, "submit": 0, "seen": 0, "small": [0, 1], "bottleneck": 0, "dynam": 0, "suitabl": 0, "pad": [0, 1], "fix": 0, "spent": 0, "trade": [0, 1], "off": 0, "consid": 0, "larg": [0, 1], "spend": 0, "fusion": 0, "long": [0, 1], "faster": [0, 1], "due": 0, "worth": [0, 1], "As": [0, 1], "could": [0, 1], "latter": 0, "job": 0, "imag": [0, 1], "nightli": 0, "runtime_vers": 0, "export": 0, "u": [0, 1], "central2": 0, "project_id": 0, "id": 0, "accelerator_typ": 0, "32": 0, "tpu_nam": 0, "your_tpu_nam": 0, "subnetwork": 0, "tpusubnet": 0, "similar": 0, "pip3": 0, "googleapi": 0, "com": 0, "cp38": 0, "linux_x86_64": 0, "whl": 0, "librari": 0, "sudo": 0, "apt": 0, "libopenbla": 0, "dev": 0, "libgl1": 0, "guidelin": 0, "modifi": 0, "replac": 0, "remov": 0, "progress": 0, "bar": 0, "rememb": 0, "someth": 0, "txt2img": 0, "prompt": 0, "photograph": 0, "astronaut": 0, "ride": 0, "hors": 0, "diff": 0, "found": 0, "relat": 0, "written": 0, "precision_scop": 0, "addition": 0, "particular": 0, "configur": 0, "frozenclipembedd": 0, "therefor": 0, "well": [0, 1], "simplic": [0, 1], "tutori": 0, "self": 0, "ddim": 0, "schedul": 0, "top": 0, "attr": 0, "statement": [0, 1], "disabl": 0, "stop": 0, "fall": 0, "especi": 0, "slow": 0, "tri": 0, "huge": [0, 1], "difficult": 0, "degrad": 0, "easier": [0, 1], "readi": 0, "reason": 0, "cover": 0, "sd": 0, "xl": 0, "git": 0, "github": 0, "text_to_imag": 0, "inference_tpu_single_devic": 0, "speed": [0, 1], "inference_tpu_multidevic": 0, "warn": 0, "watch": 0, "made": [0, 1], "lora": 0, "model_id": 0, "stabilityai": 0, "9": [0, 1], "pipelin": 0, "dpmsolvermultistepschedul": 0, "pip": 0, "r": 0, "txt": 0, "invisible_watermark": 0, "transform": 0, "safetensor": 0, "agre": 0, "licens": 0, "card": 0, "account": 0, "token": 0, "cli": 0, "login": 0, "_your_copied_token__": 0, "readm": 0, "pipe": 0, "hour": 0, "wherea": 0, "likewis": 0, "gpt": 0, "15": 0, "min": 0, "epoch": 0, "becom": [0, 1], "subsequ": 0, "mayb": 0, "try": [0, 1], "notic": 0, "piec": 0, "__call__": 0, "commit": 0, "nonzero": [0, 1], "item": 0, "desir": 0, "down": 0, "rule": 0, "thumb": 0, "fit": 0, "durat": 0, "enough": [0, 1], "least": 0, "good": [0, 1], "constantli": 0, "idl": 0, "inference_tpu_": 0, "capture_profil": 0, "give": 0, "xp": 0, "measur": 0, "net": 0, "analyz": 0, "cannot": [0, 1], "portion": 0, "busi": 0, "middl": [0, 1], "scroll": 0, "find": 0, "occupi": 0, "demonstr": 0, "displai": 0, "largest": 0, "appropri": 0, "zoom": 0, "timelin": 0, "period": 0, "examin": 0, "thing": 0, "did": 0, "pipe_watermark": 0, "closer": 0, "preced": 0, "finish": 0, "proceed": 0, "watermark": 0, "numpi": 0, "arrai": 0, "cv2": 0, "pywt": 0, "leav": 0, "broken": 0, "rewrit": [0, 1], "those": 0, "rerun": 0, "wa": [0, 1], "scale_model_input": 0, "overal": 0, "my_funct": 0, "Then": 0, "your_script": 0, "magic": [0, 1], "treat": 0, "constant": 0, "xla_no_special_scalar": 0, "subject": 0, "softwar": 0, "peculiar": 0, "implement": [0, 1], "depth": 0, "want": [0, 1], "detial": 0, "__version__": 0, "cu121": 0, "100": 0, "t2": 0, "200": 0, "300": 0, "test_train_mp_imagenet": 0, "fake_data": 0, "branch": [0, 1], "rx": 0, "r2": 0, "conclud": 0, "correctli": 0, "diagnos": 0, "counter": 0, "extrem": 0, "bug": 0, "pt_xla_debug": 0, "xla_dynamo_debug": 0, "summari": 0, "compiletim": 0, "too": [0, 1], "frequent": 0, "21": 0, "11": [0, 1], "transferfromdevicetim": 0, "aten": [0, 1], "_ctc_loss": 0, "_ctc_loss_backward": 0, "open": 0, "23": 0, "12": 0, "hash": 0, "537d4b0264b029688281412214d252e9": 0, "588": 0, "320": 0, "frame": 0, "workspac": 0, "dk2": 0, "840": 0, "broadcast_master_param": 0, "1230": 0, "train_imagenet": 0, "261": 0, "365": 0, "_intern": 0, "176": 0, "_thread_fn": 0, "70": 0, "usr": 0, "lib": 0, "concurr": 0, "57": 0, "_worker": 0, "80": 0, "manual": 0, "exit": 0, "steptrac": 0, "region": 0, "decid": [0, 1], "often": [0, 1], "expect": [0, 1], "frequenc": 0, "pair": 0, "stabil": 0, "keep": [0, 1], "met": 0, "short": [0, 1], "short_metrics_report": 0, "metrics_report": 0, "destroi": 0, "term": [0, 1], "percentil": 0, "totalsampl": 0, "202": 0, "06m09s401ms746": 0, "001u": 0, "valuer": 0, "778ms572": 0, "062u": 0, "rate": 0, "425201": 0, "001ms32": 0, "778u": 0, "001ms61": 0, "283u": 0, "001ms79": 0, "236u": 0, "001ms110": 0, "973u": 0, "001ms228": 0, "773u": 0, "001ms339": 0, "183u": 0, "90": 0, "001ms434": 0, "305u": 0, "95": 0, "002ms921": 0, "063u": 0, "99": 0, "21s102ms853": 0, "173u": 0, "track": 0, "statu": 0, "cachedsynctensor": 0, "395": 0, "context": [0, 1], "area": 0, "rout": 0, "engin": 0, "qualifi": 0, "c": [0, 1], "namespac": 0, "33": 0, "than": [0, 1], "_local_scalar_dens": 0, "feel": 0, "clear_al": 0, "resourc": 0, "offici": 0, "colab": 0, "notebook": 0, "behav": 0, "semant": [0, 1], "constraint": 0, "suggest": 0, "certain": [0, 1], "pattern": [0, 1], "bad": 0, "mind": 0, "expens": [0, 1], "speedup": 0, "rest": [0, 1], "direct": 0, "indirect": 0, "mask": [0, 1], "thu": 0, "low": 0, "variat": 0, "don": [0, 1], "nativ": 0, "translat": 0, "slowdown": 0, "ask": [0, 1], "unless": [0, 1], "most": 0, "checkout": 0, "scalar": [0, 1], "substitut": 0, "flow": 0, "applic": 0, "clip_grad": 0, "norm": 0, "problemat": 0, "impact": [0, 1], "patch": 0, "clip_grad_norm_": 0, "dramat": 0, "els": [0, 1], "total_norm": 0, "zero": 0, "p": [0, 1], "param_norm": 0, "grad": 0, "norm_typ": 0, "add_": 0, "clip_coef": 0, "max_norm": 0, "1e": 0, "mul_": 0, "data_parallel": 0, "drop": 0, "last": 0, "stride": 0, "reconstruct": 0, "shallow": 0, "ty": 0, "_xlac": [0, 1], "_get_xla_tensors_text": [0, 1], "_get_xla_tensors_hlo": 0, "prior": 0, "degre": 0, "xla_ir_debug": 0, "_xla_ir": 0, "propag": 0, "metadata": 0, "xla_save_tensors_fil": 0, "realli": [0, 1], "big": [0, 1], "left": 0, "append": 0, "clean": 0, "sheet": 0, "xla_save_tensors_fmt": 0, "store": 0, "_xla_save_tensor": 0, "text": 0, "dot": 0, "graphviz": 0, "xla_flag": 0, "xla_dump_to": 0, "dir_nam": 0, "unoptim": 0, "optimz": 0, "xla_metrics_fil": 0, "xla_save_hlo_fil": 0, "offend": 0, "xla_sync_wait": 0, "forc": [0, 1], "xla_use_eager_debug_mod": 0, "bypass": 0, "higher": 0, "optimizaiton": 0, "skip": 0, "bifloat16": 0, "arithmet": 0, "accur": 0, "accumul": 0, "4096": 0, "dtype": [0, 1], "4097": 0, "averag": 0, "stai": 0, "fp32": 0, "xla_use_f16": 0, "float16": 0, "half": 0, "tf_cpp_log_thread_id": 0, "tf": [0, 1], "tf_cpp_vmodul": 0, "vlog": 0, "form": [0, 1], "tf_cpp_min_log_level": 0, "turn": 0, "tf_vlog": 0, "xla_dump_hlo_graph": 0, "rais": 0, "xla_util": 0, "cc": 0, "save1": 0, "xla_graph_executor": 0, "pjrt_computation_cli": 0, "pr": 0, "dir": 0, "pytorch_test_with_slow": 0, "test_torch": 0, "k": 0, "test_put_xla_uint8": 0, "torch_test_devic": 0, "pytorch_test_bas": 0, "migrat": 0, "jax": 0, "public": 0, "init": 0, "renam": 0, "regist": 0, "xla_backend": 0, "torchrun": 0, "init_method": 0, "plugin": 0, "xpu": 0, "neuron": 0, "continu": 0, "xrt_tpu_config": 0, "libtpu": 0, "30": 0, "thousand": 0, "preview": 0, "On": 0, "safe": 0, "broadcast": 0, "manual_se": 0, "pjrt_backend": 0, "o": 0, "ddp": 0, "dist": 0, "init_process_group": 0, "rank": 0, "world_siz": 0, "42": 0, "128": 0, "gradient_as_bucket_view": 0, "mseloss": 0, "001": 0, "confirm": 0, "localservic": 0, "localhost": 0, "51011": 0, "master_addr": 0, "master_port": 0, "12355": 0, "overhead": 0, "grpc": 0, "torchbench": 0, "observ": 0, "35": 0, "2048": 0, "chip": 0, "learn": 0, "test_train_mp_mnist": 0, "batch_siz": 0, "256": 0, "num_epoch": 0, "tpu_process_bound": 0, "tpu_visible_chip": 0, "r1": 0, "preinstal": 0, "docker_imag": 0, "gcr": 0, "io": 0, "authent": 0, "privat": 0, "gcp": 0, "repositori": 0, "auth": 0, "rm": 0, "privileg": 0, "simpli": 0, "gpu_num_devic": 0, "nnode": 0, "num_gpu_devic": 0, "pjrt_distribut": 0, "physic": 0, "number_gpu_vm": 0, "node_rank": 0, "current_node_rank": 0, "nproc_per_nod": 0, "number_local_gpu_devic": 0, "rdzv_endpoint": 0, "internal_ip_address": 0, "port": 0, "multinode_train": 0, "endpoint": 0, "ip": 0, "address": 0, "omit": 0, "machine_0": 0, "machine_1": 0, "machine_0_internal_ip_address": 0, "ident": 0, "although": [0, 1], "mostli": 0, "interchang": 0, "subtl": 0, "importantli": 0, "around": [0, 1], "architectur": 0, "latenc": 0, "deseri": 0, "gain": 0, "interact": 0, "normal": [0, 1], "plan": 0, "simpler": 0, "xla_dist": 0, "sdk": 0, "reimplement": 0, "collect": 0, "enhanc": 0, "substanti": 0, "vari": [0, 1], "consist": 0, "servic": 0, "connect": 0, "practic": [0, 1], "unreli": 0, "inbound": 0, "entir": 0, "restart": 0, "impos": 0, "unwant": 0, "permit": 0, "subset": 0, "old": 0, "alter": 0, "html": 0, "_": 0, "all_gather_object": 0, "gloo": 0, "backend": [0, 1], "new_group": 0, "subgroup": 0, "experi": 0, "altern": [0, 1], "less": [0, 1], "reliabl": 0, "blob": 0, "md": 0, "strongli": 0, "world": [0, 1], "queri": 0, "_all_gath": 0, "No": [0, 1], "int32": 0, "zeros_lik": 0, "get_world_s": 0, "task": 0, "175": 0, "chart": 0, "breakdown": 0, "tfrt": 0, "legaci": 0, "streamexecutor": 0, "tpu_legaci": 0, "shown": [0, 1], "v": [0, 1], "comparison": 0, "jit": 0, "unmodifi": 0, "hook": 0, "biggest": 0, "right": [0, 1], "bridg": 0, "torchfx": 0, "recogn": [0, 1], "technologi": 0, "fx": 0, "openxla": 0, "a_xla": 0, "b_xla": 0, "compiled_cod": 0, "resnet18": 0, "torchvis": 0, "eval_model": 0, "xla_resnet18": 0, "eval": 0, "dynamo_resnet18": 0, "no_grad": 0, "resent18": 0, "binari": 0, "invok": 0, "bench": 0, "59": 0, "64": 0, "resnext50_32x4d": 0, "91": 0, "alexnet": 0, "28": 0, "mobilenet_v2": 0, "18": 0, "62": 0, "mnasnet1_0": 0, "68": 0, "vgg16": 0, "bert_pytorch": 0, "49": 0, "squeezenet1_1": 0, "29": 0, "timm_vision_transform": 0, "52": 0, "geomean": 0, "04": 0, "team": 0, "train_model": 0, "crossentropyloss": 0, "pred": 0, "train_model_main": 0, "dynamo_train_model": 0, "xla_optim": 0, "weight_decai": 0, "extract": 0, "07": 0, "43": 0, "19": 0, "81": 0, "87": 0, "41": 0, "fwd": 0, "bwd": 0, "e2": 0, "easili": [0, 1], "hide": 0, "cost": 0, "scenario": 0, "larger": 0, "forward": 0, "best": 0, "ideal": [0, 1], "promis": 0, "complex": 0, "tradit": 0, "hard": [0, 1], "expand": 0, "mention": [0, 1], "super": 0, "excit": 0, "upcom": 0, "heavili": 0, "invest": 0, "upstream": 0, "matur": 0, "stori": 0, "xlafullyshardeddataparallel": 0, "my_modul": 0, "adam": 0, "0001": 0, "sum": 0, "individu": 0, "layer": 0, "outer": 0, "wrapper": 0, "leftov": 0, "arxiv": 0, "ab": 0, "1910": 0, "02054": 0, "reshard_after_forward": 0, "test_train_mp_mnist_fsdp_with_ckpt": 0, "test_train_mp_imagenet_fsdp": 0, "interleav": 0, "submodul": 0, "inner": 0, "fsdpvitmodel": 0, "ronghanghu": 0, "vit_10b_fsdp_exampl": 0, "run_vit_train": 0, "checkpoint_modul": 0, "pull": 0, "3524": 0, "auto_wrap_polici": 0, "size_based_auto_wrap_polici": 0, "polici": 0, "100m": 0, "transformer_auto_wrap_polici": 0, "conv2d": 0, "partial": [0, 1], "transformer_layer_cl": 0, "auto_wrapper_cal": 0, "remateri": 0, "m": [0, 1], "kwarg": 0, "resum": 0, "get_shard_metadata": 0, "consolidate_sharded_model_checkpoint": 0, "stitch": 0, "ckpt": 0, "shard_metadata": 0, "ckpt_path": 0, "f": 0, "pth": 0, "bash": 0, "consolidate_sharded_ckpt": 0, "ckpt_prefix": 0, "your_sharded_checkpoint_fil": 0, "ckpt_suffix": 0, "_rank": 0, "inspir": 0, "fairscal": 0, "fullyshardeddataparallel": 0, "readthedoc": 0, "en": 0, "explicit": 0, "resort": 0, "approach": [0, 1], "newer": 0, "recurs": 0, "98": 0, "accuraci": 0, "drop_last": 0, "use_nested_fsdp": 0, "use_gradient_checkpoint": 0, "final_ckpt": 0, "75": 0, "1k": 0, "datadir": 0, "test_set_batch_s": 0, "eval_interv": 0, "num_warmup_epoch": 0, "lr_scheduler_divide_every_n_epoch": 0, "lr_scheduler_divisor": 0, "residu": 0, "algorithm": 0, "vision": 0, "vit": 0, "against": 0, "abil": [0, 1], "And": [0, 1], "who": 0, "eager": [0, 1], "mode": [0, 1], "know": [0, 1], "nccl": 0, "new_rank": 0, "ddp_model": 0, "final": 0, "launcher": 0, "demo_fn": 0, "everyth": [0, 1], "touch": 0, "plu": 0, "five": 0, "sy": 0, "tempfil": 0, "cleanup": 0, "destroy_process_group": 0, "toymodel": 0, "__init__": 0, "net1": 0, "1000000": 0, "relu": 0, "net2": 0, "demo_bas": 0, "assert": 0, "graident_as_bucket_view": 0, "label": 0, "run_demo": 0, "tot": 0, "statist": 0, "produc": [0, 1], "median": 0, "90th": 0, "std": 0, "cv": 0, "418": 0, "54": 0, "419": 0, "22": 0, "430": 0, "40": 0, "76": 0, "02": 0, "97": 0, "407": 0, "60": 0, "39": 0, "seem": 0, "extra": 0, "17864": 0, "20108": 0, "96": 0, "24351": 0, "74": 0, "5866": 0, "83": 0, "10701": 0, "11770": 0, "00": 0, "14313": 0, "78": 0, "3102": 0, "92": 0, "round": 0, "sens": 0, "amort": 0, "logdir": 0, "converg": 0, "achiev": 0, "48": 0, "caution": 0, "interest": 0, "enforc": [0, 1], "crash": 0, "nvidia": 0, "driver": 0, "publish": 0, "prebuilt": 0, "cuda11": 0, "config": 0, "central1": 0, "pkg": 0, "nightly_3": 0, "8_cuda_12": 0, "toolkit": 0, "datacent": 0, "latest": 0, "curl": 0, "fssl": 0, "libnvidia": 0, "gpgkei": 0, "gpg": 0, "dearmor": 0, "keyr": 0, "l": 0, "deb": 0, "sed": 0, "sign": 0, "tee": 0, "d": [0, 1], "ctk": 0, "systemctl": 0, "shm": 0, "16g": 0, "bin": 0, "exec": 0, "awk": 0, "nr": 0, "visibl": [0, 1], "smi": 0, "verifi": 0, "20ab2c7a2d06": 0, "dec": 0, "06": 0, "24": 0, "2022": 0, "510": 0, "47": 0, "03": 0, "bu": 0, "disp": 0, "volatil": 0, "uncorr": 0, "ecc": 0, "fan": 0, "temp": 0, "perf": [0, 1], "pwr": 0, "cap": 0, "mig": 0, "tesla": 0, "v100": 0, "sxm2": 0, "00000000": 0, "36c": 0, "p0": 0, "38w": 0, "300w": 0, "0mib": 0, "16384mib": 0, "gi": 0, "pid": 0, "ld_library_path": 0, "echo": 0, "link": 0, "bashrc": 0, "lib64": 0, "cp310": 0, "manylinux_2_28_x86_64": 0, "prepar": 0, "begin": 0, "38": 0, "89059": 0, "82": 0, "globalr": 0, "79297": 0, "117": 0, "45": 0, "84": 0, "36": 0, "43628": 0, "281": 0, "83108": 0, "346": 0, "88": 0, "108": 0, "99023": 0, "373": 0, "132": 0, "56": 0, "92699": 0, "384": 0, "152": 0, "14": 0, "120": 0, "68816": 0, "388": 0, "169": 0, "09": 0, "use_cuda": 0, "xla_cuda": 0, "successfulli": 0, "gspmd": 0, "annot": 0, "ml": 0, "proper": 0, "hint": 0, "figur": 0, "strategi": 0, "non": [0, 1], "suit": 0, "framework": [0, 1], "popular": 0, "shift": 0, "unlock": 0, "advanc": 0, "use_spmd": 0, "b8b484515a97f74e013dcf38125c44d53a41f011": 0, "l214": 0, "is_spmd": 0, "mechan": 0, "mark_shard": 0, "src": 0, "union": 0, "partition_spec": 0, "logic": [0, 1], "xlatensor": 0, "rfc": 0, "np": 0, "mesh_shap": 0, "num_devic": 0, "global_runtime_device_count": 0, "device_id": 0, "m1_shard": 0, "isinst": 0, "techniqu": 0, "comment": [0, 1], "simplelinear": 0, "fc1": 0, "enumer": 0, "sharidng": 0, "spatial": 0, "ouput": 0, "cluster": 0, "interconnect": 0, "deriv": 0, "sub": 0, "abstract": 0, "axis_nam": 0, "nad": 0, "get_logical_mesh": 0, "ordereddict": 0, "tile": 0, "intend": 0, "manipul": 0, "nice": 0, "arrang": 0, "involv": [0, 1], "center": 0, "dcn": 0, "hybridmesh": 0, "box": 0, "multislic": 0, "accept": 0, "ici_mesh_shap": 0, "dcn_mesh_shap": 0, "denot": 0, "device_mesh": 0, "axi": 0, "input_tensor": 0, "evenli": 0, "equival": 0, "3d": 0, "2d": [0, 1], "two_d_partially_repl": 0, "two_d_fully_shard": 0, "4d": 0, "four_d": 0, "delai": 0, "carri": [0, 1], "lazili": [0, 1], "satisfi": 0, "subclass": 0, "__torch_dispatch__": 0, "spmdpartition": 0, "global_tensor": 0, "strictli": 0, "local_shard": 0, "xlashard": 0, "4e8e5511555073ce8b6d1a436bf808c9333dcac6": 0, "xla_sharded_tensor": 0, "l12": 0, "ongo": 0, "effort": [0, 1], "distributedtensor": 0, "prototyp": 0, "proof": 0, "concept": 0, "distribute_tensor": 0, "devicemesh": 0, "big_tensor": 0, "100000": 0, "my_dtensor": 0, "tune": 0, "_input_sharding_": 0, "field": 0, "input_mesh": 0, "dedic": 0, "planner": 0, "spmdsaveplann": 0, "spmdloadplann": 0, "dist_cp": 0, "distributed_checkpoint": 0, "xc": 0, "storage_writ": 0, "filesystemwrit": 0, "checkpoint_dir": 0, "storage_read": 0, "filesystemread": 0, "manag": 0, "all_step": 0, "save_async": 0, "unblock": 0, "dispatch": 0, "preemption": 0, "detect": 0, "termin": 0, "provis": 0, "queuedresourc": 0, "autocheckpoint": 0, "chkpt_on_preempt": 0, "fsspec": 0, "filesystem": 0, "gc": 0, "chkpt_mgr": 0, "my": [0, 1], "bucket": [0, 1], "tracked_step": 0, "choos": 0, "highest": 0, "best_step": 0, "discov": 0, "_after": 0, "_the": 0, "unnecessari": 0, "forth": 0, "decis": 0, "hardcod": 0, "leverag": 0, "num_gpu_machin": 0, "rank_of_current_machin": 0, "machine_0_ip_address": 0, "training_or_inference_script_using_spmd": 0, "xla_use_spmd": 0, "test_train_spmd_imagenet": 0, "0th": 0, "w": 0, "h": 0, "highli": 0, "param": 0, "named_paramet": 0, "quick": 0, "plai": 0, "512": 0, "throughput": 0, "am": 0, "roughli": 0, "4x": 0, "placement": 0, "visualize_shard": 0, "rich": 0, "generated_t": 0, "use_color": 0, "alt_text": 0, "asset": 0, "spmd_debug_1": 0, "png": 0, "style": 0, "partial_repl": 0, "envvar": 0, "xla_spmd_auto": 0, "_tensor": 0, "auto_polici": 0, "device_count": 0, "mymodul": 0, "sharded_model": 0, "behvaior": 0, "xla_auto_use_group_shard": 0, "reshard": 0, "xla_auto_spmd_mesh": 0, "unset": 0, "fsdpv2": 0, "epxress": 0, "famou": 0, "aim": [0, 1], "offer": 0, "enjoi": 0, "bring": 0, "tabl": 0, "review": 0, "spmd_fully_sharded_data_parallel": 0, "spmdfullyshardeddataparallel": 0, "autowrap": 0, "shard_output": 0, "categori": 0, "element": [0, 1], "logit": 0, "children": 0, "infinit": 0, "compil": 1, "xla": 1, "ar": 1, "static": 1, "word": 1, "chang": 1, "hurt": 1, "understand": 1, "debug": 1, "user": 1, "pov": 1, "sai": 1, "support": 1, "assur": 1, "gone": 1, "coverag": 1, "detail": 1, "explan": 1, "common": 1, "do": 1, "get": 1, "rid": 1, "mainli": 1, "problem": 1, "beginn": 1, "propos": 1, "here": 1, "reli": 1, "impract": 1, "assumpt": 1, "ye": 1, "exampl": 1, "differ": 1, "sentenc": 1, "length": 1, "ll": 1, "new": 1, "more": 1, "kinda": 1, "anti": 1, "pytorch": 1, "frontend": 1, "lazi": 1, "matter": 1, "cpu": 1, "workaround": 1, "okai": 1, "teach": 1, "data": 1, "torch": 1, "theoret": 1, "memori": 1, "sort": 1, "devic": 1, "node": 1, "obviou": 1, "s64": 1, "8": 1, "num_output": 1, "10": 1, "mul": 1, "inde": 1, "_get_xla_tensor_dimension_s": 1, "commonli": 1, "basic": 1, "correct": 1, "wrong": 1, "wors": 1, "take": 1, "probabl": 1, "how": 1, "upper": 1, "nit": 1, "api": 1, "rand": 1, "solv": 1, "should": 1, "kept": 1, "earli": 1, "accessor": 1, "implicitli": 1, "current": 1, "overload": 1, "A": 1, "explod": 1, "cheap": 1, "ve": 1, "variant": 1, "hoc": 1, "think": 1, "verison": 1, "x": 1, "bla": 1, "blabla": 1, "execut": 1, "proce": 1, "choic": 1, "multipl": 1, "combin": 1, "wide": 1, "adopt": 1, "code": 1, "uglier": 1, "win": 1, "pars": 1, "automat": 1, "torchscript": 1, "somehow": 1, "merg": 1, "properli": 1, "haven": 1, "thought": 1, "trivial": 1, "side": 1, "That": 1, "why": 1, "hit": 1, "bandwidth": 1, "automag": 1, "gold": 1, "smart": 1, "trick": 1, "tbh": 1, "longer": 1, "y": 1, "The": 1, "unawar": 1, "hope": 1, "smash": 1, "an": 1, "blocker": 1, "ahead": 1, "nnc": 1, "symbol": 1, "optim": 1, "concret": 1, "kernel": 1, "exactli": 1, "transpos": 1, "With": 1, "brian": 1, "hirsh": 1, "bdhirsh": 1, "question": 1, "move": 1, "stick": 1, "torch_warn": 1, "yea": 1, "hei": 1, "won": 1, "run": 1, "blaze": 1, "fast": 1, "isn": 1, "model": 1, "devirtu": 1, "sound": 1, "great": 1, "truth": 1, "irvalu": 1, "discrep": 1, "followup": 1, "1000": 1, "singl": 1, "presenc": 1, "get_dimention_s": 1, "didn": 1, "condit": 1, "exponenti": 1, "blowup": 1, "fewer": 1, "opportun": 1, "feasibl": 1, "annoi": 1, "variabl": 1, "wasn": 1, "materiz": 1, "combo": 1}, "objects": {"torch_xla.core": [[0, 0, 0, "-", "functions"], [0, 0, 0, "-", "xla_model"]], "torch_xla.core.functions": [[0, 1, 1, "", "all_gather"], [0, 1, 1, "", "all_reduce"]], "torch_xla.core.xla_model": [[0, 1, 1, "", "add_step_closure"], [0, 1, 1, "", "all_gather"], [0, 1, 1, "", "all_reduce"], [0, 1, 1, "", "all_to_all"], [0, 1, 1, "", "do_on_ordinals"], [0, 1, 1, "", "get_local_ordinal"], [0, 1, 1, "", "get_memory_info"], [0, 1, 1, "", "get_ordinal"], [0, 1, 1, "", "get_rng_state"], [0, 1, 1, "", "get_stablehlo"], [0, 1, 1, "", "get_stablehlo_bytecode"], [0, 1, 1, "", "get_xla_supported_devices"], [0, 1, 1, "", "is_master_ordinal"], [0, 1, 1, "", "mesh_reduce"], [0, 1, 1, "", "optimizer_step"], [0, 1, 1, "", "rendezvous"], [0, 1, 1, "", "save"], [0, 1, 1, "", "set_rng_state"], [0, 1, 1, "", "wait_device_ops"], [0, 1, 1, "", "xla_device"], [0, 1, 1, "", "xla_device_hw"], [0, 1, 1, "", "xrt_world_size"]], "torch_xla.distributed": [[0, 0, 0, "-", "parallel_loader"], [0, 0, 0, "-", "xla_multiprocessing"]], "torch_xla.distributed.parallel_loader": [[0, 2, 1, "", "ParallelLoader"]], "torch_xla.distributed.parallel_loader.ParallelLoader": [[0, 3, 1, "", "per_device_loader"]], "torch_xla.distributed.xla_multiprocessing": [[0, 2, 1, "", "MpModelWrapper"], [0, 2, 1, "", "MpSerialExecutor"], [0, 1, 1, "", "spawn"]], "torch_xla.distributed.xla_multiprocessing.MpModelWrapper": [[0, 3, 1, "", "to"]], "torch_xla.distributed.xla_multiprocessing.MpSerialExecutor": [[0, 3, 1, "", "run"]], "torch_xla.utils": [[0, 0, 0, "-", "serialization"], [0, 0, 0, "-", "utils"]], "torch_xla.utils.serialization": [[0, 1, 1, "", "load"], [0, 1, 1, "", "save"]], "torch_xla.utils.utils": [[0, 2, 1, "", "DataWrapper"], [0, 2, 1, "", "SampleGenerator"]]}, "objtypes": {"0": "py:module", "1": "py:function", "2": "py:class", "3": "py:method"}, "objnames": {"0": ["py", "module", "Python module"], "1": ["py", "function", "Python function"], "2": ["py", "class", "Python class"], "3": ["py", "method", "Python method"]}, "titleterms": {"pytorch": 0, "xla": 0, "devic": 0, "creat": 0, "an": 0, "tensor": [0, 1], "ar": 0, "run": 0, "model": 0, "singl": 0, "multipl": 0, "multi": 0, "process": 0, "tpu": 0, "pod": 0, "deep": 0, "dive": 0, "lazi": 0, "bfloat16": 0, "memori": 0, "layout": 0, "move": 0, "from": [0, 1], "cpu": 0, "save": 0, "load": 0, "compil": 0, "cach": 0, "further": 0, "read": 0, "api": 0, "xla_model": 0, "distribut": 0, "util": 0, "test": 0, "beginn": 0, "": [0, 1], "guid": 0, "basic": 0, "high": 0, "level": 0, "understand": 0, "some": [0, 1], "detail": 0, "setup": 0, "convert": 0, "code": 0, "exampl": 0, "1": [0, 1], "stabl": 0, "diffus": 0, "infer": 0, "lightn": 0, "2": [0, 1], "hf": 0, "profil": 0, "perform": 0, "analysi": 0, "troubleshoot": 0, "saniti": 0, "check": 0, "version": 0, "A": 0, "simpl": 0, "calcul": 0, "resnet": 0, "With": 0, "fake": 0, "data": 0, "debug": 0, "tool": 0, "dynamo": 0, "auto": 0, "metric": 0, "execut": 0, "get": 0, "report": 0, "The": 0, "clear": 0, "known": 0, "caveat": 0, "quirk": 0, "more": 0, "environ": 0, "variabl": 0, "common": 0, "combin": 0, "reproduc": 0, "ci": 0, "cd": 0, "unit": 0, "failur": 0, "pjrt": 0, "runtim": 0, "tl": 0, "dr": 0, "benefit": 0, "quickstart": 0, "docker": 0, "gpu": 0, "node": 0, "train": 0, "differ": 0, "xrt": 0, "multithread": 0, "v2": 0, "v3": 0, "chang": 0, "xm": 0, "rendezv": 0, "torch": 0, "new": 0, "torchdynamo": 0, "integr": 0, "featur": 0, "gap": 0, "take": 0, "awai": 0, "fulli": 0, "shard": 0, "parallel": 0, "fsdp": 0, "script": 0, "mnist": 0, "imagenet": 0, "instal": 0, "clone": 0, "repo": 0, "8": 0, "50": 0, "10": 0, "billion": 0, "paramet": 0, "how": 0, "do": 0, "distributeddataparallel": 0, "background": 0, "motiv": 0, "us": [0, 1], "benchmark": 0, "resnet50": 0, "real": [0, 1], "disclaim": 0, "instanc": 0, "wheel": 0, "amp": 0, "automat": 0, "mix": 0, "precis": 0, "develop": 0, "build": 0, "sourc": [0, 1], "support": 0, "spmd": 0, "user": 0, "what": [0, 1], "i": [0, 1], "design": 0, "overview": 0, "aannot": 0, "mesh": 0, "hybrid": 0, "partit": 0, "spec": 0, "xlashardedtensor": 0, "dtensor": 0, "awar": 0, "host": 0, "checkpoint": 0, "checkpointmanag": 0, "group": 0, "virtual": 0, "optim": 0, "number": 0, "refer": 0, "express": 0, "here": 0, "2x2": 0, "ax": 0, "x": 0, "y": 0, "can": [0, 1], "visual": 0, "visualize_tensor_shard": 0, "method": 0, "current": 0, "should": 0, "via": 0, "distribute_modul": 0, "output": [0, 1], "gradient": 0, "huggingfac": 0, "llama": 0, "recompil": 1, "torch_xla": 1, "let": 1, "first": 1, "start": 1, "fact": 1, "constraint": 1, "input": 1, "dataset": 1, "oper": 1, "bound": 1, "dynam": 1, "shape": 1, "fix": 1, "case": 1, "when": 1, "you": 1, "without": 1, "queri": 1, "its": 1, "dimens": 1, "3": 1, "control": 1, "flow": 1, "conclus": 1, "appendix": 1}, "envversion": {"sphinx.domains.c": 2, "sphinx.domains.changeset": 1, "sphinx.domains.citation": 1, "sphinx.domains.cpp": 8, "sphinx.domains.index": 1, "sphinx.domains.javascript": 2, "sphinx.domains.math": 2, "sphinx.domains.python": 3, "sphinx.domains.rst": 2, "sphinx.domains.std": 2, "sphinx.ext.intersphinx": 1, "sphinx.ext.todo": 2, "sphinx.ext.viewcode": 1, "sphinx": 57}, "alltitles": {"PyTorch on XLA Devices": [[0, "pytorch-on-xla-devices"]], "Creating an XLA Tensor": [[0, "creating-an-xla-tensor"]], "XLA Tensors are PyTorch Tensors": [[0, "xla-tensors-are-pytorch-tensors"]], "Running Models on XLA Devices": [[0, "running-models-on-xla-devices"]], "Running on a Single XLA Device": [[0, "running-on-a-single-xla-device"]], "Running on Multiple XLA Devices with Multi-processing": [[0, "running-on-multiple-xla-devices-with-multi-processing"]], "Running on TPU Pods": [[0, "running-on-tpu-pods"]], "XLA Tensor Deep Dive": [[0, "id3"]], "XLA Tensors are Lazy": [[0, "xla-tensors-are-lazy"]], "XLA Tensors and bFloat16": [[0, "xla-tensors-and-bfloat16"]], "Memory Layout": [[0, "memory-layout"]], "Moving XLA Tensors to and from the CPU": [[0, "moving-xla-tensors-to-and-from-the-cpu"]], "Saving and Loading XLA Tensors": [[0, "saving-and-loading-xla-tensors"]], "Compilation Caching": [[0, "compilation-caching"]], "Further Reading": [[0, "further-reading"]], "PyTorch/XLA API": [[0, "pytorch-xla-api"]], "xla_model": [[0, "module-torch_xla.core.xla_model"]], "distributed": [[0, "module-torch_xla.distributed.parallel_loader"]], "utils": [[0, "module-torch_xla.utils.utils"]], "test": [[0, "test"]], "Beginner\u2019s Guide to PyTorch/XLA": [[0, "beginner-s-guide-to-pytorch-xla"]], "Basic high-level understanding of some XLA details": [[0, "basic-high-level-understanding-of-some-xla-details"]], "TPU Setup": [[0, "tpu-setup"]], "Converting code to PyTorch XLA": [[0, "converting-code-to-pytorch-xla"]], "Example 1. Stable Diffusion inference in PyTorch Lightning on a Single TPU Device": [[0, "example-1-stable-diffusion-inference-in-pytorch-lightning-on-a-single-tpu-device"]], "Example 2. HF Stable Diffusion Inference": [[0, "example-2-hf-stable-diffusion-inference"]], "Running on a Single TPU device": [[0, "running-on-a-single-tpu-device"]], "Profiling and performance analysis": [[0, "profiling-and-performance-analysis"]], "Running on Multiple TPU Devices": [[0, "running-on-multiple-tpu-devices"]], "Running on Pods": [[0, "running-on-pods"]], "Troubleshooting": [[0, "troubleshooting"]], "Sanity Check": [[0, "sanity-check"]], "Check PyTorch/XLA Version": [[0, "check-pytorch-xla-version"]], "Perform A Simple Calculation": [[0, "perform-a-simple-calculation"]], "Run Resnet With Fake Data": [[0, "run-resnet-with-fake-data"]], "Performance Debugging": [[0, "performance-debugging"]], "PyTorch/XLA Debugging Tool": [[0, "pytorch-xla-debugging-tool"]], "PyTorch/XLA + Dynamo Debugging Tool": [[0, "pytorch-xla-dynamo-debugging-tool"]], "Perform A Auto-Metrics Analysis": [[0, "perform-a-auto-metrics-analysis"]], "Compilation & Execution Analysis": [[0, "compilation-execution-analysis"]], "Get A Metrics Report": [[0, "get-a-metrics-report"]], "Understand The Metrics Report": [[0, "understand-the-metrics-report"]], "Clear The Metrics Report": [[0, "clear-the-metrics-report"]], "Performance Profiling": [[0, "performance-profiling"]], "Known Performance Caveats": [[0, "known-performance-caveats"]], "XLA Tensor Quirks": [[0, "xla-tensor-quirks"]], "More Debugging Tools": [[0, "more-debugging-tools"]], "Environment Variables": [[0, "environment-variables"]], "Common Debugging Environment Variables Combinations": [[0, "common-debugging-environment-variables-combinations"]], "Reproducing PyTorch/XLA CI/CD unit test failures.": [[0, "reproducing-pytorch-xla-ci-cd-unit-test-failures"]], "PJRT Runtime": [[0, "pjrt-runtime"]], "TL;DR": [[0, "tl-dr"]], "Benefits": [[0, "benefits"]], "Quickstart": [[0, "quickstart"]], "CPU": [[0, "cpu"]], "TPU": [[0, "tpu"]], "Pods": [[0, "pods"]], "Docker": [[0, "docker"], [0, "id21"]], "GPU": [[0, "gpu"]], "Single-node GPU training": [[0, "single-node-gpu-training"]], "Multi-node GPU training": [[0, "multi-node-gpu-training"]], "Differences from XRT": [[0, "differences-from-xrt"]], "Multithreading on TPU v2/v3": [[0, "id18"]], "Changes to xm.rendezvous": [[0, "changes-to-xm-rendezvous"]], "PJRT and torch.distributed": [[0, "pjrt-and-torch-distributed"]], "Performance": [[0, "performance"]], "New TPU runtime": [[0, "new-tpu-runtime"]], "TorchDynamo(torch.compile) integration in PyTorch XLA": [[0, "torchdynamo-torch-compile-integration-in-pytorch-xla"]], "Integration": [[0, "integration"]], "Inference": [[0, "inference"]], "Training": [[0, "training"]], "Feature gaps": [[0, "feature-gaps"]], "Take away": [[0, "take-away"]], "Fully Sharded Data Parallel (FSDP) in PyTorch XLA": [[0, "fully-sharded-data-parallel-fsdp-in-pytorch-xla"]], "Example training scripts on MNIST and ImageNet": [[0, "example-training-scripts-on-mnist-and-imagenet"]], "Installation": [[0, "installation"]], "Clone PyTorch/XLA repo": [[0, "clone-pytorch-xla-repo"]], "Train MNIST on v3-8 TPU": [[0, "train-mnist-on-v3-8-tpu"]], "Train ImageNet with ResNet-50 on v3-8 TPU": [[0, "train-imagenet-with-resnet-50-on-v3-8-tpu"]], "Example training scripts on TPU pod (with 10 billion parameters)": [[0, "example-training-scripts-on-tpu-pod-with-10-billion-parameters"]], "How to do DistributedDataParallel": [[0, "how-to-do-distributeddataparallel"]], "Background / Motivation": [[0, "background-motivation"]], "How to use DistributedDataParallel": [[0, "how-to-use-distributeddataparallel"]], "Benchmarking": [[0, "benchmarking"]], "Resnet50 with fake data": [[0, "resnet50-with-fake-data"]], "MNIST with fake data": [[0, "mnist-with-fake-data"]], "MNIST with real data": [[0, "mnist-with-real-data"]], "Disclaimer": [[0, "disclaimer"]], "How to run with PyTorch/XLA:GPU": [[0, "how-to-run-with-pytorch-xla-gpu"]], "Create a GPU instance": [[0, "create-a-gpu-instance"]], "Environment Setup": [[0, "environment-setup"]], "Check environment variable": [[0, "check-environment-variable"]], "Wheel": [[0, "wheel"]], "Run a simple model": [[0, "run-a-simple-model"]], "AMP (AUTOMATIC MIXED PRECISION)": [[0, "amp-automatic-mixed-precision"]], "Develop PyTorch/XLA on a GPU instance (build PyTorch/XLA from source with GPU support)": [[0, "develop-pytorch-xla-on-a-gpu-instance-build-pytorch-xla-from-source-with-gpu-support"]], "PyTorch/XLA SPMD User Guide": [[0, "pytorch-xla-spmd-user-guide"]], "What is PyTorch/XLA SPMD?": [[0, "what-is-pytorch-xla-spmd"]], "PyTorch/XLA SPMD Design Overview": [[0, "pytorch-xla-spmd-design-overview"]], "Simple Example & Sharding Aannotation API": [[0, "simple-example-sharding-aannotation-api"]], "Mesh": [[0, "id28"]], "Hybrid Mesh": [[0, "hybrid-mesh"]], "Partition Spec": [[0, "id29"]], "XLAShardedTensor": [[0, "xlashardedtensor"]], "DTensor Integration": [[0, "dtensor-integration"]], "Sharding-Aware Host-to-Device Data Loading": [[0, "sharding-aware-host-to-device-data-loading"]], "Distributed Checkpointing": [[0, "distributed-checkpointing"]], "CheckpointManager": [[0, "checkpointmanager"]], "Process Groups": [[0, "process-groups"]], "Virtual Device Optimization": [[0, "virtual-device-optimization"]], "Number of processes": [[0, "number-of-processes"]], "Running SPMD on TPU Pod": [[0, "running-spmd-on-tpu-pod"]], "Running SPMD on GPU": [[0, "running-spmd-on-gpu"]], "Reference Examples": [[0, "reference-examples"]], "Use SPMD to express Data Parallel": [[0, "use-spmd-to-express-data-parallel"]], "Use SPMD to express FSDP(Fully Sharded Data Parallel)": [[0, "use-spmd-to-express-fsdp-fully-sharded-data-parallel"]], "Running Resnet50 example with SPMD": [[0, "running-resnet50-example-with-spmd"]], "SPMD Debugging Tool": [[0, "spmd-debugging-tool"]], "Here, mesh is a 2x2 mesh with axes \u2018x\u2019 and \u2018y\u2019": [[0, "here-mesh-is-a-2x2-mesh-with-axes-x-and-y"]], "A tensor\u2019s sharding can be visualized using the visualize_tensor_sharding method": [[0, "a-tensor-s-sharding-can-be-visualized-using-the-visualize-tensor-sharding-method"]], "Currently, model should be loaded to xla device via distribute_module.": [[0, "currently-model-should-be-loaded-to-xla-device-via-distribute-module"]], "Fully Sharded Data Parallel via SPMD": [[0, "fully-sharded-data-parallel-via-spmd"]], "Sharding output": [[0, "sharding-output"]], "Gradient checkpointing": [[0, "id51"]], "HuggingFace Llama 2 Example": [[0, "huggingface-llama-2-example"]], "Source of recompilations in torch_xla": [[1, "source-of-recompilations-in-torch-xla"]], "Let\u2019s first start with some facts/constraints:": [[1, "lets-first-start-with-some-facts-constraints"]], "#1. From input dataset.": [[1, "from-input-dataset"]], "#2. From operator output": [[1, "from-operator-output"]], "2.1 Bounded dynamic shape can fix the case when you use the tensor with dynamic shape as a Tensor, without querying its real dimension.": [[1, "bounded-dynamic-shape-can-fix-the-case-when-you-use-the-tensor-with-dynamic-shape-as-a-tensor-without-querying-its-real-dimension"]], "2.2 what if real dimension is queried on a tensor with dynamic shape?": [[1, "what-if-real-dimension-is-queried-on-a-tensor-with-dynamic-shape"]], "#3. From control flow": [[1, "from-control-flow"]], "Conclusion:": [[1, "conclusion"]], "Appendix:": [[1, "appendix"]]}, "indexentries": {"datawrapper (class in torch_xla.utils.utils)": [[0, "torch_xla.utils.utils.DataWrapper"]], "mpmodelwrapper (class in torch_xla.distributed.xla_multiprocessing)": [[0, "torch_xla.distributed.xla_multiprocessing.MpModelWrapper"]], "mpserialexecutor (class in torch_xla.distributed.xla_multiprocessing)": [[0, "torch_xla.distributed.xla_multiprocessing.MpSerialExecutor"]], "parallelloader (class in torch_xla.distributed.parallel_loader)": [[0, "torch_xla.distributed.parallel_loader.ParallelLoader"]], "samplegenerator (class in torch_xla.utils.utils)": [[0, "torch_xla.utils.utils.SampleGenerator"]], "add_step_closure() (in module torch_xla.core.xla_model)": [[0, "torch_xla.core.xla_model.add_step_closure"]], "all_gather() (in module torch_xla.core.functions)": [[0, "torch_xla.core.functions.all_gather"]], "all_gather() (in module torch_xla.core.xla_model)": [[0, "torch_xla.core.xla_model.all_gather"]], "all_reduce() (in module torch_xla.core.functions)": [[0, "torch_xla.core.functions.all_reduce"]], "all_reduce() (in module torch_xla.core.xla_model)": [[0, "torch_xla.core.xla_model.all_reduce"]], "all_to_all() (in module torch_xla.core.xla_model)": [[0, "torch_xla.core.xla_model.all_to_all"]], "do_on_ordinals() (in module torch_xla.core.xla_model)": [[0, "torch_xla.core.xla_model.do_on_ordinals"]], "get_local_ordinal() (in module torch_xla.core.xla_model)": [[0, "torch_xla.core.xla_model.get_local_ordinal"]], "get_memory_info() (in module torch_xla.core.xla_model)": [[0, "torch_xla.core.xla_model.get_memory_info"]], "get_ordinal() (in module torch_xla.core.xla_model)": [[0, "torch_xla.core.xla_model.get_ordinal"]], "get_rng_state() (in module torch_xla.core.xla_model)": [[0, "torch_xla.core.xla_model.get_rng_state"]], "get_stablehlo() (in module torch_xla.core.xla_model)": [[0, "torch_xla.core.xla_model.get_stablehlo"]], "get_stablehlo_bytecode() (in module torch_xla.core.xla_model)": [[0, "torch_xla.core.xla_model.get_stablehlo_bytecode"]], "get_xla_supported_devices() (in module torch_xla.core.xla_model)": [[0, "torch_xla.core.xla_model.get_xla_supported_devices"]], "is_master_ordinal() (in module torch_xla.core.xla_model)": [[0, "torch_xla.core.xla_model.is_master_ordinal"]], "load() (in module torch_xla.utils.serialization)": [[0, "torch_xla.utils.serialization.load"]], "mesh_reduce() (in module torch_xla.core.xla_model)": [[0, "torch_xla.core.xla_model.mesh_reduce"]], "module": [[0, "module-torch_xla.core.functions"], [0, "module-torch_xla.core.xla_model"], [0, "module-torch_xla.distributed.parallel_loader"], [0, "module-torch_xla.distributed.xla_multiprocessing"], [0, "module-torch_xla.utils.serialization"], [0, "module-torch_xla.utils.utils"]], "optimizer_step() (in module torch_xla.core.xla_model)": [[0, "torch_xla.core.xla_model.optimizer_step"]], "per_device_loader() (torch_xla.distributed.parallel_loader.parallelloader method)": [[0, "torch_xla.distributed.parallel_loader.ParallelLoader.per_device_loader"]], "rendezvous() (in module torch_xla.core.xla_model)": [[0, "torch_xla.core.xla_model.rendezvous"]], "run() (torch_xla.distributed.xla_multiprocessing.mpserialexecutor method)": [[0, "torch_xla.distributed.xla_multiprocessing.MpSerialExecutor.run"]], "save() (in module torch_xla.core.xla_model)": [[0, "torch_xla.core.xla_model.save"]], "save() (in module torch_xla.utils.serialization)": [[0, "torch_xla.utils.serialization.save"]], "set_rng_state() (in module torch_xla.core.xla_model)": [[0, "torch_xla.core.xla_model.set_rng_state"]], "spawn() (in module torch_xla.distributed.xla_multiprocessing)": [[0, "torch_xla.distributed.xla_multiprocessing.spawn"]], "to() (torch_xla.distributed.xla_multiprocessing.mpmodelwrapper method)": [[0, "torch_xla.distributed.xla_multiprocessing.MpModelWrapper.to"]], "torch_xla.core.functions": [[0, "module-torch_xla.core.functions"]], "torch_xla.core.xla_model": [[0, "module-torch_xla.core.xla_model"]], "torch_xla.distributed.parallel_loader": [[0, "module-torch_xla.distributed.parallel_loader"]], "torch_xla.distributed.xla_multiprocessing": [[0, "module-torch_xla.distributed.xla_multiprocessing"]], "torch_xla.utils.serialization": [[0, "module-torch_xla.utils.serialization"]], "torch_xla.utils.utils": [[0, "module-torch_xla.utils.utils"]], "wait_device_ops() (in module torch_xla.core.xla_model)": [[0, "torch_xla.core.xla_model.wait_device_ops"]], "xla_device() (in module torch_xla.core.xla_model)": [[0, "torch_xla.core.xla_model.xla_device"]], "xla_device_hw() (in module torch_xla.core.xla_model)": [[0, "torch_xla.core.xla_model.xla_device_hw"]], "xrt_world_size() (in module torch_xla.core.xla_model)": [[0, "torch_xla.core.xla_model.xrt_world_size"]]}})
\ No newline at end of file