mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-01-09 15:08:02 -05:00
3888 lines
300 KiB
HTML
3888 lines
300 KiB
HTML
|
|
<!doctype html>
|
|
<html lang="en" class="no-js">
|
|
<head>
|
|
|
|
<meta charset="utf-8">
|
|
<meta name="viewport" content="width=device-width,initial-scale=1">
|
|
|
|
|
|
|
|
<link rel="canonical" href="https://docs.tinygrad.org/tensor/movement/">
|
|
|
|
|
|
<link rel="prev" href="../creation/">
|
|
|
|
|
|
<link rel="next" href="../elementwise/">
|
|
|
|
|
|
|
|
|
|
|
|
<link rel="icon" href="../../favicon.svg">
|
|
<meta name="generator" content="mkdocs-1.6.1, mkdocs-material-9.7.1">
|
|
|
|
|
|
|
|
<title>Movement - tinygrad docs</title>
|
|
|
|
|
|
|
|
<link rel="stylesheet" href="../../assets/stylesheets/main.484c7ddc.min.css">
|
|
|
|
|
|
<link rel="stylesheet" href="../../assets/stylesheets/palette.ab4e12ef.min.css">
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
<link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
|
|
<link rel="stylesheet" href="https://fonts.googleapis.com/css?family=Roboto:300,300i,400,400i,700,700i%7CRoboto+Mono:400,400i,700,700i&display=fallback">
|
|
<style>:root{--md-text-font:"Roboto";--md-code-font:"Roboto Mono"}</style>
|
|
|
|
|
|
|
|
<link rel="stylesheet" href="../../assets/_markdown_exec_pyodide.css">
|
|
|
|
<link rel="stylesheet" href="../../assets/_markdown_exec_ansi.css">
|
|
|
|
<link rel="stylesheet" href="../../assets/_mkdocstrings.css">
|
|
|
|
<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>
|
|
|
|
|
|
|
|
|
|
|
|
</head>
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
<body dir="ltr" data-md-color-scheme="default" data-md-color-primary="black" data-md-color-accent="lime">
|
|
|
|
|
|
<input class="md-toggle" data-md-toggle="drawer" type="checkbox" id="__drawer" autocomplete="off">
|
|
<input class="md-toggle" data-md-toggle="search" type="checkbox" id="__search" autocomplete="off">
|
|
<label class="md-overlay" for="__drawer"></label>
|
|
<div data-md-component="skip">
|
|
|
|
|
|
<a href="#movement-low-level" class="md-skip">
|
|
Skip to content
|
|
</a>
|
|
|
|
</div>
|
|
<div data-md-component="announce">
|
|
|
|
</div>
|
|
|
|
|
|
|
|
|
|
|
|
|
|
<header class="md-header md-header--shadow" data-md-component="header">
|
|
<nav class="md-header__inner md-grid" aria-label="Header">
|
|
<a href="../.." title="tinygrad docs" class="md-header__button md-logo" aria-label="tinygrad docs" data-md-component="logo">
|
|
|
|
<img src="../../logo_tiny_dark.svg" alt="logo">
|
|
|
|
</a>
|
|
<label class="md-header__button md-icon" for="__drawer">
|
|
|
|
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24"><path d="M3 6h18v2H3zm0 5h18v2H3zm0 5h18v2H3z"/></svg>
|
|
</label>
|
|
<div class="md-header__title" data-md-component="header-title">
|
|
<div class="md-header__ellipsis">
|
|
<div class="md-header__topic">
|
|
<span class="md-ellipsis">
|
|
tinygrad docs
|
|
</span>
|
|
</div>
|
|
<div class="md-header__topic" data-md-component="header-topic">
|
|
<span class="md-ellipsis">
|
|
|
|
Movement
|
|
|
|
</span>
|
|
</div>
|
|
</div>
|
|
</div>
|
|
|
|
|
|
<form class="md-header__option" data-md-component="palette">
|
|
|
|
|
|
|
|
|
|
<input class="md-option" data-md-color-media="(prefers-color-scheme)" data-md-color-scheme="default" data-md-color-primary="black" data-md-color-accent="lime" aria-label="Switch to light mode" type="radio" name="__palette" id="__palette_0">
|
|
|
|
<label class="md-header__button md-icon" title="Switch to light mode" for="__palette_1" hidden>
|
|
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24"><path d="m14.3 16-.7-2h-3.2l-.7 2H7.8L11 7h2l3.2 9zM20 8.69V4h-4.69L12 .69 8.69 4H4v4.69L.69 12 4 15.31V20h4.69L12 23.31 15.31 20H20v-4.69L23.31 12zm-9.15 3.96h2.3L12 9z"/></svg>
|
|
</label>
|
|
|
|
|
|
|
|
|
|
|
|
<input class="md-option" data-md-color-media="(prefers-color-scheme: light)" data-md-color-scheme="default" data-md-color-primary="black" data-md-color-accent="lime" aria-label="Switch to dark mode" type="radio" name="__palette" id="__palette_1">
|
|
|
|
<label class="md-header__button md-icon" title="Switch to dark mode" for="__palette_2" hidden>
|
|
<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>
|
|
</label>
|
|
|
|
|
|
|
|
|
|
|
|
<input class="md-option" data-md-color-media="(prefers-color-scheme: dark)" data-md-color-scheme="slate" data-md-color-primary="black" data-md-color-accent="lime" aria-label="Switch to system preference" type="radio" name="__palette" id="__palette_2">
|
|
|
|
<label class="md-header__button md-icon" title="Switch to system preference" for="__palette_0" hidden>
|
|
<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>
|
|
</label>
|
|
|
|
|
|
</form>
|
|
|
|
|
|
|
|
<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>
|
|
|
|
|
|
|
|
|
|
|
|
<label class="md-header__button md-icon" for="__search">
|
|
|
|
<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>
|
|
</label>
|
|
<div class="md-search" data-md-component="search" role="dialog">
|
|
<label class="md-search__overlay" for="__search"></label>
|
|
<div class="md-search__inner" role="search">
|
|
<form class="md-search__form" name="search">
|
|
<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>
|
|
<label class="md-search__icon md-icon" for="__search">
|
|
|
|
<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>
|
|
|
|
<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>
|
|
</label>
|
|
<nav class="md-search__options" aria-label="Search">
|
|
|
|
<button type="reset" class="md-search__icon md-icon" title="Clear" aria-label="Clear" tabindex="-1">
|
|
|
|
<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>
|
|
</button>
|
|
</nav>
|
|
|
|
<div class="md-search__suggest" data-md-component="search-suggest"></div>
|
|
|
|
</form>
|
|
<div class="md-search__output">
|
|
<div class="md-search__scrollwrap" tabindex="0" data-md-scrollfix>
|
|
<div class="md-search-result" data-md-component="search-result">
|
|
<div class="md-search-result__meta">
|
|
Initializing search
|
|
</div>
|
|
<ol class="md-search-result__list" role="presentation"></ol>
|
|
</div>
|
|
</div>
|
|
</div>
|
|
</div>
|
|
</div>
|
|
|
|
|
|
|
|
<div class="md-header__source">
|
|
<a href="https://github.com/tinygrad/tinygrad/" title="Go to repository" class="md-source" data-md-component="source">
|
|
<div class="md-source__icon md-icon">
|
|
|
|
<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>
|
|
</div>
|
|
<div class="md-source__repository">
|
|
GitHub
|
|
</div>
|
|
</a>
|
|
</div>
|
|
|
|
</nav>
|
|
|
|
</header>
|
|
|
|
<div class="md-container" data-md-component="container">
|
|
|
|
|
|
|
|
|
|
|
|
|
|
<main class="md-main" data-md-component="main">
|
|
<div class="md-main__inner md-grid">
|
|
|
|
|
|
|
|
<div class="md-sidebar md-sidebar--primary" data-md-component="sidebar" data-md-type="navigation" >
|
|
<div class="md-sidebar__scrollwrap">
|
|
<div class="md-sidebar__inner">
|
|
|
|
|
|
|
|
|
|
|
|
|
|
<nav class="md-nav md-nav--primary md-nav--integrated" aria-label="Navigation" data-md-level="0">
|
|
<label class="md-nav__title" for="__drawer">
|
|
<a href="../.." title="tinygrad docs" class="md-nav__button md-logo" aria-label="tinygrad docs" data-md-component="logo">
|
|
|
|
<img src="../../logo_tiny_dark.svg" alt="logo">
|
|
|
|
</a>
|
|
tinygrad docs
|
|
</label>
|
|
|
|
<div class="md-nav__source">
|
|
<a href="https://github.com/tinygrad/tinygrad/" title="Go to repository" class="md-source" data-md-component="source">
|
|
<div class="md-source__icon md-icon">
|
|
|
|
<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>
|
|
</div>
|
|
<div class="md-source__repository">
|
|
GitHub
|
|
</div>
|
|
</a>
|
|
</div>
|
|
|
|
<ul class="md-nav__list" data-md-scrollfix>
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
<li class="md-nav__item md-nav__item--active md-nav__item--section md-nav__item--nested">
|
|
|
|
|
|
|
|
<input class="md-nav__toggle md-toggle " type="checkbox" id="__nav_1" checked>
|
|
|
|
|
|
<div class="md-nav__link md-nav__container">
|
|
<a href="../.." class="md-nav__link ">
|
|
|
|
|
|
|
|
<span class="md-ellipsis">
|
|
|
|
|
|
Home
|
|
|
|
|
|
|
|
</span>
|
|
|
|
|
|
|
|
</a>
|
|
|
|
|
|
<label class="md-nav__link " for="__nav_1" id="__nav_1_label" tabindex="">
|
|
<span class="md-nav__icon md-icon"></span>
|
|
</label>
|
|
|
|
</div>
|
|
|
|
<nav class="md-nav" data-md-level="1" aria-labelledby="__nav_1_label" aria-expanded="true">
|
|
<label class="md-nav__title" for="__nav_1">
|
|
<span class="md-nav__icon md-icon"></span>
|
|
|
|
|
|
Home
|
|
|
|
|
|
</label>
|
|
<ul class="md-nav__list" data-md-scrollfix>
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
<li class="md-nav__item">
|
|
<a href="../../quickstart/" class="md-nav__link">
|
|
|
|
|
|
|
|
<span class="md-ellipsis">
|
|
|
|
|
|
Quickstart
|
|
|
|
|
|
|
|
</span>
|
|
|
|
|
|
|
|
</a>
|
|
</li>
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
<li class="md-nav__item">
|
|
<a href="../../showcase/" class="md-nav__link">
|
|
|
|
|
|
|
|
<span class="md-ellipsis">
|
|
|
|
|
|
Showcase
|
|
|
|
|
|
|
|
</span>
|
|
|
|
|
|
|
|
</a>
|
|
</li>
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
<li class="md-nav__item">
|
|
<a href="../../mnist/" class="md-nav__link">
|
|
|
|
|
|
|
|
<span class="md-ellipsis">
|
|
|
|
|
|
MNIST Tutorial
|
|
|
|
|
|
|
|
</span>
|
|
|
|
|
|
|
|
</a>
|
|
</li>
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
<li class="md-nav__item md-nav__item--active md-nav__item--nested">
|
|
|
|
|
|
|
|
<input class="md-nav__toggle md-toggle " type="checkbox" id="__nav_1_5" checked>
|
|
|
|
|
|
<label class="md-nav__link" for="__nav_1_5" id="__nav_1_5_label" tabindex="0">
|
|
|
|
|
|
|
|
<span class="md-ellipsis">
|
|
|
|
|
|
API Reference
|
|
|
|
|
|
|
|
</span>
|
|
|
|
|
|
|
|
<span class="md-nav__icon md-icon"></span>
|
|
</label>
|
|
|
|
<nav class="md-nav" data-md-level="2" aria-labelledby="__nav_1_5_label" aria-expanded="true">
|
|
<label class="md-nav__title" for="__nav_1_5">
|
|
<span class="md-nav__icon md-icon"></span>
|
|
|
|
|
|
API Reference
|
|
|
|
|
|
</label>
|
|
<ul class="md-nav__list" data-md-scrollfix>
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
<li class="md-nav__item md-nav__item--active md-nav__item--nested">
|
|
|
|
|
|
|
|
<input class="md-nav__toggle md-toggle " type="checkbox" id="__nav_1_5_1" checked>
|
|
|
|
|
|
<div class="md-nav__link md-nav__container">
|
|
<a href="../" class="md-nav__link ">
|
|
|
|
|
|
|
|
<span class="md-ellipsis">
|
|
|
|
|
|
Tensor
|
|
|
|
|
|
|
|
</span>
|
|
|
|
|
|
|
|
</a>
|
|
|
|
|
|
<label class="md-nav__link " for="__nav_1_5_1" id="__nav_1_5_1_label" tabindex="0">
|
|
<span class="md-nav__icon md-icon"></span>
|
|
</label>
|
|
|
|
</div>
|
|
|
|
<nav class="md-nav" data-md-level="3" aria-labelledby="__nav_1_5_1_label" aria-expanded="true">
|
|
<label class="md-nav__title" for="__nav_1_5_1">
|
|
<span class="md-nav__icon md-icon"></span>
|
|
|
|
|
|
Tensor
|
|
|
|
|
|
</label>
|
|
<ul class="md-nav__list" data-md-scrollfix>
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
<li class="md-nav__item">
|
|
<a href="../properties/" class="md-nav__link">
|
|
|
|
|
|
|
|
<span class="md-ellipsis">
|
|
|
|
|
|
Properties
|
|
|
|
|
|
|
|
</span>
|
|
|
|
|
|
|
|
</a>
|
|
</li>
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
<li class="md-nav__item">
|
|
<a href="../creation/" class="md-nav__link">
|
|
|
|
|
|
|
|
<span class="md-ellipsis">
|
|
|
|
|
|
Creation
|
|
|
|
|
|
|
|
</span>
|
|
|
|
|
|
|
|
</a>
|
|
</li>
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
<li class="md-nav__item md-nav__item--active">
|
|
|
|
<input class="md-nav__toggle md-toggle" type="checkbox" id="__toc">
|
|
|
|
|
|
|
|
<label class="md-nav__link md-nav__link--active" for="__toc">
|
|
|
|
|
|
|
|
<span class="md-ellipsis">
|
|
|
|
|
|
Movement
|
|
|
|
|
|
|
|
</span>
|
|
|
|
|
|
|
|
<span class="md-nav__icon md-icon"></span>
|
|
</label>
|
|
|
|
<a href="./" class="md-nav__link md-nav__link--active">
|
|
|
|
|
|
|
|
<span class="md-ellipsis">
|
|
|
|
|
|
Movement
|
|
|
|
|
|
|
|
</span>
|
|
|
|
|
|
|
|
</a>
|
|
|
|
|
|
|
|
<nav class="md-nav md-nav--secondary" aria-label="Table of contents">
|
|
|
|
|
|
|
|
|
|
<label class="md-nav__title" for="__toc">
|
|
<span class="md-nav__icon md-icon"></span>
|
|
Table of contents
|
|
</label>
|
|
<ul class="md-nav__list" data-md-component="toc" data-md-scrollfix>
|
|
|
|
<li class="md-nav__item">
|
|
<a href="#movement-low-level" class="md-nav__link">
|
|
<span class="md-ellipsis">
|
|
|
|
Movement (low level)
|
|
|
|
</span>
|
|
</a>
|
|
|
|
<nav class="md-nav" aria-label="Movement (low level)">
|
|
<ul class="md-nav__list">
|
|
|
|
<li class="md-nav__item">
|
|
<a href="#tinygrad.Tensor.view" class="md-nav__link">
|
|
<span class="md-ellipsis">
|
|
|
|
<code class="doc-symbol doc-symbol-toc doc-symbol-method"></code> view
|
|
|
|
</span>
|
|
</a>
|
|
|
|
</li>
|
|
|
|
<li class="md-nav__item">
|
|
<a href="#tinygrad.Tensor.reshape" class="md-nav__link">
|
|
<span class="md-ellipsis">
|
|
|
|
<code class="doc-symbol doc-symbol-toc doc-symbol-method"></code> reshape
|
|
|
|
</span>
|
|
</a>
|
|
|
|
</li>
|
|
|
|
<li class="md-nav__item">
|
|
<a href="#tinygrad.Tensor.expand" class="md-nav__link">
|
|
<span class="md-ellipsis">
|
|
|
|
<code class="doc-symbol doc-symbol-toc doc-symbol-method"></code> expand
|
|
|
|
</span>
|
|
</a>
|
|
|
|
</li>
|
|
|
|
<li class="md-nav__item">
|
|
<a href="#tinygrad.Tensor.permute" class="md-nav__link">
|
|
<span class="md-ellipsis">
|
|
|
|
<code class="doc-symbol doc-symbol-toc doc-symbol-method"></code> permute
|
|
|
|
</span>
|
|
</a>
|
|
|
|
</li>
|
|
|
|
<li class="md-nav__item">
|
|
<a href="#tinygrad.Tensor.flip" class="md-nav__link">
|
|
<span class="md-ellipsis">
|
|
|
|
<code class="doc-symbol doc-symbol-toc doc-symbol-method"></code> flip
|
|
|
|
</span>
|
|
</a>
|
|
|
|
</li>
|
|
|
|
<li class="md-nav__item">
|
|
<a href="#tinygrad.Tensor.shrink" class="md-nav__link">
|
|
<span class="md-ellipsis">
|
|
|
|
<code class="doc-symbol doc-symbol-toc doc-symbol-method"></code> shrink
|
|
|
|
</span>
|
|
</a>
|
|
|
|
</li>
|
|
|
|
<li class="md-nav__item">
|
|
<a href="#tinygrad.Tensor.pad" class="md-nav__link">
|
|
<span class="md-ellipsis">
|
|
|
|
<code class="doc-symbol doc-symbol-toc doc-symbol-method"></code> pad
|
|
|
|
</span>
|
|
</a>
|
|
|
|
</li>
|
|
|
|
</ul>
|
|
</nav>
|
|
|
|
</li>
|
|
|
|
<li class="md-nav__item">
|
|
<a href="#movement-high-level" class="md-nav__link">
|
|
<span class="md-ellipsis">
|
|
|
|
Movement (high level)
|
|
|
|
</span>
|
|
</a>
|
|
|
|
<nav class="md-nav" aria-label="Movement (high level)">
|
|
<ul class="md-nav__list">
|
|
|
|
<li class="md-nav__item">
|
|
<a href="#tinygrad.Tensor.__getitem__" class="md-nav__link">
|
|
<span class="md-ellipsis">
|
|
|
|
<code class="doc-symbol doc-symbol-toc doc-symbol-method"></code> __getitem__
|
|
|
|
</span>
|
|
</a>
|
|
|
|
</li>
|
|
|
|
<li class="md-nav__item">
|
|
<a href="#tinygrad.Tensor.gather" class="md-nav__link">
|
|
<span class="md-ellipsis">
|
|
|
|
<code class="doc-symbol doc-symbol-toc doc-symbol-method"></code> gather
|
|
|
|
</span>
|
|
</a>
|
|
|
|
</li>
|
|
|
|
<li class="md-nav__item">
|
|
<a href="#tinygrad.Tensor.cat" class="md-nav__link">
|
|
<span class="md-ellipsis">
|
|
|
|
<code class="doc-symbol doc-symbol-toc doc-symbol-method"></code> cat
|
|
|
|
</span>
|
|
</a>
|
|
|
|
</li>
|
|
|
|
<li class="md-nav__item">
|
|
<a href="#tinygrad.Tensor.stack" class="md-nav__link">
|
|
<span class="md-ellipsis">
|
|
|
|
<code class="doc-symbol doc-symbol-toc doc-symbol-method"></code> stack
|
|
|
|
</span>
|
|
</a>
|
|
|
|
</li>
|
|
|
|
<li class="md-nav__item">
|
|
<a href="#tinygrad.Tensor.repeat" class="md-nav__link">
|
|
<span class="md-ellipsis">
|
|
|
|
<code class="doc-symbol doc-symbol-toc doc-symbol-method"></code> repeat
|
|
|
|
</span>
|
|
</a>
|
|
|
|
</li>
|
|
|
|
<li class="md-nav__item">
|
|
<a href="#tinygrad.Tensor.repeat_interleave" class="md-nav__link">
|
|
<span class="md-ellipsis">
|
|
|
|
<code class="doc-symbol doc-symbol-toc doc-symbol-method"></code> repeat_interleave
|
|
|
|
</span>
|
|
</a>
|
|
|
|
</li>
|
|
|
|
<li class="md-nav__item">
|
|
<a href="#tinygrad.Tensor.split" class="md-nav__link">
|
|
<span class="md-ellipsis">
|
|
|
|
<code class="doc-symbol doc-symbol-toc doc-symbol-method"></code> split
|
|
|
|
</span>
|
|
</a>
|
|
|
|
</li>
|
|
|
|
<li class="md-nav__item">
|
|
<a href="#tinygrad.Tensor.chunk" class="md-nav__link">
|
|
<span class="md-ellipsis">
|
|
|
|
<code class="doc-symbol doc-symbol-toc doc-symbol-method"></code> chunk
|
|
|
|
</span>
|
|
</a>
|
|
|
|
</li>
|
|
|
|
<li class="md-nav__item">
|
|
<a href="#tinygrad.Tensor.unfold" class="md-nav__link">
|
|
<span class="md-ellipsis">
|
|
|
|
<code class="doc-symbol doc-symbol-toc doc-symbol-method"></code> unfold
|
|
|
|
</span>
|
|
</a>
|
|
|
|
</li>
|
|
|
|
<li class="md-nav__item">
|
|
<a href="#tinygrad.Tensor.meshgrid" class="md-nav__link">
|
|
<span class="md-ellipsis">
|
|
|
|
<code class="doc-symbol doc-symbol-toc doc-symbol-method"></code> meshgrid
|
|
|
|
</span>
|
|
</a>
|
|
|
|
</li>
|
|
|
|
<li class="md-nav__item">
|
|
<a href="#tinygrad.Tensor.squeeze" class="md-nav__link">
|
|
<span class="md-ellipsis">
|
|
|
|
<code class="doc-symbol doc-symbol-toc doc-symbol-method"></code> squeeze
|
|
|
|
</span>
|
|
</a>
|
|
|
|
</li>
|
|
|
|
<li class="md-nav__item">
|
|
<a href="#tinygrad.Tensor.unsqueeze" class="md-nav__link">
|
|
<span class="md-ellipsis">
|
|
|
|
<code class="doc-symbol doc-symbol-toc doc-symbol-method"></code> unsqueeze
|
|
|
|
</span>
|
|
</a>
|
|
|
|
</li>
|
|
|
|
<li class="md-nav__item">
|
|
<a href="#tinygrad.Tensor.T" class="md-nav__link">
|
|
<span class="md-ellipsis">
|
|
|
|
<code class="doc-symbol doc-symbol-toc doc-symbol-attribute"></code> T
|
|
|
|
</span>
|
|
</a>
|
|
|
|
</li>
|
|
|
|
<li class="md-nav__item">
|
|
<a href="#tinygrad.Tensor.transpose" class="md-nav__link">
|
|
<span class="md-ellipsis">
|
|
|
|
<code class="doc-symbol doc-symbol-toc doc-symbol-method"></code> transpose
|
|
|
|
</span>
|
|
</a>
|
|
|
|
</li>
|
|
|
|
<li class="md-nav__item">
|
|
<a href="#tinygrad.Tensor.flatten" class="md-nav__link">
|
|
<span class="md-ellipsis">
|
|
|
|
<code class="doc-symbol doc-symbol-toc doc-symbol-method"></code> flatten
|
|
|
|
</span>
|
|
</a>
|
|
|
|
</li>
|
|
|
|
<li class="md-nav__item">
|
|
<a href="#tinygrad.Tensor.unflatten" class="md-nav__link">
|
|
<span class="md-ellipsis">
|
|
|
|
<code class="doc-symbol doc-symbol-toc doc-symbol-method"></code> unflatten
|
|
|
|
</span>
|
|
</a>
|
|
|
|
</li>
|
|
|
|
<li class="md-nav__item">
|
|
<a href="#tinygrad.Tensor.diag" class="md-nav__link">
|
|
<span class="md-ellipsis">
|
|
|
|
<code class="doc-symbol doc-symbol-toc doc-symbol-method"></code> diag
|
|
|
|
</span>
|
|
</a>
|
|
|
|
</li>
|
|
|
|
<li class="md-nav__item">
|
|
<a href="#tinygrad.Tensor.roll" class="md-nav__link">
|
|
<span class="md-ellipsis">
|
|
|
|
<code class="doc-symbol doc-symbol-toc doc-symbol-method"></code> roll
|
|
|
|
</span>
|
|
</a>
|
|
|
|
</li>
|
|
|
|
<li class="md-nav__item">
|
|
<a href="#tinygrad.Tensor.rearrange" class="md-nav__link">
|
|
<span class="md-ellipsis">
|
|
|
|
<code class="doc-symbol doc-symbol-toc doc-symbol-method"></code> rearrange
|
|
|
|
</span>
|
|
</a>
|
|
|
|
</li>
|
|
|
|
</ul>
|
|
</nav>
|
|
|
|
</li>
|
|
|
|
</ul>
|
|
|
|
</nav>
|
|
|
|
</li>
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
<li class="md-nav__item">
|
|
<a href="../elementwise/" class="md-nav__link">
|
|
|
|
|
|
|
|
<span class="md-ellipsis">
|
|
|
|
|
|
Elementwise
|
|
|
|
|
|
|
|
</span>
|
|
|
|
|
|
|
|
</a>
|
|
</li>
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
<li class="md-nav__item">
|
|
<a href="../ops/" class="md-nav__link">
|
|
|
|
|
|
|
|
<span class="md-ellipsis">
|
|
|
|
|
|
Complex Ops
|
|
|
|
|
|
|
|
</span>
|
|
|
|
|
|
|
|
</a>
|
|
</li>
|
|
|
|
|
|
|
|
|
|
</ul>
|
|
</nav>
|
|
|
|
</li>
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
<li class="md-nav__item">
|
|
<a href="../../dtypes/" class="md-nav__link">
|
|
|
|
|
|
|
|
<span class="md-ellipsis">
|
|
|
|
|
|
dtypes
|
|
|
|
|
|
|
|
</span>
|
|
|
|
|
|
|
|
</a>
|
|
</li>
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
<li class="md-nav__item">
|
|
<a href="../../nn/" class="md-nav__link">
|
|
|
|
|
|
|
|
<span class="md-ellipsis">
|
|
|
|
|
|
nn (Neural Networks)
|
|
|
|
|
|
|
|
</span>
|
|
|
|
|
|
|
|
</a>
|
|
</li>
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
<li class="md-nav__item">
|
|
<a href="../../env_vars/" class="md-nav__link">
|
|
|
|
|
|
|
|
<span class="md-ellipsis">
|
|
|
|
|
|
Environment Variables
|
|
|
|
|
|
|
|
</span>
|
|
|
|
|
|
|
|
</a>
|
|
</li>
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
<li class="md-nav__item">
|
|
<a href="../../runtime/" class="md-nav__link">
|
|
|
|
|
|
|
|
<span class="md-ellipsis">
|
|
|
|
|
|
Runtime
|
|
|
|
|
|
|
|
</span>
|
|
|
|
|
|
|
|
</a>
|
|
</li>
|
|
|
|
|
|
|
|
|
|
</ul>
|
|
</nav>
|
|
|
|
</li>
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
<li class="md-nav__item md-nav__item--nested">
|
|
|
|
|
|
|
|
|
|
|
|
<input class="md-nav__toggle md-toggle md-toggle--indeterminate" type="checkbox" id="__nav_1_6" >
|
|
|
|
|
|
<label class="md-nav__link" for="__nav_1_6" id="__nav_1_6_label" tabindex="0">
|
|
|
|
|
|
|
|
<span class="md-ellipsis">
|
|
|
|
|
|
Developer
|
|
|
|
|
|
|
|
</span>
|
|
|
|
|
|
|
|
<span class="md-nav__icon md-icon"></span>
|
|
</label>
|
|
|
|
<nav class="md-nav" data-md-level="2" aria-labelledby="__nav_1_6_label" aria-expanded="false">
|
|
<label class="md-nav__title" for="__nav_1_6">
|
|
<span class="md-nav__icon md-icon"></span>
|
|
|
|
|
|
Developer
|
|
|
|
|
|
</label>
|
|
<ul class="md-nav__list" data-md-scrollfix>
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
<li class="md-nav__item">
|
|
<a href="../../developer/developer/" class="md-nav__link">
|
|
|
|
|
|
|
|
<span class="md-ellipsis">
|
|
|
|
|
|
Intro
|
|
|
|
|
|
|
|
</span>
|
|
|
|
|
|
|
|
</a>
|
|
</li>
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
<li class="md-nav__item">
|
|
<a href="../../developer/layout/" class="md-nav__link">
|
|
|
|
|
|
|
|
<span class="md-ellipsis">
|
|
|
|
|
|
Layout
|
|
|
|
|
|
|
|
</span>
|
|
|
|
|
|
|
|
</a>
|
|
</li>
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
<li class="md-nav__item">
|
|
<a href="../../developer/speed/" class="md-nav__link">
|
|
|
|
|
|
|
|
<span class="md-ellipsis">
|
|
|
|
|
|
Speed
|
|
|
|
|
|
|
|
</span>
|
|
|
|
|
|
|
|
</a>
|
|
</li>
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
<li class="md-nav__item">
|
|
<a href="../../developer/uop/" class="md-nav__link">
|
|
|
|
|
|
|
|
<span class="md-ellipsis">
|
|
|
|
|
|
UOp
|
|
|
|
|
|
|
|
</span>
|
|
|
|
|
|
|
|
</a>
|
|
</li>
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
<li class="md-nav__item md-nav__item--nested">
|
|
|
|
|
|
|
|
|
|
|
|
<input class="md-nav__toggle md-toggle md-toggle--indeterminate" type="checkbox" id="__nav_1_6_5" >
|
|
|
|
|
|
<label class="md-nav__link" for="__nav_1_6_5" id="__nav_1_6_5_label" tabindex="0">
|
|
|
|
|
|
|
|
<span class="md-ellipsis">
|
|
|
|
|
|
Runtime
|
|
|
|
|
|
|
|
</span>
|
|
|
|
|
|
|
|
<span class="md-nav__icon md-icon"></span>
|
|
</label>
|
|
|
|
<nav class="md-nav" data-md-level="3" aria-labelledby="__nav_1_6_5_label" aria-expanded="false">
|
|
<label class="md-nav__title" for="__nav_1_6_5">
|
|
<span class="md-nav__icon md-icon"></span>
|
|
|
|
|
|
Runtime
|
|
|
|
|
|
</label>
|
|
<ul class="md-nav__list" data-md-scrollfix>
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
<li class="md-nav__item">
|
|
<a href="../../developer/runtime/" class="md-nav__link">
|
|
|
|
|
|
|
|
<span class="md-ellipsis">
|
|
|
|
|
|
Runtime Overview
|
|
|
|
|
|
|
|
</span>
|
|
|
|
|
|
|
|
</a>
|
|
</li>
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
<li class="md-nav__item">
|
|
<a href="../../developer/hcq/" class="md-nav__link">
|
|
|
|
|
|
|
|
<span class="md-ellipsis">
|
|
|
|
|
|
HCQ
|
|
|
|
|
|
|
|
</span>
|
|
|
|
|
|
|
|
</a>
|
|
</li>
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
<li class="md-nav__item">
|
|
<a href="../../developer/am/" class="md-nav__link">
|
|
|
|
|
|
|
|
<span class="md-ellipsis">
|
|
|
|
|
|
AM Driver
|
|
|
|
|
|
|
|
</span>
|
|
|
|
|
|
|
|
</a>
|
|
</li>
|
|
|
|
|
|
|
|
|
|
</ul>
|
|
</nav>
|
|
|
|
</li>
|
|
|
|
|
|
|
|
|
|
</ul>
|
|
</nav>
|
|
|
|
</li>
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
<li class="md-nav__item">
|
|
<a href="../../tinybox/" class="md-nav__link">
|
|
|
|
|
|
|
|
<span class="md-ellipsis">
|
|
|
|
|
|
tinybox
|
|
|
|
|
|
|
|
</span>
|
|
|
|
|
|
|
|
</a>
|
|
</li>
|
|
|
|
|
|
|
|
|
|
</ul>
|
|
</nav>
|
|
|
|
</li>
|
|
|
|
|
|
|
|
</ul>
|
|
</nav>
|
|
</div>
|
|
</div>
|
|
</div>
|
|
|
|
|
|
|
|
|
|
<div class="md-content" data-md-component="content">
|
|
|
|
|
|
|
|
|
|
|
|
|
|
<nav class="md-path" aria-label="Navigation" >
|
|
<ol class="md-path__list">
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
<li class="md-path__item">
|
|
<a href="../.." class="md-path__link">
|
|
|
|
<span class="md-ellipsis">
|
|
Home
|
|
</span>
|
|
|
|
</a>
|
|
</li>
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
<li class="md-path__item">
|
|
<a href="../" class="md-path__link">
|
|
|
|
<span class="md-ellipsis">
|
|
API Reference
|
|
</span>
|
|
|
|
</a>
|
|
</li>
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
<li class="md-path__item">
|
|
<a href="../" class="md-path__link">
|
|
|
|
<span class="md-ellipsis">
|
|
Tensor
|
|
</span>
|
|
|
|
</a>
|
|
</li>
|
|
|
|
|
|
|
|
|
|
</ol>
|
|
</nav>
|
|
|
|
|
|
<article class="md-content__inner md-typeset">
|
|
|
|
|
|
|
|
|
|
|
|
<a href="https://github.com/tinygrad/tinygrad/edit/master/docs/tensor/movement.md" title="Edit this page" class="md-content__button md-icon" rel="edit">
|
|
|
|
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24"><path d="M10 20H6V4h7v5h5v3.1l2-2V8l-6-6H6c-1.1 0-2 .9-2 2v16c0 1.1.9 2 2 2h4zm10.2-7c.1 0 .3.1.4.2l1.3 1.3c.2.2.2.6 0 .8l-1 1-2.1-2.1 1-1c.1-.1.2-.2.4-.2m0 3.9L14.1 23H12v-2.1l6.1-6.1z"/></svg>
|
|
</a>
|
|
|
|
|
|
|
|
|
|
|
|
<a href="https://github.com/tinygrad/tinygrad/raw/master/docs/tensor/movement.md" title="View source of this page" class="md-content__button md-icon">
|
|
|
|
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24"><path d="M17 18c.56 0 1 .44 1 1s-.44 1-1 1-1-.44-1-1 .44-1 1-1m0-3c-2.73 0-5.06 1.66-6 4 .94 2.34 3.27 4 6 4s5.06-1.66 6-4c-.94-2.34-3.27-4-6-4m0 6.5a2.5 2.5 0 0 1-2.5-2.5 2.5 2.5 0 0 1 2.5-2.5 2.5 2.5 0 0 1 2.5 2.5 2.5 2.5 0 0 1-2.5 2.5M9.27 20H6V4h7v5h5v4.07c.7.08 1.36.25 2 .49V8l-6-6H6a2 2 0 0 0-2 2v16a2 2 0 0 0 2 2h4.5a8.2 8.2 0 0 1-1.23-2"/></svg>
|
|
</a>
|
|
|
|
|
|
|
|
<h1>Movement</h1>
|
|
|
|
<h2 id="movement-low-level">Movement (low level)<a class="headerlink" href="#movement-low-level" title="Permanent link">¤</a></h2>
|
|
|
|
|
|
<div class="doc doc-object doc-function">
|
|
|
|
|
|
<h3 id="tinygrad.Tensor.view" class="doc doc-heading">
|
|
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">view</span>
|
|
|
|
|
|
<a href="#tinygrad.Tensor.view" class="headerlink" title="Permanent link">¤</a></h3>
|
|
<div class="language-python doc-signature highlight"><pre><span></span><code><span class="nf">view</span><span class="p">(</span><span class="n">shape</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"><a class="autorefs autorefs-external" title="<code>typing.Self</code>" href="https://docs.python.org/3/library/typing.html#typing.Self">Self</a></span>
|
|
</code></pre></div>
|
|
|
|
<div class="doc doc-contents first">
|
|
|
|
<p><code class="language-python highlight"><span class="o">.</span><span class="n">view</span></code> is an alias for <code class="language-python highlight"><span class="o">.</span><span class="n">reshape</span></code>.</p>
|
|
|
|
|
|
<details class="mkdocstrings-source">
|
|
<summary>Source code in <code>tinygrad/mixin/movement.py</code></summary>
|
|
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal">177</span>
|
|
<span class="normal">178</span>
|
|
<span class="normal">179</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="k">def</span><span class="w"> </span><span class="nf">view</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">shape</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">Self</span><span class="p">:</span>
|
|
<span class="w"> </span><span class="sd">"""`.view` is an alias for `.reshape`."""</span>
|
|
<span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">shape</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">)</span>
|
|
</code></pre></div></td></tr></table></div>
|
|
</details>
|
|
</div>
|
|
|
|
</div>
|
|
|
|
<div class="doc doc-object doc-function">
|
|
|
|
|
|
<h3 id="tinygrad.Tensor.reshape" class="doc doc-heading">
|
|
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">reshape</span>
|
|
|
|
|
|
<a href="#tinygrad.Tensor.reshape" class="headerlink" title="Permanent link">¤</a></h3>
|
|
<div class="language-python doc-signature highlight"><pre><span></span><code><span class="nf">reshape</span><span class="p">(</span><span class="n">shape</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"><a class="autorefs autorefs-external" title="<code>typing.Self</code>" href="https://docs.python.org/3/library/typing.html#typing.Self">Self</a></span>
|
|
</code></pre></div>
|
|
|
|
<div class="doc doc-contents first">
|
|
|
|
<p>Returns a tensor with the same data as the original tensor but with a different shape.
|
|
<code class="language-python highlight"><span class="n">shape</span></code> can be passed as a tuple or as separate arguments.</p>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="n">t</span> <span class="o">=</span> <span class="n">Tensor</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="mi">6</span><span class="p">)</span>
|
|
<span class="nb">print</span><span class="p">(</span><span class="n">t</span><span class="o">.</span><span class="n">reshape</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="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">[[</span><span class="mi">0</span> <span class="mi">1</span> <span class="mi">2</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">3</span> <span class="mi">4</span> <span class="mi">5</span><span class="p">]]</span>
|
|
</code></pre></div>
|
|
|
|
|
|
<details class="mkdocstrings-source">
|
|
<summary>Source code in <code>tinygrad/mixin/movement.py</code></summary>
|
|
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal"> 86</span>
|
|
<span class="normal"> 87</span>
|
|
<span class="normal"> 88</span>
|
|
<span class="normal"> 89</span>
|
|
<span class="normal"> 90</span>
|
|
<span class="normal"> 91</span>
|
|
<span class="normal"> 92</span>
|
|
<span class="normal"> 93</span>
|
|
<span class="normal"> 94</span>
|
|
<span class="normal"> 95</span>
|
|
<span class="normal"> 96</span>
|
|
<span class="normal"> 97</span>
|
|
<span class="normal"> 98</span>
|
|
<span class="normal"> 99</span>
|
|
<span class="normal">100</span>
|
|
<span class="normal">101</span>
|
|
<span class="normal">102</span>
|
|
<span class="normal">103</span>
|
|
<span class="normal">104</span>
|
|
<span class="normal">105</span>
|
|
<span class="normal">106</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="k">def</span><span class="w"> </span><span class="nf">reshape</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">shape</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">Self</span><span class="p">:</span>
|
|
<span class="w"> </span><span class="sd">"""</span>
|
|
<span class="sd"> Returns a tensor with the same data as the original tensor but with a different shape.</span>
|
|
<span class="sd"> `shape` can be passed as a tuple or as separate arguments.</span>
|
|
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> t = Tensor.arange(6)</span>
|
|
<span class="sd"> print(t.reshape(2, 3).numpy())</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> """</span>
|
|
<span class="c1"># resolve None and args</span>
|
|
<span class="n">new_shape</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">([</span><span class="n">s</span> <span class="k">if</span> <span class="n">s</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="k">else</span> <span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">s</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">argfix</span><span class="p">(</span><span class="n">shape</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">))])</span>
|
|
<span class="c1"># resolve -1</span>
|
|
<span class="k">if</span> <span class="p">(</span><span class="n">c</span> <span class="o">:=</span> <span class="n">new_shape</span><span class="o">.</span><span class="n">count</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">))</span> <span class="o">></span> <span class="mi">1</span><span class="p">:</span>
|
|
<span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="sa">f</span><span class="s2">"only one dimension can be inferred using -1, getting </span><span class="si">{</span><span class="n">new_shape</span><span class="si">}</span><span class="s2">"</span><span class="p">)</span>
|
|
<span class="k">if</span> <span class="n">c</span><span class="p">:</span>
|
|
<span class="n">new_shape</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">([</span><span class="o">-</span><span class="n">prod</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span> <span class="o">//</span> <span class="n">prod</span><span class="p">(</span><span class="n">new_shape</span><span class="p">)</span> <span class="k">if</span> <span class="n">s</span> <span class="o">==</span> <span class="o">-</span><span class="mi">1</span> <span class="k">else</span> <span class="n">s</span> <span class="k">for</span> <span class="n">s</span> <span class="ow">in</span> <span class="n">new_shape</span><span class="p">])</span>
|
|
<span class="k">if</span> <span class="n">prod</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span> <span class="o">!=</span> <span class="n">prod</span><span class="p">(</span><span class="n">new_shape</span><span class="p">):</span>
|
|
<span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">"size mismatch, can't reshape (</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="si">}</span><span class="s2">) -> (</span><span class="si">{</span><span class="n">new_shape</span><span class="si">}</span><span class="s2">)"</span><span class="p">)</span>
|
|
<span class="n">ret</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_mop</span><span class="p">(</span><span class="n">Ops</span><span class="o">.</span><span class="n">RESHAPE</span><span class="p">,</span> <span class="n">arg</span><span class="o">=</span><span class="n">new_shape</span><span class="p">)</span>
|
|
<span class="k">return</span> <span class="bp">self</span> <span class="k">if</span> <span class="n">ret</span><span class="o">.</span><span class="n">shape</span> <span class="o">==</span> <span class="bp">self</span><span class="o">.</span><span class="n">shape</span> <span class="k">else</span> <span class="n">ret</span>
|
|
</code></pre></div></td></tr></table></div>
|
|
</details>
|
|
</div>
|
|
|
|
</div>
|
|
|
|
<div class="doc doc-object doc-function">
|
|
|
|
|
|
<h3 id="tinygrad.Tensor.expand" class="doc doc-heading">
|
|
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">expand</span>
|
|
|
|
|
|
<a href="#tinygrad.Tensor.expand" class="headerlink" title="Permanent link">¤</a></h3>
|
|
<div class="language-python doc-signature highlight"><pre><span></span><code><span class="nf">expand</span><span class="p">(</span><span class="n">shape</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"><a class="autorefs autorefs-external" title="<code>typing.Self</code>" href="https://docs.python.org/3/library/typing.html#typing.Self">Self</a></span>
|
|
</code></pre></div>
|
|
|
|
<div class="doc doc-contents first">
|
|
|
|
<p>Returns a tensor that is expanded to the shape that is specified.
|
|
Expand can also increase the number of dimensions that a tensor has.</p>
|
|
<p>Passing a <code class="language-python highlight"><span class="o">-</span><span class="mi">1</span></code> or <code class="language-python highlight"><span class="kc">None</span></code> to a dimension means that its size will not be changed.</p>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="n">t</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>
|
|
<span class="nb">print</span><span class="p">(</span><span class="n">t</span><span class="o">.</span><span class="n">expand</span><span class="p">(</span><span class="mi">4</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">[[</span><span class="mi">1</span> <span class="mi">2</span> <span class="mi">3</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">1</span> <span class="mi">2</span> <span class="mi">3</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">1</span> <span class="mi">2</span> <span class="mi">3</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">1</span> <span class="mi">2</span> <span class="mi">3</span><span class="p">]]</span>
|
|
</code></pre></div>
|
|
|
|
|
|
<details class="mkdocstrings-source">
|
|
<summary>Source code in <code>tinygrad/mixin/movement.py</code></summary>
|
|
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal">71</span>
|
|
<span class="normal">72</span>
|
|
<span class="normal">73</span>
|
|
<span class="normal">74</span>
|
|
<span class="normal">75</span>
|
|
<span class="normal">76</span>
|
|
<span class="normal">77</span>
|
|
<span class="normal">78</span>
|
|
<span class="normal">79</span>
|
|
<span class="normal">80</span>
|
|
<span class="normal">81</span>
|
|
<span class="normal">82</span>
|
|
<span class="normal">83</span>
|
|
<span class="normal">84</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="k">def</span><span class="w"> </span><span class="nf">expand</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">shape</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">Self</span><span class="p">:</span>
|
|
<span class="w"> </span><span class="sd">"""</span>
|
|
<span class="sd"> Returns a tensor that is expanded to the shape that is specified.</span>
|
|
<span class="sd"> Expand can also increase the number of dimensions that a tensor has.</span>
|
|
|
|
<span class="sd"> Passing a `-1` or `None` to a dimension means that its size will not be changed.</span>
|
|
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> t = Tensor([1, 2, 3])</span>
|
|
<span class="sd"> print(t.expand(4, -1).numpy())</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> """</span>
|
|
<span class="n">new_shape</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">(</span><span class="n">from_</span> <span class="k">if</span> <span class="n">to</span> <span class="o">==</span> <span class="o">-</span><span class="mi">1</span> <span class="ow">or</span> <span class="n">to</span> <span class="ow">is</span> <span class="kc">None</span> <span class="k">else</span> <span class="n">to</span> <span class="k">for</span> <span class="n">from_</span><span class="p">,</span> <span class="n">to</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="o">*</span><span class="p">(</span><span class="n">_align_left</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="n">argfix</span><span class="p">(</span><span class="n">shape</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">)))))</span>
|
|
<span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_broadcast_to</span><span class="p">(</span><span class="n">new_shape</span><span class="p">)</span>
|
|
</code></pre></div></td></tr></table></div>
|
|
</details>
|
|
</div>
|
|
|
|
</div>
|
|
|
|
<div class="doc doc-object doc-function">
|
|
|
|
|
|
<h3 id="tinygrad.Tensor.permute" class="doc doc-heading">
|
|
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">permute</span>
|
|
|
|
|
|
<a href="#tinygrad.Tensor.permute" class="headerlink" title="Permanent link">¤</a></h3>
|
|
<div class="language-python doc-signature highlight"><pre><span></span><code><span class="nf">permute</span><span class="p">(</span><span class="n">order</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"><a class="autorefs autorefs-external" title="<code>typing.Self</code>" href="https://docs.python.org/3/library/typing.html#typing.Self">Self</a></span>
|
|
</code></pre></div>
|
|
|
|
<div class="doc doc-contents first">
|
|
|
|
<p>Returns a tensor that is a permutation of the original tensor.
|
|
The new tensor has the same data as the original tensor but with the dimensions permuted according to the order specified.
|
|
<code class="language-python highlight"><span class="n">order</span></code> can be passed as a tuple or as separate arguments.</p>
|
|
<p><div class="language-python highlight"><pre><span></span><code><span class="n">t</span> <span class="o">=</span> <span class="n">Tensor</span><span class="o">.</span><span class="n">empty</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="mi">5</span><span class="p">)</span>
|
|
<span class="nb">print</span><span class="p">(</span><span class="n">t</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">5</span><span class="p">)</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="nb">print</span><span class="p">(</span><span class="n">t</span><span class="o">.</span><span class="n">permute</span><span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">(</span><span class="mi">5</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">)</span>
|
|
</code></pre></div></p>
|
|
|
|
|
|
<details class="mkdocstrings-source">
|
|
<summary>Source code in <code>tinygrad/mixin/movement.py</code></summary>
|
|
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal">130</span>
|
|
<span class="normal">131</span>
|
|
<span class="normal">132</span>
|
|
<span class="normal">133</span>
|
|
<span class="normal">134</span>
|
|
<span class="normal">135</span>
|
|
<span class="normal">136</span>
|
|
<span class="normal">137</span>
|
|
<span class="normal">138</span>
|
|
<span class="normal">139</span>
|
|
<span class="normal">140</span>
|
|
<span class="normal">141</span>
|
|
<span class="normal">142</span>
|
|
<span class="normal">143</span>
|
|
<span class="normal">144</span>
|
|
<span class="normal">145</span>
|
|
<span class="normal">146</span>
|
|
<span class="normal">147</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="k">def</span><span class="w"> </span><span class="nf">permute</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">order</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">Self</span><span class="p">:</span>
|
|
<span class="w"> </span><span class="sd">"""</span>
|
|
<span class="sd"> Returns a tensor that is a permutation of the original tensor.</span>
|
|
<span class="sd"> The new tensor has the same data as the original tensor but with the dimensions permuted according to the order specified.</span>
|
|
<span class="sd"> `order` can be passed as a tuple or as separate arguments.</span>
|
|
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> t = Tensor.empty(2, 3, 5)</span>
|
|
<span class="sd"> print(t.shape)</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> print(t.permute(2, 0, 1).shape)</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> """</span>
|
|
<span class="n">order_arg</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_resolve_dim</span><span class="p">(</span><span class="n">x</span><span class="p">)</span> <span class="k">for</span> <span class="n">x</span> <span class="ow">in</span> <span class="n">argfix</span><span class="p">(</span><span class="n">order</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">))</span>
|
|
<span class="k">if</span> <span class="nb">sorted</span><span class="p">(</span><span class="n">order_arg</span><span class="p">)</span> <span class="o">!=</span> <span class="nb">list</span><span class="p">(</span><span class="nb">range</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">ndim</span><span class="p">)):</span>
|
|
<span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="sa">f</span><span class="s2">"order is not a valid permutation, getting </span><span class="si">{</span><span class="n">order_arg</span><span class="si">}</span><span class="s2">"</span><span class="p">)</span>
|
|
<span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_mop</span><span class="p">(</span><span class="n">Ops</span><span class="o">.</span><span class="n">PERMUTE</span><span class="p">,</span> <span class="n">arg</span><span class="o">=</span><span class="n">order_arg</span><span class="p">)</span> <span class="k">if</span> <span class="n">order_arg</span> <span class="o">!=</span> <span class="nb">tuple</span><span class="p">(</span><span class="nb">range</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">ndim</span><span class="p">))</span> <span class="k">else</span> <span class="bp">self</span>
|
|
</code></pre></div></td></tr></table></div>
|
|
</details>
|
|
</div>
|
|
|
|
</div>
|
|
|
|
<div class="doc doc-object doc-function">
|
|
|
|
|
|
<h3 id="tinygrad.Tensor.flip" class="doc doc-heading">
|
|
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">flip</span>
|
|
|
|
|
|
<a href="#tinygrad.Tensor.flip" class="headerlink" title="Permanent link">¤</a></h3>
|
|
<div class="language-python doc-signature highlight"><pre><span></span><code><span class="nf">flip</span><span class="p">(</span><span class="n">axis</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"><a class="autorefs autorefs-external" title="<code>typing.Self</code>" href="https://docs.python.org/3/library/typing.html#typing.Self">Self</a></span>
|
|
</code></pre></div>
|
|
|
|
<div class="doc doc-contents first">
|
|
|
|
<p>Returns a tensor that reverses the order of the original tensor along given <code class="language-python highlight"><span class="n">axis</span></code>.
|
|
<code class="language-python highlight"><span class="n">axis</span></code> can be passed as a tuple or as separate arguments.</p>
|
|
<p><div class="language-python highlight"><pre><span></span><code><span class="n">t</span> <span class="o">=</span> <span class="n">Tensor</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="mi">6</span><span class="p">)</span><span class="o">.</span><span class="n">reshape</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="nb">print</span><span class="p">(</span><span class="n">t</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">[[</span><span class="mi">0</span> <span class="mi">1</span> <span class="mi">2</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">3</span> <span class="mi">4</span> <span class="mi">5</span><span class="p">]]</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="nb">print</span><span class="p">(</span><span class="n">t</span><span class="o">.</span><span class="n">flip</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">[[</span><span class="mi">3</span> <span class="mi">4</span> <span class="mi">5</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">0</span> <span class="mi">1</span> <span class="mi">2</span><span class="p">]]</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="nb">print</span><span class="p">(</span><span class="n">t</span><span class="o">.</span><span class="n">flip</span><span class="p">((</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">))</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">[[</span><span class="mi">5</span> <span class="mi">4</span> <span class="mi">3</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">2</span> <span class="mi">1</span> <span class="mi">0</span><span class="p">]]</span>
|
|
</code></pre></div></p>
|
|
|
|
|
|
<details class="mkdocstrings-source">
|
|
<summary>Source code in <code>tinygrad/mixin/movement.py</code></summary>
|
|
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal">149</span>
|
|
<span class="normal">150</span>
|
|
<span class="normal">151</span>
|
|
<span class="normal">152</span>
|
|
<span class="normal">153</span>
|
|
<span class="normal">154</span>
|
|
<span class="normal">155</span>
|
|
<span class="normal">156</span>
|
|
<span class="normal">157</span>
|
|
<span class="normal">158</span>
|
|
<span class="normal">159</span>
|
|
<span class="normal">160</span>
|
|
<span class="normal">161</span>
|
|
<span class="normal">162</span>
|
|
<span class="normal">163</span>
|
|
<span class="normal">164</span>
|
|
<span class="normal">165</span>
|
|
<span class="normal">166</span>
|
|
<span class="normal">167</span>
|
|
<span class="normal">168</span>
|
|
<span class="normal">169</span>
|
|
<span class="normal">170</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="k">def</span><span class="w"> </span><span class="nf">flip</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">axis</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">Self</span><span class="p">:</span>
|
|
<span class="w"> </span><span class="sd">"""</span>
|
|
<span class="sd"> Returns a tensor that reverses the order of the original tensor along given `axis`.</span>
|
|
<span class="sd"> `axis` can be passed as a tuple or as separate arguments.</span>
|
|
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> t = Tensor.arange(6).reshape(2, 3)</span>
|
|
<span class="sd"> print(t.numpy())</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> print(t.flip(0).numpy())</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> print(t.flip((0, 1)).numpy())</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> """</span>
|
|
<span class="n">axis_arg</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_resolve_dim</span><span class="p">(</span><span class="n">x</span><span class="p">)</span> <span class="k">for</span> <span class="n">x</span> <span class="ow">in</span> <span class="n">argfix</span><span class="p">(</span><span class="n">axis</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">))</span>
|
|
<span class="k">assert</span> <span class="nb">all</span><span class="p">(</span><span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="nb">bool</span><span class="p">)</span> <span class="ow">and</span> <span class="n">x</span> <span class="o">>=</span> <span class="mi">0</span> <span class="ow">and</span> <span class="n">x</span> <span class="o"><</span> <span class="bp">self</span><span class="o">.</span><span class="n">ndim</span> <span class="k">for</span> <span class="n">x</span> <span class="ow">in</span> <span class="n">axis_arg</span><span class="p">),</span> <span class="sa">f</span><span class="s2">"flip args must be axis ints </span><span class="si">{</span><span class="n">axis_arg</span><span class="si">}</span><span class="s2">"</span>
|
|
<span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">axis_arg</span><span class="p">)</span> <span class="o">!=</span> <span class="nb">len</span><span class="p">(</span><span class="n">dedup</span><span class="p">(</span><span class="n">axis_arg</span><span class="p">)):</span>
|
|
<span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="sa">f</span><span class="s2">"dim can appear at most once, getting </span><span class="si">{</span><span class="n">axis_arg</span><span class="si">}</span><span class="s2">"</span><span class="p">)</span>
|
|
<span class="n">flip_arg</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">([</span><span class="n">i</span> <span class="ow">in</span> <span class="n">axis_arg</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">))])</span>
|
|
<span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_mop</span><span class="p">(</span><span class="n">Ops</span><span class="o">.</span><span class="n">FLIP</span><span class="p">,</span> <span class="n">arg</span><span class="o">=</span><span class="n">flip_arg</span><span class="p">)</span> <span class="k">if</span> <span class="nb">any</span><span class="p">(</span><span class="n">flip_arg</span><span class="p">)</span> <span class="k">else</span> <span class="bp">self</span>
|
|
</code></pre></div></td></tr></table></div>
|
|
</details>
|
|
</div>
|
|
|
|
</div>
|
|
|
|
<div class="doc doc-object doc-function">
|
|
|
|
|
|
<h3 id="tinygrad.Tensor.shrink" class="doc doc-heading">
|
|
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">shrink</span>
|
|
|
|
|
|
<a href="#tinygrad.Tensor.shrink" class="headerlink" title="Permanent link">¤</a></h3>
|
|
<div class="language-python doc-signature highlight"><pre><span></span><code><span class="nf">shrink</span><span class="p">(</span><span class="n">arg</span><span class="p">:</span> <span class="n"><a class="autorefs autorefs-external" href="https://docs.python.org/3/library/stdtypes.html#tuple">tuple</a></span><span class="p">[</span><span class="n"><a class="autorefs autorefs-external" href="https://docs.python.org/3/library/stdtypes.html#tuple">tuple</a></span><span class="p">[</span><span class="n"><span title="tinygrad.mixin.movement.sint">sint</span></span><span class="p">,</span> <span class="n"><span title="tinygrad.mixin.movement.sint">sint</span></span><span class="p">]</span> <span class="o">|</span> <span class="kc">None</span><span class="p">,</span> <span class="o">...</span><span class="p">])</span> <span class="o">-></span> <span class="n"><a class="autorefs autorefs-external" title="<code>typing.Self</code>" href="https://docs.python.org/3/library/typing.html#typing.Self">Self</a></span>
|
|
</code></pre></div>
|
|
|
|
<div class="doc doc-contents first">
|
|
|
|
<p>Returns a tensor that shrinks the each axis based on input arg.
|
|
<code class="language-python highlight"><span class="n">arg</span></code> must have the same length as <code class="language-python highlight"><span class="bp">self</span><span class="o">.</span><span class="n">ndim</span></code>.
|
|
For each axis, it can be <code class="language-python highlight"><span class="kc">None</span></code>, which means no shrink, or a tuple <code class="language-python highlight"><span class="p">(</span><span class="n">start</span><span class="p">,</span> <span class="n">end</span><span class="p">)</span></code> that works the same as Python slice.</p>
|
|
<p><div class="language-python highlight"><pre><span></span><code><span class="n">t</span> <span class="o">=</span> <span class="n">Tensor</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="mi">9</span><span class="p">)</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="mi">3</span><span class="p">,</span> <span class="mi">3</span><span class="p">)</span>
|
|
<span class="nb">print</span><span class="p">(</span><span class="n">t</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">[[</span><span class="mi">0</span> <span class="mi">1</span> <span class="mi">2</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">3</span> <span class="mi">4</span> <span class="mi">5</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">6</span> <span class="mi">7</span> <span class="mi">8</span><span class="p">]]</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="nb">print</span><span class="p">(</span><span class="n">t</span><span class="o">.</span><span class="n">shrink</span><span class="p">(((</span><span class="kc">None</span><span class="p">,</span> <span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">3</span><span class="p">))))</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">[[</span><span class="mi">1</span> <span class="mi">2</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">4</span> <span class="mi">5</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">7</span> <span class="mi">8</span><span class="p">]]</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="nb">print</span><span class="p">(</span><span class="n">t</span><span class="o">.</span><span class="n">shrink</span><span class="p">((((</span><span class="mi">0</span><span class="p">,</span> <span class="mi">2</span><span class="p">),</span> <span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">2</span><span class="p">))))</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">[[</span><span class="mi">0</span> <span class="mi">1</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">3</span> <span class="mi">4</span><span class="p">]]</span>
|
|
</code></pre></div></p>
|
|
|
|
|
|
<details class="mkdocstrings-source">
|
|
<summary>Source code in <code>tinygrad/mixin/movement.py</code></summary>
|
|
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal">108</span>
|
|
<span class="normal">109</span>
|
|
<span class="normal">110</span>
|
|
<span class="normal">111</span>
|
|
<span class="normal">112</span>
|
|
<span class="normal">113</span>
|
|
<span class="normal">114</span>
|
|
<span class="normal">115</span>
|
|
<span class="normal">116</span>
|
|
<span class="normal">117</span>
|
|
<span class="normal">118</span>
|
|
<span class="normal">119</span>
|
|
<span class="normal">120</span>
|
|
<span class="normal">121</span>
|
|
<span class="normal">122</span>
|
|
<span class="normal">123</span>
|
|
<span class="normal">124</span>
|
|
<span class="normal">125</span>
|
|
<span class="normal">126</span>
|
|
<span class="normal">127</span>
|
|
<span class="normal">128</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="k">def</span><span class="w"> </span><span class="nf">shrink</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">arg</span><span class="p">:</span> <span class="nb">tuple</span><span class="p">[</span><span class="nb">tuple</span><span class="p">[</span><span class="n">sint</span><span class="p">,</span> <span class="n">sint</span><span class="p">]</span> <span class="o">|</span> <span class="kc">None</span><span class="p">,</span> <span class="o">...</span><span class="p">])</span> <span class="o">-></span> <span class="n">Self</span><span class="p">:</span>
|
|
<span class="w"> </span><span class="sd">"""</span>
|
|
<span class="sd"> Returns a tensor that shrinks the each axis based on input arg.</span>
|
|
<span class="sd"> `arg` must have the same length as `self.ndim`.</span>
|
|
<span class="sd"> For each axis, it can be `None`, which means no shrink, or a tuple `(start, end)` that works the same as Python slice.</span>
|
|
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> t = Tensor.arange(9).reshape(3, 3)</span>
|
|
<span class="sd"> print(t.numpy())</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> print(t.shrink(((None, (1, 3)))).numpy())</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> print(t.shrink((((0, 2), (0, 2)))).numpy())</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> """</span>
|
|
<span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">ndim</span> <span class="o">!=</span> <span class="nb">len</span><span class="p">(</span><span class="n">arg</span><span class="p">):</span>
|
|
<span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">"</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">ndim</span><span class="si">=}</span><span class="s2"> != </span><span class="si">{</span><span class="nb">len</span><span class="p">(</span><span class="n">arg</span><span class="p">)</span><span class="si">=}</span><span class="s2">"</span><span class="p">)</span>
|
|
<span class="n">ret</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_mop</span><span class="p">(</span><span class="n">Ops</span><span class="o">.</span><span class="n">SHRINK</span><span class="p">,</span> <span class="n">arg</span><span class="o">=</span><span class="p">[</span><span class="n">x</span> <span class="k">if</span> <span class="n">x</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="k">else</span> <span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="n">s</span><span class="p">)</span> <span class="k">for</span> <span class="n">x</span><span class="p">,</span> <span class="n">s</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">arg</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">)])</span>
|
|
<span class="k">return</span> <span class="bp">self</span> <span class="k">if</span> <span class="n">ret</span><span class="o">.</span><span class="n">shape</span> <span class="o">==</span> <span class="bp">self</span><span class="o">.</span><span class="n">shape</span> <span class="k">else</span> <span class="n">ret</span>
|
|
</code></pre></div></td></tr></table></div>
|
|
</details>
|
|
</div>
|
|
|
|
</div>
|
|
|
|
<div class="doc doc-object doc-function">
|
|
|
|
|
|
<h3 id="tinygrad.Tensor.pad" class="doc doc-heading">
|
|
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">pad</span>
|
|
|
|
|
|
<a href="#tinygrad.Tensor.pad" class="headerlink" title="Permanent link">¤</a></h3>
|
|
<div class="language-python doc-signature highlight"><pre><span></span><code><span class="nf">pad</span><span class="p">(</span>
|
|
<span class="n">padding</span><span class="p">:</span> <span class="p">(</span>
|
|
<span class="n"><a class="autorefs autorefs-external" title="<code>typing.Sequence</code>" href="https://docs.python.org/3/library/typing.html#typing.Sequence">Sequence</a></span><span class="p">[</span><span class="n"><span title="tinygrad.uop.ops.sint">sint</span></span><span class="p">]</span> <span class="o">|</span> <span class="n"><a class="autorefs autorefs-external" title="<code>typing.Sequence</code>" href="https://docs.python.org/3/library/typing.html#typing.Sequence">Sequence</a></span><span class="p">[</span><span class="n"><a class="autorefs autorefs-external" href="https://docs.python.org/3/library/stdtypes.html#tuple">tuple</a></span><span class="p">[</span><span class="n"><span title="tinygrad.uop.ops.sint">sint</span></span><span class="p">,</span> <span class="n"><span title="tinygrad.uop.ops.sint">sint</span></span><span class="p">]</span> <span class="o">|</span> <span class="kc">None</span><span class="p">]</span>
|
|
<span class="p">),</span>
|
|
<span class="n">mode</span><span class="p">:</span> <span class="n"><a class="autorefs autorefs-external" href="https://docs.python.org/3/library/stdtypes.html#str">str</a></span> <span class="o">=</span> <span class="s2">"constant"</span><span class="p">,</span>
|
|
<span class="n">value</span><span class="p">:</span> <span class="n"><a class="autorefs autorefs-internal" title="<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">float</span> (<code>tinygrad.tensor.Tensor.float</code>)" href="../elementwise/#tinygrad.Tensor.float">float</a></span> <span class="o">=</span> <span class="mf">0.0</span><span class="p">,</span>
|
|
<span class="p">)</span> <span class="o">-></span> <span class="n"><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">Tensor</span> (<code>tinygrad.tensor.Tensor</code>)" href="../#tinygrad.Tensor">Tensor</a></span>
|
|
</code></pre></div>
|
|
|
|
<div class="doc doc-contents first">
|
|
|
|
<p>Returns a tensor with padding applied based on the input <code class="language-python highlight"><span class="n">padding</span></code>.</p>
|
|
<p><code class="language-python highlight"><span class="n">padding</span></code> supports two padding structures:</p>
|
|
<ol>
|
|
<li>
|
|
<p>Flat padding: <code class="language-python highlight"><span class="p">(</span><span class="n">padding_left</span><span class="p">,</span> <span class="n">padding_right</span><span class="p">,</span> <span class="n">padding_top</span><span class="p">,</span> <span class="n">padding_bottom</span><span class="p">,</span> <span class="o">...</span><span class="p">)</span></code></p>
|
|
<ul>
|
|
<li>This structure matches PyTorch's pad.</li>
|
|
<li><code class="language-python highlight"><span class="n">padding</span></code> length must be even.</li>
|
|
</ul>
|
|
</li>
|
|
<li>
|
|
<p>Group padding: <code class="language-python highlight"><span class="p">(</span><span class="o">...</span><span class="p">,</span> <span class="p">(</span><span class="n">padding_top</span><span class="p">,</span> <span class="n">padding_bottom</span><span class="p">),</span> <span class="p">(</span><span class="n">padding_left</span><span class="p">,</span> <span class="n">padding_right</span><span class="p">))</span></code></p>
|
|
<ul>
|
|
<li>This structure matches pad for JAX, NumPy, TensorFlow, and others.</li>
|
|
<li>For each axis, padding can be <code class="language-python highlight"><span class="kc">None</span></code>, meaning no padding, or a tuple <code class="language-python highlight"><span class="p">(</span><span class="n">start</span><span class="p">,</span> <span class="n">end</span><span class="p">)</span></code>.</li>
|
|
<li><code class="language-python highlight"><span class="n">padding</span></code> must have the same length as <code class="language-python highlight"><span class="bp">self</span><span class="o">.</span><span class="n">ndim</span></code>.</li>
|
|
</ul>
|
|
</li>
|
|
</ol>
|
|
<p>Padding values can be negative, resulting in dimension shrinks that work similarly to Python negative slices.
|
|
Padding modes is selected with <code class="language-python highlight"><span class="n">mode</span></code> which supports <code class="language-python highlight"><span class="n">constant</span></code>, <code class="language-python highlight"><span class="n">reflect</span></code> and <code class="language-python highlight"><span class="n">replicate</span></code>.</p>
|
|
<p><div class="language-python highlight"><pre><span></span><code><span class="n">t</span> <span class="o">=</span> <span class="n">Tensor</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="mi">9</span><span class="p">)</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">3</span><span class="p">)</span>
|
|
<span class="nb">print</span><span class="p">(</span><span class="n">t</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">[[[[</span><span class="mi">0</span> <span class="mi">1</span> <span class="mi">2</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">3</span> <span class="mi">4</span> <span class="mi">5</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">6</span> <span class="mi">7</span> <span class="mi">8</span><span class="p">]]]]</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="nb">print</span><span class="p">(</span><span class="n">t</span><span class="o">.</span><span class="n">pad</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">0</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">))</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">[[[[</span><span class="mi">0</span> <span class="mi">0</span> <span class="mi">1</span> <span class="mi">2</span> <span class="mi">0</span> <span class="mi">0</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">0</span> <span class="mi">3</span> <span class="mi">4</span> <span class="mi">5</span> <span class="mi">0</span> <span class="mi">0</span><span class="p">]]]]</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="nb">print</span><span class="p">(</span><span class="n">t</span><span class="o">.</span><span class="n">pad</span><span class="p">(((</span><span class="kc">None</span><span class="p">,</span> <span class="kc">None</span><span class="p">,</span> <span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">),</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="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">[[[[</span><span class="mi">0</span> <span class="mi">0</span> <span class="mi">1</span> <span class="mi">2</span> <span class="mi">0</span> <span class="mi">0</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">0</span> <span class="mi">3</span> <span class="mi">4</span> <span class="mi">5</span> <span class="mi">0</span> <span class="mi">0</span><span class="p">]]]]</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="nb">print</span><span class="p">(</span><span class="n">t</span><span class="o">.</span><span class="n">pad</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">0</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">),</span> <span class="n">value</span><span class="o">=-</span><span class="nb">float</span><span class="p">(</span><span class="s1">'inf'</span><span class="p">))</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">[[[[</span><span class="o">-</span><span class="n">inf</span> <span class="mf">0.</span> <span class="mf">1.</span> <span class="mf">2.</span> <span class="o">-</span><span class="n">inf</span> <span class="o">-</span><span class="n">inf</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="o">-</span><span class="n">inf</span> <span class="mf">3.</span> <span class="mf">4.</span> <span class="mf">5.</span> <span class="o">-</span><span class="n">inf</span> <span class="o">-</span><span class="n">inf</span><span class="p">]]]]</span>
|
|
</code></pre></div></p>
|
|
|
|
|
|
<details class="mkdocstrings-source">
|
|
<summary>Source code in <code>tinygrad/tensor.py</code></summary>
|
|
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal">1038</span>
|
|
<span class="normal">1039</span>
|
|
<span class="normal">1040</span>
|
|
<span class="normal">1041</span>
|
|
<span class="normal">1042</span>
|
|
<span class="normal">1043</span>
|
|
<span class="normal">1044</span>
|
|
<span class="normal">1045</span>
|
|
<span class="normal">1046</span>
|
|
<span class="normal">1047</span>
|
|
<span class="normal">1048</span>
|
|
<span class="normal">1049</span>
|
|
<span class="normal">1050</span>
|
|
<span class="normal">1051</span>
|
|
<span class="normal">1052</span>
|
|
<span class="normal">1053</span>
|
|
<span class="normal">1054</span>
|
|
<span class="normal">1055</span>
|
|
<span class="normal">1056</span>
|
|
<span class="normal">1057</span>
|
|
<span class="normal">1058</span>
|
|
<span class="normal">1059</span>
|
|
<span class="normal">1060</span>
|
|
<span class="normal">1061</span>
|
|
<span class="normal">1062</span>
|
|
<span class="normal">1063</span>
|
|
<span class="normal">1064</span>
|
|
<span class="normal">1065</span>
|
|
<span class="normal">1066</span>
|
|
<span class="normal">1067</span>
|
|
<span class="normal">1068</span>
|
|
<span class="normal">1069</span>
|
|
<span class="normal">1070</span>
|
|
<span class="normal">1071</span>
|
|
<span class="normal">1072</span>
|
|
<span class="normal">1073</span>
|
|
<span class="normal">1074</span>
|
|
<span class="normal">1075</span>
|
|
<span class="normal">1076</span>
|
|
<span class="normal">1077</span>
|
|
<span class="normal">1078</span>
|
|
<span class="normal">1079</span>
|
|
<span class="normal">1080</span>
|
|
<span class="normal">1081</span>
|
|
<span class="normal">1082</span>
|
|
<span class="normal">1083</span>
|
|
<span class="normal">1084</span>
|
|
<span class="normal">1085</span>
|
|
<span class="normal">1086</span>
|
|
<span class="normal">1087</span>
|
|
<span class="normal">1088</span>
|
|
<span class="normal">1089</span>
|
|
<span class="normal">1090</span>
|
|
<span class="normal">1091</span>
|
|
<span class="normal">1092</span>
|
|
<span class="normal">1093</span>
|
|
<span class="normal">1094</span>
|
|
<span class="normal">1095</span>
|
|
<span class="normal">1096</span>
|
|
<span class="normal">1097</span>
|
|
<span class="normal">1098</span>
|
|
<span class="normal">1099</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="k">def</span><span class="w"> </span><span class="nf">pad</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">padding</span><span class="p">:</span><span class="n">Sequence</span><span class="p">[</span><span class="n">sint</span><span class="p">]</span><span class="o">|</span><span class="n">Sequence</span><span class="p">[</span><span class="nb">tuple</span><span class="p">[</span><span class="n">sint</span><span class="p">,</span> <span class="n">sint</span><span class="p">]</span><span class="o">|</span><span class="kc">None</span><span class="p">],</span> <span class="n">mode</span><span class="p">:</span><span class="nb">str</span><span class="o">=</span><span class="s2">"constant"</span><span class="p">,</span> <span class="n">value</span><span class="p">:</span><span class="nb">float</span><span class="o">=</span><span class="mf">0.0</span><span class="p">)</span> <span class="o">-></span> <span class="n">Tensor</span><span class="p">:</span>
|
|
<span class="w"> </span><span class="sd">"""</span>
|
|
<span class="sd"> Returns a tensor with padding applied based on the input `padding`.</span>
|
|
|
|
<span class="sd"> `padding` supports two padding structures:</span>
|
|
|
|
<span class="sd"> 1. Flat padding: `(padding_left, padding_right, padding_top, padding_bottom, ...)`</span>
|
|
<span class="sd"> - This structure matches PyTorch's pad.</span>
|
|
<span class="sd"> - `padding` length must be even.</span>
|
|
|
|
<span class="sd"> 2. Group padding: `(..., (padding_top, padding_bottom), (padding_left, padding_right))`</span>
|
|
<span class="sd"> - This structure matches pad for JAX, NumPy, TensorFlow, and others.</span>
|
|
<span class="sd"> - For each axis, padding can be `None`, meaning no padding, or a tuple `(start, end)`.</span>
|
|
<span class="sd"> - `padding` must have the same length as `self.ndim`.</span>
|
|
|
|
<span class="sd"> Padding values can be negative, resulting in dimension shrinks that work similarly to Python negative slices.</span>
|
|
<span class="sd"> Padding modes is selected with `mode` which supports `constant`, `reflect` and `replicate`.</span>
|
|
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> t = Tensor.arange(9).reshape(1, 1, 3, 3)</span>
|
|
<span class="sd"> print(t.numpy())</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> print(t.pad((1, 2, 0, -1)).numpy())</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> print(t.pad(((None, None, (0, -1), (1, 2)))).numpy())</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> print(t.pad((1, 2, 0, -1), value=-float('inf')).numpy())</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> """</span>
|
|
<span class="k">if</span> <span class="n">mode</span> <span class="ow">not</span> <span class="ow">in</span> <span class="p">{</span><span class="s2">"constant"</span><span class="p">,</span> <span class="s2">"reflect"</span><span class="p">,</span> <span class="s2">"replicate"</span><span class="p">,</span> <span class="s2">"circular"</span><span class="p">}:</span> <span class="k">raise</span> <span class="ne">NotImplementedError</span><span class="p">(</span><span class="sa">f</span><span class="s2">"</span><span class="si">{</span><span class="n">mode</span><span class="si">=}</span><span class="s2"> is not supported"</span><span class="p">)</span>
|
|
<span class="c1"># flat padding</span>
|
|
<span class="k">if</span> <span class="nb">all</span><span class="p">(</span><span class="nb">isinstance</span><span class="p">(</span><span class="n">p</span><span class="p">,</span> <span class="p">(</span><span class="nb">int</span><span class="p">,</span><span class="n">UOp</span><span class="p">))</span> <span class="k">for</span> <span class="n">p</span> <span class="ow">in</span> <span class="n">padding</span><span class="p">):</span>
|
|
<span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">padding</span><span class="p">)</span><span class="o">%</span><span class="mi">2</span> <span class="o">!=</span> <span class="mi">0</span><span class="p">:</span> <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">"Flat padding must have even number of pads"</span><span class="p">)</span>
|
|
<span class="n">pX</span> <span class="o">=</span> <span class="n">_flat_to_grouped</span><span class="p">(</span><span class="nb">tuple</span><span class="p">(</span><span class="n">cast</span><span class="p">(</span><span class="n">Sequence</span><span class="p">[</span><span class="n">sint</span><span class="p">],</span> <span class="n">padding</span><span class="p">))</span> <span class="o">+</span> <span class="p">(</span><span class="mi">0</span><span class="p">,</span><span class="mi">0</span><span class="p">)</span><span class="o">*</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">ndim</span> <span class="o">-</span> <span class="nb">len</span><span class="p">(</span><span class="n">padding</span><span class="p">)</span><span class="o">//</span><span class="mi">2</span><span class="p">))</span>
|
|
<span class="c1"># group padding</span>
|
|
<span class="k">else</span><span class="p">:</span> <span class="n">pX</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">((</span><span class="mi">0</span><span class="p">,</span><span class="mi">0</span><span class="p">)</span> <span class="k">if</span> <span class="n">p</span> <span class="ow">is</span> <span class="kc">None</span> <span class="k">else</span> <span class="n">p</span> <span class="k">for</span> <span class="n">p</span> <span class="ow">in</span> <span class="n">cast</span><span class="p">(</span><span class="n">Sequence</span><span class="p">[</span><span class="nb">tuple</span><span class="p">[</span><span class="n">sint</span><span class="p">,</span> <span class="n">sint</span><span class="p">]</span><span class="o">|</span><span class="kc">None</span><span class="p">],</span> <span class="n">padding</span><span class="p">))</span>
|
|
<span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">pX</span><span class="p">)</span> <span class="o">!=</span> <span class="bp">self</span><span class="o">.</span><span class="n">ndim</span><span class="p">:</span> <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">"padding length is improper, </span><span class="si">{</span><span class="n">padding</span><span class="si">=}</span><span class="s2"> </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">ndim</span><span class="si">=}</span><span class="s2">"</span><span class="p">)</span>
|
|
<span class="n">X</span><span class="p">,</span> <span class="n">pads</span> <span class="o">=</span> <span class="bp">self</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">((</span><span class="n">smax</span><span class="p">(</span><span class="n">pB</span><span class="p">,</span><span class="mi">0</span><span class="p">),</span> <span class="n">smax</span><span class="p">(</span><span class="n">pA</span><span class="p">,</span><span class="mi">0</span><span class="p">))</span> <span class="k">for</span> <span class="n">pB</span><span class="p">,</span><span class="n">pA</span> <span class="ow">in</span> <span class="n">pX</span><span class="p">)</span>
|
|
<span class="k">if</span> <span class="n">mode</span> <span class="o">==</span> <span class="s2">"constant"</span><span class="p">:</span>
|
|
<span class="k">def</span><span class="w"> </span><span class="nf">_constant</span><span class="p">(</span><span class="n">x</span><span class="p">:</span><span class="n">Tensor</span><span class="p">,</span><span class="n">px</span><span class="p">,</span><span class="n">v</span><span class="p">)</span> <span class="o">-></span> <span class="n">Tensor</span><span class="p">:</span>
|
|
<span class="k">return</span> <span class="n">x</span><span class="o">.</span><span class="n">_apply_uop</span><span class="p">(</span><span class="n">UOp</span><span class="o">.</span><span class="n">pad</span><span class="p">,</span> <span class="n">arg</span><span class="o">=</span><span class="n">px</span><span class="p">)</span> <span class="k">if</span> <span class="n">v</span> <span class="o">==</span> <span class="mi">0</span> <span class="k">else</span> <span class="p">(</span><span class="n">x</span><span class="o">.</span><span class="n">_apply_uop</span><span class="p">(</span><span class="n">UOp</span><span class="o">.</span><span class="n">pad</span><span class="p">,</span> <span class="n">arg</span><span class="o">=</span><span class="n">px</span><span class="p">)</span><span class="o">+</span><span class="n">Tensor</span><span class="o">.</span><span class="n">ones_like</span><span class="p">(</span><span class="n">x</span><span class="p">)</span><span class="o">.</span><span class="n">_apply_uop</span><span class="p">(</span><span class="n">UOp</span><span class="o">.</span><span class="n">pad</span><span class="p">,</span> <span class="n">arg</span><span class="o">=</span><span class="n">px</span><span class="p">)</span><span class="o">.</span><span class="n">where</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span><span class="n">v</span><span class="p">))</span>
|
|
<span class="k">return</span> <span class="n">_constant</span><span class="p">(</span><span class="n">X</span><span class="p">,</span> <span class="n">pX</span><span class="p">,</span> <span class="n">value</span><span class="p">)</span> <span class="k">if</span> <span class="nb">all</span><span class="p">(</span><span class="n">resolve</span><span class="p">(</span><span class="n">p</span> <span class="o">>=</span> <span class="mi">0</span><span class="p">)</span> <span class="k">for</span> <span class="n">p</span> <span class="ow">in</span> <span class="n">flatten</span><span class="p">(</span><span class="n">pX</span><span class="p">))</span> <span class="k">else</span> \
|
|
<span class="n">_constant</span><span class="p">(</span><span class="n">X</span><span class="o">.</span><span class="n">shrink</span><span class="p">(</span><span class="nb">tuple</span><span class="p">((</span><span class="o">-</span><span class="n">smin</span><span class="p">(</span><span class="n">pB</span><span class="p">,</span><span class="mi">0</span><span class="p">),</span><span class="n">smin</span><span class="p">(</span><span class="n">pA</span><span class="o">+</span><span class="n">s</span><span class="p">,</span><span class="n">s</span><span class="p">))</span> <span class="k">for</span> <span class="p">(</span><span class="n">pB</span><span class="p">,</span><span class="n">pA</span><span class="p">),</span><span class="n">s</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">pX</span><span class="p">,</span> <span class="n">X</span><span class="o">.</span><span class="n">shape</span><span class="p">))),</span> <span class="n">pads</span><span class="p">,</span> <span class="n">value</span><span class="p">)</span>
|
|
<span class="k">assert</span> <span class="n">all_int</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">),</span> <span class="sa">f</span><span class="s2">"does not support symbolic shape </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="si">}</span><span class="s2">"</span>
|
|
<span class="k">if</span> <span class="n">mode</span> <span class="o">==</span> <span class="s2">"circular"</span><span class="p">:</span>
|
|
<span class="k">if</span> <span class="nb">any</span><span class="p">(</span><span class="n">pB</span><span class="o">></span><span class="n">sh</span> <span class="ow">or</span> <span class="n">pA</span><span class="o">></span><span class="n">sh</span> <span class="k">for</span> <span class="p">(</span><span class="n">pB</span><span class="p">,</span><span class="n">pA</span><span class="p">),</span><span class="n">sh</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">pX</span><span class="p">,</span> <span class="n">X</span><span class="o">.</span><span class="n">shape</span><span class="p">)):</span> <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s1">'Padding value causes wrapping around more than once.'</span><span class="p">)</span>
|
|
<span class="k">if</span> <span class="nb">any</span><span class="p">(</span><span class="n">pB</span><span class="o"><</span><span class="mi">0</span> <span class="ow">or</span> <span class="n">pA</span><span class="o"><</span><span class="mi">0</span> <span class="k">for</span> <span class="n">pB</span><span class="p">,</span><span class="n">pA</span> <span class="ow">in</span> <span class="n">pX</span><span class="p">):</span> <span class="k">raise</span> <span class="ne">NotImplementedError</span><span class="p">(</span><span class="s2">"Negative pads with circular pads is not supported"</span><span class="p">)</span>
|
|
<span class="n">orig_shape</span><span class="p">,</span> <span class="n">X</span> <span class="o">=</span> <span class="n">X</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="n">X</span><span class="o">.</span><span class="n">repeat</span><span class="p">(</span><span class="nb">tuple</span><span class="p">(</span><span class="mi">1</span> <span class="o">+</span> <span class="nb">bool</span><span class="p">(</span><span class="n">pB</span><span class="p">)</span> <span class="o">+</span> <span class="nb">bool</span><span class="p">(</span><span class="n">pA</span><span class="p">)</span> <span class="k">for</span> <span class="n">pB</span><span class="p">,</span><span class="n">pA</span> <span class="ow">in</span> <span class="n">pads</span><span class="p">))</span>
|
|
<span class="k">return</span> <span class="n">X</span><span class="o">.</span><span class="n">shrink</span><span class="p">(</span><span class="nb">tuple</span><span class="p">((</span><span class="mi">0</span> <span class="k">if</span> <span class="n">pB</span> <span class="o">==</span> <span class="mi">0</span> <span class="k">else</span> <span class="n">osh</span><span class="o">-</span><span class="n">pB</span><span class="p">,</span> <span class="n">xsh</span> <span class="k">if</span> <span class="n">pA</span> <span class="o">==</span> <span class="mi">0</span> <span class="k">else</span> <span class="n">xsh</span><span class="o">-</span><span class="n">osh</span><span class="o">+</span><span class="n">pA</span><span class="p">)</span> <span class="k">for</span> <span class="p">(</span><span class="n">pB</span><span class="p">,</span><span class="n">pA</span><span class="p">),</span><span class="n">osh</span><span class="p">,</span><span class="n">xsh</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">pads</span><span class="p">,</span> <span class="n">orig_shape</span><span class="p">,</span> <span class="n">X</span><span class="o">.</span><span class="n">shape</span><span class="p">)))</span>
|
|
<span class="k">for</span> <span class="n">d</span><span class="p">,(</span><span class="n">pB</span><span class="p">,</span><span class="n">pA</span><span class="p">)</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">pads</span><span class="p">):</span>
|
|
<span class="k">if</span> <span class="n">mode</span> <span class="o">==</span> <span class="s2">"reflect"</span><span class="p">:</span>
|
|
<span class="k">if</span> <span class="n">pB</span> <span class="o">>=</span> <span class="p">(</span><span class="n">s</span><span class="o">:=</span><span class="n">X</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="n">d</span><span class="p">])</span> <span class="ow">or</span> <span class="n">pA</span><span class="o">>=</span><span class="n">s</span><span class="p">:</span> <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">"Padding (</span><span class="si">{</span><span class="n">pB</span><span class="si">}</span><span class="s2">, </span><span class="si">{</span><span class="n">pA</span><span class="si">}</span><span class="s2">) should be less than the input size=</span><span class="si">{</span><span class="n">s</span><span class="si">}</span><span class="s2"> for dim=</span><span class="si">{</span><span class="n">d</span><span class="si">}</span><span class="s2">."</span><span class="p">)</span>
|
|
<span class="n">slcB</span><span class="p">,</span> <span class="n">slcA</span><span class="p">,</span> <span class="o">=</span> <span class="nb">slice</span><span class="p">(</span><span class="n">pB</span><span class="p">,</span><span class="mi">0</span><span class="p">,</span><span class="o">-</span><span class="mi">1</span><span class="p">),</span> <span class="nb">slice</span><span class="p">(</span><span class="n">s</span><span class="o">-</span><span class="mi">2</span> <span class="k">if</span> <span class="n">s</span><span class="o">-</span><span class="mi">2</span><span class="o">>=</span><span class="mi">0</span> <span class="k">else</span> <span class="kc">None</span><span class="p">,</span> <span class="n">s</span><span class="o">-</span><span class="mi">2</span><span class="o">-</span><span class="n">pA</span> <span class="k">if</span> <span class="n">s</span><span class="o">-</span><span class="mi">2</span><span class="o">-</span><span class="n">pA</span><span class="o">>=</span><span class="mi">0</span> <span class="k">else</span> <span class="kc">None</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">)</span>
|
|
<span class="n">xB</span><span class="p">,</span> <span class="n">xA</span> <span class="o">=</span> <span class="p">(</span><span class="n">X</span><span class="p">[[</span><span class="n">slc</span> <span class="k">if</span> <span class="n">i</span> <span class="o">==</span> <span class="n">d</span> <span class="k">else</span> <span class="nb">slice</span><span class="p">(</span><span class="kc">None</span><span class="p">)</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">X</span><span class="o">.</span><span class="n">ndim</span><span class="p">)]]</span> <span class="k">if</span> <span class="n">p</span> <span class="o">></span> <span class="mi">0</span> <span class="k">else</span> <span class="kc">None</span> <span class="k">for</span> <span class="n">slc</span><span class="p">,</span> <span class="n">p</span> <span class="ow">in</span> <span class="p">((</span><span class="n">slcB</span><span class="p">,</span> <span class="n">pB</span><span class="p">),</span> <span class="p">(</span><span class="n">slcA</span><span class="p">,</span> <span class="n">pA</span><span class="p">)))</span>
|
|
<span class="k">if</span> <span class="n">mode</span> <span class="o">==</span> <span class="s2">"replicate"</span><span class="p">:</span>
|
|
<span class="n">shrB</span><span class="p">,</span> <span class="n">shrA</span><span class="p">,</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">((</span><span class="mi">0</span><span class="p">,</span><span class="mi">1</span><span class="p">)</span> <span class="k">if</span> <span class="n">i</span><span class="o">==</span><span class="n">d</span> <span class="k">else</span> <span class="kc">None</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">X</span><span class="o">.</span><span class="n">ndim</span><span class="p">)),</span> <span class="nb">tuple</span><span class="p">((</span><span class="n">X</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="n">i</span><span class="p">]</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span><span class="n">X</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="n">i</span><span class="p">])</span> <span class="k">if</span> <span class="n">i</span><span class="o">==</span><span class="n">d</span> <span class="k">else</span> <span class="kc">None</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">X</span><span class="o">.</span><span class="n">ndim</span><span class="p">))</span>
|
|
<span class="n">xB</span><span class="p">,</span> <span class="n">xA</span> <span class="o">=</span> <span class="p">(</span><span class="n">X</span><span class="o">.</span><span class="n">shrink</span><span class="p">(</span><span class="n">shr</span><span class="p">)</span><span class="o">.</span><span class="n">expand</span><span class="p">(</span><span class="nb">tuple</span><span class="p">(</span><span class="n">p</span> <span class="k">if</span> <span class="n">i</span><span class="o">==</span><span class="n">d</span> <span class="k">else</span> <span class="kc">None</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">X</span><span class="o">.</span><span class="n">ndim</span><span class="p">)))</span> <span class="k">if</span> <span class="n">p</span> <span class="o">></span> <span class="mi">0</span> <span class="k">else</span> <span class="kc">None</span> <span class="k">for</span> <span class="n">shr</span><span class="p">,</span> <span class="n">p</span> <span class="ow">in</span> <span class="p">((</span><span class="n">shrB</span><span class="p">,</span> <span class="n">pB</span><span class="p">),</span> <span class="p">(</span><span class="n">shrA</span><span class="p">,</span> <span class="n">pA</span><span class="p">)))</span>
|
|
<span class="n">X</span> <span class="o">=</span> <span class="n">Tensor</span><span class="o">.</span><span class="n">cat</span><span class="p">(</span><span class="o">*</span><span class="p">(</span><span class="n">X_</span> <span class="k">for</span> <span class="n">X_</span> <span class="ow">in</span> <span class="p">(</span><span class="n">xB</span><span class="p">,</span> <span class="n">X</span><span class="p">,</span> <span class="n">xA</span><span class="p">)</span> <span class="k">if</span> <span class="n">X_</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">),</span> <span class="n">dim</span><span class="o">=</span><span class="n">d</span><span class="p">)</span>
|
|
<span class="k">return</span> <span class="n">X</span><span class="o">.</span><span class="n">shrink</span><span class="p">(</span><span class="nb">tuple</span><span class="p">((</span><span class="o">-</span><span class="nb">min</span><span class="p">(</span><span class="n">pB</span><span class="p">,</span><span class="mi">0</span><span class="p">),</span> <span class="nb">min</span><span class="p">(</span><span class="n">pA</span><span class="o">+</span><span class="n">s</span><span class="p">,</span><span class="n">s</span><span class="p">))</span> <span class="k">for</span> <span class="p">(</span><span class="n">pB</span><span class="p">,</span><span class="n">pA</span><span class="p">),</span><span class="n">s</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">pX</span><span class="p">,</span> <span class="n">X</span><span class="o">.</span><span class="n">shape</span><span class="p">)))</span>
|
|
</code></pre></div></td></tr></table></div>
|
|
</details>
|
|
</div>
|
|
|
|
</div><h2 id="movement-high-level">Movement (high level)<a class="headerlink" href="#movement-high-level" title="Permanent link">¤</a></h2>
|
|
|
|
|
|
<div class="doc doc-object doc-function">
|
|
|
|
|
|
<h3 id="tinygrad.Tensor.__getitem__" class="doc doc-heading">
|
|
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">__getitem__</span>
|
|
|
|
|
|
<a href="#tinygrad.Tensor.__getitem__" class="headerlink" title="Permanent link">¤</a></h3>
|
|
<div class="language-python doc-signature highlight"><pre><span></span><code><span class="nf">__getitem__</span><span class="p">(</span><span class="n">indices</span><span class="p">)</span> <span class="o">-></span> <span class="n"><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">Tensor</span> (<code>tinygrad.tensor.Tensor</code>)" href="../#tinygrad.Tensor">Tensor</a></span>
|
|
</code></pre></div>
|
|
|
|
<div class="doc doc-contents first">
|
|
|
|
<p>Retrieves a sub-tensor using indexing.</p>
|
|
<p>Supported Index Types: <code class="language-python highlight"><span class="nb">int</span> <span class="o">|</span> <span class="nb">slice</span> <span class="o">|</span> <span class="n">Tensor</span> <span class="o">|</span> <span class="kc">None</span> <span class="o">|</span> <span class="nb">list</span> <span class="o">|</span> <span class="nb">tuple</span> <span class="o">|</span> <span class="bp">Ellipsis</span></code></p>
|
|
<p>Examples:
|
|
<div class="language-python highlight"><pre><span></span><code><span class="n">t</span> <span class="o">=</span> <span class="n">Tensor</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="mi">12</span><span class="p">)</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="mi">3</span><span class="p">,</span> <span class="mi">4</span><span class="p">)</span>
|
|
<span class="nb">print</span><span class="p">(</span><span class="n">t</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">[[</span> <span class="mi">0</span> <span class="mi">1</span> <span class="mi">2</span> <span class="mi">3</span><span class="p">]</span>
|
|
<span class="p">[</span> <span class="mi">4</span> <span class="mi">5</span> <span class="mi">6</span> <span class="mi">7</span><span class="p">]</span>
|
|
<span class="p">[</span> <span class="mi">8</span> <span class="mi">9</span> <span class="mi">10</span> <span class="mi">11</span><span class="p">]]</span>
|
|
</code></pre></div></p>
|
|
<ul>
|
|
<li>
|
|
<p>Int Indexing: Select an element or sub-tensor using integers for each dimension.
|
|
<div class="language-python highlight"><pre><span></span><code><span class="nb">print</span><span class="p">(</span><span class="n">t</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="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="mi">6</span>
|
|
</code></pre></div></p>
|
|
</li>
|
|
<li>
|
|
<p>Slice Indexing: Select a range of elements using slice notation (<code class="language-python highlight"><span class="n">start</span><span class="p">:</span><span class="n">end</span><span class="p">:</span><span class="n">stride</span></code>).
|
|
<div class="language-python highlight"><pre><span></span><code><span class="nb">print</span><span class="p">(</span><span class="n">t</span><span class="p">[</span><span class="mi">0</span><span class="p">:</span><span class="mi">2</span><span class="p">,</span> <span class="p">::</span><span class="mi">2</span><span class="p">]</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">[[</span><span class="mi">0</span> <span class="mi">2</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">4</span> <span class="mi">6</span><span class="p">]]</span>
|
|
</code></pre></div></p>
|
|
</li>
|
|
<li>
|
|
<p>Tensor Indexing: Use another tensor as indices for advanced indexing. Using <code class="language-python highlight"><span class="nb">tuple</span></code> or <code class="language-python highlight"><span class="nb">list</span></code> here also works.
|
|
<div class="language-python highlight"><pre><span></span><code><span class="nb">print</span><span class="p">(</span><span class="n">t</span><span class="p">[</span><span class="n">Tensor</span><span class="p">([</span><span class="mi">2</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">]),</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><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">[</span><span class="mi">9</span> <span class="mi">2</span> <span class="mi">7</span><span class="p">]</span>
|
|
</code></pre></div></p>
|
|
</li>
|
|
<li>
|
|
<p><code class="language-python highlight"><span class="kc">None</span></code> Indexing: Add a new dimension to the tensor.
|
|
<div class="language-python highlight"><pre><span></span><code><span class="nb">print</span><span class="p">(</span><span class="n">t</span><span class="p">[:,</span> <span class="kc">None</span><span class="p">]</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">(</span><span class="mi">3</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">4</span><span class="p">)</span>
|
|
</code></pre></div></p>
|
|
</li>
|
|
</ul>
|
|
<div class="admonition note">
|
|
<p class="admonition-title">Note</p>
|
|
<p>Out-of-bounds indexing results in a value of <code class="language-python highlight"><span class="mi">0</span></code>.
|
|
<div class="language-python highlight"><pre><span></span><code><span class="n">t</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>
|
|
<span class="nb">print</span><span class="p">(</span><span class="n">t</span><span class="p">[</span><span class="n">Tensor</span><span class="p">([</span><span class="mi">4</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">2</span><span class="p">])]</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">[</span><span class="mi">0</span> <span class="mi">0</span> <span class="mi">3</span><span class="p">]</span>
|
|
</code></pre></div></p>
|
|
</div>
|
|
|
|
|
|
<details class="mkdocstrings-source">
|
|
<summary>Source code in <code>tinygrad/tensor.py</code></summary>
|
|
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal">1212</span>
|
|
<span class="normal">1213</span>
|
|
<span class="normal">1214</span>
|
|
<span class="normal">1215</span>
|
|
<span class="normal">1216</span>
|
|
<span class="normal">1217</span>
|
|
<span class="normal">1218</span>
|
|
<span class="normal">1219</span>
|
|
<span class="normal">1220</span>
|
|
<span class="normal">1221</span>
|
|
<span class="normal">1222</span>
|
|
<span class="normal">1223</span>
|
|
<span class="normal">1224</span>
|
|
<span class="normal">1225</span>
|
|
<span class="normal">1226</span>
|
|
<span class="normal">1227</span>
|
|
<span class="normal">1228</span>
|
|
<span class="normal">1229</span>
|
|
<span class="normal">1230</span>
|
|
<span class="normal">1231</span>
|
|
<span class="normal">1232</span>
|
|
<span class="normal">1233</span>
|
|
<span class="normal">1234</span>
|
|
<span class="normal">1235</span>
|
|
<span class="normal">1236</span>
|
|
<span class="normal">1237</span>
|
|
<span class="normal">1238</span>
|
|
<span class="normal">1239</span>
|
|
<span class="normal">1240</span>
|
|
<span class="normal">1241</span>
|
|
<span class="normal">1242</span>
|
|
<span class="normal">1243</span>
|
|
<span class="normal">1244</span>
|
|
<span class="normal">1245</span>
|
|
<span class="normal">1246</span>
|
|
<span class="normal">1247</span>
|
|
<span class="normal">1248</span>
|
|
<span class="normal">1249</span>
|
|
<span class="normal">1250</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="k">def</span><span class="w"> </span><span class="fm">__getitem__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">indices</span><span class="p">)</span> <span class="o">-></span> <span class="n">Tensor</span><span class="p">:</span>
|
|
<span class="w"> </span><span class="sd">"""</span>
|
|
<span class="sd"> Retrieves a sub-tensor using indexing.</span>
|
|
|
|
<span class="sd"> Supported Index Types: `int | slice | Tensor | None | list | tuple | Ellipsis`</span>
|
|
|
|
<span class="sd"> Examples:</span>
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> t = Tensor.arange(12).reshape(3, 4)</span>
|
|
<span class="sd"> print(t.numpy())</span>
|
|
<span class="sd"> ```</span>
|
|
|
|
<span class="sd"> - Int Indexing: Select an element or sub-tensor using integers for each dimension.</span>
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> print(t[1, 2].numpy())</span>
|
|
<span class="sd"> ```</span>
|
|
|
|
<span class="sd"> - Slice Indexing: Select a range of elements using slice notation (`start:end:stride`).</span>
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> print(t[0:2, ::2].numpy())</span>
|
|
<span class="sd"> ```</span>
|
|
|
|
<span class="sd"> - Tensor Indexing: Use another tensor as indices for advanced indexing. Using `tuple` or `list` here also works.</span>
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> print(t[Tensor([2, 0, 1]), Tensor([1, 2, 3])].numpy())</span>
|
|
<span class="sd"> ```</span>
|
|
|
|
<span class="sd"> - `None` Indexing: Add a new dimension to the tensor.</span>
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> print(t[:, None].shape)</span>
|
|
<span class="sd"> ```</span>
|
|
|
|
<span class="sd"> NOTE: Out-of-bounds indexing results in a value of `0`.</span>
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> t = Tensor([1, 2, 3])</span>
|
|
<span class="sd"> print(t[Tensor([4, 3, 2])].numpy())</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> """</span>
|
|
<span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_getitem</span><span class="p">(</span><span class="n">indices</span><span class="p">)</span>
|
|
</code></pre></div></td></tr></table></div>
|
|
</details>
|
|
</div>
|
|
|
|
</div>
|
|
|
|
<div class="doc doc-object doc-function">
|
|
|
|
|
|
<h3 id="tinygrad.Tensor.gather" class="doc doc-heading">
|
|
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">gather</span>
|
|
|
|
|
|
<a href="#tinygrad.Tensor.gather" class="headerlink" title="Permanent link">¤</a></h3>
|
|
<div class="language-python doc-signature highlight"><pre><span></span><code><span class="nf">gather</span><span class="p">(</span><span class="n">dim</span><span class="p">:</span> <span class="n"><a class="autorefs autorefs-internal" title="<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">int</span> (<code>tinygrad.tensor.Tensor.int</code>)" href="../elementwise/#tinygrad.Tensor.int">int</a></span><span class="p">,</span> <span class="n">index</span><span class="p">:</span> <span class="n"><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">Tensor</span> (<code>tinygrad.tensor.Tensor</code>)" href="../#tinygrad.Tensor">Tensor</a></span><span class="p">)</span> <span class="o">-></span> <span class="n"><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">Tensor</span> (<code>tinygrad.tensor.Tensor</code>)" href="../#tinygrad.Tensor">Tensor</a></span>
|
|
</code></pre></div>
|
|
|
|
<div class="doc doc-contents first">
|
|
|
|
<p>Gathers values along an axis specified by <code class="language-python highlight"><span class="n">dim</span></code>.</p>
|
|
<p><div class="language-python highlight"><pre><span></span><code><span class="n">t</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="p">[</span><span class="mi">3</span><span class="p">,</span> <span class="mi">4</span><span class="p">]])</span>
|
|
<span class="nb">print</span><span class="p">(</span><span class="n">t</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">[[</span><span class="mi">1</span> <span class="mi">2</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">3</span> <span class="mi">4</span><span class="p">]]</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="nb">print</span><span class="p">(</span><span class="n">t</span><span class="o">.</span><span class="n">gather</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="n">Tensor</span><span class="p">([[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span> <span class="p">[</span><span class="mi">1</span><span class="p">,</span> <span class="mi">0</span><span class="p">]]))</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">[[</span><span class="mi">1</span> <span class="mi">1</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">4</span> <span class="mi">3</span><span class="p">]]</span>
|
|
</code></pre></div></p>
|
|
|
|
|
|
<details class="mkdocstrings-source">
|
|
<summary>Source code in <code>tinygrad/tensor.py</code></summary>
|
|
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal">1273</span>
|
|
<span class="normal">1274</span>
|
|
<span class="normal">1275</span>
|
|
<span class="normal">1276</span>
|
|
<span class="normal">1277</span>
|
|
<span class="normal">1278</span>
|
|
<span class="normal">1279</span>
|
|
<span class="normal">1280</span>
|
|
<span class="normal">1281</span>
|
|
<span class="normal">1282</span>
|
|
<span class="normal">1283</span>
|
|
<span class="normal">1284</span>
|
|
<span class="normal">1285</span>
|
|
<span class="normal">1286</span>
|
|
<span class="normal">1287</span>
|
|
<span class="normal">1288</span>
|
|
<span class="normal">1289</span>
|
|
<span class="normal">1290</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="k">def</span><span class="w"> </span><span class="nf">gather</span><span class="p">(</span><span class="bp">self</span><span class="p">:</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">dim</span><span class="p">:</span><span class="nb">int</span><span class="p">,</span> <span class="n">index</span><span class="p">:</span><span class="n">Tensor</span><span class="p">)</span> <span class="o">-></span> <span class="n">Tensor</span><span class="p">:</span>
|
|
<span class="w"> </span><span class="sd">"""</span>
|
|
<span class="sd"> Gathers values along an axis specified by `dim`.</span>
|
|
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> t = Tensor([[1, 2], [3, 4]])</span>
|
|
<span class="sd"> print(t.numpy())</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> print(t.gather(1, Tensor([[0, 0], [1, 0]])).numpy())</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> """</span>
|
|
<span class="k">assert</span> <span class="n">index</span><span class="o">.</span><span class="n">ndim</span> <span class="o">==</span> <span class="bp">self</span><span class="o">.</span><span class="n">ndim</span><span class="p">,</span> <span class="sa">f</span><span class="s2">"self.ndim must equal index.ndim, </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">ndim</span><span class="si">=}</span><span class="s2">, </span><span class="si">{</span><span class="n">index</span><span class="o">.</span><span class="n">ndim</span><span class="si">=}</span><span class="s2">"</span>
|
|
<span class="n">dim</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_resolve_dim</span><span class="p">(</span><span class="n">dim</span><span class="p">)</span>
|
|
<span class="k">assert</span> <span class="nb">all</span><span class="p">(</span><span class="n">s</span> <span class="o">>=</span> <span class="n">i</span> <span class="k">for</span> <span class="n">d</span><span class="p">,(</span><span class="n">s</span><span class="p">,</span><span class="n">i</span><span class="p">)</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="nb">zip</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="n">index</span><span class="o">.</span><span class="n">shape</span><span class="p">))</span> <span class="k">if</span> <span class="n">d</span> <span class="o">!=</span> <span class="n">dim</span><span class="p">),</span> <span class="s2">"requires self.shape[d] >= index.shape[d] for all d != dim"</span>
|
|
<span class="n">index</span> <span class="o">=</span> <span class="n">index</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">device</span><span class="p">)</span>
|
|
<span class="n">x</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">shrink</span><span class="p">(</span><span class="nb">tuple</span><span class="p">((</span><span class="mi">0</span><span class="p">,</span> <span class="n">i</span><span class="p">)</span> <span class="k">if</span> <span class="n">d</span> <span class="o">!=</span> <span class="n">dim</span> <span class="k">else</span> <span class="kc">None</span> <span class="k">for</span> <span class="n">d</span><span class="p">,</span><span class="n">i</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">index</span><span class="o">.</span><span class="n">shape</span><span class="p">)))</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">transpose</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="n">dim</span><span class="p">)</span>
|
|
<span class="k">return</span> <span class="p">(</span><span class="n">index</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">_one_hot_along_dim</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="n">dim</span><span class="p">])</span><span class="o">.</span><span class="n">where</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="mi">0</span><span class="p">))</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dtype</span><span class="p">)</span>
|
|
</code></pre></div></td></tr></table></div>
|
|
</details>
|
|
</div>
|
|
|
|
</div>
|
|
|
|
<div class="doc doc-object doc-function">
|
|
|
|
|
|
<h3 id="tinygrad.Tensor.cat" class="doc doc-heading">
|
|
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">cat</span>
|
|
|
|
|
|
<a href="#tinygrad.Tensor.cat" class="headerlink" title="Permanent link">¤</a></h3>
|
|
<div class="language-python doc-signature highlight"><pre><span></span><code><span class="nf">cat</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">:</span> <span class="n"><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">Tensor</span> (<code>tinygrad.tensor.Tensor</code>)" href="../#tinygrad.Tensor">Tensor</a></span><span class="p">,</span> <span class="n">dim</span><span class="p">:</span> <span class="n"><a class="autorefs autorefs-internal" title="<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">int</span> (<code>tinygrad.tensor.Tensor.int</code>)" href="../elementwise/#tinygrad.Tensor.int">int</a></span> <span class="o">=</span> <span class="mi">0</span><span class="p">)</span> <span class="o">-></span> <span class="n"><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">Tensor</span> (<code>tinygrad.tensor.Tensor</code>)" href="../#tinygrad.Tensor">Tensor</a></span>
|
|
</code></pre></div>
|
|
|
|
<div class="doc doc-contents first">
|
|
|
|
<p>Concatenates self with other <code class="language-python highlight"><span class="n">Tensor</span></code> in <code class="language-python highlight"><span class="n">args</span></code> along an axis specified by <code class="language-python highlight"><span class="n">dim</span></code>.
|
|
All tensors must have the same shape except in the concatenating dimension.</p>
|
|
<p><div class="language-python highlight"><pre><span></span><code><span class="n">t0</span><span class="p">,</span> <span class="n">t1</span><span class="p">,</span> <span class="n">t2</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="n">Tensor</span><span class="p">([[</span><span class="mi">3</span><span class="p">,</span> <span class="mi">4</span><span class="p">]]),</span> <span class="n">Tensor</span><span class="p">([[</span><span class="mi">5</span><span class="p">,</span> <span class="mi">6</span><span class="p">]])</span>
|
|
<span class="nb">print</span><span class="p">(</span><span class="n">t0</span><span class="o">.</span><span class="n">cat</span><span class="p">(</span><span class="n">t1</span><span class="p">,</span> <span class="n">t2</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">[[</span><span class="mi">1</span> <span class="mi">2</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">3</span> <span class="mi">4</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">5</span> <span class="mi">6</span><span class="p">]]</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="nb">print</span><span class="p">(</span><span class="n">t0</span><span class="o">.</span><span class="n">cat</span><span class="p">(</span><span class="n">t1</span><span class="p">,</span> <span class="n">t2</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">[[</span><span class="mi">1</span> <span class="mi">2</span> <span class="mi">3</span> <span class="mi">4</span> <span class="mi">5</span> <span class="mi">6</span><span class="p">]]</span>
|
|
</code></pre></div></p>
|
|
|
|
|
|
<details class="mkdocstrings-source">
|
|
<summary>Source code in <code>tinygrad/tensor.py</code></summary>
|
|
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal">1292</span>
|
|
<span class="normal">1293</span>
|
|
<span class="normal">1294</span>
|
|
<span class="normal">1295</span>
|
|
<span class="normal">1296</span>
|
|
<span class="normal">1297</span>
|
|
<span class="normal">1298</span>
|
|
<span class="normal">1299</span>
|
|
<span class="normal">1300</span>
|
|
<span class="normal">1301</span>
|
|
<span class="normal">1302</span>
|
|
<span class="normal">1303</span>
|
|
<span class="normal">1304</span>
|
|
<span class="normal">1305</span>
|
|
<span class="normal">1306</span>
|
|
<span class="normal">1307</span>
|
|
<span class="normal">1308</span>
|
|
<span class="normal">1309</span>
|
|
<span class="normal">1310</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="k">def</span><span class="w"> </span><span class="nf">cat</span><span class="p">(</span><span class="bp">self</span><span class="p">:</span><span class="n">Tensor</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">:</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">dim</span><span class="p">:</span><span class="nb">int</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span> <span class="o">-></span> <span class="n">Tensor</span><span class="p">:</span>
|
|
<span class="w"> </span><span class="sd">"""</span>
|
|
<span class="sd"> Concatenates self with other `Tensor` in `args` along an axis specified by `dim`.</span>
|
|
<span class="sd"> All tensors must have the same shape except in the concatenating dimension.</span>
|
|
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> t0, t1, t2 = Tensor([[1, 2]]), Tensor([[3, 4]]), Tensor([[5, 6]])</span>
|
|
<span class="sd"> print(t0.cat(t1, t2, dim=0).numpy())</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> print(t0.cat(t1, t2, dim=1).numpy())</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> """</span>
|
|
<span class="n">dim</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_resolve_dim</span><span class="p">(</span><span class="n">dim</span><span class="p">)</span>
|
|
<span class="k">for</span> <span class="n">arg</span> <span class="ow">in</span> <span class="n">args</span><span class="p">:</span> <span class="k">assert</span> <span class="n">arg</span><span class="o">.</span><span class="n">ndim</span><span class="o">==</span><span class="bp">self</span><span class="o">.</span><span class="n">ndim</span> <span class="ow">and</span> <span class="nb">all</span><span class="p">(</span><span class="n">ti</span><span class="o">==</span><span class="n">ai</span> <span class="k">for</span> <span class="n">i</span><span class="p">,(</span><span class="n">ti</span><span class="p">,</span><span class="n">ai</span><span class="p">)</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="nb">zip</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="n">arg</span><span class="o">.</span><span class="n">shape</span><span class="p">))</span> <span class="k">if</span> <span class="n">i</span><span class="o">!=</span><span class="n">dim</span><span class="p">)</span>
|
|
<span class="n">tensors</span> <span class="o">=</span> <span class="p">[</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">]</span>
|
|
<span class="n">dim_cumsum</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="n">itertools</span><span class="o">.</span><span class="n">accumulate</span><span class="p">([</span><span class="n">t</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="n">dim</span><span class="p">]</span> <span class="k">for</span> <span class="n">t</span> <span class="ow">in</span> <span class="n">tensors</span><span class="p">],</span> <span class="n">initial</span><span class="o">=</span><span class="mi">0</span><span class="p">))</span>
|
|
<span class="k">for</span> <span class="n">i</span><span class="p">,</span><span class="n">t</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">tensors</span><span class="p">):</span> <span class="n">tensors</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="o">=</span> <span class="n">t</span><span class="o">.</span><span class="n">pad</span><span class="p">([(</span><span class="n">dim_cumsum</span><span class="p">[</span><span class="n">i</span><span class="p">],</span> <span class="n">dim_cumsum</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span><span class="o">-</span><span class="n">dim_cumsum</span><span class="p">[</span><span class="n">i</span><span class="o">+</span><span class="mi">1</span><span class="p">])</span> <span class="k">if</span> <span class="n">j</span><span class="o">==</span><span class="n">dim</span> <span class="k">else</span> <span class="kc">None</span> <span class="k">for</span> <span class="n">j</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">t</span><span class="o">.</span><span class="n">ndim</span><span class="p">)])</span>
|
|
<span class="k">return</span> <span class="n">functools</span><span class="o">.</span><span class="n">reduce</span><span class="p">(</span><span class="n">Tensor</span><span class="o">.</span><span class="n">add</span><span class="p">,</span> <span class="n">tensors</span><span class="p">)</span>
|
|
</code></pre></div></td></tr></table></div>
|
|
</details>
|
|
</div>
|
|
|
|
</div>
|
|
|
|
<div class="doc doc-object doc-function">
|
|
|
|
|
|
<h3 id="tinygrad.Tensor.stack" class="doc doc-heading">
|
|
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">stack</span>
|
|
|
|
|
|
<a href="#tinygrad.Tensor.stack" class="headerlink" title="Permanent link">¤</a></h3>
|
|
<div class="language-python doc-signature highlight"><pre><span></span><code><span class="nf">stack</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">:</span> <span class="n"><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">Tensor</span> (<code>tinygrad.tensor.Tensor</code>)" href="../#tinygrad.Tensor">Tensor</a></span><span class="p">,</span> <span class="n">dim</span><span class="p">:</span> <span class="n"><a class="autorefs autorefs-internal" title="<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">int</span> (<code>tinygrad.tensor.Tensor.int</code>)" href="../elementwise/#tinygrad.Tensor.int">int</a></span> <span class="o">=</span> <span class="mi">0</span><span class="p">)</span> <span class="o">-></span> <span class="n"><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">Tensor</span> (<code>tinygrad.tensor.Tensor</code>)" href="../#tinygrad.Tensor">Tensor</a></span>
|
|
</code></pre></div>
|
|
|
|
<div class="doc doc-contents first">
|
|
|
|
<p>Concatenates self with other <code class="language-python highlight"><span class="n">Tensor</span></code> in <code class="language-python highlight"><span class="n">args</span></code> along a new dimension specified by <code class="language-python highlight"><span class="n">dim</span></code>.</p>
|
|
<p><div class="language-python highlight"><pre><span></span><code><span class="n">t0</span><span class="p">,</span> <span class="n">t1</span><span class="p">,</span> <span class="n">t2</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="n">Tensor</span><span class="p">([</span><span class="mi">3</span><span class="p">,</span> <span class="mi">4</span><span class="p">]),</span> <span class="n">Tensor</span><span class="p">([</span><span class="mi">5</span><span class="p">,</span> <span class="mi">6</span><span class="p">])</span>
|
|
<span class="nb">print</span><span class="p">(</span><span class="n">t0</span><span class="o">.</span><span class="n">stack</span><span class="p">(</span><span class="n">t1</span><span class="p">,</span> <span class="n">t2</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">[[</span><span class="mi">1</span> <span class="mi">2</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">3</span> <span class="mi">4</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">5</span> <span class="mi">6</span><span class="p">]]</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="nb">print</span><span class="p">(</span><span class="n">t0</span><span class="o">.</span><span class="n">stack</span><span class="p">(</span><span class="n">t1</span><span class="p">,</span> <span class="n">t2</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">[[</span><span class="mi">1</span> <span class="mi">3</span> <span class="mi">5</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">2</span> <span class="mi">4</span> <span class="mi">6</span><span class="p">]]</span>
|
|
</code></pre></div></p>
|
|
|
|
|
|
<details class="mkdocstrings-source">
|
|
<summary>Source code in <code>tinygrad/tensor.py</code></summary>
|
|
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal">1312</span>
|
|
<span class="normal">1313</span>
|
|
<span class="normal">1314</span>
|
|
<span class="normal">1315</span>
|
|
<span class="normal">1316</span>
|
|
<span class="normal">1317</span>
|
|
<span class="normal">1318</span>
|
|
<span class="normal">1319</span>
|
|
<span class="normal">1320</span>
|
|
<span class="normal">1321</span>
|
|
<span class="normal">1322</span>
|
|
<span class="normal">1323</span>
|
|
<span class="normal">1324</span>
|
|
<span class="normal">1325</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="k">def</span><span class="w"> </span><span class="nf">stack</span><span class="p">(</span><span class="bp">self</span><span class="p">:</span><span class="n">Tensor</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">:</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">dim</span><span class="p">:</span><span class="nb">int</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span> <span class="o">-></span> <span class="n">Tensor</span><span class="p">:</span>
|
|
<span class="w"> </span><span class="sd">"""</span>
|
|
<span class="sd"> Concatenates self with other `Tensor` in `args` along a new dimension specified by `dim`.</span>
|
|
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> t0, t1, t2 = Tensor([1, 2]), Tensor([3, 4]), Tensor([5, 6])</span>
|
|
<span class="sd"> print(t0.stack(t1, t2, dim=0).numpy())</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> print(t0.stack(t1, t2, dim=1).numpy())</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> """</span>
|
|
<span class="c1"># checks for shapes and number of dimensions delegated to cat</span>
|
|
<span class="k">return</span> <span class="n">Tensor</span><span class="o">.</span><span class="n">cat</span><span class="p">(</span><span class="o">*</span><span class="p">[</span><span class="n">t</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="n">dim</span><span class="p">)</span> <span class="k">for</span> <span class="n">t</span> <span class="ow">in</span> <span class="n">argfix</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">)],</span> <span class="n">dim</span><span class="o">=</span><span class="n">dim</span><span class="p">)</span>
|
|
</code></pre></div></td></tr></table></div>
|
|
</details>
|
|
</div>
|
|
|
|
</div>
|
|
|
|
<div class="doc doc-object doc-function">
|
|
|
|
|
|
<h3 id="tinygrad.Tensor.repeat" class="doc doc-heading">
|
|
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">repeat</span>
|
|
|
|
|
|
<a href="#tinygrad.Tensor.repeat" class="headerlink" title="Permanent link">¤</a></h3>
|
|
<div class="language-python doc-signature highlight"><pre><span></span><code><span class="nf">repeat</span><span class="p">(</span><span class="n">repeats</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"><a class="autorefs autorefs-external" title="<code>typing.Self</code>" href="https://docs.python.org/3/library/typing.html#typing.Self">Self</a></span>
|
|
</code></pre></div>
|
|
|
|
<div class="doc doc-contents first">
|
|
|
|
<p>Repeats tensor number of times along each dimension specified by <code class="language-python highlight"><span class="n">repeats</span></code>.
|
|
<code class="language-python highlight"><span class="n">repeats</span></code> can be passed as a tuple or as separate arguments.</p>
|
|
<p><div class="language-python highlight"><pre><span></span><code><span class="n">t</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>
|
|
<span class="nb">print</span><span class="p">(</span><span class="n">t</span><span class="o">.</span><span class="n">repeat</span><span class="p">(</span><span class="mi">4</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">[[</span><span class="mi">1</span> <span class="mi">2</span> <span class="mi">3</span> <span class="mi">1</span> <span class="mi">2</span> <span class="mi">3</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">1</span> <span class="mi">2</span> <span class="mi">3</span> <span class="mi">1</span> <span class="mi">2</span> <span class="mi">3</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">1</span> <span class="mi">2</span> <span class="mi">3</span> <span class="mi">1</span> <span class="mi">2</span> <span class="mi">3</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">1</span> <span class="mi">2</span> <span class="mi">3</span> <span class="mi">1</span> <span class="mi">2</span> <span class="mi">3</span><span class="p">]]</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="nb">print</span><span class="p">(</span><span class="n">t</span><span class="o">.</span><span class="n">repeat</span><span class="p">(</span><span class="mi">4</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">(</span><span class="mi">4</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">)</span>
|
|
</code></pre></div></p>
|
|
|
|
|
|
<details class="mkdocstrings-source">
|
|
<summary>Source code in <code>tinygrad/mixin/movement.py</code></summary>
|
|
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal">335</span>
|
|
<span class="normal">336</span>
|
|
<span class="normal">337</span>
|
|
<span class="normal">338</span>
|
|
<span class="normal">339</span>
|
|
<span class="normal">340</span>
|
|
<span class="normal">341</span>
|
|
<span class="normal">342</span>
|
|
<span class="normal">343</span>
|
|
<span class="normal">344</span>
|
|
<span class="normal">345</span>
|
|
<span class="normal">346</span>
|
|
<span class="normal">347</span>
|
|
<span class="normal">348</span>
|
|
<span class="normal">349</span>
|
|
<span class="normal">350</span>
|
|
<span class="normal">351</span>
|
|
<span class="normal">352</span>
|
|
<span class="normal">353</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="k">def</span><span class="w"> </span><span class="nf">repeat</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">repeats</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">Self</span><span class="p">:</span>
|
|
<span class="w"> </span><span class="sd">"""</span>
|
|
<span class="sd"> Repeats tensor number of times along each dimension specified by `repeats`.</span>
|
|
<span class="sd"> `repeats` can be passed as a tuple or as separate arguments.</span>
|
|
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> t = Tensor([1, 2, 3])</span>
|
|
<span class="sd"> print(t.repeat(4, 2).numpy())</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> print(t.repeat(4, 2, 1).shape)</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> """</span>
|
|
<span class="n">repeats</span> <span class="o">=</span> <span class="n">argfix</span><span class="p">(</span><span class="n">repeats</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">)</span>
|
|
<span class="n">base_shape</span> <span class="o">=</span> <span class="n">_align_left</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="n">repeats</span><span class="p">)[</span><span class="mi">0</span><span class="p">]</span>
|
|
<span class="n">unsqueezed_shape</span> <span class="o">=</span> <span class="n">flatten</span><span class="p">([[</span><span class="n">s</span><span class="p">]</span> <span class="k">if</span> <span class="n">r</span> <span class="o">==</span> <span class="mi">1</span> <span class="k">else</span> <span class="p">[</span><span class="mi">1</span><span class="p">,</span> <span class="n">s</span><span class="p">]</span> <span class="k">for</span> <span class="n">r</span><span class="p">,</span> <span class="n">s</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">repeats</span><span class="p">,</span> <span class="n">base_shape</span><span class="p">)])</span>
|
|
<span class="n">expanded_shape</span> <span class="o">=</span> <span class="n">flatten</span><span class="p">([[</span><span class="n">s</span><span class="p">]</span> <span class="k">if</span> <span class="n">r</span> <span class="o">==</span> <span class="mi">1</span> <span class="k">else</span> <span class="p">[</span><span class="n">r</span><span class="p">,</span> <span class="n">s</span><span class="p">]</span> <span class="k">for</span> <span class="n">r</span><span class="p">,</span> <span class="n">s</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">repeats</span><span class="p">,</span> <span class="n">base_shape</span><span class="p">)])</span>
|
|
<span class="n">final_shape</span> <span class="o">=</span> <span class="p">[</span><span class="n">r</span> <span class="o">*</span> <span class="n">s</span> <span class="k">for</span> <span class="n">r</span><span class="p">,</span> <span class="n">s</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">repeats</span><span class="p">,</span> <span class="n">base_shape</span><span class="p">)]</span>
|
|
<span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">unsqueezed_shape</span><span class="p">)</span><span class="o">.</span><span class="n">expand</span><span class="p">(</span><span class="n">expanded_shape</span><span class="p">)</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">final_shape</span><span class="p">)</span>
|
|
</code></pre></div></td></tr></table></div>
|
|
</details>
|
|
</div>
|
|
|
|
</div>
|
|
|
|
<div class="doc doc-object doc-function">
|
|
|
|
|
|
<h3 id="tinygrad.Tensor.repeat_interleave" class="doc doc-heading">
|
|
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">repeat_interleave</span>
|
|
|
|
|
|
<a href="#tinygrad.Tensor.repeat_interleave" class="headerlink" title="Permanent link">¤</a></h3>
|
|
<div class="language-python doc-signature highlight"><pre><span></span><code><span class="nf">repeat_interleave</span><span class="p">(</span>
|
|
<span class="n">repeats</span><span class="p">:</span> <span class="n"><a class="autorefs autorefs-external" href="https://docs.python.org/3/library/functions.html#int">int</a></span><span class="p">,</span> <span class="n">dim</span><span class="p">:</span> <span class="n"><a class="autorefs autorefs-external" href="https://docs.python.org/3/library/functions.html#int">int</a></span> <span class="o">|</span> <span class="kc">None</span> <span class="o">=</span> <span class="kc">None</span>
|
|
<span class="p">)</span> <span class="o">-></span> <span class="n"><a class="autorefs autorefs-external" title="<code>typing.Self</code>" href="https://docs.python.org/3/library/typing.html#typing.Self">Self</a></span>
|
|
</code></pre></div>
|
|
|
|
<div class="doc doc-contents first">
|
|
|
|
<p>Repeats elements of a tensor.</p>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="n">t</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>
|
|
<span class="nb">print</span><span class="p">(</span><span class="n">t</span><span class="o">.</span><span class="n">repeat_interleave</span><span class="p">(</span><span class="mi">2</span><span class="p">)</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">[</span><span class="mi">1</span> <span class="mi">1</span> <span class="mi">2</span> <span class="mi">2</span> <span class="mi">3</span> <span class="mi">3</span><span class="p">]</span>
|
|
</code></pre></div>
|
|
|
|
|
|
<details class="mkdocstrings-source">
|
|
<summary>Source code in <code>tinygrad/mixin/movement.py</code></summary>
|
|
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal">319</span>
|
|
<span class="normal">320</span>
|
|
<span class="normal">321</span>
|
|
<span class="normal">322</span>
|
|
<span class="normal">323</span>
|
|
<span class="normal">324</span>
|
|
<span class="normal">325</span>
|
|
<span class="normal">326</span>
|
|
<span class="normal">327</span>
|
|
<span class="normal">328</span>
|
|
<span class="normal">329</span>
|
|
<span class="normal">330</span>
|
|
<span class="normal">331</span>
|
|
<span class="normal">332</span>
|
|
<span class="normal">333</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="k">def</span><span class="w"> </span><span class="nf">repeat_interleave</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">repeats</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">dim</span><span class="p">:</span> <span class="nb">int</span> <span class="o">|</span> <span class="kc">None</span> <span class="o">=</span> <span class="kc">None</span><span class="p">)</span> <span class="o">-></span> <span class="n">Self</span><span class="p">:</span>
|
|
<span class="w"> </span><span class="sd">"""</span>
|
|
<span class="sd"> Repeats elements of a tensor.</span>
|
|
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> t = Tensor([1, 2, 3])</span>
|
|
<span class="sd"> print(t.repeat_interleave(2).numpy())</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> """</span>
|
|
<span class="n">x</span><span class="p">,</span> <span class="n">dim</span> <span class="o">=</span> <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">flatten</span><span class="p">(),</span> <span class="mi">0</span><span class="p">)</span> <span class="k">if</span> <span class="n">dim</span> <span class="ow">is</span> <span class="kc">None</span> <span class="k">else</span> <span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">_resolve_dim</span><span class="p">(</span><span class="n">dim</span><span class="p">))</span>
|
|
<span class="n">shp</span> <span class="o">=</span> <span class="n">x</span><span class="o">.</span><span class="n">shape</span>
|
|
<span class="n">x</span> <span class="o">=</span> <span class="n">x</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="o">*</span><span class="n">shp</span><span class="p">[:</span> <span class="n">dim</span> <span class="o">+</span> <span class="mi">1</span><span class="p">],</span> <span class="mi">1</span><span class="p">,</span> <span class="o">*</span><span class="n">shp</span><span class="p">[</span><span class="n">dim</span> <span class="o">+</span> <span class="mi">1</span> <span class="p">:])</span>
|
|
<span class="n">x</span> <span class="o">=</span> <span class="n">x</span><span class="o">.</span><span class="n">expand</span><span class="p">(</span><span class="o">*</span><span class="n">shp</span><span class="p">[:</span> <span class="n">dim</span> <span class="o">+</span> <span class="mi">1</span><span class="p">],</span> <span class="n">repeats</span><span class="p">,</span> <span class="o">*</span><span class="n">shp</span><span class="p">[</span><span class="n">dim</span> <span class="o">+</span> <span class="mi">1</span> <span class="p">:])</span>
|
|
<span class="n">x</span> <span class="o">=</span> <span class="n">x</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="o">*</span><span class="n">shp</span><span class="p">[:</span><span class="n">dim</span><span class="p">],</span> <span class="n">shp</span><span class="p">[</span><span class="n">dim</span><span class="p">]</span> <span class="o">*</span> <span class="n">repeats</span><span class="p">,</span> <span class="o">*</span><span class="n">shp</span><span class="p">[</span><span class="n">dim</span> <span class="o">+</span> <span class="mi">1</span> <span class="p">:])</span>
|
|
<span class="k">return</span> <span class="n">x</span>
|
|
</code></pre></div></td></tr></table></div>
|
|
</details>
|
|
</div>
|
|
|
|
</div>
|
|
|
|
<div class="doc doc-object doc-function">
|
|
|
|
|
|
<h3 id="tinygrad.Tensor.split" class="doc doc-heading">
|
|
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">split</span>
|
|
|
|
|
|
<a href="#tinygrad.Tensor.split" class="headerlink" title="Permanent link">¤</a></h3>
|
|
<div class="language-python doc-signature highlight"><pre><span></span><code><span class="nf">split</span><span class="p">(</span>
|
|
<span class="n">sizes</span><span class="p">:</span> <span class="n"><a class="autorefs autorefs-internal" title="<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">int</span> (<code>tinygrad.tensor.Tensor.int</code>)" href="../elementwise/#tinygrad.Tensor.int">int</a></span> <span class="o">|</span> <span class="n"><a class="autorefs autorefs-external" title="<code>typing.Sequence</code>" href="https://docs.python.org/3/library/typing.html#typing.Sequence">Sequence</a></span><span class="p">[</span><span class="n"><a class="autorefs autorefs-internal" title="<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">int</span> (<code>tinygrad.tensor.Tensor.int</code>)" href="../elementwise/#tinygrad.Tensor.int">int</a></span><span class="p">],</span> <span class="n">dim</span><span class="p">:</span> <span class="n"><a class="autorefs autorefs-internal" title="<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">int</span> (<code>tinygrad.tensor.Tensor.int</code>)" href="../elementwise/#tinygrad.Tensor.int">int</a></span> <span class="o">=</span> <span class="mi">0</span>
|
|
<span class="p">)</span> <span class="o">-></span> <span class="n"><a class="autorefs autorefs-external" href="https://docs.python.org/3/library/stdtypes.html#tuple">tuple</a></span><span class="p">[</span><span class="n"><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">Tensor</span> (<code>tinygrad.tensor.Tensor</code>)" href="../#tinygrad.Tensor">Tensor</a></span><span class="p">,</span> <span class="o">...</span><span class="p">]</span>
|
|
</code></pre></div>
|
|
|
|
<div class="doc doc-contents first">
|
|
|
|
<p>Splits the tensor into chunks along the dimension specified by <code class="language-python highlight"><span class="n">dim</span></code>.
|
|
If <code class="language-python highlight"><span class="n">sizes</span></code> is an integer, it splits into equally sized chunks if possible, otherwise the last chunk will be smaller.
|
|
If <code class="language-python highlight"><span class="n">sizes</span></code> is a list, it splits into <code class="language-python highlight"><span class="nb">len</span><span class="p">(</span><span class="n">sizes</span><span class="p">)</span></code> chunks with size in <code class="language-python highlight"><span class="n">dim</span></code> according to <code class="language-python highlight"><span class="n">size</span></code>.</p>
|
|
<p><div class="language-python highlight"><pre><span></span><code><span class="n">t</span> <span class="o">=</span> <span class="n">Tensor</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="mi">10</span><span class="p">)</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="mi">5</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span>
|
|
<span class="nb">print</span><span class="p">(</span><span class="n">t</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">[[</span><span class="mi">0</span> <span class="mi">1</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">2</span> <span class="mi">3</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">4</span> <span class="mi">5</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">6</span> <span class="mi">7</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">8</span> <span class="mi">9</span><span class="p">]]</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="n">split</span> <span class="o">=</span> <span class="n">t</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="mi">2</span><span class="p">)</span>
|
|
<span class="nb">print</span><span class="p">(</span><span class="s2">"</span><span class="se">\n</span><span class="s2">"</span><span class="o">.</span><span class="n">join</span><span class="p">([</span><span class="nb">repr</span><span class="p">(</span><span class="n">x</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span> <span class="k">for</span> <span class="n">x</span> <span class="ow">in</span> <span class="n">split</span><span class="p">]))</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="n">array</span><span class="p">([[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">],</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="n">dtype</span><span class="o">=</span><span class="n">int32</span><span class="p">)</span>
|
|
<span class="n">array</span><span class="p">([[</span><span class="mi">4</span><span class="p">,</span> <span class="mi">5</span><span class="p">],</span>
|
|
<span class="p">[</span><span class="mi">6</span><span class="p">,</span> <span class="mi">7</span><span class="p">]],</span> <span class="n">dtype</span><span class="o">=</span><span class="n">int32</span><span class="p">)</span>
|
|
<span class="n">array</span><span class="p">([[</span><span class="mi">8</span><span class="p">,</span> <span class="mi">9</span><span class="p">]],</span> <span class="n">dtype</span><span class="o">=</span><span class="n">int32</span><span class="p">)</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="n">split</span> <span class="o">=</span> <span class="n">t</span><span class="o">.</span><span class="n">split</span><span class="p">([</span><span class="mi">1</span><span class="p">,</span> <span class="mi">4</span><span class="p">])</span>
|
|
<span class="nb">print</span><span class="p">(</span><span class="s2">"</span><span class="se">\n</span><span class="s2">"</span><span class="o">.</span><span class="n">join</span><span class="p">([</span><span class="nb">repr</span><span class="p">(</span><span class="n">x</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span> <span class="k">for</span> <span class="n">x</span> <span class="ow">in</span> <span class="n">split</span><span class="p">]))</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="n">array</span><span class="p">([[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">]],</span> <span class="n">dtype</span><span class="o">=</span><span class="n">int32</span><span class="p">)</span>
|
|
<span class="n">array</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>
|
|
<span class="p">[</span><span class="mi">6</span><span class="p">,</span> <span class="mi">7</span><span class="p">],</span>
|
|
<span class="p">[</span><span class="mi">8</span><span class="p">,</span> <span class="mi">9</span><span class="p">]],</span> <span class="n">dtype</span><span class="o">=</span><span class="n">int32</span><span class="p">)</span>
|
|
</code></pre></div></p>
|
|
|
|
|
|
<details class="mkdocstrings-source">
|
|
<summary>Source code in <code>tinygrad/tensor.py</code></summary>
|
|
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal">1327</span>
|
|
<span class="normal">1328</span>
|
|
<span class="normal">1329</span>
|
|
<span class="normal">1330</span>
|
|
<span class="normal">1331</span>
|
|
<span class="normal">1332</span>
|
|
<span class="normal">1333</span>
|
|
<span class="normal">1334</span>
|
|
<span class="normal">1335</span>
|
|
<span class="normal">1336</span>
|
|
<span class="normal">1337</span>
|
|
<span class="normal">1338</span>
|
|
<span class="normal">1339</span>
|
|
<span class="normal">1340</span>
|
|
<span class="normal">1341</span>
|
|
<span class="normal">1342</span>
|
|
<span class="normal">1343</span>
|
|
<span class="normal">1344</span>
|
|
<span class="normal">1345</span>
|
|
<span class="normal">1346</span>
|
|
<span class="normal">1347</span>
|
|
<span class="normal">1348</span>
|
|
<span class="normal">1349</span>
|
|
<span class="normal">1350</span>
|
|
<span class="normal">1351</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="k">def</span><span class="w"> </span><span class="nf">split</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">sizes</span><span class="p">:</span><span class="nb">int</span><span class="o">|</span><span class="n">Sequence</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span> <span class="n">dim</span><span class="p">:</span><span class="nb">int</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span> <span class="o">-></span> <span class="nb">tuple</span><span class="p">[</span><span class="n">Tensor</span><span class="p">,</span> <span class="o">...</span><span class="p">]:</span>
|
|
<span class="w"> </span><span class="sd">"""</span>
|
|
<span class="sd"> Splits the tensor into chunks along the dimension specified by `dim`.</span>
|
|
<span class="sd"> If `sizes` is an integer, it splits into equally sized chunks if possible, otherwise the last chunk will be smaller.</span>
|
|
<span class="sd"> If `sizes` is a list, it splits into `len(sizes)` chunks with size in `dim` according to `size`.</span>
|
|
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> t = Tensor.arange(10).reshape(5, 2)</span>
|
|
<span class="sd"> print(t.numpy())</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> split = t.split(2)</span>
|
|
<span class="sd"> print("\\n".join([repr(x.numpy()) for x in split]))</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> split = t.split([1, 4])</span>
|
|
<span class="sd"> print("\\n".join([repr(x.numpy()) for x in split]))</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> """</span>
|
|
<span class="n">dim</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_resolve_dim</span><span class="p">(</span><span class="n">dim</span><span class="p">)</span>
|
|
<span class="n">dim_sz</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="n">dim</span><span class="p">]</span>
|
|
<span class="k">assert</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">dim_sz</span><span class="p">,</span> <span class="nb">int</span><span class="p">),</span> <span class="sa">f</span><span class="s2">"does not support symbolic shape in split dimension </span><span class="si">{</span><span class="n">dim</span><span class="si">}</span><span class="s2">: </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="si">}</span><span class="s2">"</span>
|
|
<span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">sizes</span><span class="p">,</span> <span class="nb">int</span><span class="p">):</span> <span class="n">sizes</span> <span class="o">=</span> <span class="p">[</span><span class="nb">min</span><span class="p">(</span><span class="n">sizes</span><span class="p">,</span> <span class="n">dim_sz</span><span class="o">-</span><span class="n">i</span><span class="p">)</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="nb">max</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="n">dim_sz</span><span class="p">),</span> <span class="nb">max</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="n">sizes</span><span class="p">))]</span>
|
|
<span class="k">assert</span> <span class="nb">sum</span><span class="p">(</span><span class="n">sizes</span><span class="p">)</span> <span class="o">==</span> <span class="n">dim_sz</span><span class="p">,</span> <span class="sa">f</span><span class="s2">"expect sizes to sum exactly to </span><span class="si">{</span><span class="n">dim_sz</span><span class="si">}</span><span class="s2">, but got </span><span class="si">{</span><span class="nb">sum</span><span class="p">(</span><span class="n">sizes</span><span class="p">)</span><span class="si">}</span><span class="s2">"</span>
|
|
<span class="k">return</span> <span class="nb">tuple</span><span class="p">(</span><span class="bp">self</span><span class="p">[</span><span class="n">sl</span><span class="p">]</span> <span class="k">for</span> <span class="n">sl</span> <span class="ow">in</span> <span class="p">[</span><span class="nb">tuple</span><span class="p">([</span><span class="nb">slice</span><span class="p">(</span><span class="kc">None</span><span class="p">)]</span><span class="o">*</span><span class="n">dim</span> <span class="o">+</span> <span class="p">[</span><span class="nb">slice</span><span class="p">(</span><span class="nb">sum</span><span class="p">(</span><span class="n">sizes</span><span class="p">[:</span><span class="n">i</span><span class="p">]),</span> <span class="nb">sum</span><span class="p">(</span><span class="n">sizes</span><span class="p">[:</span><span class="n">i</span> <span class="o">+</span> <span class="mi">1</span><span class="p">]))])</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">sizes</span><span class="p">))])</span>
|
|
</code></pre></div></td></tr></table></div>
|
|
</details>
|
|
</div>
|
|
|
|
</div>
|
|
|
|
<div class="doc doc-object doc-function">
|
|
|
|
|
|
<h3 id="tinygrad.Tensor.chunk" class="doc doc-heading">
|
|
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">chunk</span>
|
|
|
|
|
|
<a href="#tinygrad.Tensor.chunk" class="headerlink" title="Permanent link">¤</a></h3>
|
|
<div class="language-python doc-signature highlight"><pre><span></span><code><span class="nf">chunk</span><span class="p">(</span><span class="n">chunks</span><span class="p">:</span> <span class="n"><a class="autorefs autorefs-internal" title="<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">int</span> (<code>tinygrad.tensor.Tensor.int</code>)" href="../elementwise/#tinygrad.Tensor.int">int</a></span><span class="p">,</span> <span class="n">dim</span><span class="p">:</span> <span class="n"><a class="autorefs autorefs-internal" title="<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">int</span> (<code>tinygrad.tensor.Tensor.int</code>)" href="../elementwise/#tinygrad.Tensor.int">int</a></span> <span class="o">=</span> <span class="mi">0</span><span class="p">)</span> <span class="o">-></span> <span class="n"><a class="autorefs autorefs-external" href="https://docs.python.org/3/library/stdtypes.html#list">list</a></span><span class="p">[</span><span class="n"><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">Tensor</span> (<code>tinygrad.tensor.Tensor</code>)" href="../#tinygrad.Tensor">Tensor</a></span><span class="p">]</span>
|
|
</code></pre></div>
|
|
|
|
<div class="doc doc-contents first">
|
|
|
|
<p>Splits the tensor into <code class="language-python highlight"><span class="n">chunks</span></code> number of chunks along the dimension <code class="language-python highlight"><span class="n">dim</span></code>.
|
|
If the tensor size along <code class="language-python highlight"><span class="n">dim</span></code> is not divisible by <code class="language-python highlight"><span class="n">chunks</span></code>, all returned chunks will be the same size except the last one.
|
|
The function may return fewer than the specified number of chunks.</p>
|
|
<p><div class="language-python highlight"><pre><span></span><code><span class="n">chunked</span> <span class="o">=</span> <span class="n">Tensor</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="mi">11</span><span class="p">)</span><span class="o">.</span><span class="n">chunk</span><span class="p">(</span><span class="mi">6</span><span class="p">)</span>
|
|
<span class="nb">print</span><span class="p">(</span><span class="s2">"</span><span class="se">\n</span><span class="s2">"</span><span class="o">.</span><span class="n">join</span><span class="p">([</span><span class="nb">repr</span><span class="p">(</span><span class="n">x</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span> <span class="k">for</span> <span class="n">x</span> <span class="ow">in</span> <span class="n">chunked</span><span class="p">]))</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="n">array</span><span class="p">([</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">],</span> <span class="n">dtype</span><span class="o">=</span><span class="n">int32</span><span class="p">)</span>
|
|
<span class="n">array</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="n">dtype</span><span class="o">=</span><span class="n">int32</span><span class="p">)</span>
|
|
<span class="n">array</span><span class="p">([</span><span class="mi">4</span><span class="p">,</span> <span class="mi">5</span><span class="p">],</span> <span class="n">dtype</span><span class="o">=</span><span class="n">int32</span><span class="p">)</span>
|
|
<span class="n">array</span><span class="p">([</span><span class="mi">6</span><span class="p">,</span> <span class="mi">7</span><span class="p">],</span> <span class="n">dtype</span><span class="o">=</span><span class="n">int32</span><span class="p">)</span>
|
|
<span class="n">array</span><span class="p">([</span><span class="mi">8</span><span class="p">,</span> <span class="mi">9</span><span class="p">],</span> <span class="n">dtype</span><span class="o">=</span><span class="n">int32</span><span class="p">)</span>
|
|
<span class="n">array</span><span class="p">([</span><span class="mi">10</span><span class="p">],</span> <span class="n">dtype</span><span class="o">=</span><span class="n">int32</span><span class="p">)</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="n">chunked</span> <span class="o">=</span> <span class="n">Tensor</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="mi">12</span><span class="p">)</span><span class="o">.</span><span class="n">chunk</span><span class="p">(</span><span class="mi">6</span><span class="p">)</span>
|
|
<span class="nb">print</span><span class="p">(</span><span class="s2">"</span><span class="se">\n</span><span class="s2">"</span><span class="o">.</span><span class="n">join</span><span class="p">([</span><span class="nb">repr</span><span class="p">(</span><span class="n">x</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span> <span class="k">for</span> <span class="n">x</span> <span class="ow">in</span> <span class="n">chunked</span><span class="p">]))</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="n">array</span><span class="p">([</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">],</span> <span class="n">dtype</span><span class="o">=</span><span class="n">int32</span><span class="p">)</span>
|
|
<span class="n">array</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="n">dtype</span><span class="o">=</span><span class="n">int32</span><span class="p">)</span>
|
|
<span class="n">array</span><span class="p">([</span><span class="mi">4</span><span class="p">,</span> <span class="mi">5</span><span class="p">],</span> <span class="n">dtype</span><span class="o">=</span><span class="n">int32</span><span class="p">)</span>
|
|
<span class="n">array</span><span class="p">([</span><span class="mi">6</span><span class="p">,</span> <span class="mi">7</span><span class="p">],</span> <span class="n">dtype</span><span class="o">=</span><span class="n">int32</span><span class="p">)</span>
|
|
<span class="n">array</span><span class="p">([</span><span class="mi">8</span><span class="p">,</span> <span class="mi">9</span><span class="p">],</span> <span class="n">dtype</span><span class="o">=</span><span class="n">int32</span><span class="p">)</span>
|
|
<span class="n">array</span><span class="p">([</span><span class="mi">10</span><span class="p">,</span> <span class="mi">11</span><span class="p">],</span> <span class="n">dtype</span><span class="o">=</span><span class="n">int32</span><span class="p">)</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="n">chunked</span> <span class="o">=</span> <span class="n">Tensor</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="mi">13</span><span class="p">)</span><span class="o">.</span><span class="n">chunk</span><span class="p">(</span><span class="mi">6</span><span class="p">)</span>
|
|
<span class="nb">print</span><span class="p">(</span><span class="s2">"</span><span class="se">\n</span><span class="s2">"</span><span class="o">.</span><span class="n">join</span><span class="p">([</span><span class="nb">repr</span><span class="p">(</span><span class="n">x</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span> <span class="k">for</span> <span class="n">x</span> <span class="ow">in</span> <span class="n">chunked</span><span class="p">]))</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="n">array</span><span class="p">([</span><span class="mi">0</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="n">dtype</span><span class="o">=</span><span class="n">int32</span><span class="p">)</span>
|
|
<span class="n">array</span><span class="p">([</span><span class="mi">3</span><span class="p">,</span> <span class="mi">4</span><span class="p">,</span> <span class="mi">5</span><span class="p">],</span> <span class="n">dtype</span><span class="o">=</span><span class="n">int32</span><span class="p">)</span>
|
|
<span class="n">array</span><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="n">dtype</span><span class="o">=</span><span class="n">int32</span><span class="p">)</span>
|
|
<span class="n">array</span><span class="p">([</span> <span class="mi">9</span><span class="p">,</span> <span class="mi">10</span><span class="p">,</span> <span class="mi">11</span><span class="p">],</span> <span class="n">dtype</span><span class="o">=</span><span class="n">int32</span><span class="p">)</span>
|
|
<span class="n">array</span><span class="p">([</span><span class="mi">12</span><span class="p">],</span> <span class="n">dtype</span><span class="o">=</span><span class="n">int32</span><span class="p">)</span>
|
|
</code></pre></div></p>
|
|
|
|
|
|
<details class="mkdocstrings-source">
|
|
<summary>Source code in <code>tinygrad/tensor.py</code></summary>
|
|
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal">1353</span>
|
|
<span class="normal">1354</span>
|
|
<span class="normal">1355</span>
|
|
<span class="normal">1356</span>
|
|
<span class="normal">1357</span>
|
|
<span class="normal">1358</span>
|
|
<span class="normal">1359</span>
|
|
<span class="normal">1360</span>
|
|
<span class="normal">1361</span>
|
|
<span class="normal">1362</span>
|
|
<span class="normal">1363</span>
|
|
<span class="normal">1364</span>
|
|
<span class="normal">1365</span>
|
|
<span class="normal">1366</span>
|
|
<span class="normal">1367</span>
|
|
<span class="normal">1368</span>
|
|
<span class="normal">1369</span>
|
|
<span class="normal">1370</span>
|
|
<span class="normal">1371</span>
|
|
<span class="normal">1372</span>
|
|
<span class="normal">1373</span>
|
|
<span class="normal">1374</span>
|
|
<span class="normal">1375</span>
|
|
<span class="normal">1376</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="k">def</span><span class="w"> </span><span class="nf">chunk</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">chunks</span><span class="p">:</span><span class="nb">int</span><span class="p">,</span> <span class="n">dim</span><span class="p">:</span><span class="nb">int</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span> <span class="o">-></span> <span class="nb">list</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]:</span>
|
|
<span class="w"> </span><span class="sd">"""</span>
|
|
<span class="sd"> Splits the tensor into `chunks` number of chunks along the dimension `dim`.</span>
|
|
<span class="sd"> If the tensor size along `dim` is not divisible by `chunks`, all returned chunks will be the same size except the last one.</span>
|
|
<span class="sd"> The function may return fewer than the specified number of chunks.</span>
|
|
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> chunked = Tensor.arange(11).chunk(6)</span>
|
|
<span class="sd"> print("\\n".join([repr(x.numpy()) for x in chunked]))</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> chunked = Tensor.arange(12).chunk(6)</span>
|
|
<span class="sd"> print("\\n".join([repr(x.numpy()) for x in chunked]))</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> chunked = Tensor.arange(13).chunk(6)</span>
|
|
<span class="sd"> print("\\n".join([repr(x.numpy()) for x in chunked]))</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> """</span>
|
|
<span class="n">dim</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_resolve_dim</span><span class="p">(</span><span class="n">dim</span><span class="p">)</span>
|
|
<span class="n">dim_sz</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="n">dim</span><span class="p">]</span>
|
|
<span class="k">assert</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">dim_sz</span><span class="p">,</span> <span class="nb">int</span><span class="p">),</span> <span class="sa">f</span><span class="s2">"does not support symbolic shape in split dimension </span><span class="si">{</span><span class="n">dim</span><span class="si">}</span><span class="s2">: </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="si">}</span><span class="s2">"</span>
|
|
<span class="k">assert</span> <span class="n">chunks</span> <span class="o">></span> <span class="mi">0</span><span class="p">,</span> <span class="sa">f</span><span class="s2">"expect chunks to be greater than 0, got: </span><span class="si">{</span><span class="n">chunks</span><span class="si">}</span><span class="s2">"</span>
|
|
<span class="k">return</span> <span class="nb">list</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="n">ceildiv</span><span class="p">(</span><span class="n">dim_sz</span><span class="p">,</span> <span class="n">chunks</span><span class="p">)</span> <span class="k">if</span> <span class="n">dim_sz</span> <span class="k">else</span> <span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">*</span><span class="n">chunks</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="n">dim</span><span class="p">))</span>
|
|
</code></pre></div></td></tr></table></div>
|
|
</details>
|
|
</div>
|
|
|
|
</div>
|
|
|
|
<div class="doc doc-object doc-function">
|
|
|
|
|
|
<h3 id="tinygrad.Tensor.unfold" class="doc doc-heading">
|
|
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">unfold</span>
|
|
|
|
|
|
<a href="#tinygrad.Tensor.unfold" class="headerlink" title="Permanent link">¤</a></h3>
|
|
<div class="language-python doc-signature highlight"><pre><span></span><code><span class="nf">unfold</span><span class="p">(</span><span class="n">dim</span><span class="p">:</span> <span class="n"><a class="autorefs autorefs-internal" title="<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">int</span> (<code>tinygrad.tensor.Tensor.int</code>)" href="../elementwise/#tinygrad.Tensor.int">int</a></span><span class="p">,</span> <span class="n">size</span><span class="p">:</span> <span class="n"><span title="tinygrad.uop.ops.sint">sint</span></span><span class="p">,</span> <span class="n">step</span><span class="p">:</span> <span class="n"><a class="autorefs autorefs-internal" title="<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">int</span> (<code>tinygrad.tensor.Tensor.int</code>)" href="../elementwise/#tinygrad.Tensor.int">int</a></span><span class="p">)</span> <span class="o">-></span> <span class="n"><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">Tensor</span> (<code>tinygrad.tensor.Tensor</code>)" href="../#tinygrad.Tensor">Tensor</a></span>
|
|
</code></pre></div>
|
|
|
|
<div class="doc doc-contents first">
|
|
|
|
<p>Unfolds the tensor along dimension <code class="language-python highlight"><span class="n">dim</span></code> into overlapping windows.</p>
|
|
<p>Each window has length <code class="language-python highlight"><span class="n">size</span></code> and begins every <code class="language-python highlight"><span class="n">step</span></code> elements of <code class="language-python highlight"><span class="bp">self</span></code>.
|
|
Returns the input tensor with dimension <code class="language-python highlight"><span class="n">dim</span></code> replaced by dims <code class="language-python highlight"><span class="p">(</span><span class="n">n_windows</span><span class="p">,</span> <span class="n">size</span><span class="p">)</span></code>
|
|
where <code class="language-python highlight"><span class="n">n_windows</span> <span class="o">=</span> <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="n">dim</span><span class="p">]</span> <span class="o">-</span> <span class="n">size</span><span class="p">)</span> <span class="o">//</span> <span class="n">step</span> <span class="o">+</span> <span class="mi">1</span></code>.</p>
|
|
<p><div class="language-python highlight"><pre><span></span><code><span class="n">unfolded</span> <span class="o">=</span> <span class="n">Tensor</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="mi">8</span><span class="p">)</span><span class="o">.</span><span class="n">unfold</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span><span class="mi">2</span><span class="p">,</span><span class="mi">2</span><span class="p">)</span>
|
|
<span class="nb">print</span><span class="p">(</span><span class="s2">"</span><span class="se">\n</span><span class="s2">"</span><span class="o">.</span><span class="n">join</span><span class="p">([</span><span class="nb">repr</span><span class="p">(</span><span class="n">x</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span> <span class="k">for</span> <span class="n">x</span> <span class="ow">in</span> <span class="n">unfolded</span><span class="p">]))</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="n">array</span><span class="p">([</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">],</span> <span class="n">dtype</span><span class="o">=</span><span class="n">int32</span><span class="p">)</span>
|
|
<span class="n">array</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="n">dtype</span><span class="o">=</span><span class="n">int32</span><span class="p">)</span>
|
|
<span class="n">array</span><span class="p">([</span><span class="mi">4</span><span class="p">,</span> <span class="mi">5</span><span class="p">],</span> <span class="n">dtype</span><span class="o">=</span><span class="n">int32</span><span class="p">)</span>
|
|
<span class="n">array</span><span class="p">([</span><span class="mi">6</span><span class="p">,</span> <span class="mi">7</span><span class="p">],</span> <span class="n">dtype</span><span class="o">=</span><span class="n">int32</span><span class="p">)</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="n">unfolded</span> <span class="o">=</span> <span class="n">Tensor</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="mi">27</span><span class="p">)</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="mi">3</span><span class="p">,</span><span class="mi">3</span><span class="p">,</span><span class="mi">3</span><span class="p">)</span><span class="o">.</span><span class="n">unfold</span><span class="p">(</span><span class="o">-</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="nb">print</span><span class="p">(</span><span class="s2">"</span><span class="se">\n</span><span class="s2">"</span><span class="o">.</span><span class="n">join</span><span class="p">([</span><span class="nb">repr</span><span class="p">(</span><span class="n">x</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span> <span class="k">for</span> <span class="n">x</span> <span class="ow">in</span> <span class="n">unfolded</span><span class="p">]))</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="n">array</span><span class="p">([[[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">]],</span>
|
|
|
|
<span class="p">[[</span><span class="mi">3</span><span class="p">,</span> <span class="mi">4</span><span class="p">]],</span>
|
|
|
|
<span class="p">[[</span><span class="mi">6</span><span class="p">,</span> <span class="mi">7</span><span class="p">]]],</span> <span class="n">dtype</span><span class="o">=</span><span class="n">int32</span><span class="p">)</span>
|
|
<span class="n">array</span><span class="p">([[[</span> <span class="mi">9</span><span class="p">,</span> <span class="mi">10</span><span class="p">]],</span>
|
|
|
|
<span class="p">[[</span><span class="mi">12</span><span class="p">,</span> <span class="mi">13</span><span class="p">]],</span>
|
|
|
|
<span class="p">[[</span><span class="mi">15</span><span class="p">,</span> <span class="mi">16</span><span class="p">]]],</span> <span class="n">dtype</span><span class="o">=</span><span class="n">int32</span><span class="p">)</span>
|
|
<span class="n">array</span><span class="p">([[[</span><span class="mi">18</span><span class="p">,</span> <span class="mi">19</span><span class="p">]],</span>
|
|
|
|
<span class="p">[[</span><span class="mi">21</span><span class="p">,</span> <span class="mi">22</span><span class="p">]],</span>
|
|
|
|
<span class="p">[[</span><span class="mi">24</span><span class="p">,</span> <span class="mi">25</span><span class="p">]]],</span> <span class="n">dtype</span><span class="o">=</span><span class="n">int32</span><span class="p">)</span>
|
|
</code></pre></div></p>
|
|
|
|
|
|
<details class="mkdocstrings-source">
|
|
<summary>Source code in <code>tinygrad/tensor.py</code></summary>
|
|
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal">1378</span>
|
|
<span class="normal">1379</span>
|
|
<span class="normal">1380</span>
|
|
<span class="normal">1381</span>
|
|
<span class="normal">1382</span>
|
|
<span class="normal">1383</span>
|
|
<span class="normal">1384</span>
|
|
<span class="normal">1385</span>
|
|
<span class="normal">1386</span>
|
|
<span class="normal">1387</span>
|
|
<span class="normal">1388</span>
|
|
<span class="normal">1389</span>
|
|
<span class="normal">1390</span>
|
|
<span class="normal">1391</span>
|
|
<span class="normal">1392</span>
|
|
<span class="normal">1393</span>
|
|
<span class="normal">1394</span>
|
|
<span class="normal">1395</span>
|
|
<span class="normal">1396</span>
|
|
<span class="normal">1397</span>
|
|
<span class="normal">1398</span>
|
|
<span class="normal">1399</span>
|
|
<span class="normal">1400</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="k">def</span><span class="w"> </span><span class="nf">unfold</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">dim</span><span class="p">:</span><span class="nb">int</span><span class="p">,</span> <span class="n">size</span><span class="p">:</span><span class="n">sint</span><span class="p">,</span> <span class="n">step</span><span class="p">:</span><span class="nb">int</span><span class="p">)</span> <span class="o">-></span> <span class="n">Tensor</span><span class="p">:</span>
|
|
<span class="w"> </span><span class="sd">"""</span>
|
|
<span class="sd"> Unfolds the tensor along dimension `dim` into overlapping windows.</span>
|
|
|
|
<span class="sd"> Each window has length `size` and begins every `step` elements of `self`.</span>
|
|
<span class="sd"> Returns the input tensor with dimension `dim` replaced by dims `(n_windows, size)`</span>
|
|
<span class="sd"> where `n_windows = (self.shape[dim] - size) // step + 1`.</span>
|
|
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> unfolded = Tensor.arange(8).unfold(0,2,2)</span>
|
|
<span class="sd"> print("\\n".join([repr(x.numpy()) for x in unfolded]))</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> unfolded = Tensor.arange(27).reshape(3,3,3).unfold(-1,2,3)</span>
|
|
<span class="sd"> print("\\n".join([repr(x.numpy()) for x in unfolded]))</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> """</span>
|
|
<span class="k">if</span> <span class="n">size</span> <span class="o"><</span> <span class="mi">0</span><span class="p">:</span> <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="sa">f</span><span class="s1">'size must be >= 0 but got </span><span class="si">{</span><span class="n">size</span><span class="si">=}</span><span class="s1">'</span><span class="p">)</span>
|
|
<span class="k">if</span> <span class="n">step</span> <span class="o"><=</span> <span class="mi">0</span><span class="p">:</span> <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="sa">f</span><span class="s1">'step must be > 0 but got </span><span class="si">{</span><span class="n">step</span><span class="si">=}</span><span class="s1">'</span><span class="p">)</span>
|
|
<span class="k">if</span> <span class="n">size</span> <span class="o">></span> <span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="n">dim</span><span class="p">]:</span> <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="sa">f</span><span class="s1">'maximum size for tensor at dimension </span><span class="si">{</span><span class="n">dim</span><span class="si">}</span><span class="s1"> is </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="n">dim</span><span class="p">]</span><span class="si">}</span><span class="s1"> but size is </span><span class="si">{</span><span class="n">size</span><span class="si">}</span><span class="s1">'</span><span class="p">)</span>
|
|
<span class="n">dim</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_resolve_dim</span><span class="p">(</span><span class="n">dim</span><span class="p">)</span>
|
|
<span class="n">perm_to_last</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">(</span><span class="n">i</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">ndim</span><span class="p">)</span> <span class="k">if</span> <span class="n">i</span> <span class="o">!=</span> <span class="n">dim</span><span class="p">)</span> <span class="o">+</span> <span class="p">(</span><span class="n">dim</span><span class="p">,)</span>
|
|
<span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">permute</span><span class="p">(</span><span class="n">perm_to_last</span><span class="p">)</span><span class="o">.</span><span class="n">_pool</span><span class="p">((</span><span class="n">size</span><span class="p">,),</span> <span class="n">step</span><span class="p">)</span><span class="o">.</span><span class="n">permute</span><span class="p">(</span><span class="n">argsort</span><span class="p">(</span><span class="n">perm_to_last</span><span class="p">)</span> <span class="o">+</span> <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">ndim</span><span class="p">,))</span>
|
|
</code></pre></div></td></tr></table></div>
|
|
</details>
|
|
</div>
|
|
|
|
</div>
|
|
|
|
<div class="doc doc-object doc-function">
|
|
|
|
|
|
<h3 id="tinygrad.Tensor.meshgrid" class="doc doc-heading">
|
|
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">meshgrid</span>
|
|
|
|
|
|
<a href="#tinygrad.Tensor.meshgrid" class="headerlink" title="Permanent link">¤</a></h3>
|
|
<div class="language-python doc-signature highlight"><pre><span></span><code><span class="nf">meshgrid</span><span class="p">(</span>
|
|
<span class="o">*</span><span class="n">args</span><span class="p">:</span> <span class="n"><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">Tensor</span> (<code>tinygrad.tensor.Tensor</code>)" href="../#tinygrad.Tensor">Tensor</a></span><span class="p">,</span> <span class="n">indexing</span><span class="p">:</span> <span class="n"><a class="autorefs autorefs-external" title="<code>typing.Literal</code>" href="https://docs.python.org/3/library/typing.html#typing.Literal">Literal</a></span><span class="p">[</span><span class="s2">"ij"</span><span class="p">,</span> <span class="s2">"xy"</span><span class="p">]</span> <span class="o">=</span> <span class="s2">"ij"</span>
|
|
<span class="p">)</span> <span class="o">-></span> <span class="n"><a class="autorefs autorefs-external" href="https://docs.python.org/3/library/stdtypes.html#tuple">tuple</a></span><span class="p">[</span><span class="n"><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">Tensor</span> (<code>tinygrad.tensor.Tensor</code>)" href="../#tinygrad.Tensor">Tensor</a></span><span class="p">,</span> <span class="o">...</span><span class="p">]</span>
|
|
</code></pre></div>
|
|
|
|
<div class="doc doc-contents first">
|
|
|
|
<p>Generates coordinate matrices from coordinate vectors.
|
|
Input tensors can be scalars or 1D tensors.</p>
|
|
<p><code class="language-python highlight"><span class="n">indexing</span></code> determines how the output grids are aligned.
|
|
<code class="language-python highlight"><span class="n">ij</span></code> indexing follows matrix-style indexing and <code class="language-python highlight"><span class="n">xy</span></code> indexing follows Cartesian-style indexing.</p>
|
|
<p><div class="language-python highlight"><pre><span></span><code><span class="n">x</span><span class="p">,</span> <span class="n">y</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> <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> <span class="mi">6</span><span class="p">])</span>
|
|
<span class="n">grid_x</span><span class="p">,</span> <span class="n">grid_y</span> <span class="o">=</span> <span class="n">x</span><span class="o">.</span><span class="n">meshgrid</span><span class="p">(</span><span class="n">y</span><span class="p">)</span>
|
|
<span class="nb">print</span><span class="p">(</span><span class="n">grid_x</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
<span class="nb">print</span><span class="p">(</span><span class="n">grid_y</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">[[</span><span class="mi">1</span> <span class="mi">1</span> <span class="mi">1</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">2</span> <span class="mi">2</span> <span class="mi">2</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">3</span> <span class="mi">3</span> <span class="mi">3</span><span class="p">]]</span>
|
|
<span class="p">[[</span><span class="mi">4</span> <span class="mi">5</span> <span class="mi">6</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">4</span> <span class="mi">5</span> <span class="mi">6</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">4</span> <span class="mi">5</span> <span class="mi">6</span><span class="p">]]</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="n">grid_x</span><span class="p">,</span> <span class="n">grid_y</span> <span class="o">=</span> <span class="n">x</span><span class="o">.</span><span class="n">meshgrid</span><span class="p">(</span><span class="n">y</span><span class="p">,</span> <span class="n">indexing</span><span class="o">=</span><span class="s2">"xy"</span><span class="p">)</span>
|
|
<span class="nb">print</span><span class="p">(</span><span class="n">grid_x</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
<span class="nb">print</span><span class="p">(</span><span class="n">grid_y</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">[[</span><span class="mi">1</span> <span class="mi">2</span> <span class="mi">3</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">1</span> <span class="mi">2</span> <span class="mi">3</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">1</span> <span class="mi">2</span> <span class="mi">3</span><span class="p">]]</span>
|
|
<span class="p">[[</span><span class="mi">4</span> <span class="mi">4</span> <span class="mi">4</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">5</span> <span class="mi">5</span> <span class="mi">5</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">6</span> <span class="mi">6</span> <span class="mi">6</span><span class="p">]]</span>
|
|
</code></pre></div></p>
|
|
|
|
|
|
<details class="mkdocstrings-source">
|
|
<summary>Source code in <code>tinygrad/tensor.py</code></summary>
|
|
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal">1402</span>
|
|
<span class="normal">1403</span>
|
|
<span class="normal">1404</span>
|
|
<span class="normal">1405</span>
|
|
<span class="normal">1406</span>
|
|
<span class="normal">1407</span>
|
|
<span class="normal">1408</span>
|
|
<span class="normal">1409</span>
|
|
<span class="normal">1410</span>
|
|
<span class="normal">1411</span>
|
|
<span class="normal">1412</span>
|
|
<span class="normal">1413</span>
|
|
<span class="normal">1414</span>
|
|
<span class="normal">1415</span>
|
|
<span class="normal">1416</span>
|
|
<span class="normal">1417</span>
|
|
<span class="normal">1418</span>
|
|
<span class="normal">1419</span>
|
|
<span class="normal">1420</span>
|
|
<span class="normal">1421</span>
|
|
<span class="normal">1422</span>
|
|
<span class="normal">1423</span>
|
|
<span class="normal">1424</span>
|
|
<span class="normal">1425</span>
|
|
<span class="normal">1426</span>
|
|
<span class="normal">1427</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="k">def</span><span class="w"> </span><span class="nf">meshgrid</span><span class="p">(</span><span class="bp">self</span><span class="p">:</span><span class="n">Tensor</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">:</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">indexing</span><span class="p">:</span><span class="n">Literal</span><span class="p">[</span><span class="s2">"ij"</span><span class="p">,</span> <span class="s2">"xy"</span><span class="p">]</span><span class="o">=</span><span class="s2">"ij"</span><span class="p">)</span> <span class="o">-></span> <span class="nb">tuple</span><span class="p">[</span><span class="n">Tensor</span><span class="p">,</span> <span class="o">...</span><span class="p">]:</span>
|
|
<span class="w"> </span><span class="sd">"""</span>
|
|
<span class="sd"> Generates coordinate matrices from coordinate vectors.</span>
|
|
<span class="sd"> Input tensors can be scalars or 1D tensors.</span>
|
|
|
|
<span class="sd"> `indexing` determines how the output grids are aligned.</span>
|
|
<span class="sd"> `ij` indexing follows matrix-style indexing and `xy` indexing follows Cartesian-style indexing.</span>
|
|
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> x, y = Tensor([1, 2, 3]), Tensor([4, 5, 6])</span>
|
|
<span class="sd"> grid_x, grid_y = x.meshgrid(y)</span>
|
|
<span class="sd"> print(grid_x.numpy())</span>
|
|
<span class="sd"> print(grid_y.numpy())</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> grid_x, grid_y = x.meshgrid(y, indexing="xy")</span>
|
|
<span class="sd"> print(grid_x.numpy())</span>
|
|
<span class="sd"> print(grid_y.numpy())</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> """</span>
|
|
<span class="k">if</span> <span class="n">indexing</span> <span class="ow">not</span> <span class="ow">in</span> <span class="p">(</span><span class="s2">"ij"</span><span class="p">,</span> <span class="s2">"xy"</span><span class="p">):</span> <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="sa">f</span><span class="s1">'indexing must be in ("ij", "xy"), got </span><span class="si">{</span><span class="n">indexing</span><span class="si">}</span><span class="s1">'</span><span class="p">)</span>
|
|
<span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">tensors</span><span class="o">:=</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">))</span> <span class="o">==</span> <span class="mi">1</span><span class="p">:</span> <span class="k">return</span> <span class="n">tensors</span>
|
|
<span class="n">basis</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">(</span><span class="nb">range</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">tensors</span><span class="p">)))</span> <span class="k">if</span> <span class="n">indexing</span> <span class="o">==</span> <span class="s2">"ij"</span> <span class="k">else</span> <span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span> <span class="o">+</span> <span class="nb">tuple</span><span class="p">(</span><span class="nb">range</span><span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="nb">len</span><span class="p">(</span><span class="n">tensors</span><span class="p">)))</span>
|
|
<span class="n">tensors</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">(</span><span class="n">t</span><span class="o">.</span><span class="n">reshape</span><span class="p">((</span><span class="o">-</span><span class="mi">1</span><span class="p">,)</span> <span class="o">+</span> <span class="p">(</span><span class="mi">1</span><span class="p">,)</span><span class="o">*</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">args</span><span class="p">)</span> <span class="o">-</span> <span class="n">i</span><span class="p">))</span> <span class="k">for</span> <span class="n">i</span><span class="p">,</span><span class="n">t</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">basis</span><span class="p">,</span> <span class="n">tensors</span><span class="p">))</span>
|
|
<span class="n">output_shape</span> <span class="o">=</span> <span class="n">_broadcast_shape</span><span class="p">(</span><span class="o">*</span><span class="p">(</span><span class="n">t</span><span class="o">.</span><span class="n">shape</span> <span class="k">for</span> <span class="n">t</span> <span class="ow">in</span> <span class="n">tensors</span><span class="p">))</span>
|
|
<span class="k">return</span> <span class="nb">tuple</span><span class="p">(</span><span class="n">t</span><span class="o">.</span><span class="n">_broadcast_to</span><span class="p">(</span><span class="n">output_shape</span><span class="p">)</span> <span class="k">for</span> <span class="n">t</span> <span class="ow">in</span> <span class="n">tensors</span><span class="p">)</span>
|
|
</code></pre></div></td></tr></table></div>
|
|
</details>
|
|
</div>
|
|
|
|
</div>
|
|
|
|
<div class="doc doc-object doc-function">
|
|
|
|
|
|
<h3 id="tinygrad.Tensor.squeeze" class="doc doc-heading">
|
|
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">squeeze</span>
|
|
|
|
|
|
<a href="#tinygrad.Tensor.squeeze" class="headerlink" title="Permanent link">¤</a></h3>
|
|
<div class="language-python doc-signature highlight"><pre><span></span><code><span class="nf">squeeze</span><span class="p">(</span><span class="n">dim</span><span class="p">:</span> <span class="n"><a class="autorefs autorefs-external" href="https://docs.python.org/3/library/functions.html#int">int</a></span> <span class="o">|</span> <span class="kc">None</span> <span class="o">=</span> <span class="kc">None</span><span class="p">)</span> <span class="o">-></span> <span class="n"><a class="autorefs autorefs-external" title="<code>typing.Self</code>" href="https://docs.python.org/3/library/typing.html#typing.Self">Self</a></span>
|
|
</code></pre></div>
|
|
|
|
<div class="doc doc-contents first">
|
|
|
|
<p>Returns a tensor with specified dimensions of input of size 1 removed.
|
|
If <code class="language-python highlight"><span class="n">dim</span></code> is not specified, all dimensions with size 1 are removed.</p>
|
|
<p><div class="language-python highlight"><pre><span></span><code><span class="n">t</span> <span class="o">=</span> <span class="n">Tensor</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="mi">2</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">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span>
|
|
<span class="nb">print</span><span class="p">(</span><span class="n">t</span><span class="o">.</span><span class="n">squeeze</span><span class="p">()</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="nb">print</span><span class="p">(</span><span class="n">t</span><span class="o">.</span><span class="n">squeeze</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">(</span><span class="mi">2</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">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="nb">print</span><span class="p">(</span><span class="n">t</span><span class="o">.</span><span class="n">squeeze</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span>
|
|
</code></pre></div></p>
|
|
|
|
|
|
<details class="mkdocstrings-source">
|
|
<summary>Source code in <code>tinygrad/mixin/movement.py</code></summary>
|
|
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal">181</span>
|
|
<span class="normal">182</span>
|
|
<span class="normal">183</span>
|
|
<span class="normal">184</span>
|
|
<span class="normal">185</span>
|
|
<span class="normal">186</span>
|
|
<span class="normal">187</span>
|
|
<span class="normal">188</span>
|
|
<span class="normal">189</span>
|
|
<span class="normal">190</span>
|
|
<span class="normal">191</span>
|
|
<span class="normal">192</span>
|
|
<span class="normal">193</span>
|
|
<span class="normal">194</span>
|
|
<span class="normal">195</span>
|
|
<span class="normal">196</span>
|
|
<span class="normal">197</span>
|
|
<span class="normal">198</span>
|
|
<span class="normal">199</span>
|
|
<span class="normal">200</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="k">def</span><span class="w"> </span><span class="nf">squeeze</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">dim</span><span class="p">:</span> <span class="nb">int</span> <span class="o">|</span> <span class="kc">None</span> <span class="o">=</span> <span class="kc">None</span><span class="p">)</span> <span class="o">-></span> <span class="n">Self</span><span class="p">:</span>
|
|
<span class="w"> </span><span class="sd">"""</span>
|
|
<span class="sd"> Returns a tensor with specified dimensions of input of size 1 removed.</span>
|
|
<span class="sd"> If `dim` is not specified, all dimensions with size 1 are removed.</span>
|
|
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> t = Tensor.zeros(2, 1, 2, 1, 2)</span>
|
|
<span class="sd"> print(t.squeeze().shape)</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> print(t.squeeze(0).shape)</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> print(t.squeeze(1).shape)</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> """</span>
|
|
<span class="k">if</span> <span class="n">dim</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
|
|
<span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="nb">tuple</span><span class="p">(</span><span class="n">dim</span> <span class="k">for</span> <span class="n">dim</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">shape</span> <span class="k">if</span> <span class="n">dim</span> <span class="o">!=</span> <span class="mi">1</span><span class="p">))</span>
|
|
<span class="n">dim</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_resolve_dim</span><span class="p">(</span><span class="n">dim</span><span class="p">)</span>
|
|
<span class="k">return</span> <span class="bp">self</span> <span class="k">if</span> <span class="ow">not</span> <span class="bp">self</span><span class="o">.</span><span class="n">ndim</span> <span class="ow">or</span> <span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="n">dim</span><span class="p">]</span> <span class="o">!=</span> <span class="mi">1</span> <span class="k">else</span> <span class="bp">self</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">[:</span><span class="n">dim</span><span class="p">]</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="n">dim</span> <span class="o">+</span> <span class="mi">1</span> <span class="p">:])</span>
|
|
</code></pre></div></td></tr></table></div>
|
|
</details>
|
|
</div>
|
|
|
|
</div>
|
|
|
|
<div class="doc doc-object doc-function">
|
|
|
|
|
|
<h3 id="tinygrad.Tensor.unsqueeze" class="doc doc-heading">
|
|
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">unsqueeze</span>
|
|
|
|
|
|
<a href="#tinygrad.Tensor.unsqueeze" class="headerlink" title="Permanent link">¤</a></h3>
|
|
<div class="language-python doc-signature highlight"><pre><span></span><code><span class="nf">unsqueeze</span><span class="p">(</span><span class="n">dim</span><span class="p">:</span> <span class="n"><a class="autorefs autorefs-external" href="https://docs.python.org/3/library/functions.html#int">int</a></span><span class="p">)</span> <span class="o">-></span> <span class="n"><a class="autorefs autorefs-external" title="<code>typing.Self</code>" href="https://docs.python.org/3/library/typing.html#typing.Self">Self</a></span>
|
|
</code></pre></div>
|
|
|
|
<div class="doc doc-contents first">
|
|
|
|
<p>Returns a tensor with a new dimension of size 1 inserted at the specified <code class="language-python highlight"><span class="n">dim</span></code>.</p>
|
|
<p><div class="language-python highlight"><pre><span></span><code><span class="n">t</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> <span class="mi">4</span><span class="p">])</span>
|
|
<span class="nb">print</span><span class="p">(</span><span class="n">t</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">[[</span><span class="mi">1</span> <span class="mi">2</span> <span class="mi">3</span> <span class="mi">4</span><span class="p">]]</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="nb">print</span><span class="p">(</span><span class="n">t</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">[[</span><span class="mi">1</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">2</span><span class="p">]</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>
|
|
</code></pre></div></p>
|
|
|
|
|
|
<details class="mkdocstrings-source">
|
|
<summary>Source code in <code>tinygrad/mixin/movement.py</code></summary>
|
|
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal">202</span>
|
|
<span class="normal">203</span>
|
|
<span class="normal">204</span>
|
|
<span class="normal">205</span>
|
|
<span class="normal">206</span>
|
|
<span class="normal">207</span>
|
|
<span class="normal">208</span>
|
|
<span class="normal">209</span>
|
|
<span class="normal">210</span>
|
|
<span class="normal">211</span>
|
|
<span class="normal">212</span>
|
|
<span class="normal">213</span>
|
|
<span class="normal">214</span>
|
|
<span class="normal">215</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="k">def</span><span class="w"> </span><span class="nf">unsqueeze</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">dim</span><span class="p">:</span> <span class="nb">int</span><span class="p">)</span> <span class="o">-></span> <span class="n">Self</span><span class="p">:</span>
|
|
<span class="w"> </span><span class="sd">"""</span>
|
|
<span class="sd"> Returns a tensor with a new dimension of size 1 inserted at the specified `dim`.</span>
|
|
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> t = Tensor([1, 2, 3, 4])</span>
|
|
<span class="sd"> print(t.unsqueeze(0).numpy())</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> print(t.unsqueeze(1).numpy())</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> """</span>
|
|
<span class="n">dim</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_resolve_dim</span><span class="p">(</span><span class="n">dim</span><span class="p">,</span> <span class="n">extra</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
|
|
<span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">[:</span><span class="n">dim</span><span class="p">]</span> <span class="o">+</span> <span class="p">(</span><span class="mi">1</span><span class="p">,)</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="n">dim</span><span class="p">:])</span>
|
|
</code></pre></div></td></tr></table></div>
|
|
</details>
|
|
</div>
|
|
|
|
</div>
|
|
|
|
<div class="doc doc-object doc-attribute">
|
|
|
|
|
|
|
|
<h3 id="tinygrad.Tensor.T" class="doc doc-heading">
|
|
<code class="doc-symbol doc-symbol-heading doc-symbol-attribute"></code> <span class="doc doc-object-name doc-attribute-name">T</span>
|
|
|
|
|
|
<span class="doc doc-labels">
|
|
<small class="doc doc-label doc-label-property"><code>property</code></small>
|
|
</span>
|
|
|
|
<a href="#tinygrad.Tensor.T" class="headerlink" title="Permanent link">¤</a></h3>
|
|
<div class="language-python doc-signature highlight"><pre><span></span><code><span class="n">T</span><span class="p">:</span> <span class="n"><a class="autorefs autorefs-external" title="<code>typing.Self</code>" href="https://docs.python.org/3/library/typing.html#typing.Self">Self</a></span>
|
|
</code></pre></div>
|
|
|
|
<div class="doc doc-contents first">
|
|
|
|
<p><code class="language-python highlight"><span class="o">.</span><span class="n">T</span></code> is an alias for <code class="language-python highlight"><span class="o">.</span><span class="n">transpose</span><span class="p">()</span></code>.</p>
|
|
|
|
</div>
|
|
|
|
</div>
|
|
|
|
<div class="doc doc-object doc-function">
|
|
|
|
|
|
<h3 id="tinygrad.Tensor.transpose" class="doc doc-heading">
|
|
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">transpose</span>
|
|
|
|
|
|
<a href="#tinygrad.Tensor.transpose" class="headerlink" title="Permanent link">¤</a></h3>
|
|
<div class="language-python doc-signature highlight"><pre><span></span><code><span class="nf">transpose</span><span class="p">(</span><span class="n">dim0</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">dim1</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span> <span class="o">-></span> <span class="n"><a class="autorefs autorefs-external" title="<code>typing.Self</code>" href="https://docs.python.org/3/library/typing.html#typing.Self">Self</a></span>
|
|
</code></pre></div>
|
|
|
|
<div class="doc doc-contents first">
|
|
|
|
<p>Returns a tensor that is a transposed version of the original tensor.
|
|
The given dimensions <code class="language-python highlight"><span class="n">dim0</span></code> and <code class="language-python highlight"><span class="n">dim1</span></code> are swapped.</p>
|
|
<p><div class="language-python highlight"><pre><span></span><code><span class="n">t</span> <span class="o">=</span> <span class="n">Tensor</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="mi">6</span><span class="p">)</span><span class="o">.</span><span class="n">reshape</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="nb">print</span><span class="p">(</span><span class="n">t</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">[[</span><span class="mi">0</span> <span class="mi">1</span> <span class="mi">2</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">3</span> <span class="mi">4</span> <span class="mi">5</span><span class="p">]]</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="nb">print</span><span class="p">(</span><span class="n">t</span><span class="o">.</span><span class="n">transpose</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">[[</span><span class="mi">0</span> <span class="mi">3</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">1</span> <span class="mi">4</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">2</span> <span class="mi">5</span><span class="p">]]</span>
|
|
</code></pre></div></p>
|
|
|
|
|
|
<details class="mkdocstrings-source">
|
|
<summary>Source code in <code>tinygrad/mixin/movement.py</code></summary>
|
|
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal">222</span>
|
|
<span class="normal">223</span>
|
|
<span class="normal">224</span>
|
|
<span class="normal">225</span>
|
|
<span class="normal">226</span>
|
|
<span class="normal">227</span>
|
|
<span class="normal">228</span>
|
|
<span class="normal">229</span>
|
|
<span class="normal">230</span>
|
|
<span class="normal">231</span>
|
|
<span class="normal">232</span>
|
|
<span class="normal">233</span>
|
|
<span class="normal">234</span>
|
|
<span class="normal">235</span>
|
|
<span class="normal">236</span>
|
|
<span class="normal">237</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="k">def</span><span class="w"> </span><span class="nf">transpose</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">dim0</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">dim1</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span> <span class="o">-></span> <span class="n">Self</span><span class="p">:</span>
|
|
<span class="w"> </span><span class="sd">"""</span>
|
|
<span class="sd"> Returns a tensor that is a transposed version of the original tensor.</span>
|
|
<span class="sd"> The given dimensions `dim0` and `dim1` are swapped.</span>
|
|
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> t = Tensor.arange(6).reshape(2, 3)</span>
|
|
<span class="sd"> print(t.numpy())</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> print(t.transpose(0, 1).numpy())</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> """</span>
|
|
<span class="n">order</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="nb">range</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">ndim</span><span class="p">))</span>
|
|
<span class="n">order</span><span class="p">[</span><span class="n">dim0</span><span class="p">],</span> <span class="n">order</span><span class="p">[</span><span class="n">dim1</span><span class="p">]</span> <span class="o">=</span> <span class="n">order</span><span class="p">[</span><span class="n">dim1</span><span class="p">],</span> <span class="n">order</span><span class="p">[</span><span class="n">dim0</span><span class="p">]</span>
|
|
<span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">permute</span><span class="p">(</span><span class="n">order</span><span class="p">)</span>
|
|
</code></pre></div></td></tr></table></div>
|
|
</details>
|
|
</div>
|
|
|
|
</div>
|
|
|
|
<div class="doc doc-object doc-function">
|
|
|
|
|
|
<h3 id="tinygrad.Tensor.flatten" class="doc doc-heading">
|
|
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">flatten</span>
|
|
|
|
|
|
<a href="#tinygrad.Tensor.flatten" class="headerlink" title="Permanent link">¤</a></h3>
|
|
<div class="language-python doc-signature highlight"><pre><span></span><code><span class="nf">flatten</span><span class="p">(</span><span class="n">start_dim</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span> <span class="n">end_dim</span><span class="o">=-</span><span class="mi">1</span><span class="p">)</span> <span class="o">-></span> <span class="n"><a class="autorefs autorefs-external" title="<code>typing.Self</code>" href="https://docs.python.org/3/library/typing.html#typing.Self">Self</a></span>
|
|
</code></pre></div>
|
|
|
|
<div class="doc doc-contents first">
|
|
|
|
<p>Flattens the tensor by reshaping it into a one-dimensional tensor.
|
|
If <code class="language-python highlight"><span class="n">start_dim</span></code> or <code class="language-python highlight"><span class="n">end_dim</span></code> are passed, only dimensions starting with <code class="language-python highlight"><span class="n">start_dim</span></code> and ending with <code class="language-python highlight"><span class="n">end_dim</span></code> are flattened.</p>
|
|
<p><div class="language-python highlight"><pre><span></span><code><span class="n">t</span> <span class="o">=</span> <span class="n">Tensor</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="mi">8</span><span class="p">)</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span>
|
|
<span class="nb">print</span><span class="p">(</span><span class="n">t</span><span class="o">.</span><span class="n">flatten</span><span class="p">()</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">[</span><span class="mi">0</span> <span class="mi">1</span> <span class="mi">2</span> <span class="mi">3</span> <span class="mi">4</span> <span class="mi">5</span> <span class="mi">6</span> <span class="mi">7</span><span class="p">]</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="nb">print</span><span class="p">(</span><span class="n">t</span><span class="o">.</span><span class="n">flatten</span><span class="p">(</span><span class="n">start_dim</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">[[</span><span class="mi">0</span> <span class="mi">1</span> <span class="mi">2</span> <span class="mi">3</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">4</span> <span class="mi">5</span> <span class="mi">6</span> <span class="mi">7</span><span class="p">]]</span>
|
|
</code></pre></div></p>
|
|
|
|
|
|
<details class="mkdocstrings-source">
|
|
<summary>Source code in <code>tinygrad/mixin/movement.py</code></summary>
|
|
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal">239</span>
|
|
<span class="normal">240</span>
|
|
<span class="normal">241</span>
|
|
<span class="normal">242</span>
|
|
<span class="normal">243</span>
|
|
<span class="normal">244</span>
|
|
<span class="normal">245</span>
|
|
<span class="normal">246</span>
|
|
<span class="normal">247</span>
|
|
<span class="normal">248</span>
|
|
<span class="normal">249</span>
|
|
<span class="normal">250</span>
|
|
<span class="normal">251</span>
|
|
<span class="normal">252</span>
|
|
<span class="normal">253</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="k">def</span><span class="w"> </span><span class="nf">flatten</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">start_dim</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span> <span class="n">end_dim</span><span class="o">=-</span><span class="mi">1</span><span class="p">)</span> <span class="o">-></span> <span class="n">Self</span><span class="p">:</span>
|
|
<span class="w"> </span><span class="sd">"""</span>
|
|
<span class="sd"> Flattens the tensor by reshaping it into a one-dimensional tensor.</span>
|
|
<span class="sd"> If `start_dim` or `end_dim` are passed, only dimensions starting with `start_dim` and ending with `end_dim` are flattened.</span>
|
|
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> t = Tensor.arange(8).reshape(2, 2, 2)</span>
|
|
<span class="sd"> print(t.flatten().numpy())</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> print(t.flatten(start_dim=1).numpy())</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> """</span>
|
|
<span class="n">start_dim</span><span class="p">,</span> <span class="n">end_dim</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_resolve_dim</span><span class="p">(</span><span class="n">start_dim</span><span class="p">),</span> <span class="bp">self</span><span class="o">.</span><span class="n">_resolve_dim</span><span class="p">(</span><span class="n">end_dim</span><span class="p">)</span>
|
|
<span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">[:</span><span class="n">start_dim</span><span class="p">]</span> <span class="o">+</span> <span class="p">(</span><span class="n">prod</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="n">start_dim</span> <span class="p">:</span> <span class="n">end_dim</span> <span class="o">+</span> <span class="mi">1</span><span class="p">]),)</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="n">end_dim</span> <span class="o">+</span> <span class="mi">1</span> <span class="p">:])</span>
|
|
</code></pre></div></td></tr></table></div>
|
|
</details>
|
|
</div>
|
|
|
|
</div>
|
|
|
|
<div class="doc doc-object doc-function">
|
|
|
|
|
|
<h3 id="tinygrad.Tensor.unflatten" class="doc doc-heading">
|
|
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">unflatten</span>
|
|
|
|
|
|
<a href="#tinygrad.Tensor.unflatten" class="headerlink" title="Permanent link">¤</a></h3>
|
|
<div class="language-python doc-signature highlight"><pre><span></span><code><span class="nf">unflatten</span><span class="p">(</span><span class="n">dim</span><span class="p">:</span> <span class="n"><a class="autorefs autorefs-external" href="https://docs.python.org/3/library/functions.html#int">int</a></span><span class="p">,</span> <span class="n">sizes</span><span class="p">:</span> <span class="n"><a class="autorefs autorefs-external" href="https://docs.python.org/3/library/stdtypes.html#tuple">tuple</a></span><span class="p">[</span><span class="n"><a class="autorefs autorefs-external" href="https://docs.python.org/3/library/functions.html#int">int</a></span><span class="p">,</span> <span class="o">...</span><span class="p">])</span> <span class="o">-></span> <span class="n"><a class="autorefs autorefs-external" title="<code>typing.Self</code>" href="https://docs.python.org/3/library/typing.html#typing.Self">Self</a></span>
|
|
</code></pre></div>
|
|
|
|
<div class="doc doc-contents first">
|
|
|
|
<p>Unflattens dimension <code class="language-python highlight"><span class="n">dim</span></code> of the tensor into multiple dimensions specified by <code class="language-python highlight"><span class="n">sizes</span></code>. <code class="language-python highlight"><span class="n">Tensor</span><span class="o">.</span><span class="n">flatten</span><span class="p">()</span></code> is the inverse of this function.</p>
|
|
<p><div class="language-python highlight"><pre><span></span><code><span class="nb">print</span><span class="p">(</span><span class="n">Tensor</span><span class="o">.</span><span class="n">ones</span><span class="p">(</span><span class="mi">3</span><span class="p">,</span> <span class="mi">4</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">unflatten</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="mi">2</span><span class="p">))</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">(</span><span class="mi">3</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="nb">print</span><span class="p">(</span><span class="n">Tensor</span><span class="o">.</span><span class="n">ones</span><span class="p">(</span><span class="mi">3</span><span class="p">,</span> <span class="mi">4</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">unflatten</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">))</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">(</span><span class="mi">3</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="nb">print</span><span class="p">(</span><span class="n">Tensor</span><span class="o">.</span><span class="n">ones</span><span class="p">(</span><span class="mi">5</span><span class="p">,</span> <span class="mi">12</span><span class="p">,</span> <span class="mi">3</span><span class="p">)</span><span class="o">.</span><span class="n">unflatten</span><span class="p">(</span><span class="o">-</span><span class="mi">2</span><span class="p">,</span> <span class="p">(</span><span class="mi">2</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="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">))</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">(</span><span class="mi">5</span><span class="p">,</span> <span class="mi">2</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="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">3</span><span class="p">)</span>
|
|
</code></pre></div></p>
|
|
|
|
|
|
<details class="mkdocstrings-source">
|
|
<summary>Source code in <code>tinygrad/mixin/movement.py</code></summary>
|
|
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal">255</span>
|
|
<span class="normal">256</span>
|
|
<span class="normal">257</span>
|
|
<span class="normal">258</span>
|
|
<span class="normal">259</span>
|
|
<span class="normal">260</span>
|
|
<span class="normal">261</span>
|
|
<span class="normal">262</span>
|
|
<span class="normal">263</span>
|
|
<span class="normal">264</span>
|
|
<span class="normal">265</span>
|
|
<span class="normal">266</span>
|
|
<span class="normal">267</span>
|
|
<span class="normal">268</span>
|
|
<span class="normal">269</span>
|
|
<span class="normal">270</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="k">def</span><span class="w"> </span><span class="nf">unflatten</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">dim</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">sizes</span><span class="p">:</span> <span class="nb">tuple</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="o">...</span><span class="p">])</span> <span class="o">-></span> <span class="n">Self</span><span class="p">:</span>
|
|
<span class="w"> </span><span class="sd">"""</span>
|
|
<span class="sd"> Unflattens dimension `dim` of the tensor into multiple dimensions specified by `sizes`. `Tensor.flatten()` is the inverse of this function.</span>
|
|
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> print(Tensor.ones(3, 4, 1).unflatten(1, (2, 2)).shape)</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> print(Tensor.ones(3, 4, 1).unflatten(1, (-1, 2)).shape)</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> print(Tensor.ones(5, 12, 3).unflatten(-2, (2, 2, 3, 1, 1)).shape)</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> """</span>
|
|
<span class="n">dim</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_resolve_dim</span><span class="p">(</span><span class="n">dim</span><span class="p">)</span>
|
|
<span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">[:</span><span class="n">dim</span><span class="p">]</span> <span class="o">+</span> <span class="n">sizes</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="n">dim</span> <span class="o">+</span> <span class="mi">1</span> <span class="p">:])</span>
|
|
</code></pre></div></td></tr></table></div>
|
|
</details>
|
|
</div>
|
|
|
|
</div>
|
|
|
|
<div class="doc doc-object doc-function">
|
|
|
|
|
|
<h3 id="tinygrad.Tensor.diag" class="doc doc-heading">
|
|
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">diag</span>
|
|
|
|
|
|
<a href="#tinygrad.Tensor.diag" class="headerlink" title="Permanent link">¤</a></h3>
|
|
<div class="language-python doc-signature highlight"><pre><span></span><code><span class="nf">diag</span><span class="p">()</span> <span class="o">-></span> <span class="n"><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">Tensor</span> (<code>tinygrad.tensor.Tensor</code>)" href="../#tinygrad.Tensor">Tensor</a></span>
|
|
</code></pre></div>
|
|
|
|
<div class="doc doc-contents first">
|
|
|
|
<p>Returns a 2-D square tensor with the elements of input as the main diagonal.</p>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="nb">print</span><span class="p">(</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><span class="o">.</span><span class="n">diag</span><span class="p">()</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">[[</span><span class="mi">1</span> <span class="mi">0</span> <span class="mi">0</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">0</span> <span class="mi">2</span> <span class="mi">0</span><span class="p">]</span>
|
|
<span class="p">[</span><span class="mi">0</span> <span class="mi">0</span> <span class="mi">3</span><span class="p">]]</span>
|
|
</code></pre></div>
|
|
|
|
|
|
<details class="mkdocstrings-source">
|
|
<summary>Source code in <code>tinygrad/tensor.py</code></summary>
|
|
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal">1429</span>
|
|
<span class="normal">1430</span>
|
|
<span class="normal">1431</span>
|
|
<span class="normal">1432</span>
|
|
<span class="normal">1433</span>
|
|
<span class="normal">1434</span>
|
|
<span class="normal">1435</span>
|
|
<span class="normal">1436</span>
|
|
<span class="normal">1437</span>
|
|
<span class="normal">1438</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="k">def</span><span class="w"> </span><span class="nf">diag</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-></span> <span class="n">Tensor</span><span class="p">:</span>
|
|
<span class="w"> </span><span class="sd">"""</span>
|
|
<span class="sd"> Returns a 2-D square tensor with the elements of input as the main diagonal.</span>
|
|
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> print(Tensor([1, 2, 3]).diag().numpy())</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> """</span>
|
|
<span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">ndim</span> <span class="o">!=</span> <span class="mi">1</span><span class="p">:</span> <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">"expect input to be 1-D, getting </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">ndim</span><span class="si">}</span><span class="s2">-D"</span><span class="p">)</span>
|
|
<span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">pad</span><span class="p">((</span><span class="kc">None</span><span class="p">,(</span><span class="mi">0</span><span class="p">,</span><span class="n">n</span><span class="o">:=</span><span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">])))</span><span class="o">.</span><span class="n">flatten</span><span class="p">()</span><span class="o">.</span><span class="n">shrink</span><span class="p">(((</span><span class="mi">0</span><span class="p">,</span><span class="n">n</span><span class="o">*</span><span class="n">n</span><span class="p">),))</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">n</span><span class="p">,</span><span class="n">n</span><span class="p">)</span>
|
|
</code></pre></div></td></tr></table></div>
|
|
</details>
|
|
</div>
|
|
|
|
</div>
|
|
|
|
<div class="doc doc-object doc-function">
|
|
|
|
|
|
<h3 id="tinygrad.Tensor.roll" class="doc doc-heading">
|
|
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">roll</span>
|
|
|
|
|
|
<a href="#tinygrad.Tensor.roll" class="headerlink" title="Permanent link">¤</a></h3>
|
|
<div class="language-python doc-signature highlight"><pre><span></span><code><span class="nf">roll</span><span class="p">(</span>
|
|
<span class="n">shifts</span><span class="p">:</span> <span class="n"><a class="autorefs autorefs-internal" title="<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">int</span> (<code>tinygrad.tensor.Tensor.int</code>)" href="../elementwise/#tinygrad.Tensor.int">int</a></span> <span class="o">|</span> <span class="n"><a class="autorefs autorefs-external" href="https://docs.python.org/3/library/stdtypes.html#tuple">tuple</a></span><span class="p">[</span><span class="n"><a class="autorefs autorefs-internal" title="<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">int</span> (<code>tinygrad.tensor.Tensor.int</code>)" href="../elementwise/#tinygrad.Tensor.int">int</a></span><span class="p">,</span> <span class="o">...</span><span class="p">],</span>
|
|
<span class="n">dims</span><span class="p">:</span> <span class="n"><a class="autorefs autorefs-internal" title="<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">int</span> (<code>tinygrad.tensor.Tensor.int</code>)" href="../elementwise/#tinygrad.Tensor.int">int</a></span> <span class="o">|</span> <span class="n"><a class="autorefs autorefs-external" href="https://docs.python.org/3/library/stdtypes.html#tuple">tuple</a></span><span class="p">[</span><span class="n"><a class="autorefs autorefs-internal" title="<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">int</span> (<code>tinygrad.tensor.Tensor.int</code>)" href="../elementwise/#tinygrad.Tensor.int">int</a></span><span class="p">,</span> <span class="o">...</span><span class="p">]</span> <span class="o">|</span> <span class="kc">None</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
|
|
<span class="p">)</span> <span class="o">-></span> <span class="n"><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">Tensor</span> (<code>tinygrad.tensor.Tensor</code>)" href="../#tinygrad.Tensor">Tensor</a></span>
|
|
</code></pre></div>
|
|
|
|
<div class="doc doc-contents first">
|
|
|
|
<p>Rolls the tensor along specified dimension(s).
|
|
The rolling operation is circular, meaning that elements that go beyond the edge are wrapped around to the beginning of the dimension.</p>
|
|
<p><div class="language-python highlight"><pre><span></span><code><span class="n">t</span> <span class="o">=</span> <span class="n">Tensor</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="mi">4</span><span class="p">)</span>
|
|
<span class="nb">print</span><span class="p">(</span><span class="n">t</span><span class="o">.</span><span class="n">roll</span><span class="p">(</span><span class="n">shifts</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">dims</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">[</span><span class="mi">3</span> <span class="mi">0</span> <span class="mi">1</span> <span class="mi">2</span><span class="p">]</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="nb">print</span><span class="p">(</span><span class="n">t</span><span class="o">.</span><span class="n">roll</span><span class="p">(</span><span class="n">shifts</span><span class="o">=-</span><span class="mi">1</span><span class="p">,</span> <span class="n">dims</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">[</span><span class="mi">1</span> <span class="mi">2</span> <span class="mi">3</span> <span class="mi">0</span><span class="p">]</span>
|
|
</code></pre></div></p>
|
|
|
|
|
|
<details class="mkdocstrings-source">
|
|
<summary>Source code in <code>tinygrad/tensor.py</code></summary>
|
|
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal">1455</span>
|
|
<span class="normal">1456</span>
|
|
<span class="normal">1457</span>
|
|
<span class="normal">1458</span>
|
|
<span class="normal">1459</span>
|
|
<span class="normal">1460</span>
|
|
<span class="normal">1461</span>
|
|
<span class="normal">1462</span>
|
|
<span class="normal">1463</span>
|
|
<span class="normal">1464</span>
|
|
<span class="normal">1465</span>
|
|
<span class="normal">1466</span>
|
|
<span class="normal">1467</span>
|
|
<span class="normal">1468</span>
|
|
<span class="normal">1469</span>
|
|
<span class="normal">1470</span>
|
|
<span class="normal">1471</span>
|
|
<span class="normal">1472</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="k">def</span><span class="w"> </span><span class="nf">roll</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">shifts</span><span class="p">:</span><span class="nb">int</span><span class="o">|</span><span class="nb">tuple</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="o">...</span><span class="p">],</span> <span class="n">dims</span><span class="p">:</span><span class="nb">int</span><span class="o">|</span><span class="nb">tuple</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="o">...</span><span class="p">]</span><span class="o">|</span><span class="kc">None</span><span class="o">=</span><span class="kc">None</span><span class="p">)</span> <span class="o">-></span> <span class="n">Tensor</span><span class="p">:</span>
|
|
<span class="w"> </span><span class="sd">"""</span>
|
|
<span class="sd"> Rolls the tensor along specified dimension(s).</span>
|
|
<span class="sd"> The rolling operation is circular, meaning that elements that go beyond the edge are wrapped around to the beginning of the dimension.</span>
|
|
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> t = Tensor.arange(4)</span>
|
|
<span class="sd"> print(t.roll(shifts=1, dims=0).numpy())</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> print(t.roll(shifts=-1, dims=0).numpy())</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> """</span>
|
|
<span class="k">if</span> <span class="n">dims</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span> <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">flatten</span><span class="p">()</span><span class="o">.</span><span class="n">roll</span><span class="p">(</span><span class="n">shifts</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
|
|
<span class="n">dims</span><span class="p">,</span> <span class="n">shifts</span><span class="p">,</span> <span class="n">slices</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_resolve_dim</span><span class="p">(</span><span class="n">d</span><span class="p">)</span> <span class="k">for</span> <span class="n">d</span> <span class="ow">in</span> <span class="n">make_tuple</span><span class="p">(</span><span class="n">dims</span><span class="p">,</span> <span class="mi">1</span><span class="p">)),</span> <span class="n">make_tuple</span><span class="p">(</span><span class="n">shifts</span><span class="p">,</span> <span class="mi">1</span><span class="p">),</span> <span class="p">[</span><span class="nb">slice</span><span class="p">(</span><span class="kc">None</span><span class="p">)]</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">ndim</span>
|
|
<span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">dims</span><span class="p">)</span> <span class="o">!=</span> <span class="nb">len</span><span class="p">(</span><span class="n">shifts</span><span class="p">):</span> <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="sa">f</span><span class="s2">"</span><span class="si">{</span><span class="nb">len</span><span class="p">(</span><span class="n">dims</span><span class="p">)</span><span class="si">=}</span><span class="s2"> != </span><span class="si">{</span><span class="nb">len</span><span class="p">(</span><span class="n">shifts</span><span class="p">)</span><span class="si">=}</span><span class="s2">"</span><span class="p">)</span>
|
|
<span class="k">for</span> <span class="n">dim</span><span class="p">,</span> <span class="n">shift</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">dims</span><span class="p">,</span> <span class="n">shifts</span><span class="p">):</span> <span class="n">slices</span><span class="p">[</span><span class="n">dim</span><span class="p">]</span> <span class="o">=</span> <span class="nb">slice</span><span class="p">(</span><span class="n">delta</span><span class="o">:=</span><span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="n">dim</span><span class="p">]</span><span class="o">-</span><span class="n">shift</span><span class="o">%</span><span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="n">dim</span><span class="p">],</span> <span class="n">delta</span><span class="o">+</span><span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="n">dim</span><span class="p">])</span>
|
|
<span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">repeat</span><span class="p">(</span><span class="o">*</span><span class="nb">tuple</span><span class="p">(</span><span class="mi">2</span> <span class="k">if</span> <span class="n">i</span> <span class="ow">in</span> <span class="n">dims</span> <span class="k">else</span> <span class="mi">1</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">ndim</span><span class="p">)))[</span><span class="n">slices</span><span class="p">]</span>
|
|
</code></pre></div></td></tr></table></div>
|
|
</details>
|
|
</div>
|
|
|
|
</div>
|
|
|
|
<div class="doc doc-object doc-function">
|
|
|
|
|
|
<h3 id="tinygrad.Tensor.rearrange" class="doc doc-heading">
|
|
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <span class="doc doc-object-name doc-function-name">rearrange</span>
|
|
|
|
|
|
<a href="#tinygrad.Tensor.rearrange" class="headerlink" title="Permanent link">¤</a></h3>
|
|
<div class="language-python doc-signature highlight"><pre><span></span><code><span class="nf">rearrange</span><span class="p">(</span><span class="n">formula</span><span class="p">:</span> <span class="n"><a class="autorefs autorefs-external" href="https://docs.python.org/3/library/stdtypes.html#str">str</a></span><span class="p">,</span> <span class="o">**</span><span class="n">sizes</span><span class="p">)</span> <span class="o">-></span> <span class="n"><a class="autorefs autorefs-external" title="<code>typing.Self</code>" href="https://docs.python.org/3/library/typing.html#typing.Self">Self</a></span>
|
|
</code></pre></div>
|
|
|
|
<div class="doc doc-contents first">
|
|
|
|
<p>Rearranges input according to formula</p>
|
|
<p>See: <a href="https://einops.rocks/api/rearrange/">https://einops.rocks/api/rearrange/</a></p>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="n">x</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="p">[</span><span class="mi">3</span><span class="p">,</span> <span class="mi">4</span><span class="p">]])</span>
|
|
<span class="nb">print</span><span class="p">(</span><span class="n">Tensor</span><span class="o">.</span><span class="n">rearrange</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="s2">"batch channel -> (batch channel)"</span><span class="p">)</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
|
|
</code></pre></div>
|
|
<div class="language-python highlight"><pre><span></span><code><span class="p">[</span><span class="mi">1</span> <span class="mi">2</span> <span class="mi">3</span> <span class="mi">4</span><span class="p">]</span>
|
|
</code></pre></div>
|
|
|
|
|
|
<details class="mkdocstrings-source">
|
|
<summary>Source code in <code>tinygrad/mixin/movement.py</code></summary>
|
|
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal">272</span>
|
|
<span class="normal">273</span>
|
|
<span class="normal">274</span>
|
|
<span class="normal">275</span>
|
|
<span class="normal">276</span>
|
|
<span class="normal">277</span>
|
|
<span class="normal">278</span>
|
|
<span class="normal">279</span>
|
|
<span class="normal">280</span>
|
|
<span class="normal">281</span>
|
|
<span class="normal">282</span>
|
|
<span class="normal">283</span>
|
|
<span class="normal">284</span>
|
|
<span class="normal">285</span>
|
|
<span class="normal">286</span>
|
|
<span class="normal">287</span>
|
|
<span class="normal">288</span>
|
|
<span class="normal">289</span>
|
|
<span class="normal">290</span>
|
|
<span class="normal">291</span>
|
|
<span class="normal">292</span>
|
|
<span class="normal">293</span>
|
|
<span class="normal">294</span>
|
|
<span class="normal">295</span>
|
|
<span class="normal">296</span>
|
|
<span class="normal">297</span>
|
|
<span class="normal">298</span>
|
|
<span class="normal">299</span>
|
|
<span class="normal">300</span>
|
|
<span class="normal">301</span>
|
|
<span class="normal">302</span>
|
|
<span class="normal">303</span>
|
|
<span class="normal">304</span>
|
|
<span class="normal">305</span>
|
|
<span class="normal">306</span>
|
|
<span class="normal">307</span>
|
|
<span class="normal">308</span>
|
|
<span class="normal">309</span>
|
|
<span class="normal">310</span>
|
|
<span class="normal">311</span>
|
|
<span class="normal">312</span>
|
|
<span class="normal">313</span>
|
|
<span class="normal">314</span>
|
|
<span class="normal">315</span></pre></div></td><td class="code"><div><pre><span></span><code><span class="k">def</span><span class="w"> </span><span class="nf">rearrange</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">formula</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="o">**</span><span class="n">sizes</span><span class="p">)</span> <span class="o">-></span> <span class="n">Self</span><span class="p">:</span>
|
|
<span class="w"> </span><span class="sd">"""</span>
|
|
<span class="sd"> Rearranges input according to formula</span>
|
|
|
|
<span class="sd"> See: https://einops.rocks/api/rearrange/</span>
|
|
|
|
<span class="sd"> ```python exec="true" source="above" session="tensor" result="python"</span>
|
|
<span class="sd"> x = Tensor([[1, 2], [3, 4]])</span>
|
|
<span class="sd"> print(Tensor.rearrange(x, "batch channel -> (batch channel)").numpy())</span>
|
|
<span class="sd"> ```</span>
|
|
<span class="sd"> """</span>
|
|
|
|
<span class="k">def</span><span class="w"> </span><span class="nf">parse_side</span><span class="p">(</span><span class="n">s</span><span class="p">:</span> <span class="nb">str</span><span class="p">)</span> <span class="o">-></span> <span class="nb">tuple</span><span class="p">[</span><span class="nb">list</span><span class="p">[</span><span class="nb">str</span><span class="p">],</span> <span class="nb">list</span><span class="p">[</span><span class="nb">tuple</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">int</span><span class="p">]]]:</span>
|
|
<span class="w"> </span><span class="sd">"""Parse one side of formula into (axis_names, dims) where dims are (start, end) index pairs for parens."""</span>
|
|
<span class="n">tokens</span> <span class="o">=</span> <span class="sa">f</span><span class="s2">" </span><span class="si">{</span><span class="n">s</span><span class="si">}</span><span class="s2"> "</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="s2">"…"</span><span class="p">,</span> <span class="s2">"..."</span><span class="p">)</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="s2">"("</span><span class="p">,</span> <span class="s2">" ( "</span><span class="p">)</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="s2">")"</span><span class="p">,</span> <span class="s2">" ) "</span><span class="p">)</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="s2">" "</span><span class="p">,</span> <span class="s2">" "</span><span class="p">)</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="s2">" 1 "</span><span class="p">,</span> <span class="s2">" ( ) "</span><span class="p">)</span><span class="o">.</span><span class="n">split</span><span class="p">()</span>
|
|
<span class="n">lparens</span><span class="p">,</span> <span class="n">rparens</span> <span class="o">=</span> <span class="p">[</span><span class="n">i</span> <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">tok</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">tokens</span><span class="p">)</span> <span class="k">if</span> <span class="n">tok</span> <span class="o">==</span> <span class="s2">"("</span><span class="p">],</span> <span class="p">[</span><span class="n">i</span> <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">tok</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">tokens</span><span class="p">)</span> <span class="k">if</span> <span class="n">tok</span> <span class="o">==</span> <span class="s2">")"</span><span class="p">]</span>
|
|
<span class="n">pairs</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="nb">zip</span><span class="p">(</span><span class="n">lparens</span><span class="p">,</span> <span class="n">rparens</span><span class="p">))</span>
|
|
<span class="k">assert</span> <span class="nb">len</span><span class="p">(</span><span class="n">lparens</span><span class="p">)</span> <span class="o">==</span> <span class="nb">len</span><span class="p">(</span><span class="n">rparens</span><span class="p">)</span> <span class="ow">and</span> <span class="nb">sorted</span><span class="p">(</span><span class="n">flatten</span><span class="p">(</span><span class="n">pairs</span><span class="p">))</span> <span class="o">==</span> <span class="n">flatten</span><span class="p">(</span><span class="n">pairs</span><span class="p">),</span> <span class="s2">"bracket mismatch"</span>
|
|
<span class="k">return</span> <span class="p">[</span><span class="n">tok</span> <span class="k">for</span> <span class="n">tok</span> <span class="ow">in</span> <span class="n">tokens</span> <span class="k">if</span> <span class="n">tok</span> <span class="ow">not</span> <span class="ow">in</span> <span class="p">(</span><span class="s2">"("</span><span class="p">,</span> <span class="s2">")"</span><span class="p">)],</span> <span class="p">[(</span><span class="n">lp</span> <span class="o">-</span> <span class="mi">2</span><span class="o">*</span><span class="n">i</span><span class="p">,</span> <span class="n">rp</span> <span class="o">-</span> <span class="mi">1</span> <span class="o">-</span> <span class="mi">2</span><span class="o">*</span><span class="n">i</span><span class="p">)</span> <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="p">(</span><span class="n">lp</span><span class="p">,</span> <span class="n">rp</span><span class="p">)</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">pairs</span><span class="p">)]</span>
|
|
|
|
<span class="k">assert</span> <span class="n">formula</span><span class="o">.</span><span class="n">count</span><span class="p">(</span><span class="s2">"->"</span><span class="p">)</span> <span class="o">==</span> <span class="mi">1</span><span class="p">,</span> <span class="s1">'need exactly one "->" in formula'</span>
|
|
<span class="p">(</span><span class="n">lhs</span><span class="p">,</span> <span class="n">unflatten_dims</span><span class="p">),</span> <span class="p">(</span><span class="n">rhs</span><span class="p">,</span> <span class="n">flatten_dims</span><span class="p">)</span> <span class="o">=</span> <span class="nb">map</span><span class="p">(</span><span class="n">parse_side</span><span class="p">,</span> <span class="n">formula</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="s2">"->"</span><span class="p">))</span>
|
|
|
|
<span class="k">for</span> <span class="n">name</span> <span class="ow">in</span> <span class="n">sizes</span><span class="p">:</span> <span class="k">assert</span> <span class="n">name</span> <span class="ow">in</span> <span class="n">lhs</span><span class="p">,</span> <span class="sa">f</span><span class="s2">"axis </span><span class="si">{</span><span class="n">name</span><span class="si">}</span><span class="s2"> is not used in transform"</span>
|
|
<span class="k">assert</span> <span class="nb">sorted</span><span class="p">(</span><span class="n">lhs</span><span class="p">)</span> <span class="o">==</span> <span class="nb">sorted</span><span class="p">(</span><span class="n">rhs</span><span class="p">)</span> <span class="ow">and</span> <span class="nb">len</span><span class="p">(</span><span class="n">lhs</span><span class="p">)</span> <span class="o">==</span> <span class="nb">len</span><span class="p">(</span><span class="nb">set</span><span class="p">(</span><span class="n">lhs</span><span class="p">)),</span> <span class="sa">f</span><span class="s2">"name mismatch in </span><span class="si">{</span><span class="n">formula</span><span class="si">}</span><span class="s2">"</span>
|
|
<span class="k">for</span> <span class="n">name</span> <span class="ow">in</span> <span class="n">lhs</span><span class="o">+</span><span class="n">rhs</span><span class="p">:</span> <span class="k">assert</span> <span class="n">name</span> <span class="o">==</span> <span class="s2">"..."</span> <span class="ow">or</span> <span class="p">(</span><span class="n">name</span><span class="o">.</span><span class="n">isidentifier</span><span class="p">()</span> <span class="ow">and</span> <span class="s2">"_"</span> <span class="ow">not</span> <span class="ow">in</span> <span class="p">(</span><span class="n">name</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">name</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">])),</span> <span class="sa">f</span><span class="s2">"invalid axis name </span><span class="si">{</span><span class="n">name</span><span class="si">}</span><span class="s2">"</span>
|
|
<span class="k">assert</span> <span class="s2">"..."</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">flatten</span><span class="p">([</span><span class="n">lhs</span><span class="p">[</span><span class="n">s</span><span class="p">:</span><span class="n">e</span><span class="p">]</span> <span class="k">for</span> <span class="n">s</span><span class="p">,</span> <span class="n">e</span> <span class="ow">in</span> <span class="n">unflatten_dims</span><span class="p">]),</span> <span class="sa">f</span><span class="s2">"cannot have collapsed ellipsis (...) in lhs of </span><span class="si">{</span><span class="n">formula</span><span class="si">}</span><span class="s2">"</span>
|
|
<span class="k">assert</span> <span class="n">lhs</span><span class="o">.</span><span class="n">count</span><span class="p">(</span><span class="s2">"..."</span><span class="p">)</span> <span class="o"><=</span> <span class="mi">1</span><span class="p">,</span> <span class="sa">f</span><span class="s2">"too many ellipses in </span><span class="si">{</span><span class="n">formula</span><span class="si">}</span><span class="s2">"</span>
|
|
|
|
<span class="c1"># resolve ellipsis</span>
|
|
<span class="k">if</span> <span class="s2">"..."</span> <span class="ow">in</span> <span class="n">lhs</span><span class="p">:</span>
|
|
<span class="n">ell_len</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span> <span class="o">-</span> <span class="nb">len</span><span class="p">(</span><span class="n">lhs</span><span class="p">)</span> <span class="o">+</span> <span class="mi">1</span> <span class="o">+</span> <span class="nb">sum</span><span class="p">(</span><span class="n">e</span> <span class="o">-</span> <span class="n">s</span> <span class="o">-</span> <span class="mi">1</span> <span class="k">for</span> <span class="n">s</span><span class="p">,</span> <span class="n">e</span> <span class="ow">in</span> <span class="n">unflatten_dims</span><span class="p">)</span>
|
|
<span class="n">lhs</span><span class="p">,</span> <span class="n">rhs</span> <span class="o">=</span> <span class="nb">map</span><span class="p">(</span><span class="k">lambda</span> <span class="n">l</span><span class="p">:</span> <span class="n">l</span><span class="p">[:(</span><span class="n">i</span> <span class="o">:=</span> <span class="n">l</span><span class="o">.</span><span class="n">index</span><span class="p">(</span><span class="s2">"..."</span><span class="p">))]</span> <span class="o">+</span> <span class="p">[</span><span class="sa">f</span><span class="s2">"...</span><span class="si">{</span><span class="n">j</span><span class="si">}</span><span class="s2">"</span> <span class="k">for</span> <span class="n">j</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">ell_len</span><span class="p">)]</span> <span class="o">+</span> <span class="n">l</span><span class="p">[</span><span class="n">i</span> <span class="o">+</span> <span class="mi">1</span><span class="p">:]</span> <span class="k">if</span> <span class="s2">"..."</span> <span class="ow">in</span> <span class="n">l</span> <span class="k">else</span> <span class="n">l</span><span class="p">,</span> <span class="p">(</span><span class="n">lhs</span><span class="p">,</span> <span class="n">rhs</span><span class="p">))</span>
|
|
<span class="k">def</span><span class="w"> </span><span class="nf">newdims</span><span class="p">(</span><span class="n">side</span><span class="p">,</span> <span class="n">s</span><span class="p">,</span> <span class="n">e</span><span class="p">):</span> <span class="k">return</span> <span class="p">(</span><span class="n">s</span> <span class="o">+</span> <span class="p">(</span><span class="n">ell_len</span> <span class="o">-</span> <span class="mi">1</span> <span class="k">if</span> <span class="s2">"...0"</span> <span class="ow">in</span> <span class="n">side</span><span class="p">[:</span><span class="n">s</span><span class="p">]</span> <span class="k">else</span> <span class="mi">0</span><span class="p">),</span> <span class="n">e</span> <span class="o">+</span> <span class="p">(</span><span class="n">ell_len</span> <span class="o">-</span> <span class="mi">1</span> <span class="k">if</span> <span class="s2">"...0"</span> <span class="ow">in</span> <span class="n">side</span><span class="p">[:</span><span class="n">e</span><span class="p">]</span> <span class="k">else</span> <span class="mi">0</span><span class="p">))</span>
|
|
<span class="n">unflatten_dims</span><span class="p">,</span> <span class="n">flatten_dims</span> <span class="o">=</span> <span class="p">[</span><span class="n">newdims</span><span class="p">(</span><span class="n">lhs</span><span class="p">,</span> <span class="n">s</span><span class="p">,</span> <span class="n">e</span><span class="p">)</span> <span class="k">for</span> <span class="n">s</span><span class="p">,</span> <span class="n">e</span> <span class="ow">in</span> <span class="n">unflatten_dims</span><span class="p">],</span> <span class="p">[</span><span class="n">newdims</span><span class="p">(</span><span class="n">rhs</span><span class="p">,</span> <span class="n">s</span><span class="p">,</span> <span class="n">e</span><span class="p">)</span> <span class="k">for</span> <span class="n">s</span><span class="p">,</span> <span class="n">e</span> <span class="ow">in</span> <span class="n">flatten_dims</span><span class="p">]</span>
|
|
|
|
<span class="c1"># unflatten -> permute -> flatten</span>
|
|
<span class="n">t</span> <span class="o">=</span> <span class="bp">self</span>
|
|
<span class="k">for</span> <span class="n">start</span><span class="p">,</span> <span class="n">end</span> <span class="ow">in</span> <span class="n">unflatten_dims</span><span class="p">:</span> <span class="n">t</span> <span class="o">=</span> <span class="n">t</span><span class="o">.</span><span class="n">unflatten</span><span class="p">(</span><span class="n">start</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">(</span><span class="n">sizes</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">lhs</span><span class="p">[</span><span class="n">i</span><span class="p">],</span> <span class="o">-</span><span class="mi">1</span><span class="p">)</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">start</span><span class="p">,</span> <span class="n">end</span><span class="p">)))</span>
|
|
<span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">name</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">lhs</span><span class="p">):</span>
|
|
<span class="k">if</span> <span class="n">name</span> <span class="ow">in</span> <span class="n">sizes</span><span class="p">:</span> <span class="k">assert</span> <span class="n">sizes</span><span class="p">[</span><span class="n">name</span><span class="p">]</span> <span class="o">==</span> <span class="n">t</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="n">i</span><span class="p">],</span> <span class="sa">f</span><span class="s2">"size provided for dimension </span><span class="si">{</span><span class="n">name</span><span class="si">}</span><span class="s2"> incorrect"</span>
|
|
<span class="n">t</span> <span class="o">=</span> <span class="n">t</span><span class="o">.</span><span class="n">permute</span><span class="p">([</span><span class="n">lhs</span><span class="o">.</span><span class="n">index</span><span class="p">(</span><span class="n">name</span><span class="p">)</span> <span class="k">for</span> <span class="n">name</span> <span class="ow">in</span> <span class="n">rhs</span><span class="p">])</span>
|
|
<span class="k">for</span> <span class="n">start</span><span class="p">,</span> <span class="n">end</span> <span class="ow">in</span> <span class="nb">reversed</span><span class="p">(</span><span class="n">flatten_dims</span><span class="p">):</span> <span class="n">t</span> <span class="o">=</span> <span class="n">t</span><span class="o">.</span><span class="n">flatten</span><span class="p">(</span><span class="n">start</span><span class="p">,</span> <span class="n">end</span> <span class="o">-</span> <span class="mi">1</span><span class="p">)</span> <span class="k">if</span> <span class="n">start</span> <span class="o"><</span> <span class="n">end</span> <span class="k">else</span> <span class="n">t</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="n">start</span><span class="p">)</span>
|
|
<span class="k">return</span> <span class="n">t</span>
|
|
</code></pre></div></td></tr></table></div>
|
|
</details>
|
|
</div>
|
|
|
|
</div>
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
</article>
|
|
</div>
|
|
|
|
|
|
<script>var target=document.getElementById(location.hash.slice(1));target&&target.name&&(target.checked=target.name.startsWith("__tabbed_"))</script>
|
|
</div>
|
|
|
|
<button type="button" class="md-top md-icon" data-md-component="top" hidden>
|
|
|
|
<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>
|
|
Back to top
|
|
</button>
|
|
|
|
</main>
|
|
|
|
<footer class="md-footer">
|
|
|
|
|
|
|
|
<nav class="md-footer__inner md-grid" aria-label="Footer" >
|
|
|
|
|
|
<a href="../creation/" class="md-footer__link md-footer__link--prev" aria-label="Previous: Creation">
|
|
<div class="md-footer__button md-icon">
|
|
|
|
<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>
|
|
</div>
|
|
<div class="md-footer__title">
|
|
<span class="md-footer__direction">
|
|
Previous
|
|
</span>
|
|
<div class="md-ellipsis">
|
|
Creation
|
|
</div>
|
|
</div>
|
|
</a>
|
|
|
|
|
|
|
|
<a href="../elementwise/" class="md-footer__link md-footer__link--next" aria-label="Next: Elementwise">
|
|
<div class="md-footer__title">
|
|
<span class="md-footer__direction">
|
|
Next
|
|
</span>
|
|
<div class="md-ellipsis">
|
|
Elementwise
|
|
</div>
|
|
</div>
|
|
<div class="md-footer__button md-icon">
|
|
|
|
<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>
|
|
</div>
|
|
</a>
|
|
|
|
</nav>
|
|
|
|
|
|
<div class="md-footer-meta md-typeset">
|
|
<div class="md-footer-meta__inner md-grid">
|
|
<div class="md-copyright">
|
|
|
|
|
|
Made with
|
|
<a href="https://squidfunk.github.io/mkdocs-material/" target="_blank" rel="noopener">
|
|
Material for MkDocs
|
|
</a>
|
|
|
|
</div>
|
|
|
|
</div>
|
|
</div>
|
|
</footer>
|
|
|
|
</div>
|
|
<div class="md-dialog" data-md-component="dialog">
|
|
<div class="md-dialog__inner md-typeset"></div>
|
|
</div>
|
|
|
|
|
|
|
|
|
|
|
|
<script id="__config" type="application/json">{"annotate": null, "base": "../..", "features": ["announce.dismiss", "content.action.edit", "content.action.view", "content.code.annotate", "content.code.copy", "content.tooltips", "navigation.footer", "navigation.indexes", "navigation.sections", "navigation.expand", "navigation.top", "navigation.path", "search.highlight", "search.suggest", "toc.follow", "toc.integrate"], "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>
|
|
|
|
|
|
<script src="../../assets/javascripts/bundle.79ae519e.min.js"></script>
|
|
|
|
<script src="../../assets/_markdown_exec_pyodide.js"></script>
|
|
|
|
|
|
</body>
|
|
</html> |