1<!-- Google Tag Manager -->
2<script>
vendor: 415 bytes, lines 2-12
2 3 (function (w, d, s, l, i) { 4 w[l] = w[l] || []; 5 w[l].push({ "gtm.start": new Date().getTime(), event: "gtm.js" }); 6 var f = d.getElementsByTagName(s)[0], 7 j = d.createElement(s), 8 dl = l != "dataLayer" ? "&l=" + l : ""; 9 j.async = true; 10 j.src = "https://www.googletagmanager.com/gtm.js?id=" + i + dl; 11 f.parentNode.insertBefore(j, f); 12 })(window, document, "script", "dataLayer", "
12GTM-KV83JWN
vendor: 4 bytes, line 12
12");
13</script>
13 14<!-- End Google Tag Manager --> 15 16 17<!doctype html> 18<html lang="en" class="no-js"> 19 <head> 20 21 <meta charset="utf-8"> 22 <meta name="viewport" content="width=device-width,initial-scale=1"> 23 24 <meta name="description" content="Scaffold for experienced Machine Learning Researchers"> 25 26 27 <meta name="author" content="DanLing Team"> 28 29 30 <link rel="canonical" href="https://danling.org/tensors/"> 31 32 33 <link rel="prev" href="../runners/utils/"> 34 35 36 <link rel="next" href="nested_tensor/"> 37 38 39 40 <link rel="alternate" href="/" hreflang="en"> 41 42 <link rel="alternate" href="/index.zh/" hreflang="zh"> 43 44 45 46 <link rel="icon" href="../assets/images/logo.ico"> 47 <meta name="generator" content="mkdocs-1.6.1, mkdocs-material-9.7.7"> 48 49 50 51 <title>Tensors - DanLing</title> 52 53 54 55 <link rel="stylesheet" href="../assets/stylesheets/main.ec1eaa64.min.css"> 56 57 58 <link rel="stylesheet" href="../assets/stylesheets/palette.ab4e12ef.min.css"> 59 60 61 62 63 64 65 66 67 68 69 70 71 <link rel="stylesheet" href="../assets/_mkdocstrings.css"> 72 73 <link rel="stylesheet" href="../assets/css/fonts.css"> 74 75 <link rel="stylesheet" href="../assets/css/extra.css"> 76 77
77<script>__md_scope=new URL("..",location),__md_hash=e=>[...e].reduce(((e,_)=>(e<<5)-e+_.charCodeAt(0)),0),__md_get=(e,_=localStorage,t=__md_scope)=>JSON.parse(_.getItem(t.pathname+"."+e)),__md_set=(e,_,t=localStorage,a=__md_scope)=>{try{t.setItem(a.pathname+"."+e,JSON.stringify(_))}catch(e){}}</script>
77 78 79 80 81 82 83 84 85
86<script id="__analytics">function __md_analytics(){function e(){dataLayer.push(arguments)}window.dataLayer=window.dataLayer||[],e("js",new Date),e("config","G-RE7EKZ9T8D"),document.addEventListener("DOMContentLoaded",(function(){document.forms.search&&document.forms.search.query.addEventListener("blur",(function(){this.value&&e("event","search",{search_term:this.value})}));document$.subscribe((function(){var t=document.forms.feedback;if(void 0!==t)for(var a of t.querySelectorAll("[type=submit]"))a.addEventListener("click",(function(a){a.preventDefault();var n=document.location.pathname,d=this.getAttribute("data-md-value");e("event","feedback",{page:n,data:d}),t.firstElementChild.disabled=!0;var r=t.querySelector(".md-feedback__note [data-md-value='"+d+"']");r&&(r.hidden=!1)})),t.hidden=!1})),location$.subscribe((function(t){e("config","G-RE7EKZ9T8D",{page_path:t.pathname})}))}));var t=document.createElement("script");t.async=!0,t.src="https://www.googletagmanager.com/gtag/js?id=G-RE7EKZ9T8D",document.getElementById("__analytics").insertAdjacentElement("afterEnd",t)}</script>
86 87 88
88<script>if("undefined"!=typeof __md_analytics){var consent=__md_get("__consent");consent&&consent.analytics&&__md_analytics()}</script>
88 89 90 91 92 93 </head> 94 95 96 97 98 99 100 101 102 103 <body dir="ltr" data-md-color-scheme="slate" data-md-color-primary="blue-grey" data-md-color-accent="teal"> 104 105 106 <input class="md-toggle" data-md-toggle="drawer" type="checkbox" id="__drawer" autocomplete="off"> 107 <input class="md-toggle" data-md-toggle="search" type="checkbox" id="__search" autocomplete="off"> 108 <label class="md-overlay" for="__drawer"></label> 109 <div data-md-component="skip"> 110 111 112 <a href="#tensor" class="md-skip"> 113 Skip to content 114 </a> 115 116 </div> 117 <div data-md-component="announce"> 118 119 </div> 120 121 122 123 124 125 126<header class="md-header md-header--shadow md-header--lifted" data-md-component="header"> 127 <nav class="md-header__inner md-grid" aria-label="Header"> 128 <a href=".." title="DanLing" class="md-header__button md-logo" aria-label="DanLing" data-md-component="logo"> 129 130 <img src="../assets/images/logo.png" alt="logo"> 131 132 </a> 133 <label class="md-header__button md-icon" for="__drawer"> 134 135 <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24"><path d="M3 6h18v2H3zm0 5h18v2H3zm0 5h18v2H3z"/></svg> 136 </label> 137 <div class="md-header__title" data-md-component="header-title"> 138 <div class="md-header__ellipsis"> 139 <div class="md-header__topic"> 140 <span class="md-ellipsis"> 141 DanLing 142 </span> 143 </div> 144 <div class="md-header__topic" data-md-component="header-topic"> 145 <span class="md-ellipsis"> 146 147 Tensors 148 149 </span> 150 </div> 151 </div> 152 </div> 153 154 155 <form class="md-header__option" data-md-component="palette"> 156 157 158 159 160 <input class="md-option" data-md-color-media="(prefers-color-scheme: dark)" data-md-color-scheme="slate" data-md-color-primary="blue-grey" data-md-color-accent="teal" aria-label="Switch to light mode" type="radio" name="__palette" id="__palette_0"> 161 162 <label class="md-header__button md-icon" title="Switch to light mode" for="__palette_1" hidden> 163 <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24"><path d="M12 18c-.89 0-1.74-.2-2.5-.55C11.56 16.5 13 14.42 13 12s-1.44-4.5-3.5-5.45C10.26 6.2 11.11 6 12 6a6 6 0 0 1 6 6 6 6 0 0 1-6 6m8-9.31V4h-4.69L12 .69 8.69 4H4v4.69L.69 12 4 15.31V20h4.69L12 23.31 15.31 20H20v-4.69L23.31 12z"/></svg> 164 </label> 165 166 167 168 169 170 <input class="md-option" data-md-color-media="(prefers-color-scheme: light)" data-md-color-scheme="default" data-md-color-primary="blue-grey" data-md-color-accent="teal" aria-label="Switch to dark mode" type="radio" name="__palette" id="__palette_1"> 171 172 <label class="md-header__button md-icon" title="Switch to dark mode" for="__palette_0" hidden> 173 <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24"><path d="M12 8a4 4 0 0 0-4 4 4 4 0 0 0 4 4 4 4 0 0 0 4-4 4 4 0 0 0-4-4m0 10a6 6 0 0 1-6-6 6 6 0 0 1 6-6 6 6 0 0 1 6 6 6 6 0 0 1-6 6m8-9.31V4h-4.69L12 .69 8.69 4H4v4.69L.69 12 4 15.31V20h4.69L12 23.31 15.31 20H20v-4.69L23.31 12z"/></svg> 174 </label> 175 176 177</form> 178 179 180 181
181<script>var palette=__md_get("__palette");if(palette&&palette.color){if("(prefers-color-scheme)"===palette.color.media){var media=matchMedia("(prefers-color-scheme: light)"),input=document.querySelector(media.matches?"[data-md-color-media='(prefers-color-scheme: light)']":"[data-md-color-media='(prefers-color-scheme: dark)']");palette.color.media=input.getAttribute("data-md-color-media"),palette.color.scheme=input.getAttribute("data-md-color-scheme"),palette.color.primary=input.getAttribute("data-md-color-primary"),palette.color.accent=input.getAttribute("data-md-color-accent")}for(var[key,value]of Object.entries(palette.color))document.body.setAttribute("data-md-color-"+key,value)}</script>
181 182 183 184 <div class="md-header__option"> 185 <div class="md-select"> 186 187 <button class="md-header__button md-icon" aria-label="Select language"> 188 <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24"><path d="m12.87 15.07-2.54-2.51.03-.03A17.5 17.5 0 0 0 14.07 6H17V4h-7V2H8v2H1v2h11.17C11.5 7.92 10.44 9.75 9 11.35 8.07 10.32 7.3 9.19 6.69 8h-2c.73 1.63 1.73 3.17 2.98 4.56l-5.09 5.02L4 19l5-5 3.11 3.11zM18.5 10h-2L12 22h2l1.12-3h4.75L21 22h2zm-2.62 7 1.62-4.33L19.12 17z"/></svg> 189 </button> 190 <div class="md-select__inner"> 191 <ul class="md-select__list"> 192 193 <li class="md-select__item"> 194 <a href="/" hreflang="en" class="md-select__link"> 195 English 196 </a> 197 </li> 198 199 <li class="md-select__item"> 200 <a href="/index.zh/" hreflang="zh" class="md-select__link"> 201 æ±è¯ 202 </a> 203 </li> 204 205 </ul> 206 </div> 207 </div> 208</div> 209 210 211 212 213 <label class="md-header__button md-icon" for="__search"> 214 215 <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24"><path d="M9.5 3A6.5 6.5 0 0 1 16 9.5c0 1.61-.59 3.09-1.56 4.23l.27.27h.79l5 5-1.5 1.5-5-5v-.79l-.27-.27A6.52 6.52 0 0 1 9.5 16 6.5 6.5 0 0 1 3 9.5 6.5 6.5 0 0 1 9.5 3m0 2C7 5 5 7 5 9.5S7 14 9.5 14 14 12 14 9.5 12 5 9.5 5"/></svg> 216 </label> 217 <div class="md-search" data-md-component="search" role="dialog"> 218 <label class="md-search__overlay" for="__search"></label> 219 <div class="md-search__inner" role="search"> 220 <form class="md-search__form" name="search"> 221 <input type="text" class="md-search__input" name="query" aria-label="Search" placeholder="Search" autocapitalize="off" autocorrect="off" autocomplete="off" spellcheck="false" data-md-component="search-query" required> 222 <label class="md-search__icon md-icon" for="__search"> 223 224 <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24"><path d="M9.5 3A6.5 6.5 0 0 1 16 9.5c0 1.61-.59 3.09-1.56 4.23l.27.27h.79l5 5-1.5 1.5-5-5v-.79l-.27-.27A6.52 6.52 0 0 1 9.5 16 6.5 6.5 0 0 1 3 9.5 6.5 6.5 0 0 1 9.5 3m0 2C7 5 5 7 5 9.5S7 14 9.5 14 14 12 14 9.5 12 5 9.5 5"/></svg> 225 226 <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24"><path d="M20 11v2H8l5.5 5.5-1.42 1.42L4.16 12l7.92-7.92L13.5 5.5 8 11z"/></svg> 227 </label> 228 <nav class="md-search__options" aria-label="Search"> 229 230 <a href="javascript:void(0)" class="md-search__icon md-icon" title="Share" aria-label="Share" data-clipboard data-clipboard-text="" data-md-component="search-share" tabindex="-1"> 231 232 <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24"><path d="M18 16.08c-.76 0-1.44.3-1.96.77L8.91 12.7c.05-.23.09-.46.09-.7s-.04-.47-.09-.7l7.05-4.11c.54.5 1.25.81 2.04.81a3 3 0 0 0 3-3 3 3 0 0 0-3-3 3 3 0 0 0-3 3c0 .24.04.47.09.7L8.04 9.81C7.5 9.31 6.79 9 6 9a3 3 0 0 0-3 3 3 3 0 0 0 3 3c.79 0 1.5-.31 2.04-.81l7.12 4.15c-.05.21-.08.43-.08.66 0 1.61 1.31 2.91 2.92 2.91s2.92-1.3 2.92-2.91A2.92 2.92 0 0 0 18 16.08"/></svg> 233 </a> 234 235 <button type="reset" class="md-search__icon md-icon" title="Clear" aria-label="Clear" tabindex="-1"> 236 237 <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24"><path d="M19 6.41 17.59 5 12 10.59 6.41 5 5 6.41 10.59 12 5 17.59 6.41 19 12 13.41 17.59 19 19 17.59 13.41 12z"/></svg> 238 </button> 239 </nav> 240 241 <div class="md-search__suggest" data-md-component="search-suggest"></div> 242 243 </form> 244 <div class="md-search__output"> 245 <div class="md-search__scrollwrap" tabindex="0" data-md-scrollfix> 246 <div class="md-search-result" data-md-component="search-result"> 247 <div class="md-search-result__meta"> 248 Initializing search 249 </div> 250 <ol class="md-search-result__list" role="presentation"></ol> 251 </div> 252 </div> 253 </div> 254 </div> 255</div> 256 257 258 259 <div class="md-header__source"> 260 <a href="https://github.com/ZhiyuanChen/DanLing" title="Go to repository" class="md-source" data-md-component="source"> 261 <div class="md-source__icon md-icon"> 262 263 <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 448 512"><!--! Font Awesome Free 7.1.0 by @fontawesome - https://fontawesome.com License - https://fontawesome.com/license/free (Icons: CC BY 4.0, Fonts: SIL OFL 1.1, Code: MIT License) Copyright 2025 Fonticons, Inc.--><path d="M439.6 236.1 244 40.5c-5.4-5.5-12.8-8.5-20.4-8.5s-15 3-20.4 8.4L162.5 81l51.5 51.5c27.1-9.1 52.7 16.8 43.4 43.7l49.7 49.7c34.2-11.8 61.2 31 35.5 56.7-26.5 26.5-70.2-2.9-56-37.3L240.3 199v121.9c25.3 12.5 22.3 41.8 9.1 55-6.4 6.4-15.2 10.1-24.3 10.1s-17.8-3.6-24.3-10.1c-17.6-17.6-11.1-46.9 11.2-56v-123c-20.8-8.5-24.6-30.7-18.6-45L142.6 101 8.5 235.1C3 240.6 0 247.9 0 255.5s3 15 8.5 20.4l195.6 195.7c5.4 5.4 12.7 8.4 20.4 8.4s15-3 20.4-8.4l194.7-194.7c5.4-5.4 8.4-12.8 8.4-20.4s-3-15-8.4-20.4"/></svg> 264 </div> 265 <div class="md-source__repository"> 266 DanLing 267 </div> 268</a> 269 </div> 270 271 </nav> 272 273 274 275<nav class="md-tabs" aria-label="Tabs" data-md-component="tabs"> 276 <div class="md-grid"> 277 <ul class="md-tabs__list"> 278 279 280 281 282 283 284 <li class="md-tabs__item"> 285 <a href=".." class="md-tabs__link"> 286 287 288 289 290 291 DanLing 292 293 </a> 294 </li> 295 296 297 298 299 300 301 302 303 <li class="md-tabs__item"> 304 <a href="../runners/" class="md-tabs__link"> 305 306 307 308 309 310 Runners 311 312 </a> 313 </li> 314 315 316 317 318 319 320 321 322 323 324 <li class="md-tabs__item md-tabs__item--active"> 325 <a href="./" class="md-tabs__link"> 326 327 328 329 330 331 Tensors 332 333 </a> 334 </li> 335 336 337 338 339 340 341 342 343 344 345 <li class="md-tabs__item"> 346 <a href="../optim/lr_scheduler/" class="md-tabs__link"> 347 348 349 350 351 352 Optim 353 354 </a> 355 </li> 356 357 358 359 360 361 362 363 364 365 <li class="md-tabs__item"> 366 <a href="../metrics/" class="md-tabs__link"> 367 368 369 370 371 372 Metrics 373 374 </a> 375 </li> 376 377 378 379 380 381 382 383 384 385 386 <li class="md-tabs__item"> 387 <a href="../utils/decorators/" class="md-tabs__link"> 388 389 390 391 392 393 Utils 394 395 </a> 396 </li> 397 398 399 400 401 402 403 404 405 406 <li class="md-tabs__item"> 407 <a href="../package/" class="md-tabs__link"> 408 409 410 411 412 413 package 414 415 </a> 416 </li> 417 418 419 420 421 422 423 424 425 426 427 <li class="md-tabs__item"> 428 <a href="../about/" class="md-tabs__link"> 429 430 431 432 433 434 About 435 436 </a> 437 </li> 438 439 440 441 442 443 444 445 446 447 448 449 <li class="md-tabs__item"> 450 <a href="../index.zh/" class="md-tabs__link"> 451 452 453 454 455 456 䏿 457 458 </a> 459 </li> 460 461 462 463 464 </ul> 465 </div> 466</nav> 467 468 469</header> 470 471 <div class="md-container" data-md-component="container"> 472 473 474 475 476 <main class="md-main" data-md-component="main"> 477 <div class="md-main__inner md-grid"> 478 479 480 481 <div class="md-sidebar md-sidebar--primary" data-md-component="sidebar" data-md-type="navigation" > 482 <div class="md-sidebar__scrollwrap"> 483 <div class="md-sidebar__inner"> 484 485 486 487 488 489 490<nav class="md-nav md-nav--primary md-nav--lifted" aria-label="Navigation" data-md-level="0"> 491 <label class="md-nav__title" for="__drawer"> 492 <a href=".." title="DanLing" class="md-nav__button md-logo" aria-label="DanLing" data-md-component="logo"> 493 494 <img src="../assets/images/logo.png" alt="logo"> 495 496 </a> 497 DanLing 498 </label> 499 500 <div class="md-nav__source"> 501 <a href="https://github.com/ZhiyuanChen/DanLing" title="Go to repository" class="md-source" data-md-component="source"> 502 <div class="md-source__icon md-icon"> 503 504 <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 448 512"><!--! Font Awesome Free 7.1.0 by @fontawesome - https://fontawesome.com License - https://fontawesome.com/license/free (Icons: CC BY 4.0, Fonts: SIL OFL 1.1, Code: MIT License) Copyright 2025 Fonticons, Inc.--><path d="M439.6 236.1 244 40.5c-5.4-5.5-12.8-8.5-20.4-8.5s-15 3-20.4 8.4L162.5 81l51.5 51.5c27.1-9.1 52.7 16.8 43.4 43.7l49.7 49.7c34.2-11.8 61.2 31 35.5 56.7-26.5 26.5-70.2-2.9-56-37.3L240.3 199v121.9c25.3 12.5 22.3 41.8 9.1 55-6.4 6.4-15.2 10.1-24.3 10.1s-17.8-3.6-24.3-10.1c-17.6-17.6-11.1-46.9 11.2-56v-123c-20.8-8.5-24.6-30.7-18.6-45L142.6 101 8.5 235.1C3 240.6 0 247.9 0 255.5s3 15 8.5 20.4l195.6 195.7c5.4 5.4 12.7 8.4 20.4 8.4s15-3 20.4-8.4l194.7-194.7c5.4-5.4 8.4-12.8 8.4-20.4s-3-15-8.4-20.4"/></svg> 505 </div> 506 <div class="md-source__repository"> 507 DanLing 508 </div> 509</a> 510 </div> 511 512 <ul class="md-nav__list" data-md-scrollfix> 513 514 515 516 517 518 519 520 <li class="md-nav__item"> 521 <a href=".." class="md-nav__link"> 522 523 524 525 <span class="md-ellipsis"> 526 527 528 DanLing 529 530 531 532 </span> 533 534 535 536 </a> 537 </li> 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 <li class="md-nav__item md-nav__item--pruned md-nav__item--nested"> 575 576 577 578 579 580 <a href="../runners/config/" class="md-nav__link"> 581 582 583 584 <span class="md-ellipsis"> 585 586 587 Runners 588 589 590 591 </span> 592 593 594 595
596 <span class="md-nav__icon md-icon"></span> 597 598 </a> 599 600 601 602 </li> 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 <li class="md-nav__item md-nav__item--active md-nav__item--section md-nav__item--nested"> 638 639 640 641 <input class="md-nav__toggle md-toggle " type="checkbox" id="__nav_3" checked> 642 643 644 <div class="md-nav__link md-nav__container"> 645 <a href="./" class="md-nav__link md-nav__link--active"> 646 647 648 649 <span class="md-ellipsis"> 650 651 652 Tensors 653 654 655 656 </span> 657 658 659 660 </a> 661 662 663 <label class="md-nav__link md-nav__link--active" for="__nav_3" id="__nav_3_label" tabindex=""> 664 <span class="md-nav__icon md-icon"></span> 665 </label> 666 667 </div> 668 669 <nav class="md-nav" data-md-level="1" aria-labelledby="__nav_3_label" aria-expanded="true"> 670 <label class="md-nav__title" for="__nav_3"> 671 <span class="md-nav__icon md-icon"></span> 672 673 674 Tensors 675 676 677 </label> 678 <ul class="md-nav__list" data-md-scrollfix> 679 680 681 682 683 684 685 686 <li class="md-nav__item"> 687 <a href="nested_tensor/" class="md-nav__link"> 688 689 690 691 <span class="md-ellipsis"> 692 693 694 NestedTensor 695 696 697 698 </span> 699 700 701 702 </a> 703 </li> 704 705 706 707 708 709 710 711 712 713 714 <li class="md-nav__item"> 715 <a href="pn_tensor/" class="md-nav__link"> 716 717 718 719 <span class="md-ellipsis"> 720 721 722 PNTensor 723 724 725 726 </span> 727 728 729 730 </a> 731 </li> 732 733 734 735 736 737 738 739 740 741 742 <li class="md-nav__item"> 743 <a href="functions/" class="md-nav__link"> 744 745 746 747 <span class="md-ellipsis"> 748 749 750 Functions 751 752 753 754 </span> 755 756 757 758 </a> 759 </li> 760 761 762 763 764 </ul> 765 </nav> 766 767 </li> 768 769 770 771 772 773 774 775 776 777 778 779 780 781 782 783 784 785 786 787 788 789 790 791 792 793 794 795 796 <li class="md-nav__item md-nav__item--pruned md-nav__item--nested"> 797 798 799 800 801 802 <a href="../optim/lr_scheduler/" class="md-nav__link"> 803 804 805 806 <span class="md-ellipsis"> 807 808 809 Optim 810 811 812 813 </span> 814 815 816 817 818 <span class="md-nav__icon md-icon"></span> 819 820 </a> 821 822 823 824 </li> 825 826 827 828 829 830 831 832 833 834 835 836 837 838 839 840 841 842 843 844 845 846 847 848 849 850 851 852 853 854 855 856 857 858 859 860 861 862 863 <li class="md-nav__item md-nav__item--pruned md-nav__item--nested"> 864 865 866 867 868 869 <a href="../metrics/metrics/" class="md-nav__link"> 870 871 872 873 <span class="md-ellipsis"> 874 875 876 Metrics 877 878 879 880 </span> 881 882 883 884 885 <span class="md-nav__icon md-icon"></span> 886 887 </a> 888 889 890 891 </li> 892 893 894 895 896 897 898 899 900 901 902 903 904 905 906 907 908 909 910 911 912 913 914 915 916 917 918 919 920 921 922 923 924 925 926 <li class="md-nav__item md-nav__item--pruned md-nav__item--nested"> 927 928 929 930 931 932 <a href="../utils/decorators/" class="md-nav__link"> 933 934 935 936 <span class="md-ellipsis"> 937 938 939 Utils 940 941 942 943 </span> 944 945 946 947 948 <span class="md-nav__icon md-icon"></span> 949 950 </a> 951 952 953 954 </li> 955 956 957 958 959 960 961 962 963 964 <li class="md-nav__item"> 965 <a href="../package/" class="md-nav__link"> 966 967 968 969 <span class="md-ellipsis"> 970 971 972 package 973 974 975 976 </span> 977 978 979 980 </a> 981 </li> 982 983 984 985 986 987 988 989 990 991 992 993 994 995 996 997 998 999 1000 1001 1002 1003 1004 1005 1006 1007 1008 1009 1010 1011 1012 1013 1014 1015 1016 <li class="md-nav__item md-nav__item--pruned md-nav__item--nested"> 1017 1018 1019 1020 1021 1022 <a href="../about/" class="md-nav__link"> 1023 1024 1025 1026 <span class="md-ellipsis"> 1027 1028 1029 About 1030 1031 1032 1033 </span> 1034 1035 1036 1037
1038 <span class="md-nav__icon md-icon"></span> 1039 1040 </a> 1041 1042 1043 1044 </li> 1045 1046 1047 1048 1049 1050 1051 1052 1053 1054 1055 1056 1057 1058 1059 1060 1061 1062 1063 1064 1065 1066 1067 1068 1069 1070 1071 1072 1073 1074 1075 1076 1077 1078 1079 <li class="md-nav__item md-nav__item--pruned md-nav__item--nested"> 1080 1081 1082 1083 1084 1085 <a href="../index.zh/" class="md-nav__link"> 1086 1087 1088 1089 <span class="md-ellipsis"> 1090 1091 1092 䏿 1093 1094 1095 1096 </span> 1097 1098 1099 1100 1101 <span class="md-nav__icon md-icon"></span> 1102 1103 </a> 1104 1105 1106 1107 </li> 1108 1109 1110 1111 </ul> 1112</nav> 1113 </div> 1114 </div> 1115 </div> 1116 1117 1118 1119 <div class="md-sidebar md-sidebar--secondary" data-md-component="sidebar" data-md-type="toc" > 1120 <div class="md-sidebar__scrollwrap"> 1121 <div class="md-sidebar__inner"> 1122 1123 1124<nav class="md-nav md-nav--secondary" aria-label="Table of contents"> 1125 1126 1127 1128 1129 1130 1131 <label class="md-nav__title" for="__toc"> 1132 <span class="md-nav__icon md-icon"></span> 1133 Table of contents 1134 </label> 1135 <ul class="md-nav__list" data-md-component="toc" data-md-scrollfix> 1136 1137 <li class="md-nav__item"> 1138 <a href="#overview" class="md-nav__link"> 1139 <span class="md-ellipsis"> 1140 1141 Overview 1142 1143 </span> 1144 </a> 1145 1146</li> 1147 1148 <li class="md-nav__item"> 1149 <a href="#key-components" class="md-nav__link"> 1150 <span class="md-ellipsis"> 1151 1152 Key Components 1153 1154 </span> 1155 </a> 1156 1157</li> 1158 1159 <li class="md-nav__item"> 1160 <a href="#quick-start" class="md-nav__link"> 1161 <span class="md-ellipsis"> 1162 1163 Quick Start 1164 1165 </span> 1166 </a> 1167 1168 <nav class="md-nav" aria-label="Quick Start"> 1169 <ul class="md-nav__list"> 1170 1171 <li class="md-nav__item"> 1172 <a href="#creating-a-nestedtensor" class="md-nav__link"> 1173 <span class="md-ellipsis"> 1174 1175 Creating a NestedTensor 1176 1177 </span> 1178 </a> 1179 1180</li> 1181 1182 <li class="md-nav__item"> 1183 <a href="#creating-from-non-tensor-data" class="md-nav__link"> 1184 <span class="md-ellipsis"> 1185 1186 Creating from Non-Tensor Data 1187 1188 </span> 1189 </a> 1190 1191</li> 1192 1193 <li class="md-nav__item"> 1194 <a href="#converting-to-torchnested_tensor" class="md-nav__link"> 1195 <span class="md-ellipsis"> 1196 1197 Converting to torch.nested_tensor 1198 1199 </span> 1200 </a> 1201 1202</li> 1203 1204 </ul> 1205 </nav> 1206 1207</li> 1208 1209 <li class="md-nav__item"> 1210 <a href="#working-with-nestedtensor" class="md-nav__link"> 1211 <span class="md-ellipsis"> 1212 1213 Working with NestedTensor 1214 1215 </span> 1216 </a> 1217 1218 <nav class="md-nav" aria-label="Working with NestedTensor"> 1219 <ul class="md-nav__list"> 1220 1221 <li class="md-nav__item"> 1222 <a href="#operations" class="md-nav__link"> 1223 <span class="md-ellipsis"> 1224 1225 Operations 1226 1227 </span> 1228 </a> 1229 1230</li> 1231 1232 <li class="md-nav__item"> 1233 <a href="#unpacking" class="md-nav__link"> 1234 <span class="md-ellipsis"> 1235 1236 Unpacking 1237 1238 </span> 1239 </a> 1240 1241</li> 1242 1243 </ul> 1244 </nav> 1245 1246</li> 1247 1248 <li class="md-nav__item"> 1249 <a href="#architecture" class="md-nav__link"> 1250 <span class="md-ellipsis"> 1251 1252 Architecture 1253 1254 </span> 1255 </a> 1256 1257 <nav class="md-nav" aria-label="Architecture"> 1258 <ul class="md-nav__list"> 1259 1260 <li class="md-nav__item"> 1261 <a href="#dispatch-system" class="md-nav__link"> 1262 <span class="md-ellipsis"> 1263 1264 Dispatch System 1265 1266 </span> 1267 </a> 1268 1269</li> 1270 1271 <li class="md-nav__item"> 1272 <a href="#key-internal-helpers" class="md-nav__link"> 1273 <span class="md-ellipsis"> 1274 1275 Key Internal Helpers 1276 1277 </span> 1278 </a> 1279 1280</li> 1281 1282 </ul> 1283 </nav> 1284 1285</li> 1286 1287 <li class="md-nav__item"> 1288 <a href="#integration-with-pytorch-dataloader" class="md-nav__link"> 1289 <span class="md-ellipsis"> 1290 1291 Integration with PyTorch DataLoader 1292 1293 </span> 1294 </a> 1295 1296</li> 1297 1298 <li class="md-nav__item"> 1299 <a href="#advanced-usage" class="md-nav__link"> 1300 <span class="md-ellipsis"> 1301 1302 Advanced Usage 1303 1304 </span> 1305 </a> 1306 1307 <nav class="md-nav" aria-label="Advanced Usage"> 1308 <ul class="md-nav__list"> 1309 1310 <li class="md-nav__item"> 1311 <a href="#custom-collation" class="md-nav__link"> 1312 <span class="md-ellipsis"> 1313 1314 Custom Collation 1315 1316 </span> 1317 </a> 1318 1319</li> 1320 1321 <li class="md-nav__item"> 1322 <a href="#working-with-pytorch-models" class="md-nav__link"> 1323 <span class="md-ellipsis"> 1324 1325 Working with PyTorch Models 1326 1327 </span> 1328 </a> 1329 1330</li> 1331 1332 <li class="md-nav__item"> 1333 <a href="#extending-with-new-operations" class="md-nav__link"> 1334 <span class="md-ellipsis"> 1335 1336 Extending with New Operations 1337 1338 </span> 1339 </a> 1340 1341</li> 1342 1343 </ul> 1344 </nav> 1345 1346</li> 1347 1348 <li class="md-nav__item"> 1349 <a href="#benchmarks" class="md-nav__link"> 1350 <span class="md-ellipsis"> 1351 1352 Benchmarks 1353 1354 </span> 1355 </a> 1356 1357 <nav class="md-nav" aria-label="Benchmarks"> 1358 <ul class="md-nav__list"> 1359 1360 <li class="md-nav__item"> 1361 <a href="#imdb-training" class="md-nav__link"> 1362 <span class="md-ellipsis"> 1363 1364 IMDB Training 1365 1366 </span> 1367 </a> 1368 1369</li> 1370 1371 <li class="md-nav__item"> 1372 <a href="#models" class="md-nav__link"> 1373 <span class="md-ellipsis"> 1374 1375 Models 1376 1377 </span> 1378 </a> 1379 1380</li> 1381 1382 <li class="md-nav__item"> 1383 <a href="#operators" class="md-nav__link"> 1384 <span class="md-ellipsis"> 1385 1386 Operators 1387 1388 </span> 1389 </a> 1390 1391</li> 1392 1393 </ul> 1394 </nav> 1395 1396</li> 1397 1398 <li class="md-nav__item"> 1399 <a href="#danling.tensors" class="md-nav__link"> 1400 <span class="md-ellipsis"> 1401 1402 Tensor API 1403 1404 </span> 1405 </a> 1406 1407</li> 1408 1409 </ul> 1410 1411</nav> 1412 </div> 1413 </div> 1414 </div> 1415 1416 1417 1418 <div class="md-content" data-md-component="content"> 1419 1420 <article class="md-content__inner md-typeset"> 1421 1422 1423 1424 1425 1426
1427 1428 1429<h1 id="tensor">Tensor<a class="headerlink" href="#tensor" title="Permanent link">¶</a></h1> 1430<p>The <code>danling.tensors</code> module provides utilities for handling tensors with variable lengths in batched operations. 1431The core feature is the <a class="autorefs autorefs-internal" title="<code class="doc-symbol doc-symbol-heading doc-symbol-class"></code> <span class="doc doc-object-name doc-class-name">NestedTensor</span>" href="nested_tensor/#danling.tensors.nested_tensor.NestedTensor"><code>NestedTensor</code></a> class which allows efficient representation of sequences of different lengths without excessive padding.</p> 1432<h2 id="overview">Overview<a class="headerlink" href="#overview" title="Permanent link">¶</a></h2> 1433<p>In many deep learning tasks, especially those involving sequences (text, time series, etc.), each example in a batch may have a different length. Traditional approaches include:</p> 1434<ol> 1435<li><strong>Padding</strong>: Adding placeholder values to make all examples the same length (wastes computation)</li> 1436<li><strong>Bucketing</strong>: Grouping similar-length examples (complicates training)</li> 1437<li><strong>Processing one sample at a time</strong>: Slow and inefficient</li> 1438</ol> 1439<p>The <code>NestedTensor</code> solves these problems by providing:</p> 1440<ul> 1441<li>A way to store variable-length tensors in a single object</li> 1442<li>Automatic padding and mask generation for efficient computation</li> 1443<li>Transparent access to the original tensors or padded representations</li> 1444<li>Support for 615+ PyTorch operations via a multi-level dispatch system</li> 1445</ul> 1446<h2 id="key-components">Key Components<a class="headerlink" href="#key-components" title="Permanent link">¶</a></h2> 1447<ul> 1448<li><a class="autorefs autorefs-internal" title="<code class="doc-symbol doc-symbol-heading doc-symbol-class"></code> <span class="doc doc-object-name doc-class-name">NestedTensor</span>" href="nested_tensor/#danling.tensors.nested_tensor.NestedTensor"><code>NestedTensor</code></a>: Main class for handling variable-length tensors in a batch.</li> 1449<li><a class="autorefs autorefs-internal" title="<code class="doc-symbol doc-symbol-heading doc-symbol-class"></code> <span class="doc doc-object-name doc-class-name">PNTensor</span>" href="pn_tensor/#danling.tensors.pn_tensor.PNTensor"><code>PNTensor</code></a>: A tensor wrapper that can be converted to NestedTensor by PyTorch DataLoader.</li> 1450<li><code>tensor()</code>: Function to create a <a class="autorefs autorefs-internal" title="<code class="doc-symbol doc-symbol-heading doc-symbol-class"></code> <span class="doc doc-object-name doc-class-name">PNTensor</span>" href="pn_tensor/#danling.tensors.pn_tensor.PNTensor"><code>PNTensor</code></a> object (similar to <code>torch.tensor()</code>).</li> 1451<li><a class="autorefs autorefs-internal" title="<code class="doc-symbol doc-symbol-heading doc-symbol-attribute"></code> <span class="doc doc-object-name doc-attribute-name">NestedTensorFuncRegistry</span> 1452 1453 1454 <span class="doc doc-labels"> 1455 <small class="doc doc-label doc-label-module-attribute"><code>module-attribute</code></small> 1456 </span>" href="functions/#danling.tensors.ops.NestedTensorFuncRegistry"><code>NestedTensorFuncRegistry</code></a>: Registry for <code>torch.*</code> and <code>F.*</code> dispatch handlers.</li> 1457<li><a class="autorefs autorefs-internal" title="<code class="doc-symbol doc-symbol-heading doc-symbol-attribute"></code> <span class="doc doc-object-name doc-attribute-name">NestedTensorAtenRegistry</span> 1458 1459 1460 <span class="doc doc-labels"> 1461 <small class="doc doc-label doc-label-module-attribute"><code>module-attribute</code></small> 1462 </span>" href="functions/#danling.tensors.ops.NestedTensorAtenRegistry"><code>NestedTensorAtenRegistry</code></a>: Registry for <code>aten</code> dispatch handlers.</li> 1463</ul> 1464<h2 id="quick-start">Quick Start<a class="headerlink" href="#quick-start" title="Permanent link">¶</a></h2> 1465<h3 id="creating-a-nestedtensor">Creating a NestedTensor<a class="headerlink" href="#creating-a-nestedtensor" title="Permanent link">¶</a></h3> 1466<div class="language-python highlight"><table class="highlighttable"><tr><th colspan="2" class="filename"><span class="filename">Python</span></th></tr><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal"> 1</span> 1467<span class="normal"> 2</span> 1468<span class="normal"> 3</span> 1469<span class="normal"> 4</span> 1470<span class="normal"> 5</span> 1471<span class="normal"> 6</span> 1472<span class="normal"> 7</span> 1473<span class="normal">
1473 8</span> 1474<span class="normal"> 9</span> 1475<span class="normal">10</span> 1476<span class="normal">11</span> 1477<span class="normal">12</span> 1478<span class="normal">13</span> 1479<span class="normal">14</span> 1480<span class="normal">15</span> 1481<span class="normal">16</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="kn">import</span><span class="w"> </span><span class="nn">torch</span> 1482<span class="kn">from</span><span class="w"> </span><span class="nn">danling.tensors</span><span class="w"> </span><span class="kn">import</span> <span class="n">NestedTensor</span> 1483 1484<span class="c1"># Create from a list of tensors with different lengths</span> 1485<span class="n">tensor1</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">tensor</span><span class="p">([</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">])</span> 1486<span class="n">tensor2</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">tensor</span><span class="p">([</span><span class="mi">4</span><span class="p">,</span> <span class="mi">5</span><span class="p">])</span> 1487<span class="n">nested</span> <span class="o">=</span> <span class="n">NestedTensor</span><span class="p">(</span><span class="n">tensor1</span><span class="p">,</span> <span class="n">tensor2</span><span class="p">)</span> 1488 1489<span class="c1"># Access properties</span> 1490<span class="nb">print</span><span class="p">(</span><span class="n">nested</span><span class="o">.</span><span class="n">tensor</span><span class="p">)</span> <span class="c1"># Padded tensor: [[1, 2, 3], [4, 5, 0]]</span> 1491<span class="nb">print</span><span class="p">(</span><span class="n">nested</span><span class="o">.</span><span class="n">mask</span><span class="p">)</span> <span class="c1"># Mask: [[True, True, True], [True, True, False]]</span> 1492<span class="nb">print</span><span class="p">(</span><span class="n">nested</span><span class="o">.</span><span class="n">concat</span><span class="p">)</span> <span class="c1"># Concatenated: [1, 2, 3, 4, 5]</span> 1493 1494<span class="c1"># Index operations</span> 1495<span class="nb">print</span><span class="p">(</span><span class="n">nested</span><span class="p">[</span><span class="mi">0</span><span class="p">])</span> <span class="c1"># First tensor: [1, 2, 3]</span> 1496<span class="nb">print</span><span class="p">(</span><span class="n">nested</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">:])</span> <span class="c1"># Slice: NestedTensor([[2, 3], [5]])</span> 1497</code></pre></div></td></tr></table></div> 1498<h3 id="creating-from-non-tensor-data">Creating from Non-Tensor Data<a class="headerlink" href="#creating-from-non-tensor-data" title="Permanent link">¶</a></h3> 1499<div class="language-python highlight"><table class="highlighttable"><tr><th colspan="2" class="filename"><span class="filename">Python</span></th></tr><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal">1</span> 1500<span class="normal">2</span> 1501<span class="normal">3</span> 1502<span class="normal">4</span> 1503<span class="normal">5</span> 1504<span class="normal">6</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="kn">from</span><span class="w"> </span><span class="nn">danling.tensors</span><span class="w"> </span><span class="kn">import</span> <span class="n">NestedTensor</span> 1505 1506<span class="c1"># Create directly from lists</span> 1507<span class="n">nested</span> <span class="o">=</span> <span class="n">NestedTensor</span><span class="p">([</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">],</span> <span class="p">[</span><span class="mi">4</span><span class="p">,</span> <span class="mi">5</span><span class="p">])</span> 1508<span class="nb">print</span><span class="p">(</span><span class="n">nested</span><span class="p">)</span> <span class="c1"># NestedTensor([[1, 2, 3], [4, 5]])</span> 1509<span class="nb">print</span><span class="p">(</span><span class="n">nested</span><span class="o">.</span><span class="n">tolist</span><span class="p">())</span> <span class="c1"># [[1, 2, 3], [4, 5]]</span> 1510</code></pre></div></td></tr></table></div> 1511<h3 id="converting-to-torchnested_tensor">Converting to torch.nested_tensor<a class="headerlink" href="#convert
1511ing-to-torchnested_tensor" title="Permanent link">¶</a></h3> 1512<div class="language-python highlight"><table class="highlighttable"><tr><th colspan="2" class="filename"><span class="filename">Python</span></th></tr><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal">1</span> 1513<span class="normal">2</span> 1514<span class="normal">3</span> 1515<span class="normal">4</span> 1516<span class="normal">5</span> 1517<span class="normal">6</span> 1518<span class="normal">7</span> 1519<span class="normal">8</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="kn">import</span><span class="w"> </span><span class="nn">torch</span> 1520<span class="kn">from</span><span class="w"> </span><span class="nn">danling.tensors</span><span class="w"> </span><span class="kn">import</span> <span class="n">NestedTensor</span> 1521 1522<span class="n">nested</span> <span class="o">=</span> <span class="n">NestedTensor</span><span class="p">([</span><span class="mf">1.0</span><span class="p">,</span> <span class="mf">2.0</span><span class="p">,</span> <span class="mf">3.0</span><span class="p">],</span> <span class="p">[</span><span class="mf">4.0</span><span class="p">,</span> <span class="mf">5.0</span><span class="p">])</span> 1523 1524<span class="c1"># Convert to PyTorch's native nested tensor</span> 1525<span class="n">native</span> <span class="o">=</span> <span class="n">nested</span><span class="o">.</span><span class="n">to_torch_nested</span><span class="p">()</span> 1526<span class="nb">print</span><span class="p">(</span><span class="nb">type</span><span class="p">(</span><span class="n">native</span><span class="p">))</span> <span class="c1"># <class 'torch.Tensor'> (nested layout)</span> 1527</code></pre></div></td></tr></table></div> 1528<h2 id="working-with-nestedtensor">Working with NestedTensor<a class="headerlink" href="#working-with-nestedtensor" title="Permanent link">¶</a></h2> 1529<h3 id="operations">Operations<a class="headerlink" href="#operations" title="Permanent link">¶</a></h3> 1530<p>NestedTensor supports many PyTorch operations:</p> 1531<div class="language-python highlight"><table class="highlighttable"><tr><th colspan="2" class="filename"><span class="filename">Python</span></th></tr><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal"> 1</span> 1532<span class="normal"> 2</span> 1533<span class="normal"> 3</span> 1534<span class="normal"> 4</span> 1535<span class="normal"> 5</span> 1536<span class="normal"> 6</span> 1537<span class="normal"> 7</span> 1538<span class="normal"> 8</span> 1539<span class="normal"> 9</span> 1540<span class="normal">10</span> 1541<span class="normal">11</span> 1542<span class="normal">12</span> 1543<span class="normal">13</span> 1544<span class="normal">14</span> 1545<span class="normal">15</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="c1"># Arithmetic operations</span> 1546<span class="n">result</span> <span class="o">=</span> <span class="n">nested</span> <span class="o">+</span> <span class="mi">10</span> 1547<span class="n">result</span> <span class="o">=</span> <span class="n">nested</span> <span class="o">*</span> <span class="mi">2</span> 1548 1549<span class="c1"># Type conversion</span> 1550<span class="n">float_nested</span> <span class="o">=</span> <span class="n">nested</span><span class="o">.</span><span class="n">float</span><span class="p">()</span> 1551<span class="n">half_nested</span> <span class="o">=</span> <span class="n">nested</span><span class="o">.</span><span class="n">half</span><span class="p">()</span> 1552 1553<span class="c1"># Device movement</span> 1554<span class="n">gpu_nested</span> <span class="o">=</span> <span class="n">nested</span><span class="o">.</span><span class="n">cuda</span><span class="p">()</span> 1555<span class="n">cpu_nested</span> <span class="o">=</span> <span class="n">gpu_nested</span><span class="o">.</span><span class="n">cpu</span><span class="p">()</span> 1556 1557<span class="c1"># Shape operations</span> 1558<span class="nb">print</span><span class="p">(</span><span class="n">nested</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span> <span class="c1"># torch.Size([2, 3])</span> 1559<span class="nb">print</span><span class="p">(</span><span class="n">nested</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="mi">0</span><span class="p">))</span> <span class="c1">
1559# 2</span> 1560</code></pre></div></td></tr></table></div> 1561<h3 id="unpacking">Unpacking<a class="headerlink" href="#unpacking" title="Permanent link">¶</a></h3> 1562<p>You can easily convert back to original tensors:</p> 1563<div class="language-python highlight"><table class="highlighttable"><tr><th colspan="2" class="filename"><span class="filename">Python</span></th></tr><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal">1</span> 1564<span class="normal">2</span> 1565<span class="normal">3</span> 1566<span class="normal">4</span> 1567<span class="normal">5</span> 1568<span class="normal">6</span> 1569<span class="normal">7</span> 1570<span class="normal">8</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="c1"># Get as a list of lists</span> 1571<span class="n">data</span> <span class="o">=</span> <span class="n">nested</span><span class="o">.</span><span class="n">tolist</span><span class="p">()</span> 1572 1573<span class="c1"># Get as a tuple of (padded_tensor, mask)</span> 1574<span class="n">tensor</span><span class="p">,</span> <span class="n">mask</span> <span class="o">=</span> <span class="n">nested</span><span class="o">.</span><span class="n">tensor_mask</span> 1575 1576<span class="c1"># Access individual items</span> 1577<span class="n">first_item</span> <span class="o">=</span> <span class="n">nested</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="c1"># Returns the first tensor</span> 1578</code></pre></div></td></tr></table></div> 1579<h2 id="architecture">Architecture<a class="headerlink" href="#architecture" title="Permanent link">¶</a></h2> 1580<p>NestedTensor uses a <strong>packed representation</strong> that stores all variable-length elements concatenated into a single contiguous tensor, tracked by offset metadata:</p> 1581<ul> 1582<li><code>_values</code>: All element tensors concatenated along dim 0 (e.g., shape <code>[total_elements, *]</code>)</li> 1583<li><code>_offsets</code>: Cumulative element counts, shape <code>(B+1,)</code>, marking where each element starts/ends</li> 1584<li><code>_physical_shape</code>: Per-element shapes, shape <code>(B, ndim)</code>, recording each element’s original dimensions</li> 1585</ul> 1586<p>This avoids the waste of padding in the internal representation while allowing efficient batch operations.</p> 1587<h3 id="dispatch-system">Dispatch System<a class="headerlink" href="#dispatch-system" title="Permanent link">¶</a></h3> 1588<p>Operations on NestedTensor are handled by a <strong>three-tier dispatch system</strong>, ordered from fastest to most flexible:</p> 1589<p><strong>Level 1 â Aten dispatch</strong> (<code>aten_functions.py</code>, 285 ops): Operates directly on the packed <code>_values</code> tensor via <code>__torch_dispatch__</code>. This is the fastest path â no Python loops, no unpacking. Used for elementwise ops (add, mul, sin, exp, …), reductions, softmax, layer_norm, etc.</p> 1590<p><strong>Level 2 â Torch function dispatch</strong> (<code>torch_functions.py</code>, 217 ops): Intercepts <code>torch.*</code> calls via <code>__torch_function__</code>. Handles ops that need dimension translation (e.g., <code>torch.flatten</code>, <code>torch.softmax</code> with non-default dim), multi-operand dispatch (e.g., <code>torch.einsum</code>), per-element matrix ops (e.g., <code>torch.det</code>, <code>torch.linalg.svd</code>), and fused attention (<code>torch._native_multi_head_attention</code>, <code>torch._transformer_encoder_layer_fwd</code>).</p> 1591<p><strong>Level 3 â NN function dispatch</strong> (<code>nn_functions.py</code>, 113 ops): Also via <code>__torch_function__</code>, handles <code>torch.nn.functional.*</code> ops including convolutions, pooling, normalization, attention, embedding, activations (F.relu, F.gelu, F.silu, …), and loss functions. Transformer-hot ops use packed fast paths; activation handlers strip inplace flags to preserve autograd on the wrapper subclass.</p> 1592<p><strong>Fallback</strong>: Any aten op without an explicit handler falls back to <code>per_element_fallback</code>, which unpacks to individual tensors, applies the op element-by-element, and repacks. Under <code>torch.compile</code>, DanLing prefers explicit failure over silently entering those eager-only fallbacks.</p> 1593<div class="language-text highlight"><table class="highlighttable"><tr><th colspan="2" class="filename"><span class="filename">Text Only</span></th></tr><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal"> 1</span> 1594<span class="normal"> 2</span> 1595<span class="normal"> 3</span> 1596<span class="normal"> 4</span> 1597<span class="normal"> 5</span> 1598<span class="normal"> 6</span> 1599<span class="normal"> 7</span> 1600<span class="normal">
1600 8</span> 1601<span class="normal"> 9</span> 1602<span class="normal">10</span> 1603<span class="normal">11</span> 1604<span class="normal">12</span> 1605<span class="normal">13</span></pre></div></td><td class="code"><div><pre><span></span><code>torch.some_op(nested_tensor) 1606 â 1607 â¼ 1608__torch_function__ ââââ handler in TorchFuncRegistry? ââââ yes âââ dispatch 1609 â no 1610 â¼ 1611aten decomposition (PyTorch lowers to aten ops) 1612 â 1613 â¼ 1614__torch_dispatch__ ââââ handler in NestedTensorAtenRegistry? ââââ yes âââ dispatch 1615 â no 1616 â¼ 1617per_element_fallback (unpack â apply per element â repack) 1618</code></pre></div></td></tr></table></div> 1619<h3 id="key-internal-helpers">Key Internal Helpers<a class="headerlink" href="#key-internal-helpers" title="Permanent link">¶</a></h3> 1620<ul> 1621<li><code>_from_packed(values, offsets, shape_tensor, ...)</code>: Direct constructor from packed representation. Used by all aten handlers to build results without function call overhead.</li> 1622<li><code>_map_storage_serial(input, fn)</code>: Per-element slow path â applies <code>fn</code> to each element via <code>_unpack()</code>. Used when ops need individual element dimensionality.</li> 1623<li><code>_translate_non_batch_dim(nt, dim)</code>: Converts a NestedTensor dim index to the corresponding element-level dim (skipping the batch dimension).</li> 1624</ul> 1625<h2 id="integration-with-pytorch-dataloader">Integration with PyTorch DataLoader<a class="headerlink" href="#integration-with-pytorch-dataloader" title="Permanent link">¶</a></h2> 1626<p>The <code>PNTensor</code> class makes it easy to use <code>NestedTensor</code> with PyTorch’s DataLoader:</p> 1627<div class="language-python highlight"><table class="highlighttable"><tr><th colspan="2" class="filename"><span class="filename">Python</span></th></tr><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal"> 1</span> 1628<span class="normal"> 2</span> 1629<span class="normal"> 3</span> 1630<span class="normal"> 4</span> 1631<span class="normal"> 5</span> 1632<span class="normal"> 6</span> 1633<span class="normal"> 7</span> 1634<span class="normal"> 8</span> 1635<span class="normal"> 9</span> 1636<span class="normal">10</span> 1637<span class="normal">11</span> 1638<span class="normal">12</span> 1639<span class="normal">13</span> 1640<span class="normal">14</span> 1641<span class="normal">15</span> 1642<span class="normal">16</span> 1643<span class="normal">17</span> 1644<span class="normal">18</span> 1645<span class="normal">19</span> 1646<span class="normal">20</span> 1647<span class="normal">21</span> 1648<span class="normal">22</span> 1649<span class="normal">23</span> 1650<span class="normal">24</span> 1651<span class="normal">25</span> 1652<span class="normal">26</span> 1653<span class="normal">27</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="kn">from</span><span class="w"> </span><span class="nn">torch.utils.data</span><span class="w"> </span><span class="kn">import</span> <span class="n">Dataset</span><span class="p">,</span> <span class="n">DataLoader</span> 1654<span class="kn">from</span><span class="w"> </span><span class="nn">danling.tensors</span><span class="w"> </span><span class="kn">import</span> <span class="n">PNTensor</span> 1655 1656<span class="k">class</span><span class="w"> </span><span class="nc">VariableLengthDataset</span><span class="p">(</span><span class="n">Dataset</span><span class="p">):</span> 1657 <span class="k">def</span><span class="w"> </span><span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">data</span><span class="p">):</span> 1658 <span class="bp">self</span><span class="o">.</span><span class="n">data</span> <span class="o">=</span> <span class="n">data</span> 1659 1660 <span class="k">def</span><span class="w"> </span><span class="fm">__len__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span> 1661 <span class="k">return</span> <span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">data</span><span class="p">)</span> 1662 1663 <span class="k">def</span><span class="w"> </span><span class="fm">__getitem__</span><span class="p">(</span><span class="bp">
1663self</span><span class="p">,</span> <span class="n">idx</span><span class="p">):</span> 1664 <span class="c1"># Return a PNTensor; DataLoader collates it into NestedTensor.</span> 1665 <span class="k">return</span> <span class="n">PNTensor</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">data</span><span class="p">[</span><span class="n">idx</span><span class="p">])</span> 1666 1667<span class="c1"># Example usage</span> 1668<span class="n">dataset</span> <span class="o">=</span> <span class="n">VariableLengthDataset</span><span class="p">([</span> 1669 <span class="p">[</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">],</span> 1670 <span class="p">[</span><span class="mi">4</span><span class="p">,</span> <span class="mi">5</span><span class="p">],</span> 1671 <span class="p">[</span><span class="mi">6</span><span class="p">,</span> <span class="mi">7</span><span class="p">,</span> <span class="mi">8</span><span class="p">,</span> <span class="mi">9</span><span class="p">]</span> 1672<span class="p">])</span> 1673<span class="n">dataloader</span> <span class="o">=</span> <span class="n">DataLoader</span><span class="p">(</span><span class="n">dataset</span><span class="p">,</span> <span class="n">batch_size</span><span class="o">=</span><span class="mi">3</span><span class="p">)</span> 1674 1675<span class="c1"># The batches are NestedTensor objects</span> 1676<span class="k">for</span> <span class="n">batch</span> <span class="ow">in</span> <span class="n">dataloader</span><span class="p">:</span> 1677 <span class="nb">print</span><span class="p">(</span><span class="nb">type</span><span class="p">(</span><span class="n">batch</span><span class="p">))</span> <span class="c1"># <class 'danling.tensors.nested_tensor.NestedTensor'></span> 1678 <span class="nb">print</span><span class="p">(</span><span class="n">batch</span><span class="o">.</span><span class="n">tensor</span><span class="p">)</span> <span class="c1"># Padded tensor</span> 1679 <span class="nb">print</span><span class="p">(</span><span class="n">batch</span><span class="o">.</span><span class="n">mask</span><span class="p">)</span> <span class="c1"># Mask</span> 1680</code></pre></div></td></tr></table></div> 1681<h2 id="advanced-usage">Advanced Usage<a class="headerlink" href="#advanced-usage" title="Permanent link">¶</a></h2> 1682<h3 id="custom-collation">Custom Collation<a class="headerlink" href="#custom-collation" title="Permanent link">¶</a></h3> 1683<p>If you need more control over collation:</p> 1684<div class="language-python highlight"><table class="highlighttable"><tr><th colspan="2" class="filename"><span class="filename">Python</span></th></tr><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal"> 1</span> 1685<span class="normal"> 2</span> 1686<span class="normal"> 3</span> 1687<span class="normal"> 4</span> 1688<span class="normal"> 5</span> 1689<span class="normal"> 6</span> 1690<span class="normal"> 7</span> 1691<span class="normal"> 8</span> 1692<span class="normal"> 9</span> 1693<span class="normal">10</span> 1694<span class="normal">11</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="kn">from</span><span class="w"> </span><span class="nn">torch.utils.data</span><span class="w"> </span><span class="kn">import</span> <span class="n">DataLoader</span> 1695<span class="kn">from</span><span class="w"> </span><span class="nn">danling.tensors</span><span class="w"> </span><span class="kn">import</span> <span class="n">NestedTensor</span> 1696 1697<span class="k">def</span><span class="w"> </span><span class="nf">custom_collate_fn</span><span class="p">(</span><span class="n">batch</span><span class="p">):</span> 1698 <span class="k">return</span> <span class="n">NestedTensor</span><span class="p">(</span><span class="o">*</span><span class="n">batch</span><span class="p">)</span> 1699 1700<span class="n">dataloader</span> <span class="o">=</span> <span class="n">DataLoader</span><span class="p">(</span> 1701 <span class="n">dataset</span><span class="p">,</span> 1702 <span class="n">batch_size</span><span class="o">=</span><span class="mi">32</span><span class="p">,</span> 1703 <span class="n">collate_fn</span><span class="o">=</span><span class="n">custom_collate_fn</span> 1704<span class="p">)</span> 1705</code></pre></div></td></tr></table></div> 1706<h3 id="working-with-pytorch-models">Working with PyTorch Models<a class="headerlink" href="#working-with-pytorch-models" title="Permanent link">¶</a></h3> 1707<p><code>NestedTensor</code> works natively with PyTorch’s built-in transformer and vision models â no padding or masks needed:</p> 1708<div class="language-python highlight"><table class="highlighttable"><tr><th colspan="2" class="filename"><span class="filename">Python</span></th></tr><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal"> 1</span> 1709<span class="normal"> 2</span> 1710<span class="normal"> 3</span> 1711<span class="normal"> 4</span> 1712<span class="normal"> 5</span> 1713<span class="normal"> 6</span> 1714<span class="normal"> 7</span> 1715<span class="normal">
1715 8</span> 1716<span class="normal"> 9</span> 1717<span class="normal">10</span> 1718<span class="normal">11</span> 1719<span class="normal">12</span> 1720<span class="normal">13</span> 1721<span class="normal">14</span> 1722<span class="normal">15</span> 1723<span class="normal">16</span> 1724<span class="normal">17</span> 1725<span class="normal">18</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="kn">import</span><span class="w"> </span><span class="nn">torch.nn</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">nn</span> 1726<span class="kn">from</span><span class="w"> </span><span class="nn">danling.tensors</span><span class="w"> </span><span class="kn">import</span> <span class="n">NestedTensor</span> 1727 1728<span class="c1"># Variable-length sequences</span> 1729<span class="n">nt</span> <span class="o">=</span> <span class="n">NestedTensor</span><span class="p">([</span><span class="n">torch</span><span class="o">.</span><span class="n">randn</span><span class="p">(</span><span class="mi">3</span><span class="p">,</span> <span class="mi">512</span><span class="p">),</span> <span class="n">torch</span><span class="o">.</span><span class="n">randn</span><span class="p">(</span><span class="mi">7</span><span class="p">,</span> <span class="mi">512</span><span class="p">)])</span> 1730 1731<span class="c1"># Passes directly through nn.TransformerEncoder, nn.Transformer, etc.</span> 1732<span class="n">encoder</span> <span class="o">=</span> <span class="n">nn</span><span class="o">.</span><span class="n">TransformerEncoder</span><span class="p">(</span> 1733 <span class="n">nn</span><span class="o">.</span><span class="n">TransformerEncoderLayer</span><span class="p">(</span><span class="n">d_model</span><span class="o">=</span><span class="mi">512</span><span class="p">,</span> <span class="n">nhead</span><span class="o">=</span><span class="mi">8</span><span class="p">,</span> <span class="n">batch_first</span><span class="o">=</span><span class="kc">True</span><span class="p">),</span> 1734 <span class="n">num_layers</span><span class="o">=</span><span class="mi">6</span><span class="p">,</span> 1735<span class="p">)</span> 1736<span class="n">output</span> <span class="o">=</span> <span class="n">encoder</span><span class="p">(</span><span class="n">nt</span><span class="p">)</span> <span class="c1"># NestedTensor â only real tokens computed</span> 1737 1738<span class="c1"># Also works with torchvision models for variable-size images</span> 1739<span class="kn">import</span><span class="w"> </span><span class="nn">torchvision.models</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">models</span> 1740<span class="n">resnet</span> <span class="o">=</span> <span class="n">models</span><span class="o">.</span><span class="n">resnet18</span><span class="p">()</span> 1741<span class="n">images</span> <span class="o">=</span> <span class="n">NestedTensor</span><span class="p">([</span><span class="n">torch</span><span class="o">.</span><span class="n">randn</span><span class="p">(</span><span class="mi">3</span><span class="p">,</span> <span class="mi">224</span><span class="p">,</span> <span class="mi">224</span><span class="p">),</span> <span class="n">torch</span><span class="o">.</span><span class="n">randn</span><span class="p">(</span><span class="mi">3</span><span class="p">,</span> <span class="mi">160</span><span class="p">,</span> <span class="mi">160</span><span class="p">)])</span> 1742<span class="n">features</span> <span class="o">=</span> <span class="n">resnet</span><span class="p">(</span><span class="n">images</span><span class="p">)</span> 1743</code></pre></div></td></tr></table></div> 1744<p>For models that require padded input (e.g., HuggingFace transformers), materialize with <code>.tensor</code> and <code>.mask</code>:</p> 1745<div class="language-python highlight"><table class="highlighttable"><tr><th colspan="2" class="filename"><span class="filename">Python</span></th></tr><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal">1</span> 1746<span class="normal">2</span> 1747<span class="normal">3</span> 1748<span class="normal">4</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="n">outputs</span> <span class="o">=</span> <span class="n">model</span><span class="p">(</span> 1749 <span class="n">input_ids</span><span class="o">=</span><span class="n">nested_inputs</span><span class="o">.</span><span class="n">tensor</span><span class="p">,</span> 1750 <span class="n">attention_mask</span><span class="o">=</span><span class="n">nested_inputs</span><span class="o">.</span><span class="n">mask</span> 1751<span class="p">)</span> 1752</code></pre></div></td></tr></table></div> 1753<h3 id="extending-with-new-operations">Extending with New Operations<a class="headerlink" href="#extending-with-new-operations" title="Permanent link">¶</a></h3> 1754<p>You can register new <code>torch.*</code> functions to work with NestedTensor:</p> 1755<div class="language-python highlight"><table class="highlighttable"><tr><th colspan="2" class="filename"><span class="filename">Python</span></th></tr><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal">1</span> 1756<span class="normal">2</span> 1757<span class="normal">3</span> 1758<span class="normal">4</span> 1759<span class="normal">5</span> 1760<span class="normal">6</span> 1761<span class="normal">7</span> 1762<span class="normal">8</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="kn">
1762import</span><span class="w"> </span><span class="nn">torch</span> 1763<span class="kn">from</span><span class="w"> </span><span class="nn">danling.tensors.ops</span><span class="w"> </span><span class="kn">import</span> <span class="n">NestedTensorFuncRegistry</span> 1764 1765<span class="nd">@NestedTensorFuncRegistry</span><span class="o">.</span><span class="n">implement</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">my_custom_op</span><span class="p">)</span> 1766<span class="k">def</span><span class="w"> </span><span class="nf">my_custom_op</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span> 1767 <span class="c1"># Shape-preserving elementwise ops can reuse the reference structure.</span> 1768 <span class="n">values</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">my_custom_op</span><span class="p">(</span><span class="nb">input</span><span class="o">.</span><span class="n">concat</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span> 1769 <span class="k">return</span> <span class="nb">input</span><span class="o">.</span><span class="n">packed_like</span><span class="p">(</span><span class="n">values</span><span class="p">)</span> 1770</code></pre></div></td></tr></table></div> 1771<p>For ops that are purely elementwise on the packed data, register at the aten level instead:</p> 1772<div class="language-python highlight"><table class="highlighttable"><tr><th colspan="2" class="filename"><span class="filename">Python</span></th></tr><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal">1</span> 1773<span class="normal">2</span> 1774<span class="normal">3</span> 1775<span class="normal">4</span> 1776<span class="normal">5</span> 1777<span class="normal">6</span> 1778<span class="normal">7</span> 1779<span class="normal">8</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="kn">from</span><span class="w"> </span><span class="nn">danling.tensors.ops</span><span class="w"> </span><span class="kn">import</span> <span class="n">NestedTensorAtenRegistry</span> 1780<span class="n">aten</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span> 1781 1782<span class="nd">@NestedTensorAtenRegistry</span><span class="o">.</span><span class="n">implement</span><span class="p">(</span><span class="n">aten</span><span class="o">.</span><span class="n">my_op</span><span class="o">.</span><span class="n">default</span><span class="p">,</span> <span class="n">compile_safe</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span> 1783<span class="k">def</span><span class="w"> </span><span class="nf">_my_handler</span><span class="p">(</span><span class="n">func</span><span class="p">,</span> <span class="n">args</span><span class="p">,</span> <span class="n">kwargs</span><span class="p">):</span> 1784 <span class="n">source</span> <span class="o">=</span> <span class="n">args</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> 1785 <span class="n">values</span> <span class="o">=</span> <span class="n">func</span><span class="p">(</span><span class="n">source</span><span class="o">.</span><span class="n">concat</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">[</span><span class="mi">1</span><span class="p">:],</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span> 1786 <span class="k">return</span> <span class="n">source</span><span class="o">.</span><span class="n">packed_like</span><span class="p">(</span><span class="n">values</span><span class="p">)</span> 1787</code></pre></div></td></tr></table></div> 1788<p><code>packed_like</code> requires the packed output shape to remain unchanged. Operations 1789that retain the ragged lengths but replace or add a static feature tail can use 1790<code>packed_with_static_tail</code>:</p> 1791<div class="language-python highlight"><table class="highlighttable"><tr><th colspan="2" class="filename"><span class="filename">Python</span></th></tr><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal">1</span> 1792<span class="normal">2</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="n">atom_values</span> <span class="o">=</span> <span class="n">token_values</span><span class="o">.</span><span class="n">index_select</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="n">packed_atom_to_token</span><span class="p">)</span> 1793<span class="n">atoms</span> <span class="o">=</span> <span class="n">atom_mapping</span><span class="o">.</span><span class="n">packed_with_static_tail</span><span class="p">(</span><span class="n">atom_values</span><span class="p">)</span> 1794</code></pre></div></td></tr></table></div> 1795<p>For a canonical reference, the ragged dimensions are a leading logical prefix, 1796<code>packed_dim_order</code> is the identity, and the packed-value tail may have any 1797rank. The packed leading dimension is unchanged, while 1798<code>atom_values.shape[1:]</code> becomes the new static tail.</p> 1799<p>An explicit tensor-backed layout with one non-leading logical ragged dimension 1800may also replace its static dimensions, provided their packed rank does not 1801change. Tail sizes follow <code>packed_dim_order</code> and return to their original 1802logical positions. For example, sampled single representations remain
1803sample-major logically while token-major in packed storage:</p> 1804<div class="language-python highlight"><table class="highlighttable"><tr><th colspan="2" class="filename"><span class="filename">Python</span></th></tr><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal">1</span> 1805<span class="normal">2</span> 1806<span class="normal">3</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="n">single</span> <span class="o">=</span> <span class="n">NestedTensor</span><span class="p">(</span><span class="n">single_elements</span><span class="p">,</span> <span class="n">ragged_dims</span><span class="o">=</span><span class="p">(</span><span class="mi">1</span><span class="p">,))</span> <span class="c1"># each (S, N_i, C)</span> 1807<span class="n">output</span> <span class="o">=</span> <span class="n">single</span><span class="o">.</span><span class="n">packed_with_static_tail</span><span class="p">(</span><span class="n">output_values</span><span class="p">)</span> <span class="c1"># (sum N_i, S, C_out)</span> 1808<span class="c1"># each output element is (S, N_i, C_out), with packed_dim_order == (1, 0, 2)</span> 1809</code></pre></div></td></tr></table></div> 1810<p>Explicit ragged layouts retain tensor-backed row splits whenever 1811<code>packed_dim_order</code> begins with <code>ragged_dims</code>. This includes an explicit leading 1812single-ragged layout such as <code>ragged_dims=(0,)</code>; the ragged dimensions also need 1813not be a leading logical prefix. For example, elements shaped <code>(1, H, N_i, N_i)</code> 1814with <code>ragged_dims=(2, 3)</code> pack in <code>(2, 3, 0, 1)</code> order and remain tensor-backed. 1815These row splits are available through 1816<code>ragged_level_offsets(level)</code> and travel with <code>packed_like</code>, shape-preserving 1817static-tail operations, autograd transforms, and serialization. Because 1818per-sample lengths are tensor inputs rather than Python flatten metadata, 1819fixed-rank layouts such as <code>(N_i, N_i, C)</code> and <code>(S, N_i, C)</code> can reuse one 1820<code>torch.compile(dynamic=True)</code> graph across different <code>N_i</code>. This does not 1821expose the private child names or add a general packed constructor. Inferred 1822list construction without an explicit <code>ragged_dims</code> declaration keeps its 1823existing Python-metadata contract.</p> 1824<p>When an operator produces a new one-dimensional ragged topology, use a concrete 1825CPU integer lengths tensor:</p> 1826<div class="language-python highlight"><table class="highlighttable"><tr><th colspan="2" class="filename"><span class="filename">Python</span></th></tr><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal">1</span> 1827<span class="normal">2</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="n">token_values</span> <span class="o">=</span> <span class="n">segmented_mean</span><span class="p">(</span><span class="n">atom_values</span><span class="p">,</span> <span class="n">packed_atom_to_token</span><span class="p">)</span> 1828<span class="n">tokens</span> <span class="o">=</span> <span class="n">atoms</span><span class="o">.</span><span class="n">packed_with_lengths</span><span class="p">(</span><span class="n">token_values</span><span class="p">,</span> <span class="n">token_lengths</span><span class="p">)</span> 1829</code></pre></div></td></tr></table></div> 1830<p><code>packed_with_lengths</code> requires one non-negative length per batch element and 1831<code>token_lengths.sum() == token_values.shape[0]</code>. It constructs a canonical 1832leading ragged dimension and uses the remaining packed-value dimensions as the 1833static tail.</p> 1834<p>Square pair operators can rebuild two canonical ragged dimensions from the 1835same lengths without materializing Python element shapes:</p> 1836<div class="language-python highlight"><table class="highlighttable"><tr><th colspan="2" class="filename"><span class="filename">Python</span></th></tr><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal">1</span> 1837<span class="normal">2</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="n">pair</span> <span class="o">=</span> <span class="n">reference</span><span class="o">.</span><span class="n">packed_with_square_lengths</span><span class="p">(</span><span class="n">pair_values</span><span class="p">,</span> <span class="n">token_lengths</span><span class="p">)</span> 1838<span class="c1"># element i has shape (token_lengths[i], token_lengths[i], *pair_values.shape[1:])</span> 1839</code></pre></div></td></tr></table></div> 1840<p>Here <code>pair_values.shape[0]</code> must equal <code>token_lengths.square().sum()</code>. The 1841result persistently carries both CSR row-split levels as tensor metadata, so a 1842fixed batch size can reuse one dynamic compiled graph across different square 1843layouts. Rectangular pair operators use two independent length vectors:</p> 1844<div class="language-python highlight"><table class="highlighttable"><tr><th colspan="2" class="filename"><span class="filename">Python</span></th></tr><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal">1</span> 1845<span class="normal">2</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="n">distances</span> <span class="o">=</span> <span class="n">reference</span><span class="o">.</span><span class="n">packed_with_rectangular_lengths</span><span class="p">(</span><span class="n">distance_values</span><span class="p">,</span> <span class="n">query_lengths</span><span class="p">,</span> <span class="n">key_lengths</span><span class="p">)</span> 1846<span class="c1"># element i has shape (query_lengths[i], key_lengths[i], *distance_values.shape[1:])</span> 1847</code></pre></div></td></tr></table></div> 1848<p>Here the packed leading length is <code>(query_lengths * key_lengths).sum()</code> and 1849both ragged maxima remain tensor-backed graph inputs. All three length-based 1850reconstruction methods are zero-copy with respect 1851to their packed values and preserve dtype, device, strides, pinning, subclass, 1852runtime configuration, and autograd history. Reconstruction stores dynamic 1853lengths in tensor-backed offsets and shape metadata rather than Python tuples, 1854so its tracing cost does not grow with metadata rank. The current compiled 1855contract covers structural index consumers, runtime-validated same-layout 1856elementwise operations, static-tail broadcasting, and static-tail 1857normalization. Tensor-backed view/index remapping, padding, 1858<code>broadcast_tensors</code>, global-query <code>einsum</code>, and ragged-dimension softmax remain 1859staged; these paths raise an explicit compile error instead of materializing 1860Python metadata or silently assuming a layout.</p> 1861<p>For pairwise distances between explicit canonical <code>(P_i, M)</code> and <code>(R_i, M)</code> 1862elements, use <code>left.cdist(right)</code> in compiled code. A single-sample batch calls 1863native <code>cdist</code> directly; larger compiled batches use a registered segmented 1864operation that runs native <code>cdist</code> independently over the packed samples. 1865Eager mode issues the same native calls directly to avoid custom-op dispatcher 1866overhead. The result is 1867reconstructed as <code>(P_i, R_i)</code> through the rectangular metadata above. It does not 1868materialize padding, a cross-sample matrix, pair index tensors, or 1869<code>sum(P_i * R_i) x M</code> gathered operands. Eager <code>torch.cdist(left, right)</code> remains 1870supported, but the method form is the AOT/Inductor entry point because PyTorch 1871treats the built-in function as an opaque graph leaf.</p> 1872<p>For cumulative products along a canonical ragged dimension, use 1873<code>input.cumprod(dim)</code> in compiled training code. The segmented operator invokes 1874the current device’s native <code>cumprod</code> and backward kernel independently for
1875each packed sample. It therefore preserves native rounding, underflow, zero, 1876and non-finite behavior without padding or reassociating the product. Eager 1877<code>torch.cumprod(input, dim)</code> remains supported.</p> 1878<p>The result can remain a <code>NestedTensor</code> throughout a compiled model or cross a 1879compiled/eager training boundary as a wrapper-only output. DanLing preserves 1880the outer wrapper’s AOTAutograd edge and projects it back to <code>result.concat</code> 1881without padding or copying, so both wrapper outputs and direct packed outputs 1882remain differentiable. 1883They intentionally do not expose the general private packed constructor.</p> 1884<p><code>packed_offsets()</code> returns the boundaries of complete logical batch elements 1885in the flattened leading dimension of <code>concat</code>. It is the public operator 1886metadata interface for segmented kernels: for <code>(N_i, N_i, C)</code> elements it 1887returns cumulative <code>N_i * N_i</code> cell counts, whereas 1888<code>ragged_level_offsets(level)</code> returns row splits within the ragged hierarchy. 1889The same contract applies to single-ragged and non-leading-ragged layouts. 1890The canonical CPU integer tensor is returned without a copy; optional <code>device</code> 1891and <code>dtype</code> conversions are cached on the <code>NestedTensor</code> instance. An explicit 1892accelerator index is required for device caching; index-less device requests 1893retain PyTorch’s current-device semantics and are converted on every call.</p> 1894<p><code>element_sizes()</code> returns the exact logical shape of every batch element as a 1895CPU <code>torch.int64</code> tensor with shape <code>(batch_size, element_rank)</code>. Columns remain 1896in logical element-dimension order regardless of <code>batch_first</code> or 1897<code>packed_dim_order</code>, so zero-volume shapes such as <code>(0, 3)</code> and <code>(0, 7)</code> remain 1898distinguishable even when their packed offsets coincide. The method returns the 1899canonical tensor-backed metadata without a copy. It intentionally has no 1900device or dtype conversion arguments; consumers that need another placement or 1901integer width can apply <code>.to(...)</code> explicitly.</p> 1902<p><code>torch.repeat_interleave(input, repeats, dim=batch_dim)</code> and 1903<code>input.repeat_interleave(repeats, dim=batch_dim)</code> accept a non-negative integer 1904<code>repeats</code> and duplicate complete logical batch elements in their original 1905order. They repeat packed sample segments and every tensor-backed ragged 1906row-split level directly, without padding or per-element Python dispatch. 1907<code>input.repeat_batch(repeats)</code> is the equivalent explicit batch operation and 1908the canonical AOTAutograd-safe entry for compiled model code. Explicit single- 1909and multi-ragged layouts can reuse one 1910<code>torch.compile(fullgraph=True, dynamic=True)</code> graph across different ragged 1911lengths, and gradients from repeated samples accumulate into the original 1912packed values. Tensor-valued batch repeat counts are intentionally unsupported; 1913non-batch dimensions retain ordinary per-element <code>torch.repeat_interleave</code> 1914semantics.</p> 1915<p><code>packed_dim_order</code> exposes the read-only mapping from logical element dimensions 1916to physical packed-storage order when an operator needs to validate its layout.</p> 1917<p>Critical packed paths can use <code>nested_execution_guard</code> from <code>danling.tensors</code> in 1918tests or diagnostics to reject iteration, per-element fallback, padded 1919materialization, or dense repacking instead of silently accepting a slow path.</p> 1920<h2 id="benchmarks">Benchmarks<a class="headerlink" href="#benchmarks" title="Permanent link">¶</a></h2> 1921<p>Benchmarked on a single NVIDIA B200 180GB GPU with PyTorch 2.11, bfloat16.</p> 1922<p>Run with: <code>python scripts/benchmark_nested_tensor.py</code></p> 1923<h3 id="imdb-training">IMDB Training<a class="headerlink" href="#imdb-training" title="Permanent link">¶</a></h3> 1924<p>Real workload benchmark from <a href="https://github.com/ZhiyuanChen/DanLing/blob/f67c0f258da2d3ff42c6e850a250c1212f2ea5a9/examples/tensors/imdb.py"><code>examples/tensors/imdb.py</code></a>, using a BERT-large-shaped <code>torch.nn.TransformerEncoder</code> on IMDB with long variable-length sequences.</p> 1925<p>Config: <code>bert-large-uncased</code>, 2 epochs, batch size <code>32</code>, max length <code>8192</code>, <code>d_model=1024</code>, <code>nhead=16</code>, <code>num_layers=24</code></p> 1926<table> 1927<thead> 1928<tr> 1929<th style="text-align: left;">Metric</th> 1930<th style="text-align: left;">NestedTensor</th> 1931<th style="text-align: left;">Padded</th> 1932<th style="text-align: left;">Result</th> 1933</tr> 1934</thead> 1935<tbody> 1936<tr> 1937<td style="text-align: left;">Training step compute (forward + backward, all epochs)</td> 1938<td style="text-align: left;"><code>154819.4 ms</code></td> 1939<td style="text-align: left;"><code>306926.7 ms</code></td> 1940<td style="text-align: left;"><code>1.98x faster</code></td> 1941</tr> 1942<tr> 1943<td style="text-align: left;">Peak extra CUDA memory per training step</td> 1944<td style="text-align: left;"><code>12.68 GiB</code></td> 1945<td style="text-align: left;"><code>74.67 GiB</code></td> 1946<td style="text-align: left;"><code>83% lower</code></td> 1947</tr> 1948</tbody> 1949</table> 1950<p>This run measured nearly 2x faster model compute and an 83% reduction in peak extra CUDA memory for the NestedTensor path.</p> 1951<blockquote> 1952<p><strong>Note:</strong> This benchmark compares native PyTorch <code>nn.TransformerEncoder</code> execution on NestedTensor vs padded input. The timing is model forward+backward compute, not full end-to-e
1952nd wall clock including tokenization, data loading, or validation.</p> 1953</blockquote> 1954<h3 id="models">Models<a class="headerlink" href="#models" title="Permanent link">¶</a></h3> 1955<p>Synthetic model benchmarks covering <code>TransformerEncoder</code>, <code>TransformerDecoder</code>, <code>Transformer</code>, and <code>ResNet-50</code> across varying occupancy levels on a single NVIDIA B200 180GB GPU.</p> 1956<table> 1957<thead> 1958<tr> 1959<th style="text-align: left;">Model</th> 1960<th style="text-align: left;">Mode</th> 1961<th style="text-align: left;">Occ.</th> 1962<th style="text-align: left;">Padded (eager)</th> 1963<th style="text-align: left;">Padded (compiled)</th> 1964<th style="text-align: left;">DanLing (eager)</th> 1965<th style="text-align: left;">DanLing (compiled)</th> 1966<th style="text-align: left;">DL vs Padded</th> 1967<th style="text-align: left;">DL vs Compiled</th> 1968</tr> 1969</thead> 1970<tbody> 1971<tr> 1972<td style="text-align: left;">TransformerEncoder</td> 1973<td style="text-align: left;">Infer</td> 1974<td style="text-align: left;">20%</td> 1975<td style="text-align: left;">2.70 ms</td> 1976<td style="text-align: left;">39.98 ms</td> 1977<td style="text-align: left;">5.00 ms</td> 1978<td style="text-align: left;">1.10 ms</td> 1979<td style="text-align: left;">0.54x</td> 1980<td style="text-align: left;">7.99x</td> 1981</tr> 1982<tr> 1983<td style="text-align: left;">TransformerEncoder</td> 1984<td style="text-align: left;">Train</td> 1985<td style="text-align: left;">20%</td> 1986<td style="text-align: left;">26.95 ms</td> 1987<td style="text-align: left;">19.39 ms</td> 1988<td style="text-align: left;">8.60 ms</td> 1989<td style="text-align: left;">ERR ms</td> 1990<td style="text-align: left;">3.13x</td> 1991<td style="text-align: left;">2.25x</td> 1992</tr> 1993<tr> 1994<td style="text-align: left;">TransformerEncoder</td> 1995<td style="text-align: left;">Infer</td> 1996<td style="text-align: left;">35%</td> 1997<td style="text-align: left;">3.68 ms</td> 1998<td style="text-align: left;">38.93 ms</td> 1999<td style="text-align: left;">5.20 ms</td> 2000<td style="text-align: left;">1.79 ms</td> 2001<td style="text-align: left;">0.71x</td> 2002<td style="text-align: left;">7.48x</td> 2003</tr> 2004<tr> 2005<td style="text-align: left;">TransformerEncoder</td> 2006<td style="text-align: left;">Train</td> 2007<td style="text-align: left;">35%</td> 2008<td style="text-align: left;">27.05 ms</td> 2009<td style="text-align: left;">19.46 ms</td> 2010<td style="text-align: left;">11.55 ms</td> 2011<td style="text-align: left;">ERR ms</td> 2012<td style="text-align: left;">2.34x</td> 2013<td style="text-align: left;">1.68x</td> 2014</tr> 2015<tr> 2016<td style="text-align: left;">TransformerEncoder</td> 2017<td style="text-align: left;">Infer</td> 2018<td style="text-align: left;">77%</td> 2019<td style="text-align: left;">6.95 ms</td> 2020<td style="text-align: left;">36.39 ms</td> 2021<td style="text-align: left;">5.52 ms</td> 2022<td style="text-align: left;">4.45 ms</td> 2023<td style="text-align: left;">1.26x</td> 2024<td style="text-align: left;">6.59x</td> 2025</tr> 2026<tr> 2027<td style="text-align: left;">TransformerEncoder</td> 2028<td style="text-align: left;">Train</td> 2029<td style="text-align: left;">77%</td> 2030<td style="text-align: left;">27.15 ms</td> 2031<td style="text-align: left;">19.71 ms</td> 2032<td style="text-align: left;">21.37 ms</td> 2033<td style="text-align: left;">ERR ms</td> 2034<td style="text-align: left;">1.27x</td> 2035<td style="text-align: left;">0.92x</td> 2036</tr> 2037<tr> 2038<td style="text-align: left;">TransformerDecoder</td> 2039<td style="text-align: left;">Infer</td> 2040<td style="text-align: left;">20%</td> 2041<td style="text-align: left;">47.38 ms</td> 2042<td style="text-align: left;">10.86 ms</td> 2043<td style="text-align: left;">9.27 ms</td> 2044<td style="text-align: left;">1.79 ms</td> 2045<td style="text-align: left;">5.11x</td> 2046<td style="text-align: left;">1.17x</td> 2047</tr> 2048<tr> 2049<td style="text-align: left;">TransformerDecoder</td> 2050<td style="text-align: left;">Train</td> 2051<td style="text-align: left;">20%</td> 2052<td style="text-align: left;">45.86 ms</td> 2053<td style="text-align: left;">33.16 ms</td> 2054<td style="text-align: left;">15.00 ms</td> 2055<td style="text-align: left;">ERR ms</td> 2056<td style="text-align: left;">3.06x</td> 2057<td style="text-align: left;">2.21x</td> 2058</tr> 2059<tr> 2060<td style="text-align: left;">TransformerDecoder</td> 2061<td style="text-align: left;">Infer</td> 2062<td style="text-align: left;">35%</td> 2063<td style="text-align: left;">46.33 ms</td> 2064<td style="text-align: left;">10.91 ms</td> 2065<td style="text-align: left;">9.19 ms</td> 2066<td style="text-align: left;">2.91 ms</td> 2067<td style="text-align: left;">5.04x</td> 2068<td style="text-align: left;">1.19x</td> 2069</tr> 2070<tr> 2071<td style="text-align: left;">TransformerDecoder</td> 2072<td style="text-align: left;">Train</td> 2073<td style="text-align: left;">35%</td> 2074<td style="text-align: left;">45.98 ms</td> 2075<td style="text-align: left;">33.30 ms</td> 2076<td style="text-align: left;">19.67 ms</td> 2077<td style="text-align: left;">ERR ms</td> 2078<td style="text-align: left;">2.34x</td> 2079<td style="text-align: left;">1.69x</td> 2080</tr> 2081<tr> 2082<td style="text-align: left;">TransformerDecoder</td> 2083<td style="text-align: left;">Infer</td> 2084<td style="text-align: left;">77%</td> 2085<td style="text-align: left;">43.79 ms</td> 2086<td style="text-align: left;">11.02 ms</td> 2087<td style="text-align: left;">9.41 ms</td> 2088<td style="text-align: left;">7.47 ms</td> 2089<td style="text-align: left;">4.65x</td> 2090<td style="text-align: left;">1.17x</td> 2091</tr> 2092<tr> 2093<td style="text-align: left;">TransformerDecoder</td> 2094<td style="text-align: left;">Train</td> 2095<td style="text-align: left;">77%</td> 2096<td style="text-align: left;">46.18 ms</td> 2097<td style="text-align: left;">33.50 ms</td> 2098<td style="text-align: left;">36.00 ms</td> 2099<td style="text-align: left;">ERR ms</td> 2100<td style="text-align: left;">1.28x</td> 2101<td style="text-align: left;">0.93x</td> 2102</tr> 2103<tr> 2104<td style="text-align: left;">Transformer</td> 2105<td style="text-align: left;">Infer</td> 2106<td style="text-align: left;">21%</td> 2107<td style="text-align: left;">47.99 ms</td> 2108<td style="text-align: left;">48.51 ms</td> 2109<td style="text-align: left;">14.81 ms</td> 2110<td style="text-align: left;">2.93 ms</td> 2111<td style="text-align: left;">3.24x</td> 2112<td style="text-align: left;">3.27x</td> 2113</tr> 2114<tr> 2115<td style="text-align: left;">Transformer</td> 2116<td style="text-align: left;">Train</td> 2117<td style="text-align: left;">21%</td> 2118<td style="text-align: left;">68.79 ms</td> 2119<td style="text-align: left;">47.25 ms</td> 2120<td style="text-align: left;">24.53 ms</td> 2121<td style="text-align: left;">ERR ms</td> 2122<td style="text-align: left;">2.80x</td> 2123<td style="text-align: left;">1.93x</td> 2124</tr> 2125<tr> 2126<td style="text-align: left;">Transformer</td> 2127<td style="text-align: left;">Infer</td> 2128<td style="text-align: left;">40%</td> 2129<td style="text-align: left;">47.54 ms</td> 2130<td style="text-align: left;">47.52 ms</td> 2131<td style="text-align: left;">14.26 ms</td> 2132<td style="text-align: left;">5.39 ms</td> 2133<td style="text-align: left;">3.34x</td> 2134<td style="text-align: left;">3.33x</td> 2135</tr> 2136<tr> 2137<td style="text-align: left;">Transformer</td> 2138<td style="text-align: left;">Train</td> 2139<td style="text-align: left;">40%</td> 2140<td style="text-align: left;">69.03 ms</td> 2141<td style="text-align: left;">47.49 ms</td> 2142<td style="text-align: left;">32.72 ms</td> 2143<td style="text-align: left;">ERR ms</td> 2144<td style="text-align: left;">2.11x</td> 2145<td style="text-align: left;">1.45x</td> 2146</tr> 2147<tr> 2148<td style="text-align: left;">Transformer</td> 2149<td style="text-align: left;">Infer</td> 2150<td style="text-align: left;">84%</td> 2151<td style="text-align: left;">48.37 ms</td> 2152<td style="text-align: left;">45.12 ms</td> 2153<td style="text-align: left;">15.76 ms</td> 2154<td style="text-align: left;">12.51 ms</td> 2155<td style="text-align: left;">3.07x</td> 2156<td style="text-align: left;">2.86x</td> 2157</tr> 2158<tr> 2159<td style="text-align: left;">Transformer</td> 2160<td style="text-align: left;">Train</td> 2161<td style="text-align: left;">84%</td> 2162<td style="text-align: left;">69.52 ms</td> 2163<td style="text-align: left;">48.81 ms</td> 2164<td style="text-align: left;">59.88 ms</td> 2165<td style="text-align: left;">ERR ms</td> 2166<td style="text-align: left;">1.16x</td> 2167<td style="text-align: left;">0.82x</td> 2168</tr> 2169<tr> 2170<td style="text-align: left;">ResNet-50</td> 2171<td style="text-align: left;">Infer</td> 2172<td style="text-align: left;">41%</td> 2173<td style="text-align: left;">42.42 ms</td> 2174<td style="text-align: left;">ERR ms</td> 2175<td style="text-align: left;">219.49 ms</td> 2176<td style="text-align: left;">ERR ms</td> 2177<td style="text-align: left;">0.19x</td> 2178<td style="text-align: left;">N/A</td> 2179</tr> 2180<tr> 2181<td style="text-align: left;">ResNet-50</td> 2182<td style="text-align: left;">Train</td> 2183<td style="text-align: left;">41%</td> 2184<td style="text-align: left;">221.80 ms</td> 2185<td style="text-align: left;">ERR ms</td> 2186<td style="text-align: left;">513.08 ms</td> 2187<td style="text-align: left;">ERR ms</td> 2188<td style="text-align: left;">0.43x</td> 2189<td style="text-align: left;">N/A</td> 2190</tr> 2191<tr> 2192<td style="text-align: left;">ResNet-50</td> 2193<td style="text-align: left;">Infer</td> 2194<td style="text-align: left;">52%</td> 2195<td style="text-align: left;">42.56 ms</td> 2196<td style="text-align: left;">ERR ms</td> 2197<td style="text-align: left;">246.27 ms</td> 2198<td style="text-align: left;">ERR ms</td> 2199<td style="text-align: left;">0.17x</td> 2200<td style="text-align: left;">N/A</td> 2201</tr> 2202<tr> 2203<td style="text-align: left;">ResNet-50</td> 2204<td style="text-align: left;">Train</td> 2205<td style="text-align: left;">52%</td> 2206<td style="text-align: left;">221.37 ms</td> 2207<td style="text-align: left;">ERR ms</td> 2208<td style="text-align: left;">551.29 ms</td> 2209<td style="text-align: left;">ERR ms</td> 2210<td style="text-align: left;">0.40x</td> 2211<td style="text-align: left;">N/A</td> 2212</tr> 2213<tr> 2214<td style="text-align: left;">ResNet-50</td> 2215<td style="text-align: left;">Infer</td> 2216<td style="text-align: left;">81%</td> 2217<td style="text-align: left;">42.60 ms</td> 2218<td style="text-align: left;">ERR ms</td> 2219<td style="text-align: left;">318.03 ms</td> 2220<td style="text-align: left;">ERR ms</td> 2221<td style="text-align: left;">0.13x</td> 2222<td style="text-align: left;">N/A</td> 2223</tr> 2224<tr> 2225<td style="text-align: left;">ResNet-50</td> 2226<td style="text-align: left;">Train</td> 2227<td style="text-align: left;">81%</td> 2228<td style="text-align: left;">229.88 ms</td> 2229<td style="text-align: left;">ERR ms</td> 2230<td style="text-align: left;">709.83 ms</td> 2231<td style="text-align: left;">ERR ms</td> 2232<td style="text-align: left;">0.32x</td> 2233<td style="text-align: left;">N/A</td> 2234</tr> 2235</tbody> 2236</table> 2237<blockquote> 2238<p><strong>Note:</strong> ResNet-50 uses per-element dispatch (each image processed individually through conv/pool/BN layers). Inference is slower than padded due to per-element repacking overhead. BatchNorm statistics are computed correctly across all elements via concatenated storage.</p> 2239</blockquote> 2240<h3 id="operators">Operators<a class="headerlink" href="#operators" title="Permanent link">¶</a></h3> 2241<p>Synthetic operator benchmarks covering common transformer-hot ops and tensor primitives across padded tensors, DanLing <code>NestedTensor</code>, and <code>torch.nested</code>.</p> 2242<table> 2243<thead> 2244<tr> 2245<th style="text-align: left;">Operator</th> 2246<th style="text-align: left;">Occ.</th> 2247<th style="text-align: left;">Padded (eager)</th> 2248<th style="text-align: left;">Padded (compiled)</th> 2249<th style="text-align: left;">DanLing (eager)</th> 2250<th style="text-align: left;">DanLing (compiled)</th> 2251<th style="text-align: left;">torch.nested (eager)</th> 2252<th style="text-align: left;">torch.nested (compiled)</th> 2253<th style="text-align: left;">DL vs Padded</th> 2254<th style="text-align: left;">DL vs torch.nested</th> 2255</tr> 2256</thead> 2257<tbody> 2258<tr> 2259<td style="text-align: left;">F.linear</td> 2260<td style="text-align: left;">35%</td> 2261<td style="text-align: left;">0.05 ms</td> 2262<td style="text-align: left;">0.05 ms</td> 2263<td style="text-align: left;">0.13 ms</td> 2264<td style="text-align: left;">0.17 ms</td> 2265<td style="text-align: left;">0.13 ms</td> 2266<td style="text-align: left;">0.39 ms</td> 2267<td style="text-align: left;">0.27x</td> 2268<td style="text-align: left;">2.29x</td> 2269</tr> 2270<tr> 2271<td style="text-align: left;">F.layer_norm</td> 2272<td style="text-align: left;">35%</td> 2273<td style="text-align: left;">0.17 ms</td> 2274<td style="text-align: left;">0.05 ms</td> 2275<td style="text-align: left;">0.10 ms</td> 2276<td style="text-align: left;">0.16 ms</td> 2277<td style="text-align: left;">0.20 ms</td> 2278<td style="text-align: left;">0.38 ms</td> 2279<td style="text-align: left;">0.28x</td> 2280<td style="text-align: left;">2.36x</td> 2281</tr> 2282<tr> 2283<td style="text-align: left;">F.relu</td> 2284<td style="text-align: left;">35%</td> 2285<td style="text-align: left;">0.05 ms</td> 2286<td style="text-align: left;">0.05 ms</td> 2287<td style="text-align: left;">0.09 ms</td> 2288<td style="text-align: left;">0.16 ms</td> 2289<td style="text-align: left;">0.10 ms</td> 2290<td style="text-align: left;">0.37 ms</td> 2291<td style="text-align: left;">0.29x</td> 2292<td style="text-align: left;">2.38x</td> 2293</tr> 2294<tr> 2295<td style="text-align: left;">F.gelu</td> 2296<td style="text-align: left;">35%</td> 2297<td style="text-align: left;">0.08 ms</td> 2298<td style="text-align: left;">0.08 ms</td> 2299<td style="text-align: left;">0.08 ms</td> 2300<td style="text-align: left;">0.16 ms</td> 2301<td style="text-align: left;">0.10 ms</td> 2302<td style="text-align: left;">0.36 ms</td> 2303<td style="text-align: left;">0.51x</td> 2304<td style="text-align: left;">2.33x</td> 2305</tr> 2306<tr> 2307<td style="text-align: left;">F.softmax</td> 2308<td style="text-align: left;">35%</td> 2309<td style="text-align: left;">0.12 ms</td> 2310<td style="text-align: left;">0.06 ms</td> 2311<td style="text-align: left;">0.11 ms</td> 2312<td style="text-align: left;">0.16 ms</td> 2313<td style="text-align: left;">0.15 ms</td> 2314<td style="text-align: left;">0.37 ms</td> 2315<td style="text-align: left;">0.36x</td> 2316<td style="text-align: left;">2.31x</td> 2317</tr> 2318<tr> 2319<td style="text-align: left;">F.embedding</td> 2320<td style="text-align: left;">35%</td> 2321<td style="text-align: left;">0.04 ms</td> 2322<td style="text-align: left;">0.04 ms</td> 2323<td style="text-align: left;">0.15 ms</td> 2324<td style="text-align: left;">0.16 ms</td> 2325<td style="text-align: left;">0.14 ms</td> 2326<td style="text-align: left;">0.36 ms</td> 2327<td style="text-align: left;">0.25x</td> 2328<td style="text-align: left;">2.27x</td> 2329</tr> 2330<tr> 2331<td style="text-align: left;">torch.matmul</td> 2332<td style="text-align: left;">35%</td> 2333<td style="text-align: left;">0.04 ms</td> 2334<td style="text-align: left;">0.05 ms</td> 2335<td style="text-align: left;">0.17 ms</td> 2336<td style="text-align: left;">0.17 ms</td> 2337<td style="text-align: left;">0.15 ms</td> 2338<td style="text-align: left;">0.39 ms</td> 2339<td style="text-align: left;">0.28x</td> 2340<td style="text-align: left;">2.31x</td> 2341</tr> 2342<tr> 2343<td style="text-align: left;">torch.add</td> 2344<td style="text-align: left;">35%</td> 2345<td style="text-align: left;">0.05 ms</td> 2346<td style="text-align: left;">0.05 ms</td> 2347<td style="text-align: left;">0.12 ms</td> 2348<td style="text-align: left;">0.15 ms</td> 2349<td style="text-align: left;">0.11 ms</td> 2350<td style="text-align: left;">0.36 ms</td> 2351<td style="text-align: left;">0.29x</td> 2352<td style="text-align: left;">2.36x</td> 2353</tr> 2354</tbody> 2355</table> 2356<h2 id="danling.tensors">Tensor API<a class="headerlink" href="#danling.tensors" title="Permanent link">¶</a></h2> 2357<ul> 2358<li><a class="autorefs autorefs-internal" title="<
2358code class="doc-symbol doc-symbol-heading doc-symbol-class"></code> <span class="doc doc-object-name doc-class-name">NestedTensor</span>" href="nested_tensor/#danling.tensors.nested_tensor.NestedTensor">NestedTensor</a>: variable-length tensor batches and their operations.</li> 2359<li><a class="autorefs autorefs-internal" title="<code class="doc-symbol doc-symbol-heading doc-symbol-class"></code> <span class="doc doc-object-name doc-class-name">PNTensor</span>" href="pn_tensor/#danling.tensors.pn_tensor.PNTensor">PNTensor</a> and <a class="autorefs autorefs-internal" title="<code class="doc-symbol doc-symbol-heading doc-symbol-function"></code> <span class="doc doc-object-name doc-function-name">tensor</span>" href="pn_tensor/#danling.tensors.pn_tensor.tensor">tensor</a>: tensor markers for collation.</li> 2360<li><a href="functions/">Dispatch functions and registries</a>: supported dispatch layers and extension interfaces.</li> 2361</ul> 2362 2363 2364 2365 2366 2367 2368 2369 2370 2371 2372 2373 2374 2375 <aside class="md-source-file"> 2376 2377 2378 <span class="md-source-file__fact"> 2379 <span class="md-icon" title="Last update"> 2380 <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24"><path d="M21 13.1c-.1 0-.3.1-.4.2l-1 1 2.1 2.1 1-1c.2-.2.2-.6 0-.8l-1.3-1.3c-.1-.1-.2-.2-.4-.2m-1.9 1.8-6.1 6V23h2.1l6.1-6.1zM12.5 7v5.2l4 2.4-1 1L11 13V7zM11 21.9c-5.1-.5-9-4.8-9-9.9C2 6.5 6.5 2 12 2c5.3 0 9.6 4.1 10 9.3-.3-.1-.6-.2-1-.2s-.7.1-1 .2C19.6 7.2 16.2 4 12 4c-4.4 0-8 3.6-8 8 0 4.1 3.1 7.5 7.1 7.9l-.1.2z"/></svg> 2381 </span> 2382 <span class="git-revision-date-localized-plugin git-revision-date-localized-plugin-iso_datetime" title="September 9, 2026 09:46:02 UTC">2026-09-09 09:46:02</span> 2383 </span> 2384 2385 2386 2387 2388 2389 </aside> 2390 2391 2392 2393 2394 2395 2396 2397 2398 2399<!-- Giscus -->
2400<script 2401 src="https://giscus.app/client.js" 2402 data-repo="ZhiyuanChen/DanLing" 2403 data-repo-id="R_kgDOGv-RWA" 2404 data-category="General" 2405 data-category-id="DIC_kwDOGv-RWM4CTx_d" 2406 data-mapping="pathname" 2407 data-strict="0" 2408 data-reactions-enabled="1" 2409 data-emit-metadata="1" 2410 data-input-position="top" 2411 data-theme="preferred_color_scheme" 2412 data-lang="en" 2413 data-loading="lazy" 2414 crossorigin="anonymous" 2415 async 2416></script>
2416 2417 2418<!-- Reload on palette change -->
2419<script> 2420 var palette = __md_get("__palette"); 2421 if (palette && typeof palette.color === "object") 2422 if (palette.color.scheme === "slate") { 2423 var giscus = document.querySelector("script[src*=giscus]"); 2424 giscus.setAttribute("data-theme", "dark"); 2425 } 2426 2427 /* Register event handlers after documented loaded */ 2428 document.addEventListener("DOMContentLoaded", function () { 2429 var ref = document.querySelector("[data-md-component=palette]"); 2430 ref.addEventListener("change", function () { 2431 var palette = __md_get("__palette"); 2432 if (palette && typeof palette.color === "object") { 2433 var theme = palette.color.scheme === "slate" ? "dark" : "light"; 2434 2435 /* Instruct Giscus to change theme */ 2436 var frame = document.querySelector(".giscus-frame"); 2437 frame.contentWindow.postMessage( 2438 { giscus: { setConfig: { theme } } }, 2439 "https://giscus.app", 2440 ); 2441 } 2442 }); 2443 }); 2444</script>
2444 2445 2446 </article> 2447 </div> 2448 2449 2450
2450<script>var tabs=__md_get("__tabs");if(Array.isArray(tabs))e:for(var set of document.querySelectorAll(".tabbed-set")){var labels=set.querySelector(".tabbed-labels");for(var tab of tabs)for(var label of labels.getElementsByTagName("label"))if(label.innerText.trim()===tab){var input=document.getElementById(label.htmlFor);input.checked=!0;continue e}}</script>
2450 2451
2452<script>var target=document.getElementById(location.hash.slice(1));target&&target.name&&(target.checked=target.name.startsWith("__tabbed_"))</script>
2452 2453 </div> 2454 2455 <button type="button" class="md-top md-icon" data-md-component="top" hidden> 2456 2457 <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24"><path d="M13 20h-2V8l-5.5 5.5-1.42-1.42L12 4.16l7.92 7.92-1.42 1.42L13 8z"/></svg> 2458 Back to top 2459</button> 2460 2461 </main> 2462 2463 <footer class="md-footer"> 2464 2465 2466 2467 <nav class="md-footer__inner md-grid" aria-label="Footer" > 2468 2469 2470 <a href="../runners/utils/" class="md-footer__link md-footer__link--prev" aria-label="Previous: Utilities"> 2471 <div class="md-footer__button md-icon"> 2472 2473 <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24"><path d="M20 11v2H8l5.5 5.5-1.42 1.42L4.16 12l7.92-7.92L13.5 5.5 8 11z"/></svg> 2474 </div> 2475 <div class="md-footer__title"> 2476 <span class="md-footer__direction"> 2477 Previous 2478 </span> 2479 <div class="md-ellipsis"> 2480 Utilities 2481 </div> 2482 </div> 2483 </a> 2484 2485 2486 2487 <a href="nested_tensor/" class="md-footer__link md-footer__link--next" aria-label="Next: NestedTensor"> 2488 <div class="md-footer__title"> 2489 <span class="md-footer__direction"> 2490 Next 2491 </span> 2492 <div class="md-ellipsis"> 2493 NestedTensor 2494 </div> 2495 </div> 2496 <div class="md-footer__button md-icon"> 2497 2498 <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24"><path d="M4 11v2h12l-5.5 5.5 1.42 1.42L19.84 12l-7.92-7.92L10.5 5.5 16 11z"/></svg> 2499 </div> 2500 </a> 2501 2502 </nav> 2503 2504 2505 <div class="md-footer-meta md-typeset"> 2506 <div class="md-footer-meta__inner md-grid"> 2507 <div class="md-copyright"> 2508 2509 <div class="md-copyright__highlight"> 2510 All rights reserved © 2021-Present, DanLing Team 2511 </div> 2512 2513 2514</div> 2515 2516 </div> 2517 </div> 2518</footer> 2519 2520 </div> 2521 <div class="md-dialog" data-md-component="dialog"> 2522 <div class="md-dialog__inner md-typeset"></div> 2523 </div> 2524 2525 2526 <div class="md-consent" data-md-component="consent" id="__consent" hidden> 2527 <div class="md-consent__overlay"></div> 2528 <aside class="md-consent__inner"> 2529 <form class="md-consent__form md-grid md-typeset" name="consent"> 2530 2531 2532 2533 2534 2535 2536<h4>Cookie Consent</h4> 2537<p>We use cookies to recognize your repeated visits and preferences, as well as to measure the effectiveness of our documentation and whether users find what they're searching for. With your consent, you're helping us to make our documentation better. Please check out our <a href="https://danling.org/about/privacy">Privacy Notice</a> for more information.</p> 2538<input class="md-toggle" type="checkbox" id="__settings" > 2539<div class="md-consent__settings"> 2540 <ul class="task-list"> 2541 2542 2543 2544 2545 2546 2547 2548 2549 2550 <li class="task-list-item"> 2551 <label class="task-list-control"> 2552 <input type="checkbox" name="analytics" checked> 2553 <span class="task-list-indicator"></span> 2554 Google Analytics 2555 </label> 2556 </li> 2557 2558 2559 2560 2561 2562 2563 2564 2565 2566 2567 2568 <li class="task-list-item"> 2569 <label class="task-list-control"> 2570 <input type="checkbox" name="github" checked> 2571 <span class="task-list-indicator"></span> 2572 GitHub 2573 </label> 2574 </li> 2575 2576 2577 2578 2579 </ul> 2580</div> 2581<div class="md-consent__controls"> 2582 2583 2584 <button class="md-button md-button--primary">Accept</button> 2585 2586 2587 2588 2589 2590 2591 2592 <label class="md-button" for="__settings">Manage settings</label> 2593 2594 2595</div> 2596 </form> 2597 </aside> 2598 </div> 2599
2599<script>var consent=__md_get("__consent");if(consent)for(var input of document.forms.consent.elements)input.name&&(input.checked=consent[input.name]||!1);else"file:"!==location.protocol&&setTimeout((function(){document.querySelector("[data-md-component=consent]").hidden=!1}),250);var form=document.forms.consent;for(var action of["submit","reset"])form.addEventListener(action,(function(e){if(e.preventDefault(),"reset"===e.type)for(var n of document.forms.consent.elements)n.name&&(n.checked=!1);__md_set("__consent",Object.fromEntries(Array.from(new FormData(form).keys()).map((function(e){return[e,!0]})))),location.hash="",location.reload()}))</script>
2599 2600 2601 2602 2603 2604
2604<script id="__config" type="application/json">{"annotate": null, "base": "..", "features": ["announce.dismiss", "content.code.annotate", "content.code.copy", "content.code.select", "content.tabs.link", "content.tooltips", "header.autohide", "navigation.expand", "navigation.footer", "navigation.indexes", "navigation.prune", "navigation.sections", "navigation.tabs", "navigation.tabs.sticky", "navigation.top", "navigation.tracking", "search.highlight", "search.share", "search.suggest", "toc.follow"], "search": "../assets/javascripts/workers/search.2c215733.min.js", "tags": null, "translations": {"clipboard.copied": "Copied to clipboard", "clipboard.copy": "Copy to clipboard", "search.result.more.one": "1 more on this page", "search.result.more.other": "# more on this page", "search.result.none": "No matching documents", "search.result.one": "1 matching document", "search.result.other": "# matching documents", "search.result.placeholder": "Type to start searching", "search.result.term.missing": "Missing", "select.version": "Select version"}, "version": null}</script>
2604 2605 2606 2607
2607<script src="../assets/javascripts/bundle.d7400e89.min.js"></script>
2607 2608 2609
2609<script src="../javascripts/mathjax.js"></script>
2609 2610 2611
2611<script src="https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js"></script>
2611 2612 2613
2613<script src="../javascripts/shortcuts.js"></script>
2613 2614 2615 2616
2616<script type="module" src="https://static.cloudflareinsights.com/beacon.min.js/v31edd6df95cf4e85bb4c19e7a9bdbcba1788362987495" integrity="sha512-iIg7k2xntmwu6/uSb5tpc/hySgZc4eoL31yB29W6tJFo2akwjPWcEqnCEdJvGexCL0KEQwVYv5BlowfhVz26hg==" data-cf-beacon='{"version":"2024.11.0","token":"9c892eda43914cc59e37016c23ab0870","r":1,"spa":2}' crossorigin="anonymous"></script>
2616 2617</body> 2618</html>
Line numbers count LF bytes from the start of the resource, as the search results do. Vendor segments are library code the classifier recognised; they are stored but not indexed. Bytes are shown as Latin1 characters, one per byte.